trying smoothing + buffer

perception_roi_server.py

@@ -21,6 +21,7 @@ import uuid
 import time
 import math
 import threading
+from collections import deque
 import shutil
 import subprocess
 from dataclasses import dataclass, field
@@ -38,6 +39,9 @@ DEFAULT_CONF = float(os.environ.get("YOLO_CONF", "0.25"))
 DEFAULT_DEVICE = os.environ.get("YOLO_DEVICE", "auto")
 FAST_DETECT_SCALE = float(os.environ.get("FAST_DETECT_SCALE", "0.65"))
 FAST_DETECT_IMGSZ = int(os.environ.get("FAST_DETECT_IMGSZ", "512"))
+SMOOTH_ALPHA = float(os.environ.get("SMOOTH_ALPHA", "0.7"))
+MJPEG_STREAM_FPS = float(os.environ.get("MJPEG_STREAM_FPS", "8"))
+MJPEG_STARTUP_BUFFER_SEC = float(os.environ.get("MJPEG_STARTUP_BUFFER_SEC", "0.6"))
 DATA_DIR = os.environ.get("DATA_DIR", "/tmp/roi_demo")
 UPLOAD_DIR = os.path.join(DATA_DIR, "uploads")
 OUTPUT_DIR = os.path.join(DATA_DIR, "outputs")
