vine_hf/vine_pipeline.py

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": {
            "<obj_id>": {
              "top_categories": [{"label": str, "probability": float}, ...],
              "top_unary": [{"frame_id": int, "predicate": str, "probability": float}, ...],
            }
          },
          "binary_keywords": {
            "<from_id>-<to_id>": {"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