store new image in processed_image

This commit is contained in:
2025-01-07 11:09:33 +01:00
parent 5e01e25d9c
commit cb988a5260

View File

@@ -11,34 +11,36 @@ model = YOLO("yolo11n.pt") # pretrained YOLO11n model
kobuki_message = "" kobuki_message = ""
latest_image = None latest_image = None
processed_image = None
yolo_results = [] 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()) 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, processed_image, yolo_results
if message.topic == "kobuki/data": if message.topic == "kobuki/data":
kobuki_message = str(message.payload.decode("utf-8")) kobuki_message = str(message.payload.decode("utf-8"))
elif message.topic == "kobuki/cam": elif message.topic == "kobuki/cam":
latest_image = np.frombuffer(message.payload, np.uint8) latest_image = np.frombuffer(message.payload, np.uint8)
latest_image = cv2.imdecode(latest_image, cv2.IMREAD_COLOR) latest_image = cv2.imdecode(latest_image, cv2.IMREAD_COLOR)
processed_image = latest_image.copy() # Create a copy for processing
# Process the image with YOLO # Process the image with YOLO
results = model(latest_image) results = model(latest_image)
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_id = int(box.cls.item())
class_name = yolo_classes[class_id] class_name = yolo_classes[class_id]
yolo_results.append({ yolo_results.append({
"class": class_name, "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 processed 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(processed_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) 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 # Create an MQTT client instance
mqtt_client = mqtt.Client() mqtt_client = mqtt.Client()
@@ -48,20 +50,21 @@ mqtt_client.loop_start()
mqtt_client.subscribe("kobuki/data") mqtt_client.subscribe("kobuki/data")
mqtt_client.subscribe("kobuki/cam") 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('/') @app.route('/')
def index(): def index():
return render_template('index.html') return render_template('index.html')
@app.route('/control', methods=["GET","POST"]) @app.route('/control', methods=["GET", "POST"])
def control(): def control():
if request.authorization and request.authorization.username == 'ishak' and request.authorization.password == 'kobuki': if request.authorization and request.authorization.username == 'ishak' and request.authorization.password == 'kobuki':
return render_template('control.html') return render_template('control.html')
else: else:
return ('Unauthorized', 401, {'WWW-Authenticate': 'Basic realm="Login Required"'}) return ('Unauthorized', 401, {'WWW-Authenticate': 'Basic realm="Login Required"'})
@app.route('/move', methods=['POST']) @app.route('/move', methods=['POST'])
def move(): def move():
data = request.get_json() data = request.get_json()
@@ -78,19 +81,22 @@ def move():
def data(): def data():
return kobuki_message return kobuki_message
@app.route('/image') @app.route('/image')
def image(): def image():
global latest_image global processed_image
if latest_image is not None: if processed_image is not None:
_, buffer = cv2.imencode('.jpg', latest_image) _, buffer = cv2.imencode('.jpg', processed_image)
return Response(buffer.tobytes(), mimetype='image/jpeg') return Response(buffer.tobytes(), mimetype='image/jpeg')
else: else:
return "No image available", 404 return "No image available", 404
@app.route('/yolo_results', methods=['GET']) @app.route('/yolo_results', methods=['GET'])
def yolo_results_endpoint(): def yolo_results_endpoint():
global yolo_results global yolo_results
return jsonify(yolo_results) return jsonify(yolo_results)
if __name__ == '__main__': if __name__ == '__main__':
app.run(debug=True, port=5000) app.run(debug=True, port=5000)