attempt to fix broken code

This commit is contained in:
2025-01-06 16:02:17 +01:00
parent 7845feb9f8
commit 228c508012

View File

@@ -9,15 +9,13 @@ app = Flask(__name__)
# Load a model # Load a model
model = YOLO("yolo11n.pt") # pretrained YOLO11n model model = YOLO("yolo11n.pt") # pretrained YOLO11n model
#from: https://medium.com/@Mert.A/how-to-segment-objects-with-yolov11-68593eb49fa8
yolo_classes = list(model.names.values())
classes_ids = [yolo_classes.index(clas) for clas in yolo_classes]
kobuki_message = "" kobuki_message = ""
latest_image = None latest_image = None
yolo_results = [] yolo_results = []
# https://medium.com/@Mert.A/how-to-segment-objects-with-yolov11-68593eb49fa8
yolo_classes = list(model.names.values())
def on_message(client, userdata, message): def on_message(client, userdata, message):
global kobuki_message, latest_image, yolo_results global kobuki_message, latest_image, yolo_results
if message.topic == "kobuki/data": if message.topic == "kobuki/data":
@@ -30,15 +28,17 @@ def on_message(client, userdata, message):
yolo_results = [] yolo_results = []
for result in results: for result in results:
for box in result.boxes: for box in result.boxes:
class_id = int(box.cls.item()) # Convert to integer
class_name = yolo_classes[class_id]
yolo_results.append({ yolo_results.append({
"class": yolo_classes[box.cls.item()], "class": class_name,
"confidence": box.conf.item(), "confidence": box.conf.item(),
"bbox": box.xyxy.tolist() "bbox": box.xyxy.tolist()
}) })
# Draw bounding box on the image # Draw bounding box on the image
x1, y1, x2, y2 = map(int, box.xyxy[0]) x1, y1, x2, y2 = map(int, box.xyxy)
cv2.rectangle(latest_image, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.rectangle(latest_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(latest_image, f"{yolo_classes[box.cls.item()]} {box.conf.item():.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2) cv2.putText(latest_image, f"{class_name} {box.conf.item():.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
# Create an MQTT client instance # Create an MQTT client instance
mqtt_client = mqtt.Client() mqtt_client = mqtt.Client()