Fix GSAM2 GPT writer state safety and background call args

inference.py

@@ -1593,8 +1593,13 @@ def _gsam2_render_frame(
     frame_objects: Dict,
     height: int,
     width: int,
+    masks_only: bool = False,
 ) -> np.ndarray:
-    """Render a single GSAM2 tracking frame (masks + boxes). CPU-only."""
+    """Render a single GSAM2 tracking frame (masks + boxes). CPU-only.
+
+    When *masks_only* is True, skip box rendering so the writer thread can
+    draw boxes later with enriched (GPT) labels.
+    """
     from models.segmenters.grounded_sam2 import ObjectInfo
 
     frame_path = os.path.join(frame_dir, frame_names[frame_idx])
@@ -1636,8 +1641,11 @@ def _gsam2_render_frame(
             box_labels.append(label)
 
     if masks_list:
-        frame = draw_masks(frame, np.stack(masks_list), labels=mask_labels)
-    if boxes_list:
+        frame = draw_masks(
+            frame, np.stack(masks_list),
+            labels=None if masks_only else mask_labels,
+        )
+    if boxes_list and not masks_only:
         frame = draw_boxes(frame, np.array(boxes_list), label_names=box_labels)
 
     return frame
@@ -1652,6 +1660,9 @@ def run_grounded_sam2_tracking(
     job_id: Optional[str] = None,
     stream_queue: Optional[Queue] = None,
     step: int = 20,
+    enable_gpt: bool = False,
+    mission_spec=None,  # Optional[MissionSpecification]
+    first_frame_gpt_results: Optional[Dict[str, Any]] = None,
 ) -> str:
     """Run Grounded-SAM-2 video tracking pipeline.
 
@@ -1957,10 +1968,12 @@ def run_grounded_sam2_tracking(
                     frm = _gsam2_render_frame(
                         frame_dir, frame_names, fidx, fobjs,
                         height, width,
+                        masks_only=enable_gpt,
                     )
+                    payload = (fidx, frm, fobjs) if enable_gpt else (fidx, frm, {})
                     while True:
                         try:
-                            render_out.put((fidx, frm), timeout=1.0)
+                            render_out.put(payload, timeout=1.0)
                             break
                         except Full:
                             if render_done:
@@ -1969,7 +1982,7 @@ def run_grounded_sam2_tracking(
                     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)
+                        render_out.put((fidx, blank, {}), timeout=5.0)
                     except Full:
                         pass
 
@@ -1980,10 +1993,105 @@ def run_grounded_sam2_tracking(
         for t in r_workers:
             t.start()
 
+        # --- ObjectInfo → detection dict adapter ---
+        def _objectinfo_to_dets(frame_objects_dict):
+            dets = []
+            for obj_id, info in frame_objects_dict.items():
+                dets.append({
+                    "label": info.class_name,
+                    "bbox": [info.x1, info.y1, info.x2, info.y2],
+                    "score": 1.0,
+                    "track_id": f"T{obj_id:02d}",
+                    "instance_id": obj_id,
+                })
+            return dets
+
+        # --- GPT enrichment thread (when enabled) ---
+        gpt_enrichment_queue: Queue = Queue(maxsize=4)
+        gpt_data_by_track: Dict[str, Dict] = {}
+        gpt_data_lock = RLock()
+        _relevance_refined = [False]
+
+        def _gsam2_enrichment_thread_fn():
+            while True:
+                item = gpt_enrichment_queue.get()
+                if item is None:
+                    break
+                frame_idx, frame_data, gpt_dets, ms = item
+                try:
+                    # LLM post-filter (LLM_EXTRACTED mode)
+                    if ms and ms.parse_mode == "LLM_EXTRACTED":
+                        unique_labels = list({
+                            d.get("label", "").lower()
+                            for d in gpt_dets if d.get("label")
+                        })
+                        relevant_labels = evaluate_relevance_llm(
+                            unique_labels, ms.operator_text
+                        )
+                        ms.relevance_criteria.required_classes = list(relevant_labels)
+                        _relevance_refined[0] = True
+                        logging.info(
+                            "GSAM2 enrichment: LLM post-filter frame %d: relevant=%s",
+                            frame_idx, relevant_labels,
+                        )
+                        for d in gpt_dets:
+                            decision = evaluate_relevance(d, ms.relevance_criteria)
+                            d["mission_relevant"] = decision.relevant
+                        gpt_dets = [d for d in gpt_dets if d.get("mission_relevant", True)]
+
+                    # GPT threat assessment
+                    if gpt_dets:
+                        cached_gpt = first_frame_gpt_results
+                        if not cached_gpt and job_id:
+                            try:
+                                from jobs.storage import get_job_storage as _gjs
+                                _job = _gjs().get(job_id)
+                                if _job and _job.first_frame_gpt_results:
+                                    cached_gpt = _job.first_frame_gpt_results
+                            except Exception:
+                                pass
+
+                        if cached_gpt:
+                            logging.info("GSAM2 enrichment: re-using cached GPT for frame %d", frame_idx)
+                            gpt_res = cached_gpt
+                        else:
+                            logging.info("GSAM2 enrichment: running GPT for frame %d...", frame_idx)
+                            frame_b64 = encode_frame_to_b64(frame_data)
+                            gpt_res = estimate_threat_gpt(
+                                detections=gpt_dets, mission_spec=ms,
+                                image_b64=frame_b64,
+                            )
+
+                        for d in gpt_dets:
+                            tid = d.get("track_id")
+                            if tid and tid in gpt_res:
+                                merged = dict(gpt_res[tid])
+                                merged["gpt_raw"] = gpt_res[tid]
+                                merged["assessment_frame_index"] = frame_idx
+                                merged["assessment_status"] = "ASSESSED"
+                                with gpt_data_lock:
+                                    gpt_data_by_track[tid] = merged
+                        logging.info("GSAM2 enrichment: GPT results stored for %d tracks", len(gpt_data_by_track))
+
+                except Exception as e:
+                    logging.error("GSAM2 enrichment thread failed for frame %d: %s", frame_idx, e)
+
         def _writer_loop():
             nonlocal render_done
             next_idx = 0
-            buf: Dict[int, np.ndarray] = {}
+            buf: Dict[int, Tuple] = {}
+
+            # Per-track bbox history (replaces ByteTracker for GSAM2)
+            track_history: Dict[int, List] = {}
+            speed_est = SpeedEstimator(fps=fps) if enable_gpt else None
+            gpt_submitted = False
+
+            # Start enrichment thread when GPT enabled
+            enrich_thread = None
+            if enable_gpt:
+                enrich_thread = Thread(target=_gsam2_enrichment_thread_fn, daemon=True)
+                enrich_thread.start()
+
             try:
                 with StreamingVideoWriter(
                     output_video_path, fps, width, height
@@ -1998,10 +2106,104 @@ def run_grounded_sam2_tracking(
                                         len(buf), next_idx,
                                     )
                                     time.sleep(0.05)
-                                idx, frm = render_out.get(timeout=1.0)
-                                buf[idx] = frm
+                                idx, frm, fobjs = render_out.get(timeout=1.0)
+                                buf[idx] = (frm, fobjs)
+
+                            frm, fobjs = buf.pop(next_idx)
+
+                            # --- GPT enrichment path ---
+                            if enable_gpt and fobjs:
+                                dets = _objectinfo_to_dets(fobjs)
+
+                                # Maintain per-track bbox history (30-frame window)
+                                for det in dets:
+                                    iid = det["instance_id"]
+                                    track_history.setdefault(iid, []).append(det["bbox"])
+                                    if len(track_history[iid]) > 30:
+                                        track_history[iid].pop(0)
+                                    # Store an immutable per-frame snapshot.
+                                    det["history"] = list(track_history[iid])
+
+                                # Speed estimation
+                                if speed_est:
+                                    speed_est.estimate(dets)
+
+                                # Relevance gate
+                                if mission_spec:
+                                    if (mission_spec.parse_mode == "LLM_EXTRACTED"
+                                            and not _relevance_refined[0]):
+                                        for d in dets:
+                                            d["mission_relevant"] = True
+                                            d["relevance_reason"] = "pending_llm_postfilter"
+                                        gpt_dets = dets
+                                    else:
+                                        for d in dets:
+                                            decision = evaluate_relevance(d, mission_spec.relevance_criteria)
+                                            d["mission_relevant"] = decision.relevant
+                                            d["relevance_reason"] = decision.reason
+                                        gpt_dets = [d for d in dets if d.get("mission_relevant", True)]
+                                else:
+                                    for d in dets:
+                                        d["mission_relevant"] = None
+                                    gpt_dets = dets
+
+                                # GPT enrichment (one-shot, first frame with detections)
+                                if gpt_dets and not gpt_submitted:
+                                    for d in gpt_dets:
+                                        d["assessment_status"] = "PENDING_GPT"
+                                    try:
+                                        gpt_enrichment_queue.put(
+                                            (
+                                                next_idx,
+                                                frm.copy(),
+                                                copy.deepcopy(gpt_dets),
+                                                mission_spec,
+                                            ),
+                                            timeout=1.0,
+                                        )
+                                        gpt_submitted = True
+                                        logging.info("GSAM2 writer: offloaded GPT enrichment for frame %d", next_idx)
+                                    except Full:
+                                        logging.warning("GSAM2 GPT enrichment queue full, skipping")
+
+                                # Merge persistent GPT data
+                                for det in dets:
+                                    tid = det["track_id"]
+                                    with gpt_data_lock:
+                                        gpt_payload = gpt_data_by_track.get(tid)
+                                    if gpt_payload:
+                                        det.update(gpt_payload)
+                                        det["assessment_status"] = "ASSESSED"
+                                    elif "assessment_status" not in det:
+                                        det["assessment_status"] = "UNASSESSED"
+
+                                # Build enriched display labels
+                                display_labels = []
+                                for d in dets:
+                                    lbl = d.get("label", "obj")
+                                    if "track_id" in d:
+                                        lbl = f"{d['track_id']} {lbl}"
+                                    if d.get("gpt_distance_m") is not None:
+                                        try:
+                                            lbl = f"{lbl} {int(float(d['gpt_distance_m']))}m"
+                                        except (TypeError, ValueError):
+                                            pass
+                                    display_labels.append(lbl)
+
+                                # Draw boxes on mask-rendered frame
+                                if dets:
+                                    boxes = np.array([d["bbox"] for d in dets])
+                                    frm = draw_boxes(frm, boxes, label_names=display_labels)
+
+                                # Store tracks for frontend
+                                if job_id:
+                                    set_track_data(job_id, next_idx, copy.deepcopy(dets))
+
+                            elif enable_gpt:
+                                # No objects this frame — still store empty track data
+                                if job_id:
+                                    set_track_data(job_id, next_idx, [])
 
-                            frm = buf.pop(next_idx)
                             writer.write(frm)
 
                             if stream_queue:
@@ -2034,6 +2236,13 @@ def run_grounded_sam2_tracking(
                             continue
             finally:
                 render_done = True
+                # Shut down enrichment thread
+                if enrich_thread:
+                    try:
+                        gpt_enrichment_queue.put(None, timeout=5.0)
+                        enrich_thread.join(timeout=30)
+                    except Exception:
+                        logging.warning("GSAM2 enrichment thread shutdown timed out")
 
         writer_thread = Thread(target=_writer_loop, daemon=True)
         writer_thread.start()

jobs/background.py

@@ -29,13 +29,17 @@ async def process_video_async(job_id: str) -> None:
         if job.mode == "segmentation":
             detection_path = await asyncio.to_thread(
                 run_grounded_sam2_tracking,
-                job.input_video_path,
-                job.output_video_path,
-                job.queries,
-                None,
-                job.segmenter_name,
-                job_id,
-                stream_queue,
+                input_video_path=job.input_video_path,
+                output_video_path=job.output_video_path,
+                queries=job.queries,
+                max_frames=None,
+                segmenter_name=job.segmenter_name,
+                job_id=job_id,
+                stream_queue=stream_queue,
+                step=20,
+                enable_gpt=job.enable_gpt,
+                mission_spec=job.mission_spec,
+                first_frame_gpt_results=job.first_frame_gpt_results,
             )
         else:
             detections_list = None