Create app.py

app.py

@@ -0,0 +1,109 @@
+import gradio as gr
+import cv2
+import numpy as np
+import torch
+from sam2.sam2_video_predictor import SAM2VideoPredictor
+from ultralytics import YOLO
+import supervision as sv
+import os
+
+# Initialize models from Hugging Face Hub
+predictor = SAM2VideoPredictor.from_pretrained("facebook/sam2.1-hiera-tiny")
+yolo_model = YOLO("ultralytics/yolo-world-v8n")  # Lightweight YOLO-World model
+
+def detect_motorcycles(frame, prompt="motorcycle"):
+    """Detect motorcycles in a frame using YOLO-World and return bounding boxes."""
+    results = yolo_model.predict(frame, prompt=prompt, device="cpu")
+    boxes = []
+    for result in results:
+        for box in result.boxes:
+            if result.names[box.cls] == prompt:
+                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
+                boxes.append([x1, y1, x2, y2])
+    return boxes
+
+def segment_and_highlight_video(video_path, prompt="motorcycle", highlight_color="red"):
+    """Segment and highlight motorcycles in a video using SAM 2 and YOLO-World."""
+    # Create temporary directory for video frames
+    frames_dir = "video_frames"
+    os.makedirs(frames_dir, exist_ok=True)
+
+    # Extract frames
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    # Limit resolution for CPU
+    if width > 640:
+        height = int(height * 640 / width)
+        width = 640
+    frame_paths = []
+
+    # Save frames as JPEG
+    frame_idx = 0
+    with sv.ImageSink(target_dir_path=frames_dir, image_name_pattern="{:05d}.jpeg") as sink:
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frame = cv2.resize(frame, (width, height))
+            sink.save_image(frame)
+            frame_paths.append(os.path.join(frames_dir, f"{frame_idx:05d}.jpeg"))
+            frame_idx += 1
+    cap.release()
+
+    # Initialize SAM 2 inference state
+    with torch.inference_mode():
+        state = predictor.init_state(video_path=frames_dir)
+
+        # Detect motorcycles in the first frame
+        first_frame = cv2.imread(frame_paths[0])
+        boxes = detect_motorcycles(first_frame, prompt)
+
+        # Add boxes as prompts for SAM 2
+        if boxes:
+            frame_idx, obj_ids, masks = predictor.add_new_points_or_box(
+                state, frame_idx=0, obj_ids=[1], boxes=np.array(boxes)
+            )
+
+        # Create output video
+        output_path = "output.mp4"
+        out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
+
+        # Color map for highlighting
+        color_map = {"red": (0, 0, 255), "green": (0, 255, 0), "blue": (255, 0, 0)}
+        highlight_rgb = color_map.get(highlight_color.lower(), (0, 0, 255))
+
+        # Propagate masks and apply highlights
+        for frame_idx, obj_ids, masks in predictor.propagate_in_video(state):
+            frame = cv2.imread(frame_paths[frame_idx])
+            mask = masks[0].astype(np.uint8) * 255  # Assuming one object
+            mask_colored = np.zeros_like(frame)
+            mask_colored[:, :, 0] = mask * highlight_rgb[0]
+            mask_colored[:, :, 1] = mask * highlight_rgb[1]
+            mask_colored[:, :, 2] = mask * highlight_rgb[2]
+            highlighted_frame = cv2.addWeighted(frame, 0.7, mask_colored, 0.3, 0)
+            out.write(highlighted_frame)
+
+        out.release()
+
+    # Clean up
+    for frame_path in frame_paths:
+        os.remove(frame_path)
+    os.rmdir(frames_dir)
+
+    return output_path
+
+# Gradio interface
+iface = gr.Interface(
+    fn=segment_and_highlight_video,
+    inputs=[
+        gr.Video(label="Upload Video"),
+        gr.Textbox(label="Prompt", placeholder="e.g., motorcycle"),
+        gr.Dropdown(choices=["red", "green", "blue"], label="Highlight Color")
+    ],
+    outputs=gr.Video(label="Highlighted Video"),
+    title="Video Segmentation with SAM 2 and YOLO-World (CPU)",
+    description="Upload a short video, specify a text prompt (e.g., 'motorcycle'), and choose a highlight color. Optimized for CPU."
+)
+iface.launch()