import torch import numpy as np import cv2 import os from typing import Dict, List, Tuple, Optional, Any, Union from transformers import Pipeline import tempfile import uuid 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. This pipeline takes a video file or frames, along with segmentation method and keyword lists, and returns probability distributions over the keywords. Segmentation Model Configuration: The pipeline requires SAM2 and GroundingDINO models for mask generation. You can configure custom paths via constructor kwargs: - sam_config_path: Path to SAM2 config (e.g., "configs/sam2.1/sam2.1_hiera_b+.yaml") - sam_checkpoint_path: Path to SAM2 checkpoint (e.g., "checkpoints/sam2.1_hiera_base_plus.pt") - gd_config_path: Path to GroundingDINO config (e.g., "groundingdino/config/GroundingDINO_SwinT_OGC.py") - gd_checkpoint_path: Path to GroundingDINO checkpoint (e.g., "checkpoints/groundingdino_swint_ogc.pth") Old: - SAM2: ~/research/sam2/ or /home/asethi04/LASER_NEW/LASER/sam2/ - GroundingDINO: /home/asethi04/LASER_NEW/LASER/GroundingDINO/ Alternative: Use set_segmentation_models() to provide pre-initialized model instances. """ 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 ): 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) # Set default parameters from config 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") def set_segmentation_models( self, *, sam_predictor=None, mask_generator=None, grounding_model=None ): """ Set pre-initialized segmentation models, bypassing automatic initialization/current_values Args: sam_predictor: Pre-built SAM2 video predictor mask_generator: Pre-built SAM2 automatic mask generator grounding_model: Pre-built GroundingDINO model """ 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 def _sanitize_parameters(self, **kwargs): """Sanitize parameters for different pipeline stages.""" preprocess_kwargs = {} forward_kwargs = {} postprocess_kwargs = {} # Preprocess parameters 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"] # Forward parameters 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 "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"] # Postprocess parameters 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"] return preprocess_kwargs, forward_kwargs, postprocess_kwargs def preprocess( self, video_input: Union[str, np.ndarray, torch.Tensor], segmentation_method: str = None, target_fps: int = None, box_threshold: float = None, text_threshold: float = None, categorical_keywords: List[str] = None, **kwargs ) -> Dict[str, Any]: """ Preprocess video input and generate masks. Args: video_input: Path to video file, or video tensor/array segmentation_method: "sam2" or "grounding_dino_sam2" target_fps: Target FPS for video processing box_threshold: Box threshold for Grounding DINO text_threshold: Text threshold for Grounding DINO categorical_keywords: Keywords for Grounding DINO segmentation Returns: Dict containing video frames, masks, and bboxes """ # Use defaults from config if not provided if segmentation_method is None: segmentation_method = self.segmentation_method if target_fps is None: target_fps = self.target_fps if box_threshold is None: box_threshold = self.box_threshold if text_threshold is None: text_threshold = self.text_threshold if categorical_keywords is None: categorical_keywords = ["object"] # Default generic category if isinstance(video_input, str): # Video file path 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)): # Video tensor/array 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)}") # Ensure video tensor is numpy array if not isinstance(video_tensor, np.ndarray): video_tensor = np.array(video_tensor) # Ensure video tensor is in correct format if len(video_tensor.shape) != 4: raise ValueError(f"Expected video tensor shape (frames, height, width, channels), got {video_tensor.shape}") # Generate masks and bboxes based on segmentation method 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, } def _generate_sam2_masks(self, video_tensor: np.ndarray) -> Tuple[Dict, Dict, Dict[str, Any]]: """Generate masks using SAM2 automatic mask generation.""" # Initialize SAM2 models if not already done 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) height, width, _ = frame.shape 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] return masks, bboxes, {"sam_masks": masks} def _generate_grounding_dino_sam2_masks( self, video_tensor: np.ndarray, categorical_keywords: List[str], box_threshold: float, text_threshold: float, video_path: str, ) -> Tuple[Dict, Dict, Dict[str, Any]]: """Generate masks using Grounding DINO + SAM2.""" # Initialize models if not already done 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)] 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, "temp_video", out_dir=tempfile.gettempdir(), classes_ls=classes_ls, target_fps=self.target_fps, visualize=self.debug_visualizations, frames=None, max_prop_time=10 ) 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) vis_data: Dict[str, Any] = { "sam_masks": masks, "dino_labels": oid_class_pred, } return masks, bboxes, vis_data def _initialize_segmentation_models(self): """Initialize segmentation models based on the requested method and configured paths.""" 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): """Initialize SAM2 video predictor and mask generator.""" 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 # Resolve SAM2 paths config_path, checkpoint_path = self._resolve_sam2_paths() # Validate paths if custom ones were provided 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}") # Only proceed if we have valid paths 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 print(type(device)) # Video predictor self.sam_predictor = build_sam2_video_predictor( config_path, checkpoint_path, device=device ) # Mask generator 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): """Initialize GroundingDINO model.""" try: from groundingdino.util.inference import Model as gd_Model except ImportError as e: print(f"Warning: Could not import GroundingDINO: {e}") return # Resolve GroundingDINO paths config_path, checkpoint_path = self._resolve_grounding_dino_paths() # Validate paths if custom ones were provided 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}") # Only proceed if we have valid paths 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 print(type(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): """Resolve SAM2 config and checkpoint paths.""" # Use custom paths if provided 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): """Resolve GroundingDINO config and checkpoint paths.""" # Use custom paths if provided if self.gd_config_path and self.gd_checkpoint_path: return self.gd_config_path, self.gd_checkpoint_path def _prepare_visualization_dir(self, name: str, enabled: bool) -> Optional[str]: """ Ensure a directory exists for visualization artifacts and return it. If visualization is disabled, returns None. """ 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: """Create a temporary video file from video tensor.""" 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) # Use OpenCV to write video height, width = video_tensor.shape[1:3] fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(temp_path, fourcc, self.target_fps, (width, height)) for frame in video_tensor: # Convert RGB to BGR for OpenCV if len(frame.shape) == 3 and frame.shape[2] == 3: frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) else: frame_bgr = frame out.write(frame_bgr.astype(np.uint8)) out.release() return temp_path def _forward(self, model_inputs: Dict[str, Any], **forward_kwargs) -> Dict[str, Any]: """Forward pass through the model.""" 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 ) -> Dict[str, Any]: """ Postprocess model outputs into user-friendly format. Args: model_outputs: Raw model outputs return_top_k: Number of top predictions to return self.visualize: Whether to include visualization data Returns: Formatted results """ results = { "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) } 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 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)) 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 vine_frame_sets = render_vine_frame_sets( frames_np, bboxes, cat_label_lookup, unary_lookup, binary_lookup, visualization_data.get("sam_masks"), ) 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 def _generate_summary(self, model_outputs: Dict[str, Any]) -> Dict[str, Any]: """Generate a summary of the predictions.""" categorical_preds = model_outputs.get("categorical_predictions", {}) unary_preds = model_outputs.get("unary_predictions", {}) binary_preds = model_outputs.get("binary_predictions", {}) summary = { "num_objects_detected": len(categorical_preds), "num_unary_predictions": len(unary_preds), "num_binary_predictions": len(binary_preds), "top_categories": [], "top_actions": [], "top_relations": [] } # Extract top categories all_categories = [] for obj_preds in categorical_preds.values(): if obj_preds: all_categories.extend(obj_preds) if all_categories: sorted_categories = sorted(all_categories, reverse=True) summary["top_categories"] = [(cat, prob) for prob, cat in sorted_categories[:3]] # Extract top actions all_actions = [] for action_preds in unary_preds.values(): if action_preds: all_actions.extend(action_preds) if all_actions: sorted_actions = sorted(all_actions, reverse=True) summary["top_actions"] = [(act, prob) for prob, act in sorted_actions[:3]] # Extract top relations all_relations = [] for rel_preds in binary_preds.values(): if rel_preds: all_relations.extend(rel_preds) if all_relations: sorted_relations = sorted(all_relations, reverse=True) summary["top_relations"] = [(rel, prob) for prob, rel in sorted_relations[:3]] return summary