Harden GSAM2 parallel pipeline and tracking reconciliation

inference.py

@@ -1586,6 +1586,63 @@ def run_segmentation(
 
 
 
+def _gsam2_render_frame(
+    frame_dir: str,
+    frame_names: List[str],
+    frame_idx: int,
+    frame_objects: Dict,
+    height: int,
+    width: int,
+) -> np.ndarray:
+    """Render a single GSAM2 tracking frame (masks + boxes). CPU-only."""
+    from models.segmenters.grounded_sam2 import ObjectInfo
+
+    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] = []
+    boxes_list: List[List[int]] = []
+    box_labels: List[str] = []
+
+    for _obj_id, obj_info in frame_objects.items():
+        mask = obj_info.mask
+        label = f"{obj_info.instance_id} {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)
+
+        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:
+        frame = draw_masks(frame, np.stack(masks_list), labels=mask_labels)
+    if boxes_list:
+        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,
@@ -1598,14 +1655,16 @@ def run_grounded_sam2_tracking(
 ) -> str:
     """Run Grounded-SAM-2 video tracking pipeline.
 
-    Unlike per-frame segmentation, this extracts all frames to JPEG,
-    runs SAM2 video predictor for temporal mask propagation, then
-    renders the results back into a video.
+    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 models.segmenters.model_loader import load_segmenter as _load_seg
+    from models.segmenters.grounded_sam2 import MaskDictionary, ObjectInfo
 
     active_segmenter = segmenter_name or "gsam2_large"
     logging.info(
@@ -1622,92 +1681,372 @@ def run_grounded_sam2_tracking(
         total_frames = len(frame_names)
         logging.info("Extracted %d frames to %s", total_frames, frame_dir)
 
-        # 2. Load segmenter
-        segmenter = _load_seg(active_segmenter)
+        num_gpus = torch.cuda.device_count()
 
-        # 3. Run tracking pipeline
-        _check_cancellation(job_id)
-        tracking_results = segmenter.process_video(frame_dir, frame_names, queries)
+        # ==================================================================
+        # Phase 1-4: Tracking  (single-GPU fallback vs multi-GPU pipeline)
+        # ==================================================================
+        if num_gpus <= 1:
+            # ---------- Single-GPU fallback ----------
+            device_str = "cuda:0" if torch.cuda.is_available() else "cpu"
+            segmenter = load_segmenter_on_device(active_segmenter, device_str)
+            _check_cancellation(job_id)
+            tracking_results = segmenter.process_video(
+                frame_dir, frame_names, queries,
+            )
+            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)
+            segmenters = []
+            with ThreadPoolExecutor(max_workers=num_gpus) as pool:
+                futs = [
+                    pool.submit(
+                        load_segmenter_on_device,
+                        active_segmenter,
+                        f"cuda:{i}",
+                    )
+                    for i in range(num_gpus)
+                ]
+                segmenters = [f.result() for f in futs]
+            logging.info("Loaded %d segmenters", len(segmenters))
+
+            # Phase 2: Init SAM2 models/state per GPU (parallel)
+            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]
+
+            # Phase 3: Parallel segment processing (queue-based workers)
+            segments = list(range(0, total_frames, step))
+            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,
+                            )
+                            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=torch.tensor(masks).to(seg.device),
+                                box_list=(
+                                    boxes.clone()
+                                    if torch.is_tensor(boxes)
+                                    else torch.tensor(boxes)
+                                ),
+                                label_list=labels,
+                            )
+
+                            segment_results = seg.propagate_segment(
+                                state, start_idx, mask_dict, step,
+                            )
+                            seg_queue_out.put(
+                                (seg_idx, start_idx, mask_dict, segment_results)
+                            )
+                        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, {})
+                            )
 
-        # 4. Render results into output video
+            seg_workers = []
+            for i in range(num_gpus):
+                t = Thread(
+                    target=_segment_worker, args=(i,), daemon=True,
+                )
+                t.start()
+                seg_workers.append(t)
+
+            for t in seg_workers:
+                t.join()
+
+            # Collect all segment outputs
+            segment_data: Dict[int, Tuple] = {}
+            while not seg_queue_out.empty():
+                seg_idx, start_idx, mask_dict, results = seg_queue_out.get()
+                segment_data[seg_idx] = (start_idx, mask_dict, results)
+
+            # Phase 4: Sequential ID reconciliation
+            global_id_counter = 0
+            sam2_masks = MaskDictionary()
+            tracking_results: Dict[int, Dict[int, ObjectInfo]] = {}
+
+            for seg_idx in sorted(segment_data.keys()):
+                start_idx, mask_dict, segment_results = segment_data[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 {}
+                            )
+                    continue
+
+                # 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:
+                    for fi in range(
+                        start_idx, min(start_idx + step, total_frames)
+                    ):
+                        tracking_results[fi] = {}
+                    continue
+
+                # Apply remapping to every frame in this segment
+                for frame_idx, frame_objects in segment_results.items():
+                    remapped: Dict[int, ObjectInfo] = {}
+                    for local_id, obj_info in frame_objects.items():
+                        global_id = remapping.get(local_id)
+                        if global_id is None:
+                            continue
+                        remapped[global_id] = ObjectInfo(
+                            instance_id=global_id,
+                            mask=obj_info.mask,
+                            class_name=obj_info.class_name,
+                            x1=obj_info.x1, y1=obj_info.y1,
+                            x2=obj_info.x2, y2=obj_info.y2,
+                        )
+                    tracking_results[frame_idx] = remapped
+
+                # Update running tracker with last frame of this segment
+                if segment_results:
+                    last_fi = max(segment_results.keys())
+                    last_objs = tracking_results.get(last_fi, {})
+                    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
+                            )
+
+            logging.info(
+                "Multi-GPU reconciliation complete: %d frames, %d objects",
+                len(tracking_results), global_id_counter,
+            )
+
+        # ==================================================================
+        # Phase 5: Parallel rendering + sequential video writing
+        # ==================================================================
         _check_cancellation(job_id)
-        import os as _os
-
-        with StreamingVideoWriter(output_video_path, fps, width, height) as writer:
-            for frame_idx in range(total_frames):
-                _check_cancellation(job_id)
-
-                # Read original frame
-                frame_path = _os.path.join(frame_dir, frame_names[frame_idx])
-                frame = cv2.imread(frame_path)
-                if frame is None:
-                    logging.warning("Failed to read frame %d, writing blank", frame_idx)
-                    frame = np.zeros((height, width, 3), dtype=np.uint8)
-
-                frame_objects = tracking_results.get(frame_idx, {})
-
-                if frame_objects:
-                    # Collect masks, boxes, and labels for rendering
-                    masks_list = []
-                    boxes_list = []
-                    label_list = []
-
-                    for obj_id, obj_info in frame_objects.items():
-                        mask = obj_info.mask
-                        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)
-                            # Resize mask if needed
-                            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)
-
-                        label = f"{obj_info.instance_id} {obj_info.class_name}"
-                        label_list.append(label)
-
-                        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])
-
-                    # Draw masks
-                    if masks_list:
-                        masks_array = np.stack(masks_list)
-                        frame = draw_masks(frame, masks_array, labels=label_list)
-
-                    # Draw boxes
-                    if boxes_list:
-                        boxes_array = np.array(boxes_list)
-                        frame = draw_boxes(frame, boxes_array, label_names=label_list)
-
-                writer.write(frame)
-
-                # Stream frame if requested
-                if stream_queue:
+
+        render_in: Queue = Queue(maxsize=32)
+        render_out: Queue = Queue(maxsize=64)
+        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:
+                    frm = _gsam2_render_frame(
+                        frame_dir, frame_names, fidx, fobjs,
+                        height, width,
+                    )
+                    while True:
+                        try:
+                            render_out.put((fidx, frm), 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:
-                        from jobs.streaming import publish_frame as _pub
-                        if job_id:
-                            _pub(job_id, frame)
-                        else:
-                            stream_queue.put(frame, timeout=0.01)
-                    except Exception:
+                        render_out.put((fidx, blank), timeout=5.0)
+                    except Full:
                         pass
 
-                if frame_idx % 30 == 0:
-                    logging.info(
-                        "Rendered frame %d / %d", frame_idx, total_frames
-                    )
+        r_workers = [
+            Thread(target=_render_worker, daemon=True)
+            for _ in range(num_render_workers)
+        ]
+        for t in r_workers:
+            t.start()
+
+        def _writer_loop():
+            nonlocal render_done
+            next_idx = 0
+            buf: Dict[int, np.ndarray] = {}
+            try:
+                with StreamingVideoWriter(
+                    output_video_path, fps, width, height
+                ) as writer:
+                    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 = render_out.get(timeout=1.0)
+                                buf[idx] = frm
+
+                            frm = buf.pop(next_idx)
+                            writer.write(frm)
+
+                            if stream_queue:
+                                try:
+                                    from jobs.streaming import (
+                                        publish_frame as _pub,
+                                    )
+                                    if job_id:
+                                        _pub(job_id, frm)
+                                    else:
+                                        stream_queue.put(frm, timeout=0.01)
+                                except Exception:
+                                    pass
+
+                            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
+
+        writer_thread = Thread(target=_writer_loop, daemon=True)
+        writer_thread.start()
+
+        # Feed render queue
+        for fidx in range(total_frames):
+            _check_cancellation(job_id)
+            fobjs = tracking_results.get(fidx, {})
+            render_in.put((fidx, fobjs))
+
+        # Sentinels for render workers
+        for _ in r_workers:
+            render_in.put(None)
+
+        for t in r_workers:
+            t.join()
+        writer_thread.join()
 
         logging.info("Grounded-SAM-2 output written to: %s", output_video_path)
         return output_video_path
 
     finally:
-        # Cleanup temp frame directory
         try:
             shutil.rmtree(frame_dir)
             logging.info("Cleaned up temp frame dir: %s", frame_dir)

models/segmenters/grounded_sam2.py

@@ -90,18 +90,24 @@ class MaskDictionary:
     ) -> int:
         """Match current detections against tracked objects via IoU."""
         updated = {}
+        used_tracked_ids = set()
         for _seg_id, seg_info in self.labels.items():
             if seg_info.mask is None or seg_info.mask.sum() == 0:
                 continue
             matched_id = 0
+            best_iou = iou_threshold
             for _obj_id, obj_info in tracking_dict.labels.items():
+                if obj_info.instance_id in used_tracked_ids:
+                    continue
                 iou = self._iou(seg_info.mask, obj_info.mask)
-                if iou > iou_threshold:
+                if iou > best_iou:
+                    best_iou = iou
                     matched_id = obj_info.instance_id
-                    break
             if not matched_id:
                 objects_count += 1
                 matched_id = objects_count
+            else:
+                used_tracked_ids.add(matched_id)
             new_info = ObjectInfo(
                 instance_id=matched_id,
                 mask=seg_info.mask,
@@ -111,6 +117,47 @@ class MaskDictionary:
         self.labels = updated
         return objects_count
 
+    def update_masks_with_remapping(
+        self,
+        tracking_dict: "MaskDictionary",
+        iou_threshold: float = 0.5,
+        objects_count: int = 0,
+    ) -> Tuple[int, Dict[int, int]]:
+        """Match detections against tracked objects, returning ID remapping.
+
+        Same logic as ``update_masks`` but additionally returns a dict
+        mapping original (local) IDs to the assigned (global) IDs.
+        """
+        updated = {}
+        remapping: Dict[int, int] = {}
+        used_tracked_ids = set()
+        for seg_id, seg_info in self.labels.items():
+            if seg_info.mask is None or seg_info.mask.sum() == 0:
+                continue
+            matched_id = 0
+            best_iou = iou_threshold
+            for _obj_id, obj_info in tracking_dict.labels.items():
+                if obj_info.instance_id in used_tracked_ids:
+                    continue
+                iou = self._iou(seg_info.mask, obj_info.mask)
+                if iou > best_iou:
+                    best_iou = iou
+                    matched_id = obj_info.instance_id
+            if not matched_id:
+                objects_count += 1
+                matched_id = objects_count
+            else:
+                used_tracked_ids.add(matched_id)
+            new_info = ObjectInfo(
+                instance_id=matched_id,
+                mask=seg_info.mask,
+                class_name=seg_info.class_name,
+            )
+            updated[matched_id] = new_info
+            remapping[seg_id] = matched_id
+        self.labels = updated
+        return objects_count, remapping
+
     def get_target_class_name(self, instance_id: int) -> str:
         info = self.labels.get(instance_id)
         return info.class_name if info else ""
@@ -277,6 +324,122 @@ class GroundedSAM2Segmenter(Segmenter):
             boxes=det.boxes,
         )
 
+    # -- Multi-GPU helper methods -------------------------------------------
+
+    def detect_keyframe(
+        self,
+        image: "Image",
+        text_prompts: List[str],
+    ) -> Tuple[Optional[np.ndarray], Optional[torch.Tensor], List[str]]:
+        """Run GDINO + SAM2 image predictor on a single keyframe.
+
+        Args:
+            image: PIL Image in RGB mode.
+            text_prompts: Text queries for Grounding DINO.
+
+        Returns:
+            ``(masks, boxes, labels)`` where *masks* is an ``(N, H, W)``
+            numpy array, *boxes* is an ``(N, 4)`` tensor on device, and
+            *labels* is a list of strings.  Returns ``(None, None, [])``
+            when no objects are detected.
+        """
+        self._ensure_models_loaded()
+
+        prompt = self._gdino_detector._build_prompt(text_prompts)
+        gdino_processor = self._gdino_detector.processor
+        gdino_model = self._gdino_detector.model
+
+        inputs = gdino_processor(
+            images=image, text=prompt, return_tensors="pt"
+        )
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = gdino_model(**inputs)
+
+        results = self._gdino_detector._post_process(
+            outputs,
+            inputs["input_ids"],
+            target_sizes=[image.size[::-1]],
+        )
+
+        input_boxes = results[0]["boxes"]
+        det_labels = results[0].get("text_labels") or results[0].get("labels", [])
+        if torch.is_tensor(det_labels):
+            det_labels = det_labels.detach().cpu().tolist()
+        det_labels = [str(l) for l in det_labels]
+
+        if input_boxes.shape[0] == 0:
+            return None, None, []
+
+        # SAM2 image predictor
+        self._image_predictor.set_image(np.array(image))
+        masks, scores, logits = self._image_predictor.predict(
+            point_coords=None,
+            point_labels=None,
+            box=input_boxes,
+            multimask_output=False,
+        )
+
+        if masks.ndim == 2:
+            masks = masks[None]
+        elif masks.ndim == 4:
+            masks = masks.squeeze(1)
+
+        return masks, input_boxes, det_labels
+
+    def propagate_segment(
+        self,
+        inference_state: Any,
+        start_idx: int,
+        mask_dict: "MaskDictionary",
+        step: int,
+    ) -> Dict[int, Dict[int, "ObjectInfo"]]:
+        """Propagate masks for a single segment via SAM2 video predictor.
+
+        Calls ``reset_state`` first, making this safe to call independently
+        (and therefore parallelisable across GPUs).
+
+        Args:
+            inference_state: SAM2 video predictor state (from ``init_state``).
+            start_idx: Starting frame index for this segment.
+            mask_dict: MaskDictionary with object masks for the keyframe.
+            step: Maximum number of frames to propagate.
+
+        Returns:
+            Dict mapping ``frame_idx`` → ``{obj_id: ObjectInfo}`` using the
+            IDs from *mask_dict* (local, not yet reconciled).
+        """
+        self._video_predictor.reset_state(inference_state)
+
+        for obj_id, obj_info in mask_dict.labels.items():
+            self._video_predictor.add_new_mask(
+                inference_state,
+                start_idx,
+                obj_id,
+                obj_info.mask,
+            )
+
+        segment_results: Dict[int, Dict[int, ObjectInfo]] = {}
+        for out_frame_idx, out_obj_ids, out_mask_logits in self._video_predictor.propagate_in_video(
+            inference_state,
+            max_frame_num_to_track=step,
+            start_frame_idx=start_idx,
+        ):
+            frame_objects: Dict[int, ObjectInfo] = {}
+            for i, out_obj_id in enumerate(out_obj_ids):
+                out_mask = (out_mask_logits[i] > 0.0)
+                info = ObjectInfo(
+                    instance_id=out_obj_id,
+                    mask=out_mask[0],
+                    class_name=mask_dict.get_target_class_name(out_obj_id),
+                )
+                info.update_box()
+                frame_objects[out_obj_id] = info
+            segment_results[out_frame_idx] = frame_objects
+
+        return segment_results
+
     # -- Video-level tracking interface -------------------------------------
 
     def process_video(