diff --git a/garbage.py b/garbage.py index 205e00c..fe25723 100644 --- a/garbage.py +++ b/garbage.py @@ -32,9 +32,10 @@ class Garbage: self.does_din = does_din -class RecognizedGarbage(Garbage): +class RecognizedGarbage: garbage_type: GarbageType + src: Garbage def __init__(self, src: Garbage, garbage_type: GarbageType) -> None: - super().__init__(src.img) self.garbage_type = garbage_type + self.src = src diff --git a/garbageTruck.py b/garbageTruck.py index 5584adc..94f9cd9 100644 --- a/garbageTruck.py +++ b/garbageTruck.py @@ -1,6 +1,6 @@ from typing import List, Tuple from agentOrientation import AgentOrientation -from garbage import RecognizedGarbage +from garbage import GarbageType, RecognizedGarbage from gameContext import GameContext class GarbageTruck: @@ -22,19 +22,19 @@ class GarbageTruck: def sort_garbage(self, RecognizedGarbage) -> None: - if RecognizedGarbage.garbage_type == 0: + if RecognizedGarbage.garbage_type == GarbageType.PAPER: self.paper.append(RecognizedGarbage) - elif RecognizedGarbage.garbage_type == 1: + elif RecognizedGarbage.garbage_type == GarbageType.PLASTIC_AND_METAL: self.plastic_and_metal.append(RecognizedGarbage) - elif RecognizedGarbage.garbage_type == 3: + elif RecognizedGarbage.garbage_type == GarbageType.GLASS: self.glass.append(RecognizedGarbage) - elif RecognizedGarbage.garbage_type == 4: + elif RecognizedGarbage.garbage_type == GarbageType.BIO: self.bio.append(RecognizedGarbage) - elif RecognizedGarbage.garbage_type == 5: + elif RecognizedGarbage.garbage_type == GarbageType.MIXED: self.mixed.append(RecognizedGarbage) def render(self, game_context: GameContext) -> None: diff --git a/machine_learning/decisionTree.py b/machine_learning/decisionTree.py index 68a3950..a2fb6d3 100644 --- a/machine_learning/decisionTree.py +++ b/machine_learning/decisionTree.py @@ -18,32 +18,78 @@ def _read_training_data() -> TrainingData: line_attributes = values[:-1] line_class = values[-1] attributes.append(line_attributes) - classes.append(line_class) + classes.append(line_class.strip()) return TrainingData(attributes, classes) +def _attributes_to_floats(attributes: list[str]) -> list[float]: + output: list[float] = [] + if attributes[0] == 'Longitiudonal': + output.append(0) + elif attributes[0] == 'Round': + output.append(1) + elif attributes[0] == 'Flat': + output.append(2) + elif attributes[0] == 'Irregular': + output.append(3) + + if attributes[1] == 'Low': + output.append(0) + elif attributes[1] == 'Medium': + output.append(1) + elif attributes[1] == 'High': + output.append(2) + + + if attributes[2] == "Yes": + output.append(0) + else: + output.append(1) + + if attributes[3] == 'Low': + output.append(0) + elif attributes[3] == 'Medium': + output.append(1) + elif attributes[3] == 'High': + output.append(2) + + if attributes[4] == 'Low': + output.append(0) + elif attributes[4] == 'Medium': + output.append(1) + elif attributes[4] == 'High': + output.append(2) + + if attributes[5] == 'Transparent': + output.append(0) + elif attributes[5] == 'Light': + output.append(1) + elif attributes[5] == 'Dark': + output.append(2) + elif attributes[5] == "Colorful": + output.append(3) + + if attributes[6] == 'Low': + output.append(0) + elif attributes[6] == 'Medium': + output.append(1) + elif attributes[6] == 'High': + output.append(2) + + if attributes[7] == "Yes": + output.append(0) + else: + output.append(1) + return output + + trainning_data = _read_training_data() X = trainning_data.attributes Y = trainning_data.classes -le_shape = LabelEncoder() -le_flexibility = LabelEncoder() -le_color = LabelEncoder() - -le_shape.fit([x[0] for x in X]) -le_flexibility.fit([x[3] for x in X]) -le_color.fit([x[4] for x in X]) - -X_encoded = np.array([ - [le_shape.transform([x[0]])[0], x[1], x[2], le_flexibility.transform([x[3]])[0], le_color.transform([x[4]])[0]] - for x in X -]) - -encoder = OneHotEncoder(categories='auto', sparse=False) -X_encoded = encoder.fit_transform(X_encoded) model = tree.DecisionTreeClassifier() -model.fit(X_encoded, Y) - +encoded = [_attributes_to_floats(x) for x in X] +dtc = model.fit(encoded, Y) joblib.dump(model, 'model.pkl') \ No newline at end of file diff --git a/machine_learning/model.pkl b/machine_learning/model.pkl index f79ca17..457f45c 100644 Binary files a/machine_learning/model.pkl and b/machine_learning/model.pkl differ diff --git a/model.pkl b/model.pkl deleted file mode 100644 index 062ea22..0000000 Binary files a/model.pkl and /dev/null differ diff --git a/movement.py b/movement.py index 0060c0e..e726f90 100644 --- a/movement.py +++ b/movement.py @@ -1,5 +1,9 @@ +import joblib +from sklearn.calibration import LabelEncoder from agentActionType import AgentActionType import time +from garbage import GarbageType, RecognizedGarbage +from garbageCan import GarbageCan from turnCar import turn_left_orientation, turn_right_orientation from garbageTruck import GarbageTruck from typing import Tuple, Dict @@ -20,9 +24,92 @@ def collect_garbage(game_context: GameContext) -> None: move_dust_car(path, game_context) next_position = calculate_next_position(game_context.dust_car) game_context.grid[next_position] = GridCellType.VISITED_GARBAGE_CAN - game_context.city.cans_dict[next_position].is_visited = True + can = game_context.city.cans_dict[next_position] + can.is_visited = True + _recognize_garbage(game_context.dust_car, can) pass +def _recognize_garbage(dust_car: GarbageTruck, can: GarbageCan) -> None: + loaded_model = joblib.load('machine_learning/model.pkl') + for garbage in can.garbage: + attributes = [garbage.shape, garbage.flexibility, garbage.does_smell, garbage.weight, garbage.size, garbage.color, garbage.softness, garbage.does_din] + encoded = attributes_to_floats(attributes) + predicted_class = loaded_model.predict([encoded])[0] + garbage_type: GarbageType = None + if predicted_class == 'PAPER': + garbage_type = GarbageType.PAPER + elif predicted_class == 'PLASTIC_AND_METAL': + garbage_type = GarbageType.PLASTIC_AND_METAL + elif garbage_type == 'GLASS': + garbage_type = GarbageType.GLASS + elif predicted_class == 'BIO' : + garbage_type = GarbageType.BIO + elif predicted_class == 'MIXED': + garbage_type = GarbageType.MIXED + print(predicted_class) + recognized_garbage = RecognizedGarbage(garbage, garbage_type) + dust_car.sort_garbage(recognized_garbage) + +def attributes_to_floats(attributes: list[str]) -> list[float]: + output: list[float] = [] + if attributes[0] == 'Longitiudonal': + output.append(0) + elif attributes[0] == 'Round': + output.append(1) + elif attributes[0] == 'Flat': + output.append(2) + elif attributes[0] == 'Irregular': + output.append(3) + + if attributes[1] == 'Low': + output.append(0) + elif attributes[1] == 'Medium': + output.append(1) + elif attributes[1] == 'High': + output.append(2) + + + if attributes[2] == "Yes": + output.append(0) + else: + output.append(1) + + if attributes[3] == 'Low': + output.append(0) + elif attributes[3] == 'Medium': + output.append(1) + elif attributes[3] == 'High': + output.append(2) + + if attributes[4] == 'Low': + output.append(0) + elif attributes[4] == 'Medium': + output.append(1) + elif attributes[4] == 'High': + output.append(2) + + if attributes[5] == 'Transparent': + output.append(0) + elif attributes[5] == 'Light': + output.append(1) + elif attributes[5] == 'Dark': + output.append(2) + elif attributes[5] == "Colorful": + output.append(3) + + if attributes[6] == 'Low': + output.append(0) + elif attributes[6] == 'Medium': + output.append(1) + elif attributes[6] == 'High': + output.append(2) + + if attributes[7] == "Yes": + output.append(0) + else: + output.append(1) + return output + def move_dust_car(actions: list[AgentActionType], game_context: GameContext) -> None: for action in actions: diff --git a/startup.py b/startup.py index 0605095..6cc7222 100644 --- a/startup.py +++ b/startup.py @@ -57,7 +57,7 @@ def create_garbage_pieces() -> List[Garbage]: for line in lines[1:]: param = line.strip().split(',') garbage_pieces.append( - Garbage('img', param[0], param[1], param[2], param[3], param[4], param[5], param[6], param[7])) + Garbage('img', param[0], param[1], param[2], param[3], param[4], param[5], param[6], param[7].strip())) return garbage_pieces