@@ -144,7 +148,7 @@ def _yolo_detect_frame(
 def _draw_boxes(frame_bgr: np.ndarray, dets: List[Dict[str, Any]]) -> np.ndarray:
     out = frame_bgr.copy()
     for d in dets:
-        b = d.get("bbox_xyxy")
+        b = d.get("bbox_xyxy_smooth") or d.get("bbox_xyxy")
         if not (isinstance(b, (list, tuple)) and len(b) == 4):
             continue
         x1, y1, x2, y2 = [int(max(0, v)) for v in b]
@@ -159,6 +163,32 @@ def _draw_boxes(frame_bgr: np.ndarray, dets: List[Dict[str, Any]]) -> np.ndarray
     return out
 
 
+def _smooth_tracks(dets: List[Dict[str, Any]], tracker: Dict[str, Any], alpha: float = SMOOTH_ALPHA) -> List[Dict[str, Any]]:
+    if not dets:
+        return dets
+    smooth = tracker.setdefault("smooth", {})
+    out: List[Dict[str, Any]] = []
+    for d in dets:
+        b = d.get("bbox_xyxy")
+        tid = d.get("track_id")
+        if not (isinstance(b, (list, tuple)) and len(b) == 4) or tid is None:
+            out.append(d)
+            continue
+        prev = smooth.get(int(tid))
+        if isinstance(prev, (list, tuple)) and len(prev) == 4:
+            sb = [
+                float(alpha) * float(prev[i]) + (1.0 - float(alpha)) * float(b[i])
+                for i in range(4)
+            ]
+        else:
+            sb = [float(v) for v in b]
+        smooth[int(tid)] = sb
+        d2 = dict(d)
+        d2["bbox_xyxy_smooth"] = sb
+        out.append(d2)
+    return out
+
+
 def _iou_xyxy(a: List[float], b: List[float]) -> float:
     ax1, ay1, ax2, ay2 = a
     bx1, by1, bx2, by2 = b
@@ -305,7 +335,7 @@ def _apply_roi_overlay(frame_bgr: np.ndarray, dets: List[Dict[str, Any]], target
     out = bg.copy()
     pad = max(2, int(min(w, h) * 0.005))
     for d in dets:
-        b = d.get("bbox_xyxy")
+        b = d.get("bbox_xyxy_smooth") or d.get("bbox_xyxy")
         if not (isinstance(b, (list, tuple)) and len(b) == 4):
             continue
         x1, y1, x2, y2 = [int(v) for v in b]
@@ -346,6 +376,9 @@ class Job:
     latest_jpeg: Optional[bytes] = None
     latest_compressed_jpeg: Optional[bytes] = None
     latest_roi_jpeg: Optional[bytes] = None
+    overlay_buffer: deque = field(default_factory=lambda: deque(maxlen=180))
+    compressed_buffer: deque = field(default_factory=lambda: deque(maxlen=180))
+    roi_buffer: deque = field(default_factory=lambda: deque(maxlen=180))
     lock: threading.Lock = field(default_factory=threading.Lock)
     tracker_state: Dict[str, Any] = field(default_factory=lambda: {"next_id": 1, "tracks": []})
 
@@ -407,6 +440,7 @@ def _process_job(job: Job):
                     ]
                     max_id = max((int(d.get("track_id", 0)) for d in dets), default=0)
                     tracker["next_id"] = max(tracker.get("next_id", 1), max_id + 1)
+                dets = _smooth_tracks(dets, tracker)
                 with job.lock:
                     job.det_by_frame[int(frame_idx)] = dets
                 last_dets = dets
@@ -418,6 +452,7 @@ def _process_job(job: Job):
             if ok2:
                 with job.lock:
                     job.latest_jpeg = jpg.tobytes()
+                    job.overlay_buffer.append(job.latest_jpeg)
 
             if overlay_writer is not None:
                 overlay_writer.write(overlay)
@@ -515,11 +550,13 @@ def _compress_job(job: Job, bandwidth_kbps: int, target_fps: int, target_w: int,
                     if okc:
                         with job.lock:
                             job.latest_compressed_jpeg = jpgc.tobytes()
+                            job.compressed_buffer.append(job.latest_compressed_jpeg)
                 if roi_frame is not None:
                     okr, jpgr = cv2.imencode(".jpg", roi_frame, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
                     if okr:
                         with job.lock:
                             job.latest_roi_jpeg = jpgr.tobytes()
+                            job.roi_buffer.append(job.latest_roi_jpeg)
             except Exception:
                 pass
 
@@ -598,6 +635,7 @@ async def track_async(
             model = get_model(job.weights)
             det0 = _yolo_detect_frame(model, frame0, conf=job.conf, queries=job.queries, device=job.device, fast_mode=job.fast_mode)
             det0 = _assign_tracks(det0, job.tracker_state)
+            det0 = _smooth_tracks(det0, job.tracker_state)
             with job.lock:
                 job.det_by_frame[0] = det0
             vis0 = _draw_boxes(frame0, det0)
@@ -605,6 +643,7 @@ async def track_async(
             if ok2:
                 with job.lock:
                     job.latest_jpeg = jpg.tobytes()
+                    job.overlay_buffer.append(job.latest_jpeg)
     except Exception:
         pass
 
@@ -665,13 +704,27 @@ def process_status(job_id: str):
         }
 
 
-def _mjpeg_generator(job: Job):
+def _mjpeg_generator_from_buffer(job: Job, buffer_attr: str, latest_attr: str):
     boundary = b"--frame"
+    stream_fps = max(2.0, min(30.0, float(MJPEG_STREAM_FPS or 8.0)))
+    frame_interval = 1.0 / stream_fps
+    min_frames = int(math.ceil(stream_fps * max(0.0, float(MJPEG_STARTUP_BUFFER_SEC or 0.0))))
+    started = False
+    next_time = time.time()
     while True:
         with job.lock:
-            jpg = job.latest_jpeg
+            buf = getattr(job, buffer_attr, None)
             status = job.status
             err = job.error
+            if buf is not None and len(buf) > 0:
+                if not started and len(buf) < min_frames:
+                    jpg = None
+                else:
+                    started = True
+                    jpg = buf.popleft()
+            else:
+                jpg = getattr(job, latest_attr, None)
+            remaining = len(buf) if buf is not None else 0
         if err:
             break
         if jpg:
@@ -679,50 +732,25 @@ def _mjpeg_generator(job: Job):
             yield b"Content-Type: image/jpeg\r\n"
             yield f"Content-Length: {len(jpg)}\r\n\r\n".encode("ascii")
             yield jpg + b"\r\n"
-        time.sleep(0.15)
-        if status in ("completed", "error"):
-            time.sleep(0.5)
+        now = time.time()
+        if now < next_time:
+            time.sleep(next_time - now)
+        next_time = max(next_time + frame_interval, now)
+        if status in ("completed", "error") and remaining == 0:
+            time.sleep(0.2)
             break
 
 
+def _mjpeg_generator(job: Job):
+    return _mjpeg_generator_from_buffer(job, "overlay_buffer", "latest_jpeg")
+
+
 def _mjpeg_generator_compressed(job: Job):
-    boundary = b"--frame"
-    while True:
-        with job.lock:
-            jpg = job.latest_compressed_jpeg
-            status = job.status
-            err = job.error
-        if err:
-            break
-        if jpg:
-            yield boundary + b"\r\n"
-            yield b"Content-Type: image/jpeg\r\n"
-            yield f"Content-Length: {len(jpg)}\r\n\r\n".encode("ascii")
-            yield jpg + b"\r\n"
-        time.sleep(0.15)
-        if status in ("completed", "error"):
-            time.sleep(0.5)
-            break
+    return _mjpeg_generator_from_buffer(job, "compressed_buffer", "latest_compressed_jpeg")
 
 
 def _mjpeg_generator_roi(job: Job):
-    boundary = b"--frame"
-    while True:
-        with job.lock:
-            jpg = job.latest_roi_jpeg
-            status = job.status
-            err = job.error
-        if err:
-            break
-        if jpg:
-            yield boundary + b"\r\n"
-            yield b"Content-Type: image/jpeg\r\n"
-            yield f"Content-Length: {len(jpg)}\r\n\r\n".encode("ascii")
-            yield jpg + b"\r\n"
-        time.sleep(0.15)
-        if status in ("completed", "error"):
-            time.sleep(0.5)
-            break
+    return _mjpeg_generator_from_buffer(job, "roi_buffer", "latest_roi_jpeg")
 
 
 @app.get("/detect/stream/{job_id}")