inference.py

# 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 jobs.storage import set_track_data, store_latest_frame
from inspection.masks import rle_encode
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 = (128, 128, 128) if not label else _color_for_label(label)
        cv2.rectangle(output, (x1, y1), (x2, y2), color, thickness=2)
        if label:
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.6
            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 = 0.6
                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")

        display_labels = [det["label"] for det in detections]

    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 for a detection."""
    return det["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 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,
    stream_queue: Optional[Queue] = None,
    first_frame_detections: Optional[List[Dict[str, Any]]] = None,
    _ttfs_t0: Optional[float] = None,
) -> Tuple[str, List[List[Dict[str, Any]]]]:

    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_inference")

    # 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)

    _ttfs(f"video_opened ({total_frames} frames, {width}x{height}, {fps:.1f}fps)")

    # 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)
        # Pre-download weights before parallel GPU init to avoid race conditions
        from models.model_loader import prefetch_weights
        prefetch_weights(active_detector)

        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)

    _ttfs(f"models_loaded ({active_detector}, {num_gpus} GPUs, depth={'yes' if depth_estimator_name else 'no'})")

    # 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)
    for i in range(num_workers):
        t = Thread(target=worker_task, args=(i,), daemon=True)
        t.start()
        workers.append(t)

    _ttfs(f"workers_started ({num_workers} GPU workers)")

    # 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 

    
    _first_frame_published = False

    def writer_loop():
        nonlocal writer_finished, _first_frame_published
        next_idx = 0
        buffer = {}

        # Initialize Tracker & Speed Estimator
        tracker = ByteTracker(frame_rate=fps)
        speed_est = SpeedEstimator(fps=fps)

        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)

                        if next_idx == 0:
                            _ttfs("first_frame_dequeued_from_gpu")

                        # Write next_idx
                        p_frame, dets = buffer.pop(next_idx)

                        # --- SEQUENTIAL TRACKING ---
                        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)

                        # --- RENDER BOXES & OVERLAYS ---
                        if dets:
                            display_boxes = np.array([d['bbox'] for d in dets])
                            display_labels = [d.get('label', 'obj') for d in dets]
                            p_frame = draw_boxes(p_frame, display_boxes, label_names=display_labels)

                        writer.write(p_frame)

                        if stream_queue:
                            from jobs.streaming import publish_frame as _publish
                            if job_id:
                                _publish(job_id, p_frame)
                            else:
                                stream_queue.put(p_frame)
                            if not _first_frame_published:
                                _first_frame_published = True
                                _ttfs("first_frame_published_to_stream")

                        all_detections_map[next_idx] = dets

                        # Store tracks for frontend access
                        if job_id:
                            set_track_data(job_id, next_idx, dets)
                            store_latest_frame(job_id, p_frame)

                        next_idx += 1

                        if next_idx % 30 == 0:
                            logging.debug("Wrote frame %d/%d", next_idx, total_frames)

                    except Empty:
                        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)
            writer_finished = True

    writer_thread = Thread(target=writer_loop, daemon=True)
    writer_thread.start()

    _ttfs("writer_thread_started")

    # 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
            if frames_fed == 0:
                _ttfs("first_frame_fed_to_gpu")
            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,
    frame_store=None,
) -> np.ndarray:
    """Render a single GSAM2 tracking frame (masks only, no bboxes). CPU-only."""
    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] = []
    mask_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)
            mask_labels.append(label)

    if masks_list:
        # Draw masks with labels — no bboxes for segmentation mode
        frame = draw_masks(frame, np.stack(masks_list), labels=mask_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,
    _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,
                        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)
                    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 = []
            masks_rle = {}  # {instance_id: rle_dict}
            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": str(obj_id),
                    "instance_id": obj_id,
                })
                # RLE-encode the mask if present
                if info.mask is not None:
                    try:
                        if isinstance(info.mask, torch.Tensor):
                            mask_np = info.mask.cpu().numpy().astype(bool)
                        else:
                            mask_np = np.asarray(info.mask).astype(bool)
                        masks_rle[obj_id] = rle_encode(mask_np)
                    except Exception:
                        logging.debug("Failed to RLE-encode mask for obj %d", obj_id)
            return dets, masks_rle

        # Shared streaming state (publisher ↔ writer)
        _stream_deque: collections.deque = collections.deque()  # unbounded — publisher drains at its own pace
        _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)

            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)

                            # Build detection records and track history
                            if fobjs:
                                dets, masks_rle = _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)
                                    det["history"] = list(track_history[iid])

                                speed_est.estimate(dets)

                                # Store tracks for frontend
                                if job_id:
                                    set_track_data(job_id, next_idx, copy.deepcopy(dets))
                                    store_latest_frame(job_id, frm)
                                    # Store masks
                                    from jobs.storage import set_mask_data as _set_mask
                                    for iid, rle in masks_rle.items():
                                        _set_mask(job_id, next_idx, iid, rle)
                            else:
                                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()

        def _stream_publisher_thread():
            """Forward every frame from deque to stream queue. Consumer handles pacing."""
            from jobs.streaming import publish_frame as _pub

            published = 0

            while True:
                frame = None
                with _stream_lock:
                    if _stream_deque:
                        frame = _stream_deque.popleft()

                if frame is not None:
                    if job_id:
                        _pub(job_id, frame)
                    elif stream_queue:
                        stream_queue.put(frame)
                    if published == 0:
                        _ttfs("first_publish_frame")
                    published += 1
                else:
                    # Deque empty — check termination
                    if _stream_writer_done.is_set():
                        with _stream_lock:
                            if not _stream_deque:
                                break
                        continue
                    time.sleep(0.01)  # brief wait for next frame

            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