Update app.py

app.py

@@ -1,15 +1,14 @@
 import os
 import cv2
-import math
 import torch
 import numpy as np
-from flask import Flask, request, jsonify, render_template_string
+from flask import Flask, Response, request, jsonify
 from ultralytics import YOLO
 
-# ------------------------------
-# CONFIG
-# ------------------------------
-MODEL_PATH = "best.pt"  # your 5MB lightweight model
+app = Flask(__name__)
+
+# ---------------- CONFIG ----------------
+MODEL_PATH = "best.pt"
 FOCAL_LENGTH_PX = 615
 KNOWN_WIDTHS_M = {
     "person": 0.5, "car": 1.8, "truck": 2.3, "bus": 2.5,
@@ -18,17 +17,13 @@ KNOWN_WIDTHS_M = {
 }
 THRESHOLDS = {"CRITICAL": 1.0, "WARNING": 2.0, "CAUTION": 3.0}
 
-# ------------------------------
-# Initialize Flask & YOLO
-# ------------------------------
-app = Flask(__name__)
 device = 0 if torch.cuda.is_available() else "cpu"
 model = YOLO(MODEL_PATH)
 model.to(device)
+model.fuse()  # small speed boost
 
-# ------------------------------
-# Utility Functions
-# ------------------------------
+
+# ---------------- UTIL FUNCTIONS ----------------
 def estimate_distance(bbox_width_px, class_name):
     if bbox_width_px <= 1:
         return None
@@ -48,19 +43,34 @@ def get_alert_level(distance_m):
         return "CAUTION"
     return "SAFE"
 
-# ------------------------------
-# API ENDPOINT: /detect
-# ------------------------------
-@app.route("/detect", methods=["POST"])
-def detect_objects():
+def annotate_frame(frame, detections):
+    for det in detections:
+        x1, y1, x2, y2 = det["bbox"]
+        color = (0, 255, 0)
+        if det["alert_level"] == "CRITICAL":
+            color = (0, 0, 255)
+        elif det["alert_level"] == "WARNING":
+            color = (0, 165, 255)
+        elif det["alert_level"] == "CAUTION":
+            color = (0, 255, 255)
+        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
+        label = f"{det['class']} {det['distance_m']}m"
+        cv2.putText(frame, label, (x1, y1 - 10),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    return frame
+
+
+# ---------------- REALTIME VIDEO STREAM ----------------
+@app.route("/stream", methods=["POST"])
+def process_frame():
     """
-    Accepts an image file (from Unity or mobile camera)
-    and returns detections + alert level JSON.
+    Accepts a single video frame (JPEG bytes from Unity or camera),
+    returns detection data + optionally annotated frame.
     """
-    if "file" not in request.files:
-        return jsonify({"error": "No file uploaded"}), 400
+    if "frame" not in request.files:
+        return jsonify({"error": "No frame uploaded"}), 400
 
-    file = request.files["file"]
+    file = request.files["frame"]
     img_bytes = np.frombuffer(file.read(), np.uint8)
     frame = cv2.imdecode(img_bytes, cv2.IMREAD_COLOR)
 
@@ -72,49 +82,42 @@ def detect_objects():
     order = {"SAFE": 0, "CAUTION": 1, "WARNING": 2, "CRITICAL": 3}
 
     for b in boxes:
-        x1, y1, x2, y2 = b.xyxy[0].tolist()
+        x1, y1, x2, y2 = map(int, b.xyxy[0].tolist())
         cls_id = int(b.cls[0].item())
         conf = float(b.conf[0].item())
         class_name = model.names.get(cls_id, str(cls_id))
         bbox_w = x2 - x1
         distance_m = estimate_distance(bbox_w, class_name)
         level = get_alert_level(distance_m)
-
         if order[level] > order[max_level]:
             max_level = level
-
         detections.append({
             "class": class_name,
             "confidence": round(conf, 3),
             "distance_m": round(distance_m, 2) if distance_m else None,
             "alert_level": level,
-            "bbox": [int(x1), int(y1), int(x2), int(y2)]
+            "bbox": [x1, y1, x2, y2]
         })
 
-    return jsonify({
-        "alert_level": max_level,
-        "detections": detections
-    })
+    annotated = annotate_frame(frame, detections)
+    _, buffer = cv2.imencode('.jpg', annotated)
+    encoded = buffer.tobytes()
+
+    return Response(encoded, mimetype='image/jpeg',
+                    headers={
+                        "X-Alert-Level": max_level,
+                        "X-Detections": str(detections)
+                    })
+
 
-# ------------------------------
-# SIMPLE FRONTEND PAGE
-# ------------------------------
 @app.route("/")
 def index():
-    return render_template_string("""
-    <html>
-      <head><title>YOLOv8 Real-Time Alerts</title></head>
-      <body style="font-family:sans-serif;text-align:center;">
-        <h2>YOLOv8 Object Detection API</h2>
-        <p>Send a POST /detect request with an image to get JSON alerts.</p>
-        <p>Example (curl):</p>
-        <code>
-        curl -X POST -F "file=@image.jpg" https://YOUR_SPACE_URL.hf.space/detect
-        </code>
-      </body>
-    </html>
-    """)
+    return """
+    <h2>YOLOv8 Real-Time Detection Stream</h2>
+    <p>POST /stream with 'frame' (JPEG) from Unity camera feed.</p>
+    """
+
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
-    app.run(host="0.0.0.0", port=port)
+    app.run(host="0.0.0.0", port=port, threaded=True)