Update app.py

app.py

@@ -1,14 +1,13 @@
 import os
 import cv2
+import json
 import torch
 import numpy as np
 from flask import Flask, Response, request, jsonify
 from ultralytics import YOLO
 
-app = Flask(__name__)
-
-# ---------------- CONFIG ----------------
-MODEL_PATH = "best.pt"
+# -------------- CONFIG --------------
+MODEL_PATH = os.environ.get("MODEL_PATH", "best.pt")
 FOCAL_LENGTH_PX = 615
 KNOWN_WIDTHS_M = {
     "person": 0.5, "car": 1.8, "truck": 2.3, "bus": 2.5,
@@ -17,15 +16,25 @@ KNOWN_WIDTHS_M = {
 }
 THRESHOLDS = {"CRITICAL": 1.0, "WARNING": 2.0, "CAUTION": 3.0}
 
+# -------------- APP INIT --------------
+app = Flask(__name__)
+
+# Prefer GPU if available
 device = 0 if torch.cuda.is_available() else "cpu"
+
+# Load model
 model = YOLO(MODEL_PATH)
 model.to(device)
-model.fuse()  # small speed boost
-
+# Fuse for a small speed boost; ignore if unsupported by your build
+try:
+    model.fuse()
+except Exception:
+    pass
 
-# ---------------- UTIL FUNCTIONS ----------------
+# -------------- UTILS --------------
 def estimate_distance(bbox_width_px, class_name):
-    if bbox_width_px <= 1:
+    """Approx distance using pinhole model D = (W * f) / w"""
+    if bbox_width_px is None or bbox_width_px <= 1:
         return None
     known_width = KNOWN_WIDTHS_M.get(class_name)
     if not known_width:
@@ -44,80 +53,102 @@ def get_alert_level(distance_m):
     return "SAFE"
 
 def annotate_frame(frame, detections):
+    """Draw boxes and labels colored by alert level."""
     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)
+        level = det["alert_level"]
+        # Color by severity
+        if level == "CRITICAL":
+            color = (0, 0, 255)    # Red
+        elif level == "WARNING":
+            color = (0, 165, 255)  # Orange
+        elif level == "CAUTION":
+            color = (0, 255, 255)  # Yellow
+        else:
+            color = (0, 255, 0)    # Green
+
         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)
+
+        dist_str = f"{det['distance_m']}m" if det["distance_m"] is not None else "n/a"
+        label = f"{det['class']} {dist_str}"
+        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
+        y_top = max(y1 - th - 6, 0)
+        cv2.rectangle(frame, (x1, y_top), (x1 + tw + 8, y_top + th + 6), color, -1)
+        cv2.putText(frame, label, (x1 + 4, y_top + th),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2, cv2.LINE_AA)
     return frame
 
+# -------------- ROUTES --------------
+@app.route("/ping")
+def ping():
+    return jsonify({"ok": True}), 200
 
-# ---------------- REALTIME VIDEO STREAM ----------------
 @app.route("/stream", methods=["POST"])
 def process_frame():
     """
-    Accepts a single video frame (JPEG bytes from Unity or camera),
-    returns detection data + optionally annotated frame.
+    Accepts a single video frame via multipart/form-data field 'frame' (JPEG bytes),
+    returns annotated JPEG as body with alert metadata in headers.
     """
     if "frame" not in request.files:
         return jsonify({"error": "No frame uploaded"}), 400
 
+    # Decode image
     file = request.files["frame"]
-    img_bytes = np.frombuffer(file.read(), np.uint8)
+    file_bytes = file.read()
+    img_bytes = np.frombuffer(file_bytes, np.uint8)
     frame = cv2.imdecode(img_bytes, cv2.IMREAD_COLOR)
+    if frame is None:
+        return jsonify({"error": "Invalid image"}), 400
 
-    results = model(frame, conf=0.25, iou=0.5, verbose=False)
+    # Run inference
+    results = model(frame, conf=0.25, iou=0.5, verbose=False, device=device)
     boxes = results[0].boxes
 
     detections = []
     max_level = "SAFE"
     order = {"SAFE": 0, "CAUTION": 1, "WARNING": 2, "CRITICAL": 3}
 
-    for b in boxes:
-        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": [x1, y1, x2, y2]
-        })
-
+    if boxes is not None and len(boxes) > 0:
+        for b in boxes:
+            x1, y1, x2, y2 = 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 = int(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)]
+            })
+
+    # Annotate and encode JPEG
     annotated = annotate_frame(frame, detections)
-    _, buffer = cv2.imencode('.jpg', annotated)
+    ok, buffer = cv2.imencode(".jpg", annotated, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
+    if not ok:
+        return jsonify({"error": "Encode failed"}), 500
     encoded = buffer.tobytes()
 
-    return Response(encoded, mimetype='image/jpeg',
-                    headers={
-                        "X-Alert-Level": max_level,
-                        "X-Detections": str(detections)
-                    })
-
-
-@app.route("/")
-def index():
-    return """
-    <h2>YOLOv8 Real-Time Detection Stream</h2>
-    <p>POST /stream with 'frame' (JPEG) from Unity camera feed.</p>
-    """
+    # Headers with metadata
+    hdr_alert = max_level
+    hdr_count = str(len(detections))
+    headers = {
+        "Content-Length": str(len(encoded)),  # some clients use it for streaming/decoding
+        "X-Alert-Level": hdr_alert,
+        "X-Detections-Count": hdr_count
+    }
 
+    return Response(encoded, mimetype="image/jpeg", headers=headers)
 
+# -------------- MAIN --------------
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
+    # threaded=True allows concurrent requests from multiple clients
     app.run(host="0.0.0.0", port=port, threaded=True)