app.py

from pathlib import Path
from collections.abc import Mapping, Sequence
from functools import lru_cache
import inspect
import shutil
import tempfile
import os
import sys

# Add src/ to sys.path so LASER, video-sam2, GroundingDINO are importable
current_dir = Path(__file__).resolve().parent
src_dir = current_dir / "src"
if src_dir.is_dir() and str(src_dir) not in sys.path:
    sys.path.insert(0, str(src_dir))

import spaces  # <-- ZeroGPU integration
import gradio as gr
import torch
from transformers import pipeline  # not strictly necessary, but fine


# -----------------------------
# Environment / diagnostics
# -----------------------------
os.environ["GRADIO_TEMP_DIR"] = str(Path(__file__).parent / "gradio_temp")
os.environ["OPENAI_API_KEY"] = "test"
os.environ["OMP_NUM_THREADS"] = "4"

print("All imports finished")
print(f"Python version: {sys.version}")
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"CUDA version: {torch.version.cuda}")
print(f"cuDNN version: {torch.backends.cudnn.version()}")
print(f"Number of GPUs: {torch.cuda.device_count()}")

if torch.cuda.is_available():
    for i in range(torch.cuda.device_count()):
        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
        print(
            f"  Memory: {torch.cuda.get_device_properties(i).total_memory / 1e9:.2f} GB"
        )

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
os.environ["TORCH_DTYPE"] = "float32"
torch.set_default_dtype(torch.float32)

current_dir = Path(__file__).resolve().parent
# For Spaces, assume checkpoints live alongside app.py or in a "checkpoints" subdir.
# If you keep them next to app.py locally, this still works.
# NOTE: SAM2 config uses Hydra, so we use just the filename (it searches in sam2/configs/)
sam_config_path = "sam2_hiera_t.yaml"  # Hydra will find this in sam2/configs/
sam_checkpoint_path = str(current_dir / "sam2_hiera_tiny.pt")
gd_config_path = str(current_dir / "GroundingDINO_SwinT_OGC.py")
gd_checkpoint_path = str(current_dir / "groundingdino_swint_ogc.pth")
visualization_dir = str(current_dir / "outputs")
print(
    f"Setting up paths: {sam_config_path}, {sam_checkpoint_path}, {gd_config_path}, {gd_checkpoint_path}"
)




@lru_cache(maxsize=1)
def _load_vine_pipeline():
    """
    Lazy-load and cache the Vine pipeline so we don't re-download/rebuild it on every request.
    """
    from vine_hf import VineConfig, VineModel, VinePipeline

    config = VineConfig(
        segmentation_method="grounding_dino_sam2",
        model_name="openai/clip-vit-base-patch32",
        use_hf_repo=True,
        model_repo="KevinX-Penn28/testing",
        box_threshold=0.35,
        text_threshold=0.25,
        target_fps=1,  # default 1 FPS
        topk_cate=5,
        white_alpha=0.3,
        visualization_dir=visualization_dir,
        visualize=True,
        debug_visualizations=False,
        device="cuda",
        categorical_pool="max",
    )
    model = VineModel(config)
    return VinePipeline(
        model=model,
        tokenizer=None,
        sam_config_path=sam_config_path,
        sam_checkpoint_path=sam_checkpoint_path,
        gd_config_path=gd_config_path,
        gd_checkpoint_path=gd_checkpoint_path,
        device="cuda",
        trust_remote_code=True,
    )


