Update app.py

app.py

@@ -1,23 +1,48 @@
-from flask import Flask, render_template, jsonify, request, send_file
+# Flask Backend (app.py)
+from flask import Flask, render_template, Response, jsonify
 import psutil
-import json
-from ultralytics import YOLO
 import cv2
 import numpy as np
-from PIL import Image
-import io
-import base64
+from ultralytics import YOLO
+import threading
+import queue
 
 app = Flask(__name__)
 modelName = "yolov9c-seg.pt"
 model = YOLO(modelName)
 
+# Global variables for frame handling
+frame_queue = queue.Queue(maxsize=2)
+processed_frame_queue = queue.Queue(maxsize=2)
+current_frame = None
+processing_active = False
+
+def process_frames():
+    global processing_active
+    while processing_active:
+        if not frame_queue.empty():
+            frame = frame_queue.get()
+            results = model(
+                frame,
+                show=False,
+                save=False,
+                show_boxes=True,
+                show_labels=True,
+                imgsz=640,
+                iou=0.1,
+                max_det=20,
+            )
+            processed_frame = results[0].plot()
+            processed_frame = cv2.cvtColor(np.array(processed_frame), cv2.COLOR_RGB2BGR)
+            
+            if processed_frame_queue.full():
+                processed_frame_queue.get()  # Remove old frame
+            processed_frame_queue.put(processed_frame)
 
 @app.route("/")
 def home():
     return render_template("index.html")
 
-
 @app.route("/sysInfo")
 def sysInfo():
     ram = psutil.virtual_memory()
@@ -26,35 +51,46 @@ def sysInfo():
     data = {"ram": ram_usage, "cpu": cpu_usage}
     return jsonify(data)
 
+def generate_frames():
+    global processing_active
+    while True:
+        if not processed_frame_queue.empty():
+            frame = processed_frame_queue.get()
+            ret, buffer = cv2.imencode('.jpg', frame)
+            frame = buffer.tobytes()
+            yield (b'--frame\r\n'
+                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
+
+@app.route('/video_feed')
+def video_feed():
+    return Response(generate_frames(),
+                    mimetype='multipart/x-mixed-replace; boundary=frame')
+
+@app.route('/start_processing')
+def start_processing():
+    global processing_active
+    if not processing_active:
+        processing_active = True
+        threading.Thread(target=process_frames, daemon=True).start()
+    return jsonify({"status": "started"})
+
+@app.route('/stop_processing')
+def stop_processing():
+    global processing_active
+    processing_active = False
+    return jsonify({"status": "stopped"})
+
+@app.route('/update_frame', methods=['POST'])
+def update_frame():
+    global current_frame
+    data = request.get_json()
+    frame_data = np.array(data['frame'])
+    
+    if frame_queue.full():
+        frame_queue.get()  # Remove old frame
+    frame_queue.put(frame_data)
+    
+    return jsonify({"status": "success"})
 
-@app.route("/processor", methods=["POST"])
-def processor():
-    global modelName, model
-    image = request.form.get("image")
-    modelNameForm = request.form.get("model", modelName)
-    if modelNameForm != modelName:
-        modelName = modelNameForm
-        model = YOLO(modelName)
-    image = image.split(",")[1]
-    image = base64.b64decode(image)
-    image = Image.open(io.BytesIO(image))
-    image = np.array(image)
-    results = model(
-        image,
-        show=False,
-        save=False,
-        show_boxes=True,
-        show_labels=True,
-        imgsz=640,
-        iou=0.1,
-        max_det=20,
-    )
-    image = results[0].plot()
-    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    image = cv2.imencode(".jpg", image)[1].tobytes()
-    image = base64.b64encode(image).decode("utf-8")
-    return jsonify({"image": image})
-
-
-app.run(debug=True, host="0.0.0.0", port=7860)
-# app.run(debug=True)
+if __name__ == "__main__":
+    app.run(debug=True, host="0.0.0.0", port=7860)