feat(inspection): persist RLE masks during segmentation inference

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

inference.py

@@ -22,6 +22,7 @@ from models.segmenters.model_loader import load_segmenter, load_segmenter_on_dev
 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
 
@@ -1253,6 +1254,7 @@ def run_grounded_sam2_tracking(
         # --- 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,
@@ -1261,7 +1263,17 @@ def run_grounded_sam2_tracking(
                     "track_id": f"T{obj_id:02d}",
                     "instance_id": obj_id,
                 })
-            return dets
+                # 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
@@ -1300,7 +1312,7 @@ def run_grounded_sam2_tracking(
 
                             # Build detection records and track history
                             if fobjs:
-                                dets = _objectinfo_to_dets(fobjs)
+                                dets, masks_rle = _objectinfo_to_dets(fobjs)
 
                                 # Maintain per-track bbox history (30-frame window)
                                 for det in dets:
@@ -1316,6 +1328,10 @@ def run_grounded_sam2_tracking(
                                 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, [])

jobs/storage.py

@@ -40,6 +40,7 @@ class JobStorage:
         self._jobs: Dict[str, JobInfo] = {}
         self._tracks: Dict[str, Dict[int, list]] = {} # job_id -> {frame_idx -> tracks}
         self._latest_frames: Dict[str, any] = {}  # job_id -> np.ndarray
+        self._mask_data: Dict[str, Dict[str, any]] = {}  # job_id -> {f"{frame_idx}:{track_id}" -> rle_dict}
         self._lock = RLock()
 
     def create(self, job: JobInfo) -> None:
@@ -78,6 +79,31 @@ class JobStorage:
         """Get the most recent frame for ISR cropping."""
         return self._latest_frames.get(job_id)
 
+    def set_mask_data(self, job_id: str, frame_idx: int, track_id: int, rle: dict) -> None:
+        """Store RLE mask for a specific object at a specific frame."""
+        with self._lock:
+            if job_id not in self._mask_data:
+                self._mask_data[job_id] = {}
+            key = f"{frame_idx}:{track_id}"
+            self._mask_data[job_id][key] = rle
+
+    def get_mask_data(self, job_id: str, frame_idx: int, track_id: int) -> dict | None:
+        """Retrieve RLE mask for a specific object at a specific frame."""
+        with self._lock:
+            key = f"{frame_idx}:{track_id}"
+            return self._mask_data.get(job_id, {}).get(key)
+
+    def get_all_masks_for_frame(self, job_id: str, frame_idx: int) -> dict:
+        """Return {track_id: rle_dict} for all objects in a frame."""
+        with self._lock:
+            prefix = f"{frame_idx}:"
+            result = {}
+            for key, rle in self._mask_data.get(job_id, {}).items():
+                if key.startswith(prefix):
+                    tid = int(key.split(":")[1])
+                    result[tid] = rle
+            return result
+
     def get(self, job_id: str) -> Optional[JobInfo]:
         with self._lock:
             return self._jobs.get(job_id)
@@ -95,6 +121,7 @@ class JobStorage:
             self._jobs.pop(job_id, None)
             self._tracks.pop(job_id, None)
             self._latest_frames.pop(job_id, None)
+            self._mask_data.pop(job_id, None)
         shutil.rmtree(get_job_directory(job_id), ignore_errors=True)
 
     def cleanup_expired(self, max_age: timedelta) -> None:
@@ -131,3 +158,12 @@ def store_latest_frame(job_id: str, frame) -> None:
 
 def get_latest_frame(job_id: str):
     return get_job_storage().get_latest_frame(job_id)
+
+def set_mask_data(job_id: str, frame_idx: int, track_id: int, rle: dict) -> None:
+    get_job_storage().set_mask_data(job_id, frame_idx, track_id, rle)
+
+def get_mask_data(job_id: str, frame_idx: int, track_id: int) -> dict | None:
+    return get_job_storage().get_mask_data(job_id, frame_idx, track_id)
+
+def get_all_masks_for_frame(job_id: str, frame_idx: int) -> dict:
+    return get_job_storage().get_all_masks_for_frame(job_id, frame_idx)