@spaces.GPU(duration=120)  # Up to ~5 minutes of H200 ZeroGPU time per call
def process_video(
    video_file,
    categorical_keywords,
    unary_keywords,
    binary_keywords,
    output_fps,
    box_threshold,
    text_threshold,
    binary_confidence_threshold,
):
    vine_pipe = _load_vine_pipeline()

    # Normalize incoming video input to a file path
    if isinstance(video_file, dict):
        video_file = (
            video_file.get("name")
            or video_file.get("filepath")
            or video_file.get("data")
        )
    if not isinstance(video_file, (str, Path)):
        raise ValueError(f"Unsupported video input type: {type(video_file)}")

    categorical_keywords = (
        [kw.strip() for kw in categorical_keywords.split(",")]
        if categorical_keywords
        else []
    )
    unary_keywords = (
        [kw.strip() for kw in unary_keywords.split(",")] if unary_keywords else []
    )
    binary_keywords = (
        [kw.strip() for kw in binary_keywords.split(",")] if binary_keywords else []
    )

    # Debug: Print what we're sending to the pipeline
    print("\n" + "=" * 80)
    print("INPUT TO VINE PIPELINE:")
    print(f"  categorical_keywords: {categorical_keywords}")
    print(f"  unary_keywords: {unary_keywords}")
    print(f"  binary_keywords: {binary_keywords}")
    print("=" * 80 + "\n")

    # Object pairs is now optional - empty list will auto-generate all pairs in vine_model.py
    object_pairs = []

    results = vine_pipe(
        inputs=video_file,
        categorical_keywords=categorical_keywords,
        unary_keywords=unary_keywords,
        binary_keywords=binary_keywords,
        object_pairs=object_pairs,
        segmentation_method="grounding_dino_sam2",
        return_top_k=5,
        include_visualizations=True,
        debug_visualizations=False,
        device="cuda",
        box_threshold=box_threshold,
        text_threshold=text_threshold,
        target_fps=output_fps,
        binary_confidence_threshold=binary_confidence_threshold,
    )

    # Debug: Print what the pipeline returned
    print("\n" + "=" * 80)
    print("PIPELINE RESULTS DEBUG:")
    print(f"  results type: {type(results)}")
    if isinstance(results, dict):
        print(f"  results keys: {list(results.keys())}")
    print("=" * 80 + "\n")

    vine_pipe.box_threshold = box_threshold
    vine_pipe.text_threshold = text_threshold
    vine_pipe.target_fps = output_fps

    if isinstance(results, Mapping):
        results_dict = results
    elif isinstance(results, Sequence) and results and isinstance(results[0], Mapping):
        results_dict = results[0]
    else:
        results_dict = {}

    visualizations = results_dict.get("visualizations") or {}
    vine = visualizations.get("vine") or {}
    all_vis = vine.get("all") or {}
    result_video_path = all_vis.get("video_path")
    if not result_video_path:
        candidates = sorted(
            Path(visualization_dir).rglob("*.mp4"),
            key=lambda p: p.stat().st_mtime,
            reverse=True,
        )
        result_video_path = str(candidates[0]) if candidates else None
    summary = results_dict.get("summary") or {}

    if result_video_path and os.path.exists(result_video_path):
        gradio_tmp = (
            Path(os.environ.get("GRADIO_TEMP_DIR", tempfile.gettempdir()))
            / "vine_outputs"
        )
        gradio_tmp.mkdir(parents=True, exist_ok=True)
        dest_path = gradio_tmp / Path(result_video_path).name
        try:
            shutil.copyfile(result_video_path, dest_path)
            video_path_for_ui = str(dest_path)
        except Exception as e:
            print(f"Warning: failed to copy video to Gradio temp dir: {e}")
            video_path_for_ui = str(result_video_path)
    else:
        video_path_for_ui = None
        print(
            "Warning: annotated video not found or empty; check visualization settings."
        )

    return video_path_for_ui, summary


def _video_component(label: str, *, is_output: bool = False):
    """
    Build a Gradio Video component that is compatible with older Gradio versions
    (no `type`/`sources`/`format` kwargs) and newer ones when available.
    """
    kwargs = {"label": label}
    sig = inspect.signature(gr.Video.__init__)

    # Only set format for OUTPUT components
    if is_output and "format" in sig.parameters:
        kwargs["format"] = "mp4"

    if not is_output:
        if "type" in sig.parameters:
            kwargs["type"] = "filepath"
        if "sources" in sig.parameters:
            kwargs["sources"] = ["upload"]
        # Restrict to MP4 files only
        if "file_types" in sig.parameters:
            kwargs["file_types"] = [".mp4"]

    if is_output and "autoplay" in sig.parameters:
        kwargs["autoplay"] = True

    return gr.Video(**kwargs)


