app.py

import os
import torch
import numpy as np
import gradio as gr
from PIL import Image
from transformers import AutoProcessor, AutoModelForCausalLM
from sam2.build_sam import build_sam2_video_predictor, build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
import cv2
import traceback
import matplotlib.pyplot as plt
import ffmpeg
from utils import load_model_without_flash_attn
import spaces


# CUDA optimizations
torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
if torch.cuda.get_device_properties(0).major >= 8:
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

# Initialize models
sam2_checkpoint = "./checkpoints/sam2_hiera_large.pt"
model_cfg = "sam2_hiera_l.yaml"

video_predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
image_predictor = SAM2ImagePredictor(sam2_model)

model_id = 'microsoft/Florence-2-large'
device = "cuda" if torch.cuda.is_available() else "cpu"

def load_florence_model():
    return AutoModelForCausalLM.from_pretrained(
        model_id, 
        trust_remote_code=True, 
        torch_dtype=torch.float16 if device == "cuda" else torch.float32
    ).eval().to(device)

florence_model = load_model_without_flash_attn(load_florence_model)
florence_processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)




def apply_color_mask(frame, mask, obj_id):
    cmap = plt.get_cmap("tab10")
    color = np.array(cmap(obj_id % 10)[:3])  # Use modulo 10 to cycle through colors
    
    # Ensure mask has the correct shape
    if mask.ndim == 4:
        mask = mask.squeeze()  # Remove singleton dimensions
    if mask.ndim == 3 and mask.shape[0] == 1:
        mask = mask[0]  # Take the first channel if it's a single-channel 3D array
    
    # Reshape mask to match frame dimensions
    mask = cv2.resize(mask.astype(np.float32), (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_LINEAR)
    
    # Expand dimensions of mask and color for broadcasting
    mask = np.expand_dims(mask, axis=2)
    color = color.reshape(1, 1, 3)
    
    colored_mask = mask * color
    return frame * (1 - mask) + colored_mask * 255

@spaces.GPU
@torch.inference_mode()
@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def run_florence(image, text_input):
    task_prompt = '<OPEN_VOCABULARY_DETECTION>'
    prompt = task_prompt + text_input
    inputs = florence_processor(text=prompt, images=image, return_tensors="pt").to('cuda', torch.bfloat16)
    generated_ids = florence_model.generate(
        input_ids=inputs["input_ids"].cuda(),
        pixel_values=inputs["pixel_values"].cuda(),
        max_new_tokens=1024,
        early_stopping=False,
        do_sample=False,
        num_beams=3,
    )
    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
    parsed_answer = florence_processor.post_process_generation(
        generated_text, 
        task=task_prompt, 
        image_size=(image.width, image.height)
    )
    
    bboxes = parsed_answer[task_prompt]['bboxes']
    if not bboxes:
        print(f"No objects detected for prompt: '{text_input}'. Trying with a default bounding box.")
        # Return a default bounding box covering the entire image
        return [0, 0, image.width, image.height]
    
    return bboxes[0]

def remove_directory_contents(directory):
    for root, dirs, files in os.walk(directory, topdown=False):
        for name in files:
            os.remove(os.path.join(root, name))
        for name in dirs:
            os.rmdir(os.path.join(root, name))

@spaces.GPU
@torch.inference_mode()
@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def process_video(video_path, prompt):
    try:
        # Get video info
        probe = ffmpeg.probe(video_path)
        video_info = next(s for s in probe['streams'] if s['codec_type'] == 'video')
        width = int(video_info['width'])
        height = int(video_info['height'])
        num_frames = int(video_info['nb_frames'])
        fps = eval(video_info['r_frame_rate'])

        print(f"Video info: {width}x{height}, {num_frames} frames, {fps} fps")

        # Read frames
        out, _ = (
            ffmpeg
            .input(video_path)
            .output('pipe:', format='rawvideo', pix_fmt='rgb24')
            .run(capture_stdout=True)
        )
        frames = np.frombuffer(out, np.uint8).reshape([-1, height, width, 3])

        print(f"Read {len(frames)} frames")

        # Florence detection on first frame
        first_frame = Image.fromarray(frames[0])
        mask_box = run_florence(first_frame, prompt)
        print("Original mask box:", mask_box)
        
        # Convert mask_box to numpy array
        mask_box = np.array(mask_box)
        print("Reshaped mask box:", mask_box)
        
        # SAM2 segmentation on first frame
        image_predictor.set_image(first_frame)
        masks, _, _ = image_predictor.predict(
            point_coords=None,
            point_labels=None,
            box=mask_box[None, :],
            multimask_output=False,
        )
        print("masks.shape", masks.shape)
        
        mask = masks.squeeze().astype(bool)
        print("Mask shape:", mask.shape)
        print("Frame shape:", frames[0].shape)
        
        # SAM2 video propagation
        temp_dir = "temp_frames"
        os.makedirs(temp_dir, exist_ok=True)
        for i, frame in enumerate(frames):
            Image.fromarray(frame).save(os.path.join(temp_dir, f"{i:04d}.jpg"))
        
        print(f"Saved {len(frames)} temporary frames")
        
        inference_state = video_predictor.init_state(video_path=temp_dir)
        _, _, _ = video_predictor.add_new_mask(
            inference_state=inference_state,
            frame_idx=0,
            obj_id=1,
            mask=mask
        )
        
        video_segments = {}
        for out_frame_idx, out_obj_ids, out_mask_logits in video_predictor.propagate_in_video(inference_state):
            video_segments[out_frame_idx] = {
                out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
                for i, out_obj_id in enumerate(out_obj_ids)
            }

        print('Segmenting for main vid done')
        print(f"Number of segmented frames: {len(video_segments)}")
        
        # Apply segmentation masks to frames
        all_segmented_frames = []
        for i, frame in enumerate(frames):
            if i in video_segments:
                for out_obj_id, mask in video_segments[i].items():
                    frame = apply_color_mask(frame, mask, out_obj_id)
                all_segmented_frames.append(frame.astype(np.uint8))
            else:
                all_segmented_frames.append(frame)
        
        print(f"Applied masks to {len(all_segmented_frames)} frames")
        
        # Clean up temporary files
        remove_directory_contents(temp_dir)
        os.rmdir(temp_dir)
        
        # Write output video using ffmpeg
        output_path = "segmented_video.mp4"
        process = (
            ffmpeg
            .input('pipe:', format='rawvideo', pix_fmt='rgb24', s=f'{width}x{height}', r=fps)
            .output(output_path, pix_fmt='yuv420p')
            .overwrite_output()
            .run_async(pipe_stdin=True)
        )
        
        for frame in all_segmented_frames:
            process.stdin.write(frame.tobytes())
        
        process.stdin.close()
        process.wait()
        
        if not os.path.exists(output_path):
            raise ValueError(f"Output video file was not created: {output_path}")
        
        print(f"Successfully created output video: {output_path}")
        return output_path

    except Exception as e:
        print(f"Error in process_video: {str(e)}")
        print(traceback.format_exc())  # This will print the full stack trace
        return None
    
@spaces.GPU(duration=600)
def segment_video(video_file, prompt):
    if video_file is None:
        return None
    output_video = process_video(video_file, prompt)
    return output_video

demo = gr.Interface(
    fn=segment_video,
    inputs=[
        gr.Video(label="Upload Video (Keep it under 10 seconds for this demo)"),
        gr.Textbox(label="Enter text prompt for object detection (eg - Gymnast , Car ) ")
    ],
    outputs=gr.Video(label="Segmented Video"),
    title="Text-Prompted Video Object Segmentation with SAMv2",
    description="""
    This demo uses [Florence-2](https://huggingface.co/microsoft/Florence-2-large), to enable text-prompted object detection for [SAM2](https://github.com/facebookresearch/segment-anything).

    1. Upload a short video (< 10 sec , you can fork this space on larger GPU for longer vids) 
    2. Describe the object to segment (The object should be visible in the first frame).
    3. Get your segmented video.
    """
)

demo.launch()