From 228c5080126f9002057b059b7c201932a5f07e81 Mon Sep 17 00:00:00 2001 From: Sam Hos Date: Mon, 6 Jan 2025 16:02:17 +0100 Subject: [PATCH] attempt to fix broken code --- src/Python/flask/web/app.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/Python/flask/web/app.py b/src/Python/flask/web/app.py index 1a66899..e9c414c 100644 --- a/src/Python/flask/web/app.py +++ b/src/Python/flask/web/app.py @@ -9,15 +9,13 @@ app = Flask(__name__) # Load a 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 = "" latest_image = None 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): global kobuki_message, latest_image, yolo_results if message.topic == "kobuki/data": @@ -30,15 +28,17 @@ def on_message(client, userdata, message): yolo_results = [] for result in results: for box in result.boxes: + class_id = int(box.cls.item()) # Convert to integer + class_name = yolo_classes[class_id] yolo_results.append({ - "class": yolo_classes[box.cls.item()], + "class": class_name, "confidence": box.conf.item(), "bbox": box.xyxy.tolist() }) # 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.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 mqtt_client = mqtt.Client()