def _create_blocks():
    """
    Build a Blocks context that works across Gradio versions.
    """
    blocks_kwargs = {"title": "VINE Demo"}
    soft_theme = None

    if hasattr(gr, "themes") and hasattr(gr.themes, "Soft"):
        try:
            soft_theme = gr.themes.Soft()
        except Exception:
            soft_theme = None

    if "theme" in inspect.signature(gr.Blocks).parameters and soft_theme is not None:
        blocks_kwargs["theme"] = soft_theme

    return gr.Blocks(**blocks_kwargs)


# Create Gradio interface with two-column layout
with _create_blocks() as demo:
    gr.Markdown(
        """
        # 🎬 VINE: Video-based Interaction and Event Detection

        Upload an MP4 video and specify keywords to detect objects, actions, and interactions in your video.
        """
    )

    with gr.Row():
        # Left column: Inputs
        with gr.Column(scale=1):
            gr.Markdown("### Input Configuration")

            video_input = _video_component("Upload Video (MP4 only)", is_output=False)
            gr.Markdown("*Note: Only MP4 format is currently supported*")

            gr.Markdown("#### Detection Keywords")
            categorical_input = gr.Textbox(
                label="Categorical Keywords",
                placeholder="e.g., person, car, dog",
                value="person, car, dog",
                info="Objects to detect in the video (comma-separated)",
            )
            unary_input = gr.Textbox(
                label="Unary Keywords",
                placeholder="e.g., walking, running, standing",
                value="walking, running, standing",
                info="Single-object actions to detect (comma-separated)",
            )
            binary_input = gr.Textbox(
                label="Binary Keywords",
                placeholder="e.g., chasing, carrying",
                info="Object-to-object interactions to detect (comma-separated)",
            )

            gr.Markdown("#### Processing Settings")
            fps_input = gr.Number(
                label="Output FPS",
                value=1,
                info="Frames per second for processing (lower = faster)",
            )

            with gr.Accordion("Advanced Settings", open=False):
                box_threshold_input = gr.Slider(
                    label="Box Threshold",
                    minimum=0.1,
                    maximum=0.9,
                    value=0.35,
                    step=0.05,
                    info="Confidence threshold for object detection",
                )
                text_threshold_input = gr.Slider(
                    label="Text Threshold",
                    minimum=0.1,
                    maximum=0.9,
                    value=0.25,
                    step=0.05,
                    info="Confidence threshold for text-based detection",
                )
                binary_confidence_input = gr.Slider(
                    label="Binary Relation Confidence Threshold",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.8,
                    step=0.05,
                    info="Minimum confidence to show binary relations and object pairs",
                )

            submit_btn = gr.Button("🚀 Process Video", variant="primary", size="lg")

        # Right column: Outputs
        with gr.Column(scale=1):
            gr.Markdown("### Results")

            video_output = _video_component("Annotated Video Output", is_output=True)

            gr.Markdown("### Detection Summary")
            summary_output = gr.JSON(label="Summary of Detected Events")

    gr.Markdown(
        """
        ---
        ### How to Use
        1. Upload an MP4 video file
        2. Specify the objects, actions, and interactions you want to detect
        3. Adjust processing settings if needed (including binary relation confidence threshold)
        4. Click "Process Video" to analyze

        The system will automatically detect all binary relations between detected objects
        and show only those with confidence above the threshold (default: 0.8).
        """
    )

    submit_btn.click(
        fn=process_video,
        inputs=[
            video_input,
            categorical_input,
            unary_input,
            binary_input,
            fps_input,
            box_threshold_input,
            text_threshold_input,
            binary_confidence_input,
        ],
        outputs=[video_output, summary_output],
    )

if __name__ == "__main__":
    print("Got to main")
    demo.launch(share=True, debug=True)