recognizing garbage

This commit is contained in:
Pawel Felcyn 2023-05-29 09:57:52 +02:00
parent 26bcc9ddf1
commit 43c13acdee
6 changed files with 172 additions and 23 deletions

View File

@ -32,9 +32,10 @@ class Garbage:
self.does_din = does_din self.does_din = does_din
class RecognizedGarbage(Garbage): class RecognizedGarbage:
garbage_type: GarbageType garbage_type: GarbageType
src: Garbage
def __init__(self, src: Garbage, garbage_type: GarbageType) -> None: def __init__(self, src: Garbage, garbage_type: GarbageType) -> None:
super().__init__(src.img)
self.garbage_type = garbage_type self.garbage_type = garbage_type
self.src = src

View File

@ -1,6 +1,6 @@
from typing import List, Tuple from typing import List, Tuple
from agentOrientation import AgentOrientation from agentOrientation import AgentOrientation
from garbage import RecognizedGarbage from garbage import GarbageType, RecognizedGarbage
from gameContext import GameContext from gameContext import GameContext
class GarbageTruck: class GarbageTruck:
@ -22,19 +22,19 @@ class GarbageTruck:
def sort_garbage(self, RecognizedGarbage) -> None: def sort_garbage(self, RecognizedGarbage) -> None:
if RecognizedGarbage.garbage_type == 0: if RecognizedGarbage.garbage_type == GarbageType.PAPER:
self.paper.append(RecognizedGarbage) self.paper.append(RecognizedGarbage)
elif RecognizedGarbage.garbage_type == 1: elif RecognizedGarbage.garbage_type == GarbageType.PLASTIC_AND_METAL:
self.plastic_and_metal.append(RecognizedGarbage) self.plastic_and_metal.append(RecognizedGarbage)
elif RecognizedGarbage.garbage_type == 3: elif RecognizedGarbage.garbage_type == GarbageType.GLASS:
self.glass.append(RecognizedGarbage) self.glass.append(RecognizedGarbage)
elif RecognizedGarbage.garbage_type == 4: elif RecognizedGarbage.garbage_type == GarbageType.BIO:
self.bio.append(RecognizedGarbage) self.bio.append(RecognizedGarbage)
elif RecognizedGarbage.garbage_type == 5: elif RecognizedGarbage.garbage_type == GarbageType.MIXED:
self.mixed.append(RecognizedGarbage) self.mixed.append(RecognizedGarbage)
def render(self, game_context: GameContext) -> None: def render(self, game_context: GameContext) -> None:

View File

@ -21,29 +21,90 @@ def _read_training_data() -> TrainingData:
classes.append(line_class) classes.append(line_class)
return TrainingData(attributes, classes) 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() trainning_data = _read_training_data()
X = trainning_data.attributes X = trainning_data.attributes
Y = trainning_data.classes Y = trainning_data.classes
le_shape = LabelEncoder() # le_shape = LabelEncoder()
le_flexibility = LabelEncoder() # le_flexibility = LabelEncoder()
le_color = LabelEncoder() # le_color = LabelEncoder()
le_shape.fit([x[0] for x in X]) # le_shape.fit([x[0] for x in X])
le_flexibility.fit([x[3] for x in X]) # le_flexibility.fit([x[3] for x in X])
le_color.fit([x[4] for x in X]) # le_color.fit([x[4] for x in X])
X_encoded = np.array([ # 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]] # [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 # for x in X
]) # ])
encoder = OneHotEncoder(categories='auto', sparse=False) # encoder = OneHotEncoder(categories='auto', sparse=False)
X_encoded = encoder.fit_transform(X_encoded) # X_encoded = encoder.fit_transform(X_encoded)
model = tree.DecisionTreeClassifier() 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') joblib.dump(model, 'model.pkl')

Binary file not shown.

BIN
model.pkl

Binary file not shown.

View File

@ -1,5 +1,9 @@
import joblib
from sklearn.calibration import LabelEncoder
from agentActionType import AgentActionType from agentActionType import AgentActionType
import time import time
from garbage import GarbageType, RecognizedGarbage
from garbageCan import GarbageCan
from turnCar import turn_left_orientation, turn_right_orientation from turnCar import turn_left_orientation, turn_right_orientation
from garbageTruck import GarbageTruck from garbageTruck import GarbageTruck
from typing import Tuple, Dict from typing import Tuple, Dict
@ -20,9 +24,92 @@ def collect_garbage(game_context: GameContext) -> None:
move_dust_car(path, game_context) move_dust_car(path, game_context)
next_position = calculate_next_position(game_context.dust_car) next_position = calculate_next_position(game_context.dust_car)
game_context.grid[next_position] = GridCellType.VISITED_GARBAGE_CAN 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 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: def move_dust_car(actions: list[AgentActionType], game_context: GameContext) -> None:
for action in actions: for action in actions: