recognizing garbage
This commit is contained in:
Pawel Felcyn
parent
26bcc9ddf1
commit
43c13acdee
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -21,29 +21,90 @@ def _read_training_data() -> TrainingData:
|
|||
classes.append(line_class)
|
||||
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 = 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])
|
||||
# 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
|
||||
])
|
||||
# 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)
|
||||
# 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')
|
||||
Binary file not shown.
89
movement.py
89
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:
|
||||
|
|
|
|||
Loading…
Reference in New Issue