{ "cells": [ { "cell_type": "markdown", "id": "c2a2297e", "metadata": {}, "source": [ "# 🎬 Video Object Removal Pipeline\n", "## SAM2 + ProPainter | Colab Version\n", "\n", "**Requirements:** GPU runtime (T4 or better)\n", "\n", "Run Cell 1 (setup), then Cell 2 (launch). You'll get a public URL with the full UI built in — no separate HTML file needed!" ] }, { "cell_type": "code", "execution_count": null, "id": "c44b31bc", "metadata": {}, "outputs": [], "source": [ "# ============================================================\n", "# 📦 CELL 1: SETUP — Install everything (run once)\n", "# ============================================================\n", "import os, sys\n", "\n", "!git clone https://github.com/facebookresearch/sam2.git /content/sam2 -q\n", "os.chdir('/content/sam2')\n", "!pip install -e . -q\n", "!wget -q -P /content/sam2/checkpoints/ https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt\n", "!git clone https://github.com/sczhou/ProPainter.git /content/ProPainter -q\n", "!pip install -r /content/ProPainter/requirements.txt -q\n", "!pip install gradio -q\n", "\n", "print(\"✅ ALL SETUP DONE!\")" ] }, { "cell_type": "markdown", "id": "5be7dcf6", "metadata": {}, "source": [ "### 🚀 Launch the App\n", "Run this cell to get a public URL with the complete UI." ] }, { "cell_type": "code", "execution_count": null, "id": "93daccf6", "metadata": {}, "outputs": [], "source": [ "# ============================================================\n", "# 🚀 CELL 2: LAUNCH APP — Full UI with video upload + object selection\n", "# ============================================================\n", "\n", "import gradio as gr\n", "import os, sys, json, shutil, subprocess, tempfile\n", "import torch\n", "import numpy as np\n", "import cv2\n", "\n", "def extract_frame(video_path):\n", " \"\"\"Extract first frame from video for object selection.\"\"\"\n", " if not video_path:\n", " return None\n", " cap = cv2.VideoCapture(video_path)\n", " ret, frame = cap.read()\n", " cap.release()\n", " if ret:\n", " return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n", " return None\n", "\n", "def remove_objects(video_path, coords_json):\n", " \"\"\"Main pipeline: video + coordinates → result video.\"\"\"\n", " try:\n", " if not video_path:\n", " return None, \"❌ No video provided!\"\n", "\n", " coords = json.loads(coords_json) if coords_json else []\n", " if not coords:\n", " return None, \"❌ No coordinates provided! Click on the frame to select objects.\"\n", "\n", " work_dir = tempfile.mkdtemp(prefix=\"objremover_\")\n", " frames_dir = os.path.join(work_dir, \"frames\")\n", " masks_dir = os.path.join(work_dir, \"masks\")\n", " output_dir = os.path.join(work_dir, \"output\")\n", "\n", " for d in [frames_dir, masks_dir, output_dir]:\n", " os.makedirs(d, exist_ok=True)\n", "\n", " # Extract frames\n", " subprocess.run([\n", " 'ffmpeg', '-i', video_path, '-q:v', '2',\n", " os.path.join(frames_dir, '%05d.jpg'),\n", " '-hide_banner', '-loglevel', 'quiet'\n", " ], check=True)\n", "\n", " num_frames = len([f for f in os.listdir(frames_dir) if f.endswith('.jpg')])\n", " print(f\"📁 Extracted {num_frames} frames\")\n", " if num_frames == 0:\n", " return None, \"❌ Could not extract frames!\"\n", "\n", " # SAM2\n", " from hydra.core.global_hydra import GlobalHydra\n", " from hydra import initialize_config_dir\n", " from sam2.build_sam import build_sam2_video_predictor\n", "\n", " os.chdir('/content/sam2')\n", " if '/content/sam2' not in sys.path:\n", " sys.path.insert(0, '/content/sam2')\n", " GlobalHydra.instance().clear()\n", "\n", " with initialize_config_dir(config_dir=\"/content/sam2/sam2/configs/sam2.1\", version_base=None):\n", " predictor = build_sam2_video_predictor(\n", " \"sam2.1_hiera_l.yaml\",\n", " \"/content/sam2/checkpoints/sam2.1_hiera_large.pt\",\n", " device=\"cuda\"\n", " )\n", "\n", " inference_state = predictor.init_state(video_path=frames_dir)\n", "\n", " obj_id = 1\n", " for coord in coords:\n", " if coord.get('type') == 'point':\n", " predictor.add_new_points_or_box(\n", " inference_state=inference_state, frame_idx=0, obj_id=obj_id,\n", " points=np.array([[coord['x'], coord['y']]], dtype=np.float32),\n", " labels=np.array([1], dtype=np.int32)\n", " )\n", " print(f\"📍 Point ({coord['x']}, {coord['y']}) → obj {obj_id}\")\n", " elif coord.get('type') == 'box':\n", " predictor.add_new_points_or_box(\n", " inference_state=inference_state, frame_idx=0, obj_id=obj_id,\n", " box=np.array([coord['x1'], coord['y1'], coord['x2'], coord['y2']], dtype=np.float32)\n", " )\n", " print(f\"▭ Box ({coord['x1']},{coord['y1']})→({coord['x2']},{coord['y2']}) → obj {obj_id}\")\n", " obj_id += 1\n", "\n", " video_segments = {}\n", " for frame_idx, object_ids, masks in predictor.propagate_in_video(inference_state):\n", " video_segments[frame_idx] = {\n", " oid: masks[i].cpu().numpy() for i, oid in enumerate(object_ids)\n", " }\n", "\n", " print(f\"🔍 Tracked {len(video_segments)} frames\")\n", "\n", " for frame_idx, segments in video_segments.items():\n", " combined = None\n", " for oid, mask_data in segments.items():\n", " binary = (mask_data[0] > 0).astype(np.uint8) * 255\n", " combined = binary if combined is None else np.maximum(combined, binary)\n", " cv2.imwrite(os.path.join(masks_dir, f\"{frame_idx+1:05d}.png\"), combined)\n", "\n", " # ProPainter\n", " os.chdir('/content/ProPainter')\n", " sample = cv2.imread(os.path.join(frames_dir, '00001.jpg'))\n", " h, w = sample.shape[:2]\n", "\n", " subprocess.run([\n", " 'python', 'inference_propainter.py',\n", " '--video', frames_dir, '--mask', masks_dir,\n", " '--output', output_dir, '--height', str(h), '--width', str(w)\n", " ], check=True)\n", "\n", " # Find output\n", " inpaint_path = None\n", " for root, dirs, files in os.walk(output_dir):\n", " for f in files:\n", " if 'inpaint' in f and f.endswith('.mp4'):\n", " inpaint_path = os.path.join(root, f)\n", " break\n", "\n", " if not inpaint_path:\n", " return None, \"❌ ProPainter produced no output!\"\n", "\n", " final_path = os.path.join(tempfile.gettempdir(), \"result_video.mp4\")\n", " subprocess.run([\n", " 'ffmpeg', '-i', inpaint_path,\n", " '-c:v', 'libx264', '-pix_fmt', 'yuv420p',\n", " final_path, '-y', '-hide_banner', '-loglevel', 'quiet'\n", " ], check=True)\n", "\n", " size_mb = os.path.getsize(final_path) / (1024*1024)\n", " msg = f\"✅ Removed {len(coords)} object(s) across {num_frames} frames ({size_mb:.1f} MB)\"\n", " print(msg)\n", " return final_path, msg\n", "\n", " except Exception as e:\n", " import traceback\n", " traceback.print_exc()\n", " return None, f\"❌ Error: {str(e)}\"\n", "\n", "\n", "# ─── BUILD GRADIO UI ───\n", "with gr.Blocks(title=\"Video Object Remover\", theme=gr.themes.Base(primary_hue=\"emerald\")) as demo:\n", " gr.Markdown(\"# 🎬 Video Object Remover\")\n", " gr.Markdown(\"Upload a video → click on the frame to select objects → hit Remove!\")\n", "\n", " with gr.Row():\n", " with gr.Column(scale=1):\n", " video_input = gr.Video(label=\"📁 Upload Video\")\n", " frame_preview = gr.Image(label=\"👆 Click on the object to remove\", interactive=True, type=\"numpy\")\n", " coords_display = gr.Textbox(label=\"📍 Coordinates (auto-filled when you click)\", lines=3, interactive=True)\n", "\n", " with gr.Column(scale=1):\n", " result_video = gr.Video(label=\"🎬 Result Video\")\n", " status_text = gr.Textbox(label=\"Status\", interactive=False)\n", "\n", " remove_btn = gr.Button(\"🗑 Remove Objects\", variant=\"primary\", size=\"lg\")\n", "\n", " # State to track coordinates\n", " coord_state = gr.State([])\n", "\n", " def on_video_upload(video):\n", " frame = extract_frame(video)\n", " return frame, \"[]\", []\n", "\n", " def on_frame_click(frame, evt: gr.SelectData, current_coords):\n", " if frame is None:\n", " return frame, \"[]\", []\n", "\n", " x, y = evt.index[0], evt.index[1]\n", " coord = {\"type\": \"point\", \"x\": int(x), \"y\": int(y), \"frame_time\": 0.0}\n", "\n", " if current_coords is None:\n", " current_coords = []\n", " current_coords.append(coord)\n", "\n", " # Draw marker on frame\n", " marked_frame = frame.copy()\n", " cv2.circle(marked_frame, (x, y), 8, (0, 255, 100), 2)\n", " cv2.circle(marked_frame, (x, y), 3, (0, 255, 100), -1)\n", "\n", " # Draw all previous markers too\n", " for c in current_coords[:-1]:\n", " cx, cy = c['x'], c['y']\n", " cv2.circle(marked_frame, (cx, cy), 8, (0, 255, 100), 2)\n", " cv2.circle(marked_frame, (cx, cy), 3, (0, 255, 100), -1)\n", "\n", " return marked_frame, json.dumps(current_coords, indent=2), current_coords\n", "\n", " def on_remove(video, coords_json):\n", " return remove_objects(video, coords_json)\n", "\n", " video_input.change(on_video_upload, [video_input], [frame_preview, coords_display, coord_state])\n", " frame_preview.select(on_frame_click, [frame_preview, coord_state], [frame_preview, coords_display, coord_state])\n", " remove_btn.click(on_remove, [video_input, coords_display], [result_video, status_text])\n", "\n", "demo.launch(share=True, debug=True, allowed_paths=[\"/tmp\", \"/content\"])" ] } ], "metadata": { "accelerator": "GPU", "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 5 }