# CRITICAL: Clear CUDA_VISIBLE_DEVICES BEFORE any imports # HF Spaces may set this to "0" dynamically, locking us to a single GPU # import os # if "CUDA_VISIBLE_DEVICES" in os.environ: # del os.environ["CUDA_VISIBLE_DEVICES"] import os import collections import logging import time from threading import Event, RLock, Thread from queue import Queue, Full, Empty from typing import Any, Dict, List, Optional, Sequence, Tuple import cv2 import numpy as np import torch from concurrent.futures import ThreadPoolExecutor from models.detectors.base import ObjectDetector from models.model_loader import load_detector, load_detector_on_device from models.segmenters.model_loader import load_segmenter, load_segmenter_on_device from models.depth_estimators.model_loader import load_depth_estimator, load_depth_estimator_on_device from utils.video import StreamingVideoWriter from utils.relevance import evaluate_relevance from utils.enrichment import run_enrichment from utils.schemas import AssessmentStatus from jobs.storage import set_track_data import tempfile import json as json_module class AsyncVideoReader: """ Async video reader that decodes frames in a background thread. This prevents GPU starvation on multi-GPU systems by prefetching frames while the main thread is busy dispatching work to GPUs. """ def __init__(self, video_path: str, prefetch_size: int = 32): """ Initialize async video reader. Args: video_path: Path to video file prefetch_size: Number of frames to prefetch (default 32) """ from queue import Queue from threading import Thread self.video_path = video_path self.prefetch_size = prefetch_size # Open video to get metadata self._cap = cv2.VideoCapture(video_path) if not self._cap.isOpened(): raise ValueError(f"Unable to open video: {video_path}") self.fps = self._cap.get(cv2.CAP_PROP_FPS) or 30.0 self.width = int(self._cap.get(cv2.CAP_PROP_FRAME_WIDTH)) self.height = int(self._cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) self.total_frames = int(self._cap.get(cv2.CAP_PROP_FRAME_COUNT)) # Prefetch queue self._queue: Queue = Queue(maxsize=prefetch_size) self._error: Exception = None self._finished = False # Start decoder thread self._thread = Thread(target=self._decode_loop, daemon=True) self._thread.start() def _decode_loop(self): """Background thread that continuously decodes frames.""" try: while True: success, frame = self._cap.read() if not success: break self._queue.put(frame) # Blocks when queue is full (backpressure) except Exception as e: self._error = e logging.error(f"AsyncVideoReader decode error: {e}") finally: self._cap.release() self._queue.put(None) # Sentinel to signal end self._finished = True def __iter__(self): return self def __next__(self) -> np.ndarray: if self._error: raise self._error frame = self._queue.get() if frame is None: raise StopIteration return frame def close(self): """Stop the decoder thread and release resources.""" # Signal thread to stop by releasing cap (if not already done) if self._cap.isOpened(): self._cap.release() # Drain queue to unblock thread if it's waiting on put() while not self._queue.empty(): try: self._queue.get_nowait() except: break def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.close() def _check_cancellation(job_id: Optional[str]) -> None: """Check if job has been cancelled and raise exception if so.""" if job_id is None: return from jobs.storage import get_job_storage from jobs.models import JobStatus job = get_job_storage().get(job_id) if job and job.status == JobStatus.CANCELLED: raise RuntimeError("Job cancelled by user") def _color_for_label(label: str) -> Tuple[int, int, int]: # Deterministic BGR color from label text. value = abs(hash(label)) % 0xFFFFFF blue = value & 0xFF green = (value >> 8) & 0xFF red = (value >> 16) & 0xFF return (blue, green, red) def draw_boxes( frame: np.ndarray, boxes: np.ndarray, labels: Optional[Sequence[int]] = None, queries: Optional[Sequence[str]] = None, label_names: Optional[Sequence[str]] = None, ) -> np.ndarray: output = frame.copy() if boxes is None: return output for idx, box in enumerate(boxes): x1, y1, x2, y2 = [int(coord) for coord in box] if label_names is not None and idx < len(label_names): label = label_names[idx] elif labels is not None and idx < len(labels) and queries is not None: label_idx = int(labels[idx]) if 0 <= label_idx < len(queries): label = queries[label_idx] else: label = f"label_{label_idx}" else: label = f"label_{idx}" color = _color_for_label(label) cv2.rectangle(output, (x1, y1), (x2, y2), color, thickness=2) if label: font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 1.0 thickness = 2 text_size, baseline = cv2.getTextSize(label, font, font_scale, thickness) text_w, text_h = text_size pad = 4 text_x = x1 text_y = max(y1 - 6, text_h + pad) box_top_left = (text_x, text_y - text_h - pad) box_bottom_right = (text_x + text_w + pad, text_y + baseline) cv2.rectangle(output, box_top_left, box_bottom_right, color, thickness=-1) cv2.putText( output, label, (text_x + pad // 2, text_y - 2), font, font_scale, (255, 255, 255), thickness, lineType=cv2.LINE_AA, ) return output def draw_masks( frame: np.ndarray, masks: np.ndarray, alpha: float = 0.65, labels: Optional[Sequence[str]] = None, ) -> np.ndarray: output = frame.copy() if masks is None or len(masks) == 0: return output for idx, mask in enumerate(masks): if mask is None: continue if mask.ndim == 3: mask = mask[0] if mask.shape[:2] != output.shape[:2]: mask = cv2.resize(mask, (output.shape[1], output.shape[0]), interpolation=cv2.INTER_NEAREST) mask_bool = mask.astype(bool) overlay = np.zeros_like(output, dtype=np.uint8) label = None if labels and idx < len(labels): label = labels[idx] # Use a fallback key for consistent color even when no label text color_key = label if label else f"object_{idx}" color = _color_for_label(color_key) overlay[mask_bool] = color output = cv2.addWeighted(output, 1.0, overlay, alpha, 0) contours, _ = cv2.findContours( mask_bool.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE ) if contours: cv2.drawContours(output, contours, -1, color, thickness=2) # Only draw label text when explicit labels were provided if label: coords = np.column_stack(np.where(mask_bool)) if coords.size: y, x = coords[0] font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 1.0 thickness = 2 text_size, baseline = cv2.getTextSize(label, font, font_scale, thickness) text_w, text_h = text_size pad = 4 text_x = int(x) text_y = max(int(y) - 6, text_h + pad) box_top_left = (text_x, text_y - text_h - pad) box_bottom_right = (text_x + text_w + pad, text_y + baseline) cv2.rectangle(output, box_top_left, box_bottom_right, color, thickness=-1) cv2.putText( output, label, (text_x + pad // 2, text_y - 2), font, font_scale, (255, 255, 255), thickness, lineType=cv2.LINE_AA, ) return output def _build_detection_records( boxes: np.ndarray, scores: Sequence[float], labels: Sequence[int], queries: Sequence[str], label_names: Optional[Sequence[str]] = None, ) -> List[Dict[str, Any]]: detections: List[Dict[str, Any]] = [] for idx, box in enumerate(boxes): if label_names is not None and idx < len(label_names): label = label_names[idx] else: label_idx = int(labels[idx]) if idx < len(labels) else -1 if 0 <= label_idx < len(queries): label = queries[label_idx] else: label = f"label_{label_idx}" detections.append( { "label": label, "score": float(scores[idx]) if idx < len(scores) else 0.0, "bbox": [int(coord) for coord in box.tolist()], } ) return detections from utils.tracker import ByteTracker class SpeedEstimator: def __init__(self, fps: float = 30.0): self.fps = fps self.pixel_scale_map = {} # label -> pixels_per_meter (heuristic) def estimate(self, detections: List[Dict[str, Any]]): for det in detections: history = det.get('history', []) if len(history) < 5: continue # Simple heuristic: Speed based on pixel movement # We assume a base depth or size. # Delta over last 5 frames curr = history[-1] prev = history[-5] # Centroids cx1 = (curr[0] + curr[2]) / 2 cy1 = (curr[1] + curr[3]) / 2 cx2 = (prev[0] + prev[2]) / 2 cy2 = (prev[1] + prev[3]) / 2 dist_px = np.sqrt((cx1-cx2)**2 + (cy1-cy2)**2) # Heuristic scale: Assume car is ~4m long? Or just arbitrary pixel scale # If we had GPT distance, we could calibrate. # For now, let's use a dummy scale: 50px = 1m (very rough) # Speed = (dist_px / 50) meters / (5 frames / 30 fps) seconds # = (dist_px / 50) / (0.166) m/s # = (dist_px * 0.12) m/s # = * 3.6 km/h scale = 50.0 dt = 5.0 / self.fps speed_mps = (dist_px / scale) / dt speed_kph = speed_mps * 3.6 # Smoothing det['speed_kph'] = speed_kph # Direction dx = cx1 - cx2 dy = cy1 - cy2 angle = np.degrees(np.arctan2(dy, dx)) # 0 is right, 90 is down # Map to clock direction (12 is up = -90 deg) # -90 (up) -> 12 # 0 (right) -> 3 # 90 (down) -> 6 # 180 (left) -> 9 # Adjust so 12 is up (negative Y) # angle -90 is 12 clock_hour = ((angle + 90) / 30 + 12) % 12 if clock_hour == 0: clock_hour = 12.0 det['direction_clock'] = f"{int(round(clock_hour))} o'clock" det['angle_deg'] = angle # 0 is right, 90 is down (screen space) class IncrementalDepthStats: """Thread-safe incremental depth range estimator. Collects depth statistics frame-by-frame so the expensive pre-scan (opening a second video reader) can be eliminated. Before ``warmup_frames`` updates the range defaults to (0.0, 1.0). """ def __init__(self, warmup_frames: int = 30): self._lock = RLock() self._warmup = warmup_frames self._count = 0 self._global_min = float("inf") self._global_max = float("-inf") def update(self, depth_map: np.ndarray) -> None: if depth_map is None or depth_map.size == 0: return finite = depth_map[np.isfinite(depth_map)] if finite.size == 0: return lo = float(np.percentile(finite, 1)) hi = float(np.percentile(finite, 99)) with self._lock: self._global_min = min(self._global_min, lo) self._global_max = max(self._global_max, hi) self._count += 1 @property def range(self) -> Tuple[float, float]: with self._lock: if self._count < self._warmup: # Not enough data yet — use default range if self._count == 0: return (0.0, 1.0) # Use what we have but may be less stable lo, hi = self._global_min, self._global_max else: lo, hi = self._global_min, self._global_max if abs(hi - lo) < 1e-6: hi = lo + 1.0 return (lo, hi) _MODEL_LOCKS: Dict[str, RLock] = {} _MODEL_LOCKS_GUARD = RLock() _DEPTH_SCALE = float(os.getenv("DEPTH_SCALE", "25.0")) def _get_model_lock(kind: str, name: str) -> RLock: key = f"{kind}:{name}" with _MODEL_LOCKS_GUARD: lock = _MODEL_LOCKS.get(key) if lock is None: lock = RLock() _MODEL_LOCKS[key] = lock return lock def _attach_depth_metrics( frame: np.ndarray, detections: List[Dict[str, Any]], depth_estimator_name: Optional[str], depth_scale: float, # No longer used for distance calculation estimator_instance: Optional[Any] = None, ) -> None: """Attach relative depth values for visualization only. GPT handles distance estimation.""" if not detections or (not depth_estimator_name and not estimator_instance): return from models.depth_estimators.model_loader import load_depth_estimator if estimator_instance: estimator = estimator_instance # Use instance lock if available, or create one if hasattr(estimator, "lock"): lock = estimator.lock else: # Fallback (shouldn't happen with our new setup but safe) lock = _get_model_lock("depth", estimator.name) else: estimator = load_depth_estimator(depth_estimator_name) lock = _get_model_lock("depth", estimator.name) with lock: depth_result = estimator.predict(frame) depth_map = depth_result.depth_map if depth_map is None or depth_map.size == 0: return height, width = depth_map.shape[:2] raw_depths: List[Tuple[Dict[str, Any], float]] = [] for det in detections: det["depth_rel"] = None # Relative depth for visualization only bbox = det.get("bbox") if not bbox or len(bbox) < 4: continue x1, y1, x2, y2 = [int(coord) for coord in bbox[:4]] x1 = max(0, min(width - 1, x1)) y1 = max(0, min(height - 1, y1)) x2 = max(x1 + 1, min(width, x2)) y2 = max(y1 + 1, min(height, y2)) patch = depth_map[y1:y2, x1:x2] if patch.size == 0: continue # Center crop (50%) to avoid background h_p, w_p = patch.shape cy, cx = h_p // 2, w_p // 2 dy, dx = h_p // 4, w_p // 4 center_patch = patch[cy - dy : cy + dy, cx - dx : cx + dx] # Fallback to full patch if center is empty (unlikely) if center_patch.size == 0: center_patch = patch finite = center_patch[np.isfinite(center_patch)] if finite.size == 0: continue depth_raw = float(np.median(finite)) if depth_raw > 1e-6: raw_depths.append((det, depth_raw)) if not raw_depths: return # Compute relative depth (0-1) for visualization only all_raw = [d[1] for d in raw_depths] min_raw, max_raw = min(all_raw), max(all_raw) denom = max(max_raw - min_raw, 1e-6) for det, depth_raw in raw_depths: # Inverted: higher raw = closer = lower rel value (0=close, 1=far) det["depth_rel"] = 1.0 - ((depth_raw - min_raw) / denom) def infer_frame( frame: np.ndarray, queries: Sequence[str], detector_name: Optional[str] = None, depth_estimator_name: Optional[str] = None, depth_scale: float = 1.0, detector_instance: Optional[ObjectDetector] = None, depth_estimator_instance: Optional[Any] = None, ) -> Tuple[np.ndarray, List[Dict[str, Any]]]: if detector_instance: detector = detector_instance else: detector = load_detector(detector_name) text_queries = list(queries) or ["object"] try: if hasattr(detector, "lock"): lock = detector.lock else: lock = _get_model_lock("detector", detector.name) with lock: result = detector.predict(frame, text_queries) detections = _build_detection_records( result.boxes, result.scores, result.labels, text_queries, result.label_names ) if depth_estimator_name or depth_estimator_instance: try: _attach_depth_metrics( frame, detections, depth_estimator_name, depth_scale, estimator_instance=depth_estimator_instance ) except Exception: logging.exception("Depth estimation failed for frame") # Re-build display labels to include GPT distance if available display_labels = [] for i, det in enumerate(detections): label = det["label"] if det.get("gpt_distance_m") is not None: # Add GPT distance to label, e.g. "car 12m" depth_str = f"{int(det['gpt_distance_m'])}m" label = f"{label} {depth_str}" logging.debug("Object '%s' at %s (bbox: %s)", label, depth_str, det['bbox']) display_labels.append(label) except Exception: logging.exception("Inference failed for queries %s", text_queries) raise return draw_boxes( frame, result.boxes, labels=None, # Use custom labels queries=None, label_names=display_labels, ), detections def _build_display_label(det): """Build display label with GPT distance if available.""" label = det["label"] if det.get("gpt_distance_m") is not None: label = f"{label} {int(det['gpt_distance_m'])}m" return label def _attach_depth_from_result(detections, depth_result, depth_scale): """Attach relative depth values for visualization only. GPT handles distance estimation.""" depth_map = depth_result.depth_map if depth_map is None or depth_map.size == 0: return height, width = depth_map.shape[:2] raw_depths = [] for det in detections: det["depth_rel"] = None # Relative depth for visualization only bbox = det.get("bbox") if not bbox or len(bbox) < 4: continue x1, y1, x2, y2 = [int(coord) for coord in bbox[:4]] x1 = max(0, min(width - 1, x1)) y1 = max(0, min(height - 1, y1)) x2 = max(x1 + 1, min(width, x2)) y2 = max(y1 + 1, min(height, y2)) patch = depth_map[y1:y2, x1:x2] if patch.size == 0: continue h_p, w_p = patch.shape cy, cx = h_p // 2, w_p // 2 dy, dx = h_p // 4, w_p // 4 center_patch = patch[cy - dy : cy + dy, cx - dx : cx + dx] if center_patch.size == 0: center_patch = patch finite = center_patch[np.isfinite(center_patch)] if finite.size == 0: continue depth_raw = float(np.median(finite)) if depth_raw > 1e-6: raw_depths.append((det, depth_raw)) if not raw_depths: return # Compute relative depth (0-1) for visualization only all_raw = [d[1] for d in raw_depths] min_raw, max_raw = min(all_raw), max(all_raw) denom = max(max_raw - min_raw, 1e-6) for det, depth_raw in raw_depths: # Inverted: higher raw = closer = lower rel value (0=close, 1=far) det["depth_rel"] = 1.0 - ((depth_raw - min_raw) / denom) def infer_segmentation_frame( frame: np.ndarray, text_queries: Optional[List[str]] = None, segmenter_name: Optional[str] = None, segmenter_instance: Optional[Any] = None, ) -> Tuple[np.ndarray, Any]: if segmenter_instance: segmenter = segmenter_instance # Use instance lock if available if hasattr(segmenter, "lock"): lock = segmenter.lock else: lock = _get_model_lock("segmenter", segmenter.name) else: segmenter = load_segmenter(segmenter_name) lock = _get_model_lock("segmenter", segmenter.name) with lock: result = segmenter.predict(frame, text_prompts=text_queries) labels = text_queries or [] if len(labels) == 1: masks = result.masks if result.masks is not None else [] labels = [labels[0] for _ in range(len(masks))] return draw_masks(frame, result.masks, labels=labels), result def extract_first_frame(video_path: str) -> Tuple[np.ndarray, float, int, int]: cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise ValueError("Unable to open video.") fps = cap.get(cv2.CAP_PROP_FPS) or 0.0 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) success, frame = cap.read() cap.release() if not success or frame is None: raise ValueError("Video decode produced zero frames.") return frame, fps, width, height def process_first_frame( video_path: str, queries: List[str], mode: str, detector_name: Optional[str] = None, segmenter_name: Optional[str] = None, ) -> Tuple[np.ndarray, List[Dict[str, Any]]]: """Lightweight first-frame processing: detection + rendering only. GPT, depth, and LLM relevance are handled later in the async pipeline (writer enrichment thread), avoiding 2-8s synchronous startup delay. Returns: (processed_frame, detections) — all detections tagged UNASSESSED. """ frame, _, _, _ = extract_first_frame(video_path) if mode == "segmentation": processed, seg_result = infer_segmentation_frame( frame, text_queries=queries, segmenter_name=segmenter_name ) detections = [] if seg_result.boxes is not None and len(seg_result.boxes) > 0: labels = seg_result.label_names or queries or [] for idx, box in enumerate(seg_result.boxes): label = labels[idx] if idx < len(labels) else "object" detections.append({ "label": label, "bbox": [int(c) for c in box], "score": float(seg_result.scores[idx]) if seg_result.scores is not None and idx < len(seg_result.scores) else 1.0, "track_id": f"T{idx + 1:02d}", "assessment_status": AssessmentStatus.UNASSESSED, }) return processed, detections processed, detections = infer_frame( frame, queries, detector_name=detector_name ) # Tag all detections as unassessed — GPT runs later in enrichment thread for det in detections: det["assessment_status"] = AssessmentStatus.UNASSESSED return processed, detections def run_inference( input_video_path: str, output_video_path: str, queries: List[str], max_frames: Optional[int] = None, detector_name: Optional[str] = None, job_id: Optional[str] = None, depth_estimator_name: Optional[str] = None, depth_scale: float = 1.0, enable_gpt: bool = True, stream_queue: Optional[Queue] = None, mission_spec=None, # Optional[MissionSpecification] first_frame_gpt_results: Optional[Dict[str, Any]] = None, first_frame_detections: Optional[List[Dict[str, Any]]] = None, ) -> Tuple[str, List[List[Dict[str, Any]]]]: # 1. Setup Video Reader try: reader = AsyncVideoReader(input_video_path) except ValueError: logging.exception("Failed to open video at %s", input_video_path) raise fps = reader.fps width = reader.width height = reader.height total_frames = reader.total_frames if max_frames is not None: total_frames = min(total_frames, max_frames) # 2. Defaults and Config if not queries: queries = ["person", "car", "truck", "motorcycle", "bicycle", "bus", "train", "airplane"] logging.info("No queries provided, using defaults: %s", queries) logging.info("Detection queries: %s", queries) active_detector = detector_name or "yolo11" # Parallel Model Loading num_gpus = torch.cuda.device_count() detectors = [] depth_estimators = [] if num_gpus > 0: logging.info("Detected %d GPUs. Loading models in parallel...", num_gpus) def load_models_on_gpu(gpu_id: int): device_str = f"cuda:{gpu_id}" try: det = load_detector_on_device(active_detector, device_str) det.lock = RLock() depth = None if depth_estimator_name: depth = load_depth_estimator_on_device(depth_estimator_name, device_str) depth.lock = RLock() return (gpu_id, det, depth) except Exception as e: logging.error(f"Failed to load models on GPU {gpu_id}: {e}") raise with ThreadPoolExecutor(max_workers=num_gpus) as loader_pool: futures = [loader_pool.submit(load_models_on_gpu, i) for i in range(num_gpus)] results = [f.result() for f in futures] # Sort by GPU ID to ensure consistent indexing results.sort(key=lambda x: x[0]) for _, det, depth in results: detectors.append(det) depth_estimators.append(depth) else: logging.info("No GPUs detected. Loading CPU models...") det = load_detector(active_detector) det.lock = RLock() detectors.append(det) if depth_estimator_name: depth = load_depth_estimator(depth_estimator_name) depth.lock = RLock() depth_estimators.append(depth) else: depth_estimators.append(None) # 4. Incremental Depth Stats (replaces expensive pre-scan) depth_stats = IncrementalDepthStats(warmup_frames=30) if depth_estimator_name else None # queue_in: (frame_idx, frame_data) # queue_out: (frame_idx, processed_frame, detections) queue_in = Queue(maxsize=16) # Tuning for A10: buffer at least 32 frames per GPU (batch size) # GPT Latency Buffer: GPT takes ~3s. At 30fps, that's 90 frames. We need to absorb this burst. queue_out_max = max(128, (len(detectors) if detectors else 1) * 32) queue_out = Queue(maxsize=queue_out_max) # 6. Worker Function (Unified) # Robustness: Define flag early so workers can see it writer_finished = False def worker_task(gpu_idx: int): logging.info(f"Worker {gpu_idx} started. PID: {os.getpid()}") detector_instance = detectors[gpu_idx] depth_instance = depth_estimators[gpu_idx] if gpu_idx < len(depth_estimators) else None # Handle mismatched lists safely batch_size = detector_instance.max_batch_size if detector_instance.supports_batch else 1 batch_accum = [] # List[Tuple[idx, frame]] def flush_batch(): if not batch_accum: return logging.info(f"Worker {gpu_idx} flushing batch of {len(batch_accum)} frames") indices = [item[0] for item in batch_accum] frames = [item[1] for item in batch_accum] # --- UNIFIED INFERENCE --- # Separate frame 0 if we have cached detections (avoid re-detecting) cached_frame0 = None detect_indices = indices detect_frames = frames if first_frame_detections is not None and 0 in indices: f0_pos = indices.index(0) cached_frame0 = (indices[f0_pos], frames[f0_pos]) detect_indices = indices[:f0_pos] + indices[f0_pos+1:] detect_frames = frames[:f0_pos] + frames[f0_pos+1:] logging.info("Worker %d: reusing cached detections for frame 0", gpu_idx) # Run detection batch (excluding frame 0 if cached) det_results_map = {} if detect_frames: try: if detector_instance.supports_batch: with detector_instance.lock: raw_results = detector_instance.predict_batch(detect_frames, queries) else: with detector_instance.lock: raw_results = [detector_instance.predict(f, queries) for f in detect_frames] for di, dr in zip(detect_indices, raw_results): det_results_map[di] = dr except BaseException as e: logging.exception("Batch detection crashed with critical error") for di in detect_indices: det_results_map[di] = None # Run depth batch (if enabled) — always for all frames depth_results = [None] * len(frames) if depth_instance and depth_estimator_name: try: with depth_instance.lock: if depth_instance.supports_batch: depth_results = depth_instance.predict_batch(frames) else: depth_results = [depth_instance.predict(f) for f in frames] except BaseException as e: logging.exception("Batch depth crashed with critical error") # Update incremental depth stats if depth_stats is not None: for dep_res in depth_results: if dep_res and dep_res.depth_map is not None: depth_stats.update(dep_res.depth_map) # --- POST PROCESSING --- batch_det_summary = [] for i, (idx, frame, dep_res) in enumerate(zip(indices, frames, depth_results)): # 1. Detections — use cached for frame 0 if available detections = [] if cached_frame0 is not None and idx == 0: detections = [d.copy() for d in first_frame_detections] else: d_res = det_results_map.get(idx) if d_res: detections = _build_detection_records( d_res.boxes, d_res.scores, d_res.labels, queries, d_res.label_names ) batch_det_summary.append((idx, len(detections))) # 2. Frame Rendering processed = frame.copy() # A. Render Depth Heatmap (if enabled) if dep_res and dep_res.depth_map is not None: ds_min, ds_max = depth_stats.range if depth_stats else (0.0, 1.0) processed = colorize_depth_map(dep_res.depth_map, ds_min, ds_max) try: _attach_depth_from_result(detections, dep_res, depth_scale) except: pass # 3. Output while True: try: queue_out.put((idx, processed, detections), timeout=1.0) break except Full: # Robustness: Check if writer is dead if writer_finished: raise RuntimeError("Writer thread died unexpectedly") if job_id: _check_cancellation(job_id) total_dets = sum(c for _, c in batch_det_summary) if total_dets == 0 or indices[0] % 90 == 0: logging.info("Worker %d batch [frames %s]: %d total detections %s", gpu_idx, f"{indices[0]}-{indices[-1]}", total_dets, [(idx, cnt) for idx, cnt in batch_det_summary if cnt > 0]) batch_accum.clear() logging.info(f"Worker {gpu_idx} finished flushing batch") while True: try: item = queue_in.get(timeout=2.0) except Empty: # Periodic check for cancellation if main thread is slow if job_id: _check_cancellation(job_id) continue try: if item is None: logging.info(f"Worker {gpu_idx} received sentinel. Flushing and exiting.") flush_batch() break frame_idx, frame_data = item # logging.info(f"Worker {gpu_idx} got frame {frame_idx}") # Verbose if frame_idx % 30 == 0: logging.info("Processing frame %d on device %s", frame_idx, "cpu" if num_gpus==0 else f"cuda:{gpu_idx}") batch_accum.append((frame_idx, frame_data)) if len(batch_accum) >= batch_size: flush_batch() except BaseException as e: logging.exception(f"Worker {gpu_idx} CRASHED processing frame. Recovering...") # Emit empty/failed frames for the batch to keep sequence alive for idx, frm in batch_accum: try: # Fallback: Return original frame with empty detections queue_out.put((idx, frm, []), timeout=5.0) logging.info(f"Emitted fallback frame {idx}") except: pass batch_accum.clear() finally: queue_in.task_done() logging.info(f"Worker {gpu_idx} thread exiting normally.") # 6. Start Workers workers = [] num_workers = len(detectors) # If using CPU, maybe use more threads? No, CPU models usually multithread internally. # If using GPU, 1 thread per GPU is efficient. for i in range(num_workers): t = Thread(target=worker_task, args=(i,), daemon=True) t.start() workers.append(t) # 7. Start Writer / Output Collection (Main Thread or separate) # We will run writer logic in the main thread after feeding is done? # No, we must write continuously. all_detections_map = {} # writer_finished initialized earlier # writer_finished = False # --- GPT Enrichment Thread (non-blocking) --- # Runs LLM relevance + GPT threat assessment off the writer's critical path. gpt_enrichment_queue = Queue(maxsize=4) _relevance_refined = Event() def enrichment_thread_fn(tracker_ref): """Dedicated thread for GPT/LLM calls. Receives work from writer, injects results via tracker.""" while True: item = gpt_enrichment_queue.get() if item is None: break # Sentinel — shutdown frame_idx, frame_data, gpt_dets, ms = item try: gpt_res = run_enrichment( frame_idx, frame_data, gpt_dets, ms, first_frame_gpt_results=first_frame_gpt_results, job_id=job_id, relevance_refined_event=_relevance_refined, ) if gpt_res: tracker_ref.inject_metadata(gpt_dets) logging.info("Enrichment: GPT results injected into tracker for frame %d", frame_idx) except Exception as e: logging.error("Enrichment thread failed for frame %d: %s", frame_idx, e) def writer_loop(): nonlocal writer_finished next_idx = 0 buffer = {} # Initialize Tracker & Speed Estimator tracker = ByteTracker(frame_rate=fps) speed_est = SpeedEstimator(fps=fps) gpt_submitted = False # GPT enrichment submitted once for frame 0 # Start enrichment thread enrich_thread = Thread(target=enrichment_thread_fn, args=(tracker,), daemon=True) enrich_thread.start() try: with StreamingVideoWriter(output_video_path, fps, width, height) as writer: while next_idx < total_frames: # Fetch from queue try: while next_idx not in buffer: # Backpressure: bound the reorder buffer to prevent memory blowup if len(buffer) > 128: logging.warning("Writer reorder buffer too large (%d items), applying backpressure (waiting for frame %d)...", len(buffer), next_idx) time.sleep(0.05) item = queue_out.get(timeout=1.0) # wait idx, p_frame, dets = item buffer[idx] = (p_frame, dets) # Write next_idx p_frame, dets = buffer.pop(next_idx) # --- SEQUENTIAL TRACKING --- # Run tracker FIRST so detections get real track_id from ByteTracker pre_track_count = len(dets) dets = tracker.update(dets) if (next_idx % 30 == 0) or (pre_track_count > 0 and len(dets) == 0): logging.info("Writer frame %d: %d detections in -> %d tracked out", next_idx, pre_track_count, len(dets)) speed_est.estimate(dets) # --- RELEVANCE GATE (deterministic, fast — stays in writer) --- if mission_spec: if (mission_spec.parse_mode == "LLM_EXTRACTED" and not _relevance_refined.is_set()): # LLM post-filter hasn't run yet — pass all through for d in dets: d["mission_relevant"] = True d["relevance_reason"] = "pending_llm_postfilter" gpt_dets = dets else: # Normal deterministic gate (with refined or FAST_PATH classes) for d in dets: decision = evaluate_relevance(d, mission_spec.relevance_criteria) d["mission_relevant"] = decision.relevant d["relevance_reason"] = decision.reason if not decision.relevant: logging.info( json_module.dumps({ "event": "relevance_decision", "track_id": d.get("track_id"), "label": d.get("label"), "relevant": False, "reason": decision.reason, "required_classes": mission_spec.relevance_criteria.required_classes, "frame": next_idx, }) ) gpt_dets = [d for d in dets if d.get("mission_relevant", True)] else: for d in dets: d["mission_relevant"] = None gpt_dets = dets # --- GPT ENRICHMENT (non-blocking, offloaded to enrichment thread) --- if enable_gpt and gpt_dets and not gpt_submitted: # Tag as pending — enrichment thread will update to ASSESSED later for d in gpt_dets: d["assessment_status"] = AssessmentStatus.PENDING_GPT try: gpt_enrichment_queue.put( (next_idx, p_frame.copy(), gpt_dets, mission_spec), timeout=1.0, ) gpt_submitted = True logging.info("Writer: offloaded GPT enrichment for frame %d", next_idx) except Full: logging.warning("GPT enrichment queue full, skipping frame 0 GPT") # Tag unassessed detections (INV-6) for d in dets: if "assessment_status" not in d: d["assessment_status"] = AssessmentStatus.UNASSESSED # --- RENDER BOXES & OVERLAYS --- if dets: display_boxes = np.array([d['bbox'] for d in dets]) display_labels = [] for d in dets: if d.get("mission_relevant") is False: display_labels.append("") continue lbl = d.get('label', 'obj') display_labels.append(lbl) p_frame = draw_boxes(p_frame, display_boxes, label_names=display_labels) writer.write(p_frame) if stream_queue: try: from jobs.streaming import publish_frame as _publish if job_id: _publish(job_id, p_frame) else: stream_queue.put(p_frame, timeout=0.01) except: pass all_detections_map[next_idx] = dets # Store tracks for frontend access if job_id: set_track_data(job_id, next_idx, dets) next_idx += 1 if next_idx % 30 == 0: logging.debug("Wrote frame %d/%d", next_idx, total_frames) except Empty: # Normal when waiting for out-of-order worker output. if job_id: _check_cancellation(job_id) if not any(w.is_alive() for w in workers) and queue_out.empty(): logging.error( "Workers stopped unexpectedly while waiting for frame %d.", next_idx, ) break continue except Exception: logging.exception("Writer loop processing error at index %d", next_idx) if job_id: _check_cancellation(job_id) if not any(w.is_alive() for w in workers) and queue_out.empty(): logging.error( "Workers stopped unexpectedly while writer handled frame %d.", next_idx, ) break continue except Exception as e: logging.exception("Writer loop failed") finally: logging.info("Writer loop finished. Wrote %d frames (target %d)", next_idx, total_frames) # Shut down enrichment thread try: gpt_enrichment_queue.put(None, timeout=5.0) enrich_thread.join(timeout=30) except Exception: logging.warning("Enrichment thread shutdown timed out") writer_finished = True writer_thread = Thread(target=writer_loop, daemon=True) writer_thread.start() # 8. Feed Frames (Main Thread) # 8. Feed Frames (Main Thread) try: frames_fed = 0 reader_iter = iter(reader) while True: _check_cancellation(job_id) if max_frames is not None and frames_fed >= max_frames: break try: frame = next(reader_iter) except StopIteration: break queue_in.put((frames_fed, frame)) # Blocks if full frames_fed += 1 logging.info("Feeder finished. Fed %d frames (expected %d)", frames_fed, total_frames) # Update total_frames to actual count so writer knows when to stop if frames_fed != total_frames: logging.info("Updating total_frames from %d to %d (actual fed)", total_frames, frames_fed) total_frames = frames_fed # Signal workers to stop for _ in range(num_workers): try: queue_in.put(None, timeout=5.0) # Using timeout to prevent infinite block except Full: logging.warning("Failed to send stop signal to a worker (queue full)") # Wait for queue to process queue_in.join() except Exception as e: logging.exception("Feeding frames failed") # Ensure we try to signal workers even on error for _ in range(num_workers): try: queue_in.put_nowait(None) except Full: pass raise finally: reader.close() # Wait for writer writer_thread.join() # Sort detections sorted_detections = [] # If we crashed early, we return what we have max_key = max(all_detections_map.keys()) if all_detections_map else -1 for i in range(max_key + 1): sorted_detections.append(all_detections_map.get(i, [])) logging.info("Inference complete. Output: %s", output_video_path) return output_video_path, sorted_detections def _gsam2_render_frame( frame_dir: str, frame_names: List[str], frame_idx: int, frame_objects: Dict, height: int, width: int, masks_only: bool = False, frame_store=None, ) -> np.ndarray: """Render a single GSAM2 tracking frame (masks + boxes). CPU-only. When *masks_only* is True, skip box rendering so the writer thread can draw boxes later with enriched (GPT) labels. """ if frame_store is not None: frame = frame_store.get_bgr(frame_idx).copy() # .copy() — render mutates else: frame_path = os.path.join(frame_dir, frame_names[frame_idx]) frame = cv2.imread(frame_path) if frame is None: return np.zeros((height, width, 3), dtype=np.uint8) if not frame_objects: return frame masks_list: List[np.ndarray] = [] boxes_list: List[List[int]] = [] box_labels: List[str] = [] for _obj_id, obj_info in frame_objects.items(): mask = obj_info.mask label = obj_info.class_name if mask is not None: if isinstance(mask, torch.Tensor): mask_np = mask.cpu().numpy().astype(bool) else: mask_np = np.asarray(mask).astype(bool) if mask_np.shape[:2] != (height, width): mask_np = cv2.resize( mask_np.astype(np.uint8), (width, height), interpolation=cv2.INTER_NEAREST, ).astype(bool) masks_list.append(mask_np) has_box = not ( obj_info.x1 == 0 and obj_info.y1 == 0 and obj_info.x2 == 0 and obj_info.y2 == 0 ) if has_box: boxes_list.append([obj_info.x1, obj_info.y1, obj_info.x2, obj_info.y2]) box_labels.append(label) if masks_list: # Always pass labels=None here; label text is drawn by draw_boxes # below to avoid duplicate label rendering. frame = draw_masks(frame, np.stack(masks_list), labels=None) if boxes_list and not masks_only: frame = draw_boxes(frame, np.array(boxes_list), label_names=box_labels) return frame def run_grounded_sam2_tracking( input_video_path: str, output_video_path: str, queries: List[str], max_frames: Optional[int] = None, segmenter_name: Optional[str] = None, job_id: Optional[str] = None, stream_queue: Optional[Queue] = None, step: int = 20, enable_gpt: bool = False, mission_spec=None, # Optional[MissionSpecification] first_frame_gpt_results: Optional[Dict[str, Any]] = None, _perf_metrics: Optional[Dict[str, float]] = None, _perf_lock=None, num_maskmem: Optional[int] = None, detector_name: Optional[str] = None, _ttfs_t0: Optional[float] = None, ) -> str: """Run Grounded-SAM-2 video tracking pipeline. Uses multi-GPU data parallelism when multiple GPUs are available. Falls back to single-GPU ``process_video`` otherwise. """ import copy import shutil from contextlib import nullcontext from PIL import Image as PILImage from utils.video import extract_frames_to_jpeg_dir from utils.frame_store import SharedFrameStore, MemoryBudgetExceeded from models.segmenters.grounded_sam2 import MaskDictionary, ObjectInfo, LazyFrameObjects active_segmenter = segmenter_name or "GSAM2-L" def _ttfs(msg): if _ttfs_t0 is not None: logging.info("[TTFS:%s] +%.1fs %s", job_id, time.perf_counter() - _ttfs_t0, msg) _ttfs("enter run_grounded_sam2_tracking") logging.info( "Grounded-SAM-2 tracking: segmenter=%s, queries=%s, step=%d", active_segmenter, queries, step, ) # 1. Load frames — prefer in-memory SharedFrameStore, fall back to JPEG dir _use_frame_store = True frame_store = None _t_ext = time.perf_counter() try: frame_store = SharedFrameStore(input_video_path, max_frames=max_frames) fps, width, height = frame_store.fps, frame_store.width, frame_store.height total_frames = len(frame_store) frame_names = [f"{i:06d}.jpg" for i in range(total_frames)] # Write single dummy JPEG for init_state bootstrapping dummy_frame_dir = tempfile.mkdtemp(prefix="gsam2_dummy_") cv2.imwrite(os.path.join(dummy_frame_dir, "000000.jpg"), frame_store.get_bgr(0)) frame_dir = dummy_frame_dir logging.info("SharedFrameStore: %d frames in memory (dummy dir: %s)", total_frames, frame_dir) except MemoryBudgetExceeded: logging.info("Memory budget exceeded, falling back to JPEG extraction") _use_frame_store = False frame_store = None frame_dir = tempfile.mkdtemp(prefix="gsam2_frames_") frame_names, fps, width, height = extract_frames_to_jpeg_dir( input_video_path, frame_dir, max_frames=max_frames, ) total_frames = len(frame_names) try: if _perf_metrics is not None: _t_e2e = time.perf_counter() if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() _perf_metrics["frame_extraction_ms"] = (time.perf_counter() - _t_ext) * 1000.0 _ttfs(f"frame_extraction done ({total_frames} frames, in_memory={_use_frame_store})") logging.info("Loaded %d frames (in_memory=%s)", total_frames, _use_frame_store) num_gpus = torch.cuda.device_count() # ================================================================== # Phase 5: Parallel rendering + sequential video writing # (Hoisted above tracking so render pipeline starts before tracking # completes — segments are fed incrementally via callback / queue.) # ================================================================== _check_cancellation(job_id) render_in: Queue = Queue(maxsize=32) render_out: Queue = Queue(maxsize=128) render_done = False num_render_workers = min(4, os.cpu_count() or 1) def _render_worker(): while True: item = render_in.get() if item is None: break fidx, fobjs = item try: # Deferred GPU->CPU: materialize lazy objects in render thread if isinstance(fobjs, LazyFrameObjects): fobjs = fobjs.materialize() if _perf_metrics is not None: _t_r = time.perf_counter() frm = _gsam2_render_frame( frame_dir, frame_names, fidx, fobjs, height, width, masks_only=enable_gpt, frame_store=frame_store, ) if _perf_metrics is not None: _r_ms = (time.perf_counter() - _t_r) * 1000.0 if _perf_lock: with _perf_lock: _perf_metrics["render_total_ms"] += _r_ms else: _perf_metrics["render_total_ms"] += _r_ms payload = (fidx, frm, fobjs) if enable_gpt else (fidx, frm, {}) while True: try: render_out.put(payload, timeout=1.0) break except Full: if render_done: return except Exception: logging.exception("Render failed for frame %d", fidx) blank = np.zeros((height, width, 3), dtype=np.uint8) try: render_out.put((fidx, blank, {}), timeout=5.0) except Full: pass r_workers = [ Thread(target=_render_worker, daemon=True) for _ in range(num_render_workers) ] for t in r_workers: t.start() # --- ObjectInfo → detection dict adapter --- def _objectinfo_to_dets(frame_objects_dict): dets = [] for obj_id, info in frame_objects_dict.items(): dets.append({ "label": info.class_name, "bbox": [info.x1, info.y1, info.x2, info.y2], "score": 1.0, "track_id": f"T{obj_id:02d}", "instance_id": obj_id, }) return dets # --- GPT enrichment thread (when enabled) --- gpt_enrichment_queue: Queue = Queue(maxsize=4) gpt_data_by_track: Dict[str, Dict] = {} gpt_data_lock = RLock() _relevance_refined = Event() def _gsam2_enrichment_thread_fn(): while True: item = gpt_enrichment_queue.get() if item is None: break frame_idx, frame_data, gpt_dets, ms = item try: gpt_res = run_enrichment( frame_idx, frame_data, gpt_dets, ms, first_frame_gpt_results=first_frame_gpt_results, job_id=job_id, relevance_refined_event=_relevance_refined, ) # GSAM2-specific: store results in per-track dict and persist to job storage if gpt_res: for d in gpt_dets: tid = d.get("track_id") if tid and tid in gpt_res: merged = dict(gpt_res[tid]) merged["gpt_raw"] = gpt_res[tid] merged["assessment_frame_index"] = frame_idx merged["assessment_status"] = merged.get( "assessment_status", AssessmentStatus.ASSESSED ) with gpt_data_lock: gpt_data_by_track[tid] = merged logging.info("GSAM2 enrichment: GPT results stored for %d tracks", len(gpt_data_by_track)) # Persist GPT-enriched detections to job storage so # frontend polling (/detect/status) picks them up. if job_id: try: from jobs.storage import get_job_storage as _gjs _st = _gjs().get(job_id) if _st and _st.first_frame_detections: for det in _st.first_frame_detections: tid = det.get("track_id") with gpt_data_lock: payload = gpt_data_by_track.get(tid) if payload: det.update(payload) # Also sync relevance from gpt_dets src = next((d for d in gpt_dets if d.get("track_id") == tid), None) if src: if "mission_relevant" in src: det["mission_relevant"] = src["mission_relevant"] if "relevance_reason" in src: det["relevance_reason"] = src["relevance_reason"] from jobs.storage import get_job_storage as _gjs2 _gjs2().update( job_id, first_frame_detections=_st.first_frame_detections, first_frame_gpt_results=gpt_res, ) logging.info( "GSAM2 enrichment: updated first_frame_detections in job storage for %s", job_id, ) except Exception: logging.exception( "GSAM2 enrichment: failed to update job storage for %s", job_id ) except Exception as e: logging.error("GSAM2 enrichment thread failed for frame %d: %s", frame_idx, e) # Shared streaming state (publisher ↔ writer) _stream_deque: collections.deque = collections.deque(maxlen=200) _stream_lock = RLock() _stream_writer_done = Event() def _writer_loop(): nonlocal render_done _first_deposit = False next_idx = 0 buf: Dict[int, Tuple] = {} # Per-track bbox history (replaces ByteTracker for GSAM2) track_history: Dict[int, List] = {} speed_est = SpeedEstimator(fps=fps) if enable_gpt else None gpt_submitted = False # Start enrichment thread when GPT enabled enrich_thread = None if enable_gpt: enrich_thread = Thread(target=_gsam2_enrichment_thread_fn, daemon=True) enrich_thread.start() try: with StreamingVideoWriter( output_video_path, fps, width, height ) as writer: # --- Write + stream (publisher handles pacing) --- while next_idx < total_frames: try: while next_idx not in buf: if len(buf) > 128: logging.warning( "Render reorder buffer large (%d), " "waiting for frame %d", len(buf), next_idx, ) time.sleep(0.05) idx, frm, fobjs = render_out.get(timeout=1.0) buf[idx] = (frm, fobjs) frm, fobjs = buf.pop(next_idx) # --- GPT enrichment path --- if enable_gpt and fobjs: dets = _objectinfo_to_dets(fobjs) # Maintain per-track bbox history (30-frame window) for det in dets: iid = det["instance_id"] track_history.setdefault(iid, []).append(det["bbox"]) if len(track_history[iid]) > 30: track_history[iid].pop(0) # Store an immutable per-frame snapshot. det["history"] = list(track_history[iid]) # Speed estimation if speed_est: speed_est.estimate(dets) # Relevance gate if mission_spec: if (mission_spec.parse_mode == "LLM_EXTRACTED" and not _relevance_refined.is_set()): for d in dets: d["mission_relevant"] = True d["relevance_reason"] = "pending_llm_postfilter" gpt_dets = dets else: for d in dets: decision = evaluate_relevance(d, mission_spec.relevance_criteria) d["mission_relevant"] = decision.relevant d["relevance_reason"] = decision.reason gpt_dets = [d for d in dets if d.get("mission_relevant", True)] else: for d in dets: d["mission_relevant"] = None gpt_dets = dets # GPT enrichment (one-shot, first frame with detections) if gpt_dets and not gpt_submitted: for d in gpt_dets: d["assessment_status"] = AssessmentStatus.PENDING_GPT try: gpt_enrichment_queue.put( ( next_idx, frm.copy(), copy.deepcopy(gpt_dets), mission_spec, ), timeout=1.0, ) gpt_submitted = True logging.info("GSAM2 writer: offloaded GPT enrichment for frame %d", next_idx) except Full: logging.warning("GSAM2 GPT enrichment queue full, skipping") # Merge persistent GPT data for det in dets: tid = det["track_id"] with gpt_data_lock: gpt_payload = gpt_data_by_track.get(tid) if gpt_payload: det.update(gpt_payload) det["assessment_status"] = AssessmentStatus.ASSESSED elif "assessment_status" not in det: det["assessment_status"] = AssessmentStatus.UNASSESSED # Build enriched display labels display_labels = [] for d in dets: if d.get("mission_relevant") is False: display_labels.append("") continue lbl = d.get("label", "obj") if d.get("gpt_distance_m") is not None: try: lbl = f"{lbl} {int(float(d['gpt_distance_m']))}m" except (TypeError, ValueError): pass display_labels.append(lbl) # Draw boxes on mask-rendered frame if dets: boxes = np.array([d["bbox"] for d in dets]) frm = draw_boxes(frm, boxes, label_names=display_labels) # Store tracks for frontend if job_id: set_track_data(job_id, next_idx, copy.deepcopy(dets)) elif enable_gpt: # No objects this frame — still store empty track data if job_id: set_track_data(job_id, next_idx, []) if _perf_metrics is not None: _t_w = time.perf_counter() # Write to video file (always, single copy) writer.write(frm) if _perf_metrics is not None: _perf_metrics["writer_total_ms"] += (time.perf_counter() - _t_w) * 1000.0 # --- Deposit frame for stream publisher --- if stream_queue or job_id: with _stream_lock: _stream_deque.append(frm) if not _first_deposit: _first_deposit = True _ttfs("first_frame_deposited_to_deque") next_idx += 1 if next_idx % 30 == 0: logging.info( "Rendered frame %d / %d", next_idx, total_frames, ) except Empty: if job_id: _check_cancellation(job_id) if not any(t.is_alive() for t in r_workers) and render_out.empty(): logging.error( "Render workers stopped while waiting " "for frame %d", next_idx, ) break continue finally: render_done = True _stream_writer_done.set() # Shut down enrichment thread if enrich_thread: try: gpt_enrichment_queue.put(None, timeout=5.0) enrich_thread.join(timeout=30) except Exception: logging.warning("GSAM2 enrichment thread shutdown timed out") def _stream_publisher_thread(): """Adaptive-rate publisher: reads from _stream_deque, publishes at measured pace.""" from jobs.streaming import publish_frame as _pub STARTUP_WAIT = 5.0 # max seconds to accumulate before streaming MIN_FPS = 2.0 MAX_FPS = 30.0 HEARTBEAT_INTERVAL = 0.5 # re-publish last frame if deque empty LOW_WATER = 10 HIGH_WATER = 50 ADJUST_INTERVAL = 1.0 last_frame = None published = 0 # --- Phase 1: startup accumulation --- t_start = time.perf_counter() while True: elapsed = time.perf_counter() - t_start if elapsed >= STARTUP_WAIT: break if _stream_writer_done.is_set(): break time.sleep(0.1) with _stream_lock: accumulated = len(_stream_deque) elapsed = max(time.perf_counter() - t_start, 0.1) r_prod = accumulated / elapsed if accumulated > 0 else 10.0 r_stream = max(MIN_FPS, min(MAX_FPS, 0.85 * r_prod)) logging.info( "Stream publisher started: R_prod=%.1f fps, R_stream=%.1f fps, " "accumulated=%d frames in %.1fs", r_prod, r_stream, accumulated, elapsed, ) _ttfs(f"publisher: startup_wait done ({accumulated} frames in {elapsed:.1f}s)") # --- Phase 2: adaptive streaming --- last_adjust = time.perf_counter() last_publish_time = 0.0 while True: frame_interval = 1.0 / r_stream # Try to pop a frame frame = None with _stream_lock: if _stream_deque: frame = _stream_deque.popleft() if frame is not None: last_frame = frame if job_id: _pub(job_id, frame) elif stream_queue: try: stream_queue.put(frame, timeout=0.01) except Exception: pass if published == 0: _ttfs("first_publish_frame") published += 1 last_publish_time = time.perf_counter() time.sleep(frame_interval) else: # Deque empty — check termination if _stream_writer_done.is_set(): with _stream_lock: if not _stream_deque: break continue # Heartbeat: re-publish last frame to keep MJPEG alive now = time.perf_counter() if last_frame is not None and (now - last_publish_time) >= HEARTBEAT_INTERVAL: if job_id: _pub(job_id, last_frame) elif stream_queue: try: stream_queue.put(last_frame, timeout=0.01) except Exception: pass last_publish_time = now time.sleep(0.05) # Adaptive rate adjustment (every ~1s) now = time.perf_counter() if now - last_adjust >= ADJUST_INTERVAL: with _stream_lock: level = len(_stream_deque) if level < LOW_WATER: r_stream = max(MIN_FPS, r_stream * 0.9) elif level > HIGH_WATER: r_stream = min(MAX_FPS, r_stream * 1.05) last_adjust = now # Publish final frame if last_frame is not None: if job_id: _pub(job_id, last_frame) elif stream_queue: try: stream_queue.put(last_frame, timeout=0.01) except Exception: pass logging.info("Stream publisher finished: published %d frames", published) writer_thread = Thread(target=_writer_loop, daemon=True) writer_thread.start() _publisher_thread = None if stream_queue or job_id: _publisher_thread = Thread(target=_stream_publisher_thread, daemon=True) _publisher_thread.start() _ttfs("writer+publisher threads started") # ================================================================== # Phase 1-4: Tracking (single-GPU fallback vs multi-GPU pipeline) # Segments are fed incrementally to render_in as they complete. # ================================================================== try: if num_gpus <= 1: # ---------- Single-GPU fallback ---------- device_str = "cuda:0" if torch.cuda.is_available() else "cpu" _seg_kw = {"num_maskmem": num_maskmem} if num_maskmem is not None else {} if detector_name is not None: _seg_kw["detector_name"] = detector_name if _perf_metrics is not None: _t_load = time.perf_counter() segmenter = load_segmenter_on_device(active_segmenter, device_str, **_seg_kw) _check_cancellation(job_id) if _perf_metrics is not None: _perf_metrics["model_load_ms"] = (time.perf_counter() - _t_load) * 1000.0 segmenter._perf_metrics = _perf_metrics segmenter._perf_lock = None _ttfs(f"model loaded ({active_segmenter})") if _perf_metrics is not None: _t_track = time.perf_counter() def _feed_segment(seg_frames): """Fallback for empty/carry-forward segments (already CPU).""" for fidx in sorted(seg_frames.keys()): render_in.put((fidx, seg_frames[fidx])) def _feed_segment_gpu(segment_output): """Feed LazyFrameObjects into render_in (GPU->CPU deferred).""" # Deduplicate: frame_indices has one entry per (frame, obj) seen = set() for fi in segment_output.frame_indices: if fi not in seen: seen.add(fi) render_in.put((fi, LazyFrameObjects(segment_output, fi))) _ttfs("process_video started") tracking_results = segmenter.process_video( frame_dir, frame_names, queries, on_segment=_feed_segment, on_segment_output=_feed_segment_gpu, _ttfs_t0=_ttfs_t0, _ttfs_job_id=job_id, frame_store=frame_store, ) if _perf_metrics is not None: _perf_metrics["tracking_total_ms"] = (time.perf_counter() - _t_track) * 1000.0 logging.info( "Single-GPU tracking complete: %d frames", len(tracking_results), ) else: # ---------- Multi-GPU pipeline ---------- logging.info( "Multi-GPU GSAM2 tracking: %d GPUs, %d frames, step=%d", num_gpus, total_frames, step, ) # Phase 1: Load one segmenter per GPU (parallel) if _perf_metrics is not None: _t_load = time.perf_counter() segmenters = [] with ThreadPoolExecutor(max_workers=num_gpus) as pool: _seg_kw_multi = {"num_maskmem": num_maskmem} if num_maskmem is not None else {} if detector_name is not None: _seg_kw_multi["detector_name"] = detector_name futs = [ pool.submit( load_segmenter_on_device, active_segmenter, f"cuda:{i}", **_seg_kw_multi, ) for i in range(num_gpus) ] segmenters = [f.result() for f in futs] logging.info("Loaded %d segmenters", len(segmenters)) if _perf_metrics is not None: _perf_metrics["model_load_ms"] = (time.perf_counter() - _t_load) * 1000.0 import threading as _th _actual_lock = _perf_lock or _th.Lock() for seg in segmenters: seg._perf_metrics = _perf_metrics seg._perf_lock = _actual_lock _ttfs(f"model loaded ({active_segmenter}, {num_gpus} GPUs)") # Phase 2: Init SAM2 models/state per GPU (parallel) if _perf_metrics is not None: _t_init = time.perf_counter() if frame_store is not None: # Models are lazy-loaded; ensure at least one is ready so we # can read image_size. Phase 1 (load_segmenter_on_device) # only constructs the object — _video_predictor is still None. segmenters[0]._ensure_models_loaded() sam2_img_size = segmenters[0]._video_predictor.image_size # Pre-create the shared adapter (validates memory budget) shared_adapter = frame_store.sam2_adapter(image_size=sam2_img_size) _REQUIRED_KEYS = {"images", "num_frames", "video_height", "video_width", "cached_features"} def _init_seg_state(seg): seg._ensure_models_loaded() state = seg._video_predictor.init_state( video_path=frame_dir, # dummy dir with 1 JPEG offload_video_to_cpu=True, async_loading_frames=False, # 1 dummy frame, instant ) # Validate expected keys exist before patching missing = _REQUIRED_KEYS - set(state.keys()) if missing: raise RuntimeError(f"SAM2 init_state missing expected keys: {missing}") # CRITICAL: Clear cached_features BEFORE patching images # init_state caches dummy frame 0's backbone features — must evict state["cached_features"] = {} # Patch in real frame data state["images"] = shared_adapter state["num_frames"] = total_frames state["video_height"] = height state["video_width"] = width return state else: def _init_seg_state(seg): seg._ensure_models_loaded() return seg._video_predictor.init_state( video_path=frame_dir, offload_video_to_cpu=True, async_loading_frames=True, ) with ThreadPoolExecutor(max_workers=len(segmenters)) as pool: futs = [pool.submit(_init_seg_state, seg) for seg in segmenters] inference_states = [f.result() for f in futs] if _perf_metrics is not None: _perf_metrics["init_state_ms"] = (time.perf_counter() - _t_init) * 1000.0 _t_track = time.perf_counter() _ttfs("multi-GPU tracking started") # Phase 3: Parallel segment processing (queue-based workers) segments = list(range(0, total_frames, step)) num_total_segments = len(segments) seg_queue_in: Queue = Queue() seg_queue_out: Queue = Queue() for i, start_idx in enumerate(segments): seg_queue_in.put((i, start_idx)) for _ in segmenters: seg_queue_in.put(None) # sentinel iou_thresh = segmenters[0].iou_threshold def _segment_worker(gpu_idx: int): seg = segmenters[gpu_idx] state = inference_states[gpu_idx] device_type = seg.device.split(":")[0] ac = ( torch.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext() ) with ac: while True: if job_id: try: _check_cancellation(job_id) except RuntimeError as e: if "cancelled" in str(e).lower(): logging.info( "Segment worker %d cancelled.", gpu_idx, ) break raise item = seg_queue_in.get() if item is None: break seg_idx, start_idx = item try: logging.info( "GPU %d processing segment %d (frame %d)", gpu_idx, seg_idx, start_idx, ) if frame_store is not None: image = frame_store.get_pil_rgb(start_idx) else: img_path = os.path.join( frame_dir, frame_names[start_idx] ) with PILImage.open(img_path) as pil_img: image = pil_img.convert("RGB") if job_id: _check_cancellation(job_id) masks, boxes, labels = seg.detect_keyframe( image, queries, ) if masks is None: seg_queue_out.put( (seg_idx, start_idx, None, {}) ) continue mask_dict = MaskDictionary() mask_dict.add_new_frame_annotation( mask_list=masks, box_list=( boxes.clone() if torch.is_tensor(boxes) else torch.tensor(boxes) ), label_list=labels, ) segment_output = seg.propagate_segment( state, start_idx, mask_dict, step, ) seg_queue_out.put( (seg_idx, start_idx, mask_dict, segment_output) ) except RuntimeError as e: if "cancelled" in str(e).lower(): logging.info( "Segment worker %d cancelled.", gpu_idx, ) break raise except Exception: logging.exception( "Segment %d failed on GPU %d", seg_idx, gpu_idx, ) seg_queue_out.put( (seg_idx, start_idx, None, {}) ) seg_workers = [] for i in range(num_gpus): t = Thread( target=_segment_worker, args=(i,), daemon=True, ) t.start() seg_workers.append(t) # Phase 4: Streaming reconciliation — process segments in order # as they arrive, feeding render_in incrementally. _recon_accum_ms = 0.0 global_id_counter = 0 sam2_masks = MaskDictionary() tracking_results: Dict[int, Dict[int, ObjectInfo]] = {} def _mask_to_cpu(mask): """Normalize mask to CPU tensor (still used for keyframe mask_dict).""" if torch.is_tensor(mask): return mask.detach().cpu() return mask next_seg_idx = 0 segment_buffer: Dict[int, Tuple] = {} while next_seg_idx < num_total_segments: try: seg_idx, start_idx, mask_dict, segment_output = seg_queue_out.get(timeout=1.0) except Empty: if job_id: _check_cancellation(job_id) # Check if all segment workers are still alive if not any(t.is_alive() for t in seg_workers) and seg_queue_out.empty(): logging.error( "All segment workers stopped while waiting for segment %d", next_seg_idx, ) break continue segment_buffer[seg_idx] = (start_idx, mask_dict, segment_output) # Process contiguous ready segments in order while next_seg_idx in segment_buffer: start_idx, mask_dict, segment_output = segment_buffer.pop(next_seg_idx) if mask_dict is None or not mask_dict.labels: # No detections — carry forward previous masks for fi in range( start_idx, min(start_idx + step, total_frames) ): if fi not in tracking_results: tracking_results[fi] = ( { k: ObjectInfo( instance_id=v.instance_id, mask=v.mask, class_name=v.class_name, x1=v.x1, y1=v.y1, x2=v.x2, y2=v.y2, ) for k, v in sam2_masks.labels.items() } if sam2_masks.labels else {} ) render_in.put((fi, tracking_results.get(fi, {}))) if next_seg_idx == 0: _ttfs("first_segment_reconciled (multi-GPU, no detections)") next_seg_idx += 1 continue # Normalize keyframe masks to CPU before cross-GPU IoU matching. if _perf_metrics is not None: _t_rc = time.perf_counter() for info in mask_dict.labels.values(): info.mask = _mask_to_cpu(info.mask) # IoU match + get local→global remapping global_id_counter, remapping = ( mask_dict.update_masks_with_remapping( tracking_dict=sam2_masks, iou_threshold=iou_thresh, objects_count=global_id_counter, ) ) if not mask_dict.labels: if _perf_metrics is not None: _recon_accum_ms += (time.perf_counter() - _t_rc) * 1000.0 for fi in range( start_idx, min(start_idx + step, total_frames) ): tracking_results[fi] = {} render_in.put((fi, {})) if next_seg_idx == 0: _ttfs("first_segment_reconciled (multi-GPU, empty masks)") next_seg_idx += 1 continue # Materialize ONLY the last frame for IoU tracking continuity last_fi = segment_output.last_frame_idx() if last_fi is not None: last_objs = segment_output.frame_to_object_dict( last_fi, remapping=remapping, to_cpu=True, ) tracking_results[last_fi] = last_objs sam2_masks = MaskDictionary() sam2_masks.labels = copy.deepcopy(last_objs) if last_objs: first_info = next(iter(last_objs.values())) if first_info.mask is not None: m = first_info.mask sam2_masks.mask_height = ( m.shape[-2] if m.ndim >= 2 else 0 ) sam2_masks.mask_width = ( m.shape[-1] if m.ndim >= 2 else 0 ) if _perf_metrics is not None: _recon_accum_ms += (time.perf_counter() - _t_rc) * 1000.0 # Feed LazyFrameObjects to render — GPU->CPU deferred to render workers seen_fi: set = set() for fi in segment_output.frame_indices: if fi not in seen_fi: seen_fi.add(fi) render_in.put(( fi, LazyFrameObjects(segment_output, fi, remapping), )) if next_seg_idx == 0: _ttfs("first_segment_reconciled (multi-GPU)") next_seg_idx += 1 for t in seg_workers: t.join() if _perf_metrics is not None: _perf_metrics["id_reconciliation_ms"] = _recon_accum_ms _perf_metrics["tracking_total_ms"] = (time.perf_counter() - _t_track) * 1000.0 logging.info( "Multi-GPU reconciliation complete: %d frames, %d objects", len(tracking_results), global_id_counter, ) finally: # Sentinels for render workers — always sent even on error/cancellation for _ in r_workers: try: render_in.put(None, timeout=5.0) except Full: pass for t in r_workers: t.join() writer_thread.join() if _publisher_thread is not None: _publisher_thread.join(timeout=15) if _perf_metrics is not None: _perf_metrics["end_to_end_ms"] = (time.perf_counter() - _t_e2e) * 1000.0 if torch.cuda.is_available(): _perf_metrics["gpu_peak_mem_mb"] = torch.cuda.max_memory_allocated() / (1024 * 1024) logging.info("Grounded-SAM-2 output written to: %s", output_video_path) return output_video_path finally: try: shutil.rmtree(frame_dir) logging.info("Cleaned up temp frame dir: %s", frame_dir) except Exception: logging.warning("Failed to clean up temp frame dir: %s", frame_dir) def colorize_depth_map( depth_map: np.ndarray, global_min: float, global_max: float, ) -> np.ndarray: """ Convert depth map to RGB visualization using TURBO colormap. Args: depth_map: HxW float32 depth in meters global_min: Minimum depth across entire video (for stable normalization) global_max: Maximum depth across entire video (for stable normalization) Returns: HxWx3 uint8 RGB image """ import cv2 depth_clean = np.copy(depth_map) finite_mask = np.isfinite(depth_clean) if not np.isfinite(global_min) or not np.isfinite(global_max): if finite_mask.any(): local_depths = depth_clean[finite_mask].astype(np.float64, copy=False) global_min = float(np.percentile(local_depths, 1)) global_max = float(np.percentile(local_depths, 99)) else: global_min = 0.0 global_max = 1.0 # Replace NaN/inf with min value for visualization depth_clean[~finite_mask] = global_min if global_max - global_min < 1e-6: # Handle uniform depth depth_norm = np.zeros_like(depth_clean, dtype=np.uint8) else: # Clip to global range to handle outliers depth_clipped = np.clip(depth_clean, global_min, global_max) depth_norm = ((depth_clipped - global_min) / (global_max - global_min) * 255).astype(np.uint8) # Apply TURBO colormap for vibrant, perceptually uniform visualization colored = cv2.applyColorMap(depth_norm, cv2.COLORMAP_TURBO) return colored