diff --git a/src/Python/flask/web/app.py b/src/Python/flask/web/app.py index 334e643..9507142 100644 --- a/src/Python/flask/web/app.py +++ b/src/Python/flask/web/app.py @@ -11,34 +11,36 @@ model = YOLO("yolo11n.pt") # pretrained YOLO11n model kobuki_message = "" latest_image = None +processed_image = None yolo_results = [] -# https://medium.com/@Mert.A/how-to-segment-objects-with-yolov11-68593eb49fa8 +# List of class names (example for COCO dataset) yolo_classes = list(model.names.values()) def on_message(client, userdata, message): - global kobuki_message, latest_image, yolo_results + global kobuki_message, latest_image, processed_image, yolo_results if message.topic == "kobuki/data": kobuki_message = str(message.payload.decode("utf-8")) elif message.topic == "kobuki/cam": latest_image = np.frombuffer(message.payload, np.uint8) latest_image = cv2.imdecode(latest_image, cv2.IMREAD_COLOR) + processed_image = latest_image.copy() # Create a copy for processing # Process the image with YOLO results = model(latest_image) yolo_results = [] for result in results: for box in result.boxes: - class_id = int(box.cls.item()) # Convert to integer + class_id = int(box.cls.item()) class_name = yolo_classes[class_id] yolo_results.append({ "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]) - cv2.rectangle(latest_image, (x1, y1), (x2, y2), (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) + # Draw bounding box on the processed image + x1, y1, x2, y2 = map(int, box.xyxy) + cv2.rectangle(processed_image, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText(processed_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() @@ -48,20 +50,21 @@ mqtt_client.loop_start() mqtt_client.subscribe("kobuki/data") mqtt_client.subscribe("kobuki/cam") -mqtt_client.on_message = on_message # this lines needs to be under the function definition otherwise it cant find which function it needs to use +mqtt_client.on_message = on_message # this line needs to be under the function definition otherwise it can't find which function it needs to use @app.route('/') def index(): return render_template('index.html') -@app.route('/control', methods=["GET","POST"]) +@app.route('/control', methods=["GET", "POST"]) def control(): if request.authorization and request.authorization.username == 'ishak' and request.authorization.password == 'kobuki': return render_template('control.html') else: return ('Unauthorized', 401, {'WWW-Authenticate': 'Basic realm="Login Required"'}) + @app.route('/move', methods=['POST']) def move(): data = request.get_json() @@ -78,19 +81,22 @@ def move(): def data(): return kobuki_message + @app.route('/image') def image(): - global latest_image - if latest_image is not None: - _, buffer = cv2.imencode('.jpg', latest_image) + global processed_image + if processed_image is not None: + _, buffer = cv2.imencode('.jpg', processed_image) return Response(buffer.tobytes(), mimetype='image/jpeg') else: return "No image available", 404 + @app.route('/yolo_results', methods=['GET']) def yolo_results_endpoint(): global yolo_results return jsonify(yolo_results) + if __name__ == '__main__': app.run(debug=True, port=5000) \ No newline at end of file