import os import sys import uuid import hashlib import tempfile from pathlib import Path from typing import Dict, List, Tuple, Optional, Any, Union # Add src/ to sys.path so LASER, video-sam2, GroundingDINO are importable current_dir = Path(__file__).resolve().parent src_dir = current_dir.parent / "src" if src_dir.is_dir() and str(src_dir) not in sys.path: sys.path.insert(0, str(src_dir)) import cv2 import numpy as np import torch from transformers import Pipeline from .vine_config import VineConfig from .vine_model import VineModel from .vis_utils import render_dino_frames, render_sam_frames, render_vine_frame_sets from laser.loading import load_video from laser.preprocess.mask_generation_grounding_dino import generate_masks_grounding_dino class VinePipeline(Pipeline): """ Pipeline for VINE model that handles end-to-end video understanding. """ def __init__( self, sam_config_path: Optional[str] = None, sam_checkpoint_path: Optional[str] = None, gd_config_path: Optional[str] = None, gd_checkpoint_path: Optional[str] = None, **kwargs: Any, ): self.grounding_model = None self.sam_predictor = None self.mask_generator = None self.sam_config_path = sam_config_path self.sam_checkpoint_path = sam_checkpoint_path self.gd_config_path = gd_config_path self.gd_checkpoint_path = gd_checkpoint_path super().__init__(**kwargs) self.segmentation_method = getattr( self.model.config, "segmentation_method", "grounding_dino_sam2" ) self.box_threshold = getattr(self.model.config, "box_threshold", 0.35) self.text_threshold = getattr(self.model.config, "text_threshold", 0.25) self.target_fps = getattr(self.model.config, "target_fps", 1) self.visualize = getattr(self.model.config, "visualize", False) self.visualization_dir = getattr(self.model.config, "visualization_dir", None) self.debug_visualizations = getattr( self.model.config, "debug_visualizations", False ) self._device = getattr(self.model.config, "_device") if kwargs.get("device") is not None: self._device = kwargs.get("device") # ------------------------------------------------------------------ # # Segmentation model injection # ------------------------------------------------------------------ # def set_segmentation_models( self, *, sam_predictor=None, mask_generator=None, grounding_model=None, ): if sam_predictor is not None: self.sam_predictor = sam_predictor if mask_generator is not None: self.mask_generator = mask_generator if grounding_model is not None: self.grounding_model = grounding_model # ------------------------------------------------------------------ # # Pipeline protocol # ------------------------------------------------------------------ # def _sanitize_parameters(self, **kwargs: Any): preprocess_kwargs: Dict[str, Any] = {} forward_kwargs: Dict[str, Any] = {} postprocess_kwargs: Dict[str, Any] = {} if "segmentation_method" in kwargs: preprocess_kwargs["segmentation_method"] = kwargs["segmentation_method"] if "target_fps" in kwargs: preprocess_kwargs["target_fps"] = kwargs["target_fps"] if "box_threshold" in kwargs: preprocess_kwargs["box_threshold"] = kwargs["box_threshold"] if "text_threshold" in kwargs: preprocess_kwargs["text_threshold"] = kwargs["text_threshold"] if "categorical_keywords" in kwargs: preprocess_kwargs["categorical_keywords"] = kwargs["categorical_keywords"] if "categorical_keywords" in kwargs: forward_kwargs["categorical_keywords"] = kwargs["categorical_keywords"] if "unary_keywords" in kwargs: forward_kwargs["unary_keywords"] = kwargs["unary_keywords"] if "binary_keywords" in kwargs: forward_kwargs["binary_keywords"] = kwargs["binary_keywords"] if "object_pairs" in kwargs: forward_kwargs["object_pairs"] = kwargs["object_pairs"] if "batched_binary_predicates" in kwargs: # New: per-video (rel, from_cat, to_cat) triples for binary filtering forward_kwargs["batched_binary_predicates"] = kwargs["batched_binary_predicates"] if "topk_cate" in kwargs: # New: override topk_cate when binary filtering is requested forward_kwargs["topk_cate"] = kwargs["topk_cate"] if "auto_add_not_unary" in kwargs: forward_kwargs["auto_add_not_unary"] = kwargs["auto_add_not_unary"] if "return_flattened_segments" in kwargs: forward_kwargs["return_flattened_segments"] = kwargs[ "return_flattened_segments" ] if "return_valid_pairs" in kwargs: forward_kwargs["return_valid_pairs"] = kwargs["return_valid_pairs"] if "interested_object_pairs" in kwargs: forward_kwargs["interested_object_pairs"] = kwargs[ "interested_object_pairs" ] if "debug_visualizations" in kwargs: forward_kwargs["debug_visualizations"] = kwargs["debug_visualizations"] postprocess_kwargs["debug_visualizations"] = kwargs["debug_visualizations"] if "return_top_k" in kwargs: postprocess_kwargs["return_top_k"] = kwargs["return_top_k"] if "self.visualize" in kwargs: postprocess_kwargs["self.visualize"] = kwargs["self.visualize"] if "binary_confidence_threshold" in kwargs: postprocess_kwargs["binary_confidence_threshold"] = kwargs[ "binary_confidence_threshold" ] return preprocess_kwargs, forward_kwargs, postprocess_kwargs # ------------------------------------------------------------------ # # Preprocess: video + segmentation # ------------------------------------------------------------------ # def preprocess( self, video_input: Union[str, np.ndarray, torch.Tensor], segmentation_method: Optional[str] = None, target_fps: Optional[int] = None, box_threshold: Optional[float] = None, text_threshold: Optional[float] = None, categorical_keywords: Optional[List[str]] = None, **kwargs: Any, ) -> Dict[str, Any]: if segmentation_method is None: segmentation_method = self.segmentation_method if target_fps is None: target_fps = self.target_fps else: self.target_fps = target_fps if box_threshold is None: box_threshold = self.box_threshold else: self.box_threshold = box_threshold if text_threshold is None: text_threshold = self.text_threshold else: self.text_threshold = text_threshold if categorical_keywords is None: categorical_keywords = ["object"] if isinstance(video_input, str): video_tensor = load_video(video_input, target_fps=target_fps) if isinstance(video_tensor, list): video_tensor = np.array(video_tensor) elif isinstance(video_tensor, torch.Tensor): video_tensor = video_tensor.cpu().numpy() elif isinstance(video_input, (np.ndarray, torch.Tensor)): if isinstance(video_input, torch.Tensor): video_tensor = video_input.numpy() else: video_tensor = video_input else: raise ValueError(f"Unsupported video input type: {type(video_input)}") if not isinstance(video_tensor, np.ndarray): video_tensor = np.array(video_tensor) if len(video_tensor.shape) != 4: raise ValueError( f"Expected video tensor shape (frames, height, width, channels), got {video_tensor.shape}" ) visualization_data: Dict[str, Any] = {} print(f"Segmentation method: {segmentation_method}") if segmentation_method == "sam2": masks, bboxes, vis_data = self._generate_sam2_masks(video_tensor) elif segmentation_method == "grounding_dino_sam2": masks, bboxes, vis_data = self._generate_grounding_dino_sam2_masks( video_tensor, categorical_keywords, box_threshold, text_threshold, video_input, ) else: raise ValueError(f"Unsupported segmentation method: {segmentation_method}") if vis_data: visualization_data.update(vis_data) visualization_data.setdefault("sam_masks", masks) return { "video_frames": torch.tensor(video_tensor), "masks": masks, "bboxes": bboxes, "num_frames": len(video_tensor), "visualization_data": visualization_data, } # ------------------------------------------------------------------ # # Segmentation helpers # ------------------------------------------------------------------ # def _generate_sam2_masks( self, video_tensor: np.ndarray ) -> Tuple[Dict[int, Dict[int, torch.Tensor]], Dict[int, Dict[int, List[int]]], Dict[str, Any]]: print("Generating SAM2 masks...") if self.mask_generator is None: self._initialize_segmentation_models() if self.mask_generator is None: raise ValueError("SAM2 mask generator not available") masks: Dict[int, Dict[int, torch.Tensor]] = {} bboxes: Dict[int, Dict[int, List[int]]] = {} for frame_id, frame in enumerate(video_tensor): if isinstance(frame, np.ndarray) and frame.dtype != np.uint8: frame = ( (frame * 255).astype(np.uint8) if frame.max() <= 1 else frame.astype(np.uint8) ) frame_masks = self.mask_generator.generate(frame) masks[frame_id] = {} bboxes[frame_id] = {} for obj_id, mask_data in enumerate(frame_masks): mask = mask_data["segmentation"] if isinstance(mask, np.ndarray): mask = torch.from_numpy(mask) if len(mask.shape) == 2: mask = mask.unsqueeze(-1) elif len(mask.shape) == 3 and mask.shape[0] == 1: mask = mask.permute(1, 2, 0) wrapped_id = obj_id + 1 masks[frame_id][wrapped_id] = mask mask_np = ( mask.squeeze().numpy() if isinstance(mask, torch.Tensor) else mask.squeeze() ) coords = np.where(mask_np > 0) if len(coords[0]) > 0: y1, y2 = coords[0].min(), coords[0].max() x1, x2 = coords[1].min(), coords[1].max() bboxes[frame_id][wrapped_id] = [x1, y1, x2, y2] tracked_masks, tracked_bboxes = self._track_ids_across_frames(masks, bboxes) return tracked_masks, tracked_bboxes, {"sam_masks": tracked_masks} def _generate_grounding_dino_sam2_masks( self, video_tensor: np.ndarray, categorical_keywords: List[str], box_threshold: float, text_threshold: float, video_path: Union[str, None], ) -> Tuple[Dict[int, Dict[int, torch.Tensor]], Dict[int, Dict[int, List[int]]], Dict[str, Any]]: print("Generating Grounding DINO + SAM2 masks...") if self.grounding_model is None or self.sam_predictor is None: self._initialize_segmentation_models() if self.grounding_model is None or self.sam_predictor is None: raise ValueError("GroundingDINO or SAM2 models not available") temp_video_path = None if video_path is None or not isinstance(video_path, str): temp_video_path = self._create_temp_video(video_tensor) video_path = temp_video_path CHUNK = 5 classes_ls = [ categorical_keywords[i : i + CHUNK] for i in range(0, len(categorical_keywords), CHUNK) ] base_name = Path(video_path).stem fps_tag = f"fps{int(self.target_fps)}" path_hash = hashlib.md5(video_path.encode("utf-8")).hexdigest()[:8] video_cache_name = f"{base_name}_{fps_tag}_{path_hash}" video_segments, oid_class_pred, _ = generate_masks_grounding_dino( self.grounding_model, box_threshold, text_threshold, self.sam_predictor, self.mask_generator, video_tensor, video_path, video_cache_name, out_dir=tempfile.gettempdir(), classes_ls=classes_ls, target_fps=self.target_fps, visualize=self.debug_visualizations, frames=None, max_prop_time=2, ) masks: Dict[int, Dict[int, torch.Tensor]] = {} bboxes: Dict[int, Dict[int, List[int]]] = {} for frame_id, frame_masks in video_segments.items(): masks[frame_id] = {} bboxes[frame_id] = {} for obj_id, mask in frame_masks.items(): if not isinstance(mask, torch.Tensor): mask = torch.tensor(mask) masks[frame_id][obj_id] = mask mask_np = mask.numpy() if mask_np.ndim == 3 and mask_np.shape[0] == 1: mask_np = np.squeeze(mask_np, axis=0) coords = np.where(mask_np > 0) if len(coords[0]) > 0: y1, y2 = coords[0].min(), coords[0].max() x1, x2 = coords[1].min(), coords[1].max() bboxes[frame_id][obj_id] = [x1, y1, x2, y2] if temp_video_path and os.path.exists(temp_video_path): os.remove(temp_video_path) tracked_masks, tracked_bboxes = self._track_ids_across_frames(masks, bboxes) vis_data: Dict[str, Any] = { "sam_masks": tracked_masks, "dino_labels": oid_class_pred, } return tracked_masks, tracked_bboxes, vis_data # ------------------------------------------------------------------ # # ID tracking across frames # ------------------------------------------------------------------ # def _bbox_iou(self, box1: List[int], box2: List[int]) -> float: x1, y1, x2, y2 = box1 x1b, y1b, x2b, y2b = box2 ix1 = max(x1, x1b) iy1 = max(y1, y1b) ix2 = min(x2, x2b) iy2 = min(y2, y2b) iw = max(0, ix2 - ix1) ih = max(0, iy2 - iy1) inter = iw * ih if inter <= 0: return 0.0 area1 = max(0, x2 - x1) * max(0, y2 - y1) area2 = max(0, x2b - x1b) * max(0, y2b - y1b) union = area1 + area2 - inter if union <= 0: return 0.0 return inter / union def _track_ids_across_frames( self, masks: Dict[int, Dict[int, torch.Tensor]], bboxes: Dict[int, Dict[int, List[int]]], iou_threshold: float = 0.3, ) -> Tuple[Dict[int, Dict[int, torch.Tensor]], Dict[int, Dict[int, List[int]]]]: frame_ids = sorted(masks.keys()) tracked_masks: Dict[int, Dict[int, torch.Tensor]] = {} tracked_bboxes: Dict[int, Dict[int, List[int]]] = {} next_track_id = 0 prev_tracks: Dict[int, List[int]] = {} for frame_id in frame_ids: frame_masks = masks.get(frame_id, {}) frame_boxes = bboxes.get(frame_id, {}) tracked_masks[frame_id] = {} tracked_bboxes[frame_id] = {} if not frame_boxes: prev_tracks = {} continue det_ids = list(frame_boxes.keys()) prev_ids = list(prev_tracks.keys()) candidates: List[Tuple[float, int, int]] = [] for tid in prev_ids: prev_box = prev_tracks[tid] for det_id in det_ids: iou = self._bbox_iou(prev_box, frame_boxes[det_id]) if iou > iou_threshold: candidates.append((iou, tid, det_id)) candidates.sort(reverse=True) matched_prev = set() matched_det = set() for iou, tid, det_id in candidates: if tid in matched_prev or det_id in matched_det: continue matched_prev.add(tid) matched_det.add(det_id) tracked_masks[frame_id][tid] = frame_masks[det_id] tracked_bboxes[frame_id][tid] = frame_boxes[det_id] for det_id in det_ids: if det_id in matched_det: continue tid = next_track_id next_track_id += 1 tracked_masks[frame_id][tid] = frame_masks[det_id] tracked_bboxes[frame_id][tid] = frame_boxes[det_id] prev_tracks = { tid: tracked_bboxes[frame_id][tid] for tid in tracked_bboxes[frame_id].keys() } return tracked_masks, tracked_bboxes # ------------------------------------------------------------------ # # Segmentation model initialization # ------------------------------------------------------------------ # def _initialize_segmentation_models(self): if self.sam_predictor is None or self.mask_generator is None: self._initialize_sam2_models() if self.grounding_model is None: self._initialize_grounding_dino_model() def _initialize_sam2_models(self): try: from sam2.build_sam import build_sam2_video_predictor, build_sam2 from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator except ImportError as e: print(f"Warning: Could not import SAM2: {e}") return config_path, checkpoint_path = self._resolve_sam2_paths() if self.sam_config_path is not None and not os.path.exists(config_path): raise ValueError(f"SAM2 config path not found: {config_path}") if self.sam_checkpoint_path is not None and not os.path.exists(checkpoint_path): raise ValueError(f"SAM2 checkpoint path not found: {checkpoint_path}") if not os.path.exists(checkpoint_path): print(f"Warning: SAM2 checkpoint not found at {checkpoint_path}") print("SAM2 functionality will be unavailable") return try: device = self._device self.sam_predictor = build_sam2_video_predictor( config_path, checkpoint_path, device=device ) sam2_model = build_sam2( config_path, checkpoint_path, device=device, apply_postprocessing=False, ) self.mask_generator = SAM2AutomaticMaskGenerator( model=sam2_model, points_per_side=32, points_per_batch=32, pred_iou_thresh=0.7, stability_score_thresh=0.8, crop_n_layers=2, box_nms_thresh=0.6, crop_n_points_downscale_factor=2, min_mask_region_area=100, use_m2m=True, ) print("✓ SAM2 models initialized successfully") except Exception as e: raise ValueError(f"Failed to initialize SAM2 with custom paths: {e}") def _initialize_grounding_dino_model(self): try: from groundingdino.util.inference import Model as gd_Model except ImportError as e: print(f"Warning: Could not import GroundingDINO: {e}") return config_path, checkpoint_path = self._resolve_grounding_dino_paths() if self.gd_config_path is not None and not os.path.exists(config_path): raise ValueError(f"GroundingDINO config path not found: {config_path}") if self.gd_checkpoint_path is not None and not os.path.exists(checkpoint_path): raise ValueError( f"GroundingDINO checkpoint path not found: {checkpoint_path}" ) if not (os.path.exists(config_path) and os.path.exists(checkpoint_path)): print( f"Warning: GroundingDINO models not found at {config_path} / {checkpoint_path}" ) print("GroundingDINO functionality will be unavailable") return try: device = self._device self.grounding_model = gd_Model( model_config_path=config_path, model_checkpoint_path=checkpoint_path, device=device, ) print("✓ GroundingDINO model initialized successfully") except Exception as e: raise ValueError(f"Failed to initialize GroundingDINO with custom paths: {e}") def _resolve_sam2_paths(self): if self.sam_config_path and self.sam_checkpoint_path: return self.sam_config_path, self.sam_checkpoint_path def _resolve_grounding_dino_paths(self): if self.gd_config_path and self.gd_checkpoint_path: return self.gd_config_path, self.gd_checkpoint_path # ------------------------------------------------------------------ # # Video writing helpers # ------------------------------------------------------------------ # def _prepare_visualization_dir(self, name: str, enabled: bool) -> Optional[str]: if not enabled: return None if self.visualization_dir: target_dir = ( os.path.join(self.visualization_dir, name) if name else self.visualization_dir ) os.makedirs(target_dir, exist_ok=True) return target_dir return tempfile.mkdtemp(prefix=f"vine_{name}_") def _create_temp_video( self, video_tensor: np.ndarray, base_dir: Optional[str] = None, prefix: str = "temp_video", ) -> str: import subprocess if base_dir is None: base_dir = tempfile.mkdtemp(prefix=f"vine_{prefix}_") else: os.makedirs(base_dir, exist_ok=True) file_name = f"{prefix}_{uuid.uuid4().hex}.mp4" temp_path = os.path.join(base_dir, file_name) height, width = video_tensor.shape[1:3] processing_fps = max(1, self.target_fps) output_fps = processing_fps video_tensor_for_output = video_tensor ffmpeg_success = False try: ffmpeg_success = self._create_video_with_ffmpeg( video_tensor_for_output, temp_path, output_fps, width, height ) except Exception as e: print(f"FFmpeg method failed: {e}") if not ffmpeg_success: print("Using OpenCV fallback") self._create_temp_video_opencv( video_tensor_for_output, temp_path, output_fps, width, height ) return temp_path def _create_video_with_ffmpeg( self, video_tensor: np.ndarray, output_path: str, fps: int, width: int, height: int ) -> bool: import subprocess try: # Try to get FFmpeg from imageio-ffmpeg first, then fall back to system FFmpeg try: import imageio_ffmpeg ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe() print(f"Using FFmpeg from imageio-ffmpeg: {ffmpeg_exe}") except ImportError: ffmpeg_exe = "ffmpeg" print("Using system FFmpeg") ffmpeg_cmd = [ ffmpeg_exe, "-y", "-f", "rawvideo", "-vcodec", "rawvideo", "-s", f"{width}x{height}", "-pix_fmt", "rgb24", "-r", str(fps), "-i", "pipe:0", "-c:v", "libx264", "-preset", "fast", "-crf", "23", "-pix_fmt", "yuv420p", "-movflags", "+faststart", "-loglevel", "error", output_path, ] process = subprocess.Popen( ffmpeg_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) frame_data = b"" for frame in video_tensor: if frame.dtype != np.uint8: frame = ( (frame * 255).astype(np.uint8) if frame.max() <= 1 else frame.astype(np.uint8) ) frame_data += frame.tobytes() stdout, stderr = process.communicate(input=frame_data, timeout=60) if process.returncode == 0: print(f"Video created with FFmpeg (H.264) at {fps} FPS") return True else: error_msg = stderr.decode() if stderr else "Unknown error" print(f"FFmpeg error: {error_msg}") return False except FileNotFoundError: print("FFmpeg not found in PATH") return False except Exception as e: print(f"FFmpeg exception: {e}") return False def _create_temp_video_opencv( self, video_tensor: np.ndarray, temp_path: str, fps: int, width: int, height: int ) -> str: codecs_to_try = ["avc1", "X264", "mp4v"] out = None used_codec = None # Debug: Print video tensor info print(f"DEBUG: video_tensor shape: {video_tensor.shape}, dtype: {video_tensor.dtype}") print(f"DEBUG: Expected dimensions - width: {width}, height: {height}, fps: {fps}") for codec in codecs_to_try: try: fourcc = cv2.VideoWriter_fourcc(*codec) temp_out = cv2.VideoWriter(temp_path, fourcc, fps, (width, height)) if temp_out.isOpened(): out = temp_out used_codec = codec break else: temp_out.release() except Exception as e: print(f"Warning: Codec {codec} not available: {e}") continue if out is None or not out.isOpened(): raise RuntimeError( f"Failed to initialize VideoWriter with any codec. Tried: {codecs_to_try}" ) print(f"Using OpenCV with codec: {used_codec}") frame_count = 0 for frame in video_tensor: # Debug: Print first frame info if frame_count == 0: print(f"DEBUG: First frame shape: {frame.shape}, dtype: {frame.dtype}") print(f"DEBUG: First frame min: {frame.min()}, max: {frame.max()}, mean: {frame.mean()}") if len(frame.shape) == 3 and frame.shape[2] == 3: frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) else: frame_bgr = frame if frame_bgr.dtype != np.uint8: frame_bgr = ( (frame_bgr * 255).astype(np.uint8) if frame_bgr.max() <= 1 else frame_bgr.astype(np.uint8) ) # Debug: Check if frame dimensions match VideoWriter expectations if frame_count == 0: print(f"DEBUG: After conversion - frame_bgr shape: {frame_bgr.shape}, dtype: {frame_bgr.dtype}") print(f"DEBUG: After conversion - min: {frame_bgr.min()}, max: {frame_bgr.max()}") actual_height, actual_width = frame_bgr.shape[:2] if actual_height != height or actual_width != width: print(f"WARNING: Frame size mismatch! Expected ({height}, {width}), got ({actual_height}, {actual_width})") out.write(frame_bgr) frame_count += 1 print(f"DEBUG: Wrote {frame_count} frames to video") out.release() return temp_path # ------------------------------------------------------------------ # # Forward + postprocess # ------------------------------------------------------------------ # def _forward(self, model_inputs: Dict[str, Any], **forward_kwargs: Any) -> Dict[str, Any]: outputs = self.model.predict( video_frames=model_inputs["video_frames"], masks=model_inputs["masks"], bboxes=model_inputs["bboxes"], **forward_kwargs, ) outputs.setdefault("video_frames", model_inputs.get("video_frames")) outputs.setdefault("bboxes", model_inputs.get("bboxes")) outputs.setdefault("masks", model_inputs.get("masks")) outputs.setdefault("visualization_data", model_inputs.get("visualization_data")) return outputs def postprocess( self, model_outputs: Dict[str, Any], return_top_k: int = 3, visualize: Optional[bool] = None, **kwargs: Any, ) -> Dict[str, Any]: results: Dict[str, Any] = { "categorical_predictions": model_outputs.get("categorical_predictions", {}), "unary_predictions": model_outputs.get("unary_predictions", {}), "binary_predictions": model_outputs.get("binary_predictions", {}), "confidence_scores": model_outputs.get("confidence_scores", {}), "summary": self._generate_summary(model_outputs), } print("\n" + "=" * 50) print("DEBUG: Raw Model Outputs - Categorical Predictions") cat_preds = model_outputs.get("categorical_predictions", {}) for obj_id, preds in cat_preds.items(): print(f"Object {obj_id}: {preds}") print("=" * 50 + "\n") if "flattened_segments" in model_outputs: results["flattened_segments"] = model_outputs["flattened_segments"] if "valid_pairs" in model_outputs: results["valid_pairs"] = model_outputs["valid_pairs"] if "valid_pairs_metadata" in model_outputs: results["valid_pairs_metadata"] = model_outputs["valid_pairs_metadata"] if "visualization_data" in model_outputs: results["visualization_data"] = model_outputs["visualization_data"] if self.visualize and "video_frames" in model_outputs and "bboxes" in model_outputs: frames_tensor = model_outputs["video_frames"] if isinstance(frames_tensor, torch.Tensor): frames_np = frames_tensor.detach().cpu().numpy() else: frames_np = np.asarray(frames_tensor) if frames_np.dtype != np.uint8: if np.issubdtype(frames_np.dtype, np.floating): max_val = frames_np.max() if frames_np.size else 0.0 scale = 255.0 if max_val <= 1.0 else 1.0 frames_np = (frames_np * scale).clip(0, 255).astype(np.uint8) else: frames_np = frames_np.clip(0, 255).astype(np.uint8) cat_label_lookup: Dict[int, Tuple[str, float]] = {} for obj_id, preds in model_outputs.get("categorical_predictions", {}).items(): if preds: prob, label = preds[0] cat_label_lookup[obj_id] = (label, prob) unary_preds = model_outputs.get("unary_predictions", {}) unary_lookup: Dict[int, Dict[int, List[Tuple[float, str]]]] = {} for (frame_id, obj_id), preds in unary_preds.items(): if preds: unary_lookup.setdefault(frame_id, {})[obj_id] = preds[:1] binary_preds = model_outputs.get("binary_predictions", {}) binary_lookup: Dict[ int, List[Tuple[Tuple[int, int], List[Tuple[float, str]]]] ] = {} for (frame_id, obj_pair), preds in binary_preds.items(): if preds: binary_lookup.setdefault(frame_id, []).append((obj_pair, preds[:1])) bboxes = model_outputs["bboxes"] visualization_data = model_outputs.get("visualization_data", {}) visualizations: Dict[str, Dict[str, Any]] = {} debug_visualizations = kwargs.get("debug_visualizations") if debug_visualizations is None: debug_visualizations = self.debug_visualizations # Get binary confidence threshold from kwargs (default 0.0 means show all) binary_confidence_threshold = kwargs.get("binary_confidence_threshold", 0.0) vine_frame_sets = render_vine_frame_sets( frames_np, bboxes, cat_label_lookup, unary_lookup, binary_lookup, visualization_data.get("sam_masks"), binary_confidence_threshold, ) vine_visuals: Dict[str, Dict[str, Any]] = {} final_frames = vine_frame_sets.get("all", []) if final_frames: final_entry: Dict[str, Any] = {"frames": final_frames, "video_path": None} final_dir = self._prepare_visualization_dir( "all", enabled=self.visualize ) final_entry["video_path"] = self._create_temp_video( np.stack(final_frames, axis=0), base_dir=final_dir, prefix="all_visualization", ) vine_visuals["all"] = final_entry if debug_visualizations: sam_masks = visualization_data.get("sam_masks") if sam_masks: sam_frames = render_sam_frames( frames_np, sam_masks, visualization_data.get("dino_labels") ) sam_entry = {"frames": sam_frames, "video_path": None} if sam_frames: sam_dir = self._prepare_visualization_dir( "sam", enabled=self.visualize ) sam_entry["video_path"] = self._create_temp_video( np.stack(sam_frames, axis=0), base_dir=sam_dir, prefix="sam_visualization", ) visualizations["sam"] = sam_entry dino_labels = visualization_data.get("dino_labels") if dino_labels: dino_frames = render_dino_frames(frames_np, bboxes, dino_labels) dino_entry = {"frames": dino_frames, "video_path": None} if dino_frames: dino_dir = self._prepare_visualization_dir( "dino", enabled=self.visualize ) dino_entry["video_path"] = self._create_temp_video( np.stack(dino_frames, axis=0), base_dir=dino_dir, prefix="dino_visualization", ) visualizations["dino"] = dino_entry for name in ("object", "unary", "binary"): frames_list = vine_frame_sets.get(name, []) entry: Dict[str, Any] = {"frames": frames_list, "video_path": None} if frames_list: vine_dir = self._prepare_visualization_dir( name, enabled=self.visualize ) entry["video_path"] = self._create_temp_video( np.stack(frames_list, axis=0), base_dir=vine_dir, prefix=f"{name}_visualization", ) vine_visuals[name] = entry if vine_visuals: visualizations["vine"] = vine_visuals if visualizations: results["visualizations"] = visualizations return results # ------------------------------------------------------------------ # # Summary JSON # ------------------------------------------------------------------ # def _generate_summary(self, model_outputs: Dict[str, Any]) -> Dict[str, Any]: """ Per-object summary: { "num_objects_detected": N, "objects": { "": { "top_categories": [{"label": str, "probability": float}, ...], "top_unary": [{"frame_id": int, "predicate": str, "probability": float}, ...], } }, "binary_keywords": { "-": {"predicate": str, "confidence": float, "frame_id": int} } } """ categorical_preds = model_outputs.get("categorical_predictions", {}) unary_preds = model_outputs.get("unary_predictions", {}) binary_preds = model_outputs.get("binary_predictions", {}) # Debug: Print binary predictions print("\n" + "=" * 80) print("DEBUG _generate_summary: Binary predictions from model") print(f" Type: {type(binary_preds)}") print(f" Length: {len(binary_preds) if isinstance(binary_preds, dict) else 'N/A'}") print(f" Keys (first 20): {list(binary_preds.keys())[:20] if isinstance(binary_preds, dict) else 'N/A'}") if isinstance(binary_preds, dict) and len(binary_preds) > 0: print(f" Sample entries:") for i, (key, val) in enumerate(list(binary_preds.items())[:5]): print(f" {key}: {val}") print("=" * 80 + "\n") unary_by_obj: Dict[int, List[Tuple[float, str, int]]] = {} for (frame_id, obj_id), preds in unary_preds.items(): for prob, predicate in preds: prob_val = ( float(prob.detach().cpu()) if torch.is_tensor(prob) else float(prob) ) unary_by_obj.setdefault(obj_id, []).append((prob_val, predicate, frame_id)) # Process binary predictions binary_keywords: Dict[str, Dict[str, Any]] = {} for (frame_id, (from_id, to_id)), preds in binary_preds.items(): for prob, predicate in preds: prob_val = ( float(prob.detach().cpu()) if torch.is_tensor(prob) else float(prob) ) pair_key = f"{from_id}-{to_id}" # Keep only the highest confidence prediction for each pair if pair_key not in binary_keywords or prob_val > binary_keywords[pair_key]["confidence"]: binary_keywords[pair_key] = { "predicate": predicate, "confidence": prob_val, "frame_id": int(frame_id), "from_id": int(from_id), "to_id": int(to_id), } objects_summary: Dict[str, Dict[str, Any]] = {} all_obj_ids = set(categorical_preds.keys()) | set(unary_by_obj.keys()) for obj_id in sorted(all_obj_ids): cat_list = categorical_preds.get(obj_id, []) cat_sorted = sorted( [ ( float(p.detach().cpu()) if torch.is_tensor(p) else float(p), label, ) for p, label in cat_list ], key=lambda x: x[0], reverse=True, )[:3] top_categories = [ {"label": label, "probability": prob} for prob, label in cat_sorted ] unary_list = unary_by_obj.get(obj_id, []) unary_sorted = sorted(unary_list, key=lambda x: x[0], reverse=True)[:3] top_unary = [ { "frame_id": int(frame_id), "predicate": predicate, "probability": prob, } for (prob, predicate, frame_id) in unary_sorted ] objects_summary[str(obj_id)] = { "top_categories": top_categories, "top_unary": top_unary, } summary = { "num_objects_detected": len(objects_summary), "objects": objects_summary, } # Add binary keywords to summary if any exist if binary_keywords: summary["binary_keywords"] = binary_keywords print(f"\nDEBUG: Added {len(binary_keywords)} binary keywords to summary") return summary