cleanup

app.py

@@ -219,17 +219,16 @@ 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")
+        gr.Textbox(label="Enter text prompt for object detection (eg - Gymnast , Car ) ")
     ],
     outputs=gr.Video(label="Segmented Video"),
-    title="Text-Prompted Video Object Segmentation",
+    title="Text-Prompted Video Object Segmentation with SAMv2",
     description="""
-    This demo uses [Florence-2](https://huggingface.co/microsoft/Florence-2-large), a vision-language model, to enable text-prompted object detection for [SAM2](https://github.com/facebookresearch/segment-anything).
-    Florence-2 interprets your text prompt, allowing SAM2 to segment the described object in the video.
-    
-    1. Upload a short video (< 10 sec)
-    2. Describe the object to segment
-    3. Get your segmented video!
+    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.
+    3. Get your segmented video.
     """
 )
 

myapp2.py

@@ -1,204 +0,0 @@
-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
-
-# 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'
-florence_model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.bfloat16).eval().cuda()
-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
-
-def run_florence(image, text_input):
-    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
-        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)
-        )
-    return parsed_answer[task_prompt]['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))
-
-def process_video(video_path, prompt, chunk_size=30):
-    try:
-        video = cv2.VideoCapture(video_path)
-        if not video.isOpened():
-            raise ValueError("Unable to open video file")
-        
-        fps = video.get(cv2.CAP_PROP_FPS)
-        frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
-        
-        # Process video in chunks
-        all_segmented_frames = []
-        for chunk_start in range(0, frame_count, chunk_size):
-            chunk_end = min(chunk_start + chunk_size, frame_count)
-            
-            frames = []
-            video.set(cv2.CAP_PROP_POS_FRAMES, chunk_start)
-            for _ in range(chunk_end - chunk_start):
-                ret, frame = video.read()
-                if not ret:
-                    break
-                frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-            
-            if not frames:
-                print(f"No frames extracted for chunk starting at {chunk_start}")
-                continue
-            
-            # Florence detection on first frame of the chunk
-            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 and ensure it's in the correct format
-            mask_box = np.array(mask_box)
-            print("Reshaped mask box:", mask_box)
-            
-            # SAM2 segmentation on first frame
-            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
-                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 = f"temp_frames_{chunk_start}"
-            os.makedirs(temp_dir, exist_ok=True)
-            for i, frame in enumerate(frames):
-                cv2.imwrite(os.path.join(temp_dir, f"{i:04d}.jpg"), cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-            
-            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
-                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')
-            
-            # Apply segmentation masks to 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)
-            
-            # Clean up temporary files
-            remove_directory_contents(temp_dir)
-            os.rmdir(temp_dir)
-        
-        video.release()
-        
-        if not all_segmented_frames:
-            raise ValueError("No frames were processed successfully")
-
-        # Create video from segmented frames
-        output_path = "segmented_video.mp4"
-        out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, 
-                              (all_segmented_frames[0].shape[1], all_segmented_frames[0].shape[0]))
-        for frame in all_segmented_frames:
-            out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-        out.release()
-        
-        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
-
-def segment_video(video_file, prompt, chunk_size):
-    if video_file is None:
-        return None
-    output_video = process_video(video_file, prompt, int(chunk_size))
-    return output_video
-
-demo = gr.Interface(
-    fn=segment_video,
-    inputs=[
-        gr.Video(label="Upload Video"),
-        gr.Textbox(label="Enter prompt (e.g., 'a gymnast')"),
-        gr.Slider(minimum=10, maximum=100, step=10, value=30, label="Chunk Size (frames)")
-    ],
-    outputs=gr.Video(label="Segmented Video"),
-    title="Video Object Segmentation with Florence and SAM2",
-    description="Upload a video and provide a text prompt to segment a specific object throughout the video."
-)
-
-demo.launch()