Devteamdl/VOID-Quadmask-Reasoner-api
0
1"""2VOID VLM-Mask-Reasoner — Quadmask Generation Demo3Generates 4-level semantic masks for interaction-aware video inpainting.4 5Pipeline from https://github.com/Netflix/void-model:6 Stage 1: SAM2 segmentation → black mask (transformers Sam2Model)7 Stage 2: Gemini VLM scene analysis → affected objects JSON (repo code)8 Stage 3: SAM3 text-prompted segmentation → grey mask (transformers Sam3Model)9 Stage 4: Combine black + grey → quadmask (0/63/127/255) (repo code)10"""11 12import os13import sys14import json15import tempfile16import shutil17import subprocess18from pathlib import Path19 20import cv221import numpy as np22import torch23import gradio as gr24import spaces25import imageio26from PIL import Image, ImageDraw27from huggingface_hub import hf_hub_download28import openai29 30# ── Add repo modules to path ─────────────────────────────────────────────────31sys.path.insert(0, os.path.join(os.path.dirname(__file__), "VLM-MASK-REASONER"))32 33# ── Repo imports: Stage 2 (VLM) and Stage 4 (combine) ────────────────────────34from stage2_vlm_analysis import (35 process_video as vlm_process_video,36 calculate_square_grid,37)38from stage4_combine_masks import process_video as combine_process_video39# Stage 3 helpers (grid logic, mask combination — not the SegmentationModel)40from stage3a_generate_grey_masks_v2 import (41 calculate_square_grid as calc_grid_3a,42 gridify_masks,43 filter_masks_by_proximity,44 segment_object_all_frames as _repo_segment_all_frames,45 process_video_grey_masks,46)47 48# ── Constants ─────────────────────────────────────────────────────────────────49SAM2_MODEL_ID = "facebook/sam2.1-hiera-large"50SAM3_MODEL_ID = "jetjodh/sam3"51DEFAULT_VLM_MODEL = "gemini-3-flash-preview"52MAX_FRAMES = 19753FPS_DEFAULT = 1254FRAME_STRIDE = 4 # Process every Nth frame for SAM2 tracking55 56# ── Load transformers SAM2 (video model with propagation support) ─────────────57print("Loading SAM2 video model (transformers)...")58from transformers import Sam2VideoModel, Sam2VideoProcessor59from transformers.models.sam2_video.modeling_sam2_video import Sam2VideoInferenceSession60sam2_model = Sam2VideoModel.from_pretrained(SAM2_MODEL_ID).to("cuda")61sam2_processor = Sam2VideoProcessor.from_pretrained(SAM2_MODEL_ID)62print("SAM2 video model ready.")63 64# ── Load transformers SAM3 ───────────────────────────────────────────────────65print("Loading SAM3 model (transformers)...")66from transformers import Sam3Model, Sam3Processor67sam3_model = Sam3Model.from_pretrained(SAM3_MODEL_ID).to("cuda")68sam3_processor = Sam3Processor.from_pretrained(SAM3_MODEL_ID)69print("SAM3 ready.")70 71 72# ══════════════════════════════════════════════════════════════════════════════73# STAGE 1: SAM2 VIDEO SEGMENTATION (transformers Sam2VideoModel)74# Uses proper video propagation with memory — matches repo's propagate_in_video75# ══════════════════════════════════════════════════════════════════════════════76 77def stage1_segment_video(frames: list, points: list, **kwargs) -> list:78 """Segment primary object across all video frames using SAM2 video propagation.79 Matches repo: point prompts + bounding box on frame 0, propagate through video.80 Returns list of uint8 masks (0=object, 255=background)."""81 total = len(frames)82 h, w = frames[0].shape[:2]83 84 # Preprocess all frames85 pil_frames = [Image.fromarray(f) for f in frames]86 inputs = sam2_processor(images=pil_frames, return_tensors="pt").to(sam2_model.device)87 88 # Create inference session with all frames89 session = Sam2VideoInferenceSession(90 video=inputs["pixel_values"],91 video_height=h,92 video_width=w,93 inference_device=sam2_model.device,94 inference_state_device=sam2_model.device,95 dtype=torch.float32,96 )97 98 # Add point prompts + bounding box on frame 0 via processor99 # (handles normalization, object registration, and obj_with_new_inputs)100 pts = np.array(points, dtype=np.float32)101 x_min, x_max = pts[:, 0].min(), pts[:, 0].max()102 y_min, y_max = pts[:, 1].min(), pts[:, 1].max()103 x_margin = max((x_max - x_min) * 0.1, 10)104 y_margin = max((y_max - y_min) * 0.1, 10)105 box = [106 max(0, x_min - x_margin),107 max(0, y_min - y_margin),108 min(w, x_max + x_margin),109 min(h, y_max + y_margin),110 ]111 112 sam2_processor.process_new_points_or_boxes_for_video_frame(113 inference_session=session,114 frame_idx=0,115 obj_ids=[1],116 input_points=[[[[float(p[0]), float(p[1])] for p in points]]],117 input_labels=[[[1] * len(points)]],118 input_boxes=[[[float(box[0]), float(box[1]), float(box[2]), float(box[3])]]],119 )120 121 # Run forward on the prompted frame first (populates cond_frame_outputs)122 with torch.no_grad():123 sam2_model(session, frame_idx=0)124 125 # Propagate through all frames (matches repo's propagate_in_video)126 video_segments = {}127 original_sizes = [[h, w]]128 with torch.no_grad():129 for output in sam2_model.propagate_in_video_iterator(session):130 frame_idx = output.frame_idx131 # pred_masks shape varies — get the raw logits and resize to original132 mask_logits = output.pred_masks[0].cpu().float() # first object133 # Ensure 4D for interpolation: (1, 1, H_model, W_model)134 while mask_logits.dim() < 4:135 mask_logits = mask_logits.unsqueeze(0)136 mask_resized = torch.nn.functional.interpolate(137 mask_logits, size=(h, w), mode="bilinear", align_corners=False138 )139 mask = (mask_resized.squeeze() > 0.0).numpy()140 video_segments[frame_idx] = mask141 142 # Convert to uint8 masks (0=object, 255=background)143 all_masks = []144 for idx in range(total):145 if idx in video_segments:146 mask_bool = video_segments[idx]147 else:148 nearest = min(video_segments.keys(), key=lambda k: abs(k - idx))149 mask_bool = video_segments[nearest]150 mask_uint8 = np.where(mask_bool, 0, 255).astype(np.uint8)151 all_masks.append(mask_uint8)152 153 return all_masks154 155 156def write_mask_video(masks: list, fps: float, output_path: str):157 """Write list of uint8 grayscale masks to lossless MP4."""158 h, w = masks[0].shape[:2]159 temp_avi = str(Path(output_path).with_suffix('.avi'))160 fourcc = cv2.VideoWriter_fourcc(*'FFV1')161 out = cv2.VideoWriter(temp_avi, fourcc, fps, (w, h), isColor=False)162 for mask in masks:163 out.write(mask)164 out.release()165 166 cmd = [167 'ffmpeg', '-y', '-i', temp_avi,168 '-c:v', 'libx264', '-qp', '0', '-preset', 'ultrafast',169 '-pix_fmt', 'yuv444p', str(output_path),170 ]171 subprocess.run(cmd, capture_output=True)172 if os.path.exists(temp_avi):173 os.unlink(temp_avi)174 175 176# ══════════════════════════════════════════════════════════════════════════════177# STAGE 3: SAM3 TEXT-PROMPTED SEGMENTATION (transformers)178# — Drop-in replacement for repo's SegmentationModel.segment()179# ══════════════════════════════════════════════════════════════════════════════180 181class TransformersSam3Segmenter:182 """Matches the interface of the repo's SegmentationModel for stage3a."""183 model_type = "sam3"184 185 def segment(self, image_pil: Image.Image, prompt: str) -> np.ndarray:186 """Segment object by text prompt. Returns boolean mask."""187 h, w = image_pil.height, image_pil.width188 union = np.zeros((h, w), dtype=bool)189 190 try:191 inputs = sam3_processor(192 images=image_pil, text=prompt, return_tensors="pt"193 ).to(sam3_model.device)194 195 with torch.no_grad():196 outputs = sam3_model(**inputs)197 198 results = sam3_processor.post_process_instance_segmentation(199 outputs,200 threshold=0.3,201 mask_threshold=0.5,202 target_sizes=inputs.get("original_sizes").tolist(),203 )[0]204 205 masks = results.get("masks")206 if masks is not None and len(masks) > 0:207 if torch.is_tensor(masks):208 masks = masks.cpu().numpy()209 if masks.ndim == 2:210 union = masks.astype(bool)211 elif masks.ndim == 3:212 union = masks.any(axis=0).astype(bool)213 elif masks.ndim == 4:214 union = masks.any(axis=(0, 1)).astype(bool)215 except Exception as e:216 print(f" Warning: SAM3 segmentation failed for '{prompt}': {e}")217 218 return union219 220 221seg_model = TransformersSam3Segmenter()222 223 224# ══════════════════════════════════════════════════════════════════════════════225# HELPERS226# ══════════════════════════════════════════════════════════════════════════════227 228def extract_frames(video_path: str, max_frames: int = MAX_FRAMES):229 """Extract frames from video. Returns (frames_rgb_list, fps)."""230 cap = cv2.VideoCapture(video_path)231 fps = cap.get(cv2.CAP_PROP_FPS) or FPS_DEFAULT232 frames = []233 while len(frames) < max_frames:234 ret, frame = cap.read()235 if not ret:236 break237 frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))238 cap.release()239 return frames, fps240 241 242def draw_points_on_image(image: np.ndarray, points: list, radius: int = 6) -> np.ndarray:243 pil_img = Image.fromarray(image.copy())244 draw = ImageDraw.Draw(pil_img)245 for i, (x, y) in enumerate(points):246 r = radius247 draw.ellipse([x - r, y - r, x + r, y + r], fill="red", outline="white", width=2)248 draw.text((x + r + 2, y - r), str(i + 1), fill="white")249 return np.array(pil_img)250 251 252def frames_to_video(frames: list, fps: float) -> str:253 tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)254 tmp_path = tmp.name255 tmp.close()256 writer = imageio.get_writer(tmp_path, fps=fps, codec='libx264',257 output_params=['-crf', '18', '-pix_fmt', 'yuv420p'])258 for frame in frames:259 writer.append_data(frame)260 writer.close()261 return tmp_path262 263 264def create_quadmask_visualization(video_path: str, quadmask_path: str) -> str:265 cap_vid = cv2.VideoCapture(video_path)266 cap_qm = cv2.VideoCapture(quadmask_path)267 fps = cap_vid.get(cv2.CAP_PROP_FPS) or FPS_DEFAULT268 269 vis_frames = []270 while True:271 ret_v, frame = cap_vid.read()272 ret_q, qm_frame = cap_qm.read()273 if not ret_v or not ret_q:274 break275 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)276 qm = cv2.cvtColor(qm_frame, cv2.COLOR_BGR2GRAY) if len(qm_frame.shape) == 3 else qm_frame277 278 qm = np.where(qm <= 31, 0, qm)279 qm = np.where((qm > 31) & (qm <= 95), 63, qm)280 qm = np.where((qm > 95) & (qm <= 191), 127, qm)281 qm = np.where(qm > 191, 255, qm)282 283 overlay = frame_rgb.copy()284 overlay[qm == 0] = [255, 50, 50]285 overlay[qm == 63] = [255, 200, 0]286 overlay[qm == 127] = [50, 255, 50]287 result = cv2.addWeighted(frame_rgb, 0.5, overlay, 0.5, 0)288 result[qm == 255] = frame_rgb[qm == 255]289 vis_frames.append(result)290 291 cap_vid.release()292 cap_qm.release()293 return frames_to_video(vis_frames, fps) if vis_frames else None294 295 296def create_quadmask_color_video(quadmask_path: str) -> str:297 cap = cv2.VideoCapture(quadmask_path)298 fps = cap.get(cv2.CAP_PROP_FPS) or FPS_DEFAULT299 color_frames = []300 while True:301 ret, frame = cap.read()302 if not ret:303 break304 qm = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if len(frame.shape) == 3 else frame305 qm = np.where(qm <= 31, 0, qm)306 qm = np.where((qm > 31) & (qm <= 95), 63, qm)307 qm = np.where((qm > 95) & (qm <= 191), 127, qm)308 qm = np.where(qm > 191, 255, qm)309 h, w = qm.shape310 color = np.full((h, w, 3), 255, dtype=np.uint8)311 color[qm == 0] = [0, 0, 0]312 color[qm == 63] = [80, 80, 80]313 color[qm == 127] = [160, 160, 160]314 color_frames.append(color)315 cap.release()316 return frames_to_video(color_frames, fps) if color_frames else None317 318 319# ══════════════════════════════════════════════════════════════════════════════320# MAIN PIPELINE321# ══════════════════════════════════════════════════════════════════════════════322 323@spaces.GPU(duration=300)324def run_pipeline(video_path: str, points_json: str, instruction: str,325 progress=gr.Progress(track_tqdm=False)):326 """Run the full VLM-Mask-Reasoner pipeline."""327 if not video_path:328 raise gr.Error("Please upload a video.")329 if not points_json or points_json == "[]":330 raise gr.Error("Please click on the image to select at least one point on the primary object.")331 if not instruction.strip():332 raise gr.Error("Please enter an edit instruction.")333 334 points = json.loads(points_json)335 if len(points) == 0:336 raise gr.Error("Please select at least one point on the primary object.")337 338 api_key = os.environ.get("GEMINI_API_KEY", "")339 340 # Create temp output directory341 output_dir = Path(tempfile.mkdtemp(prefix="void_quadmask_"))342 input_video_path = output_dir / "input_video.mp4"343 shutil.copy2(video_path, input_video_path)344 345 # ── Stage 1: SAM2 Segmentation ──────────────────────────────────────────346 progress(0.05, desc="Stage 1: SAM2 segmentation...")347 frames, fps = extract_frames(str(input_video_path))348 if len(frames) < 2:349 raise gr.Error("Video must have at least 2 frames.")350 351 black_masks = stage1_segment_video(frames, points, stride=FRAME_STRIDE)352 black_mask_path = output_dir / "black_mask.mp4"353 write_mask_video(black_masks, fps, str(black_mask_path))354 355 # Save first frame for VLM analysis356 first_frame_path = output_dir / "first_frame.jpg"357 cv2.imwrite(str(first_frame_path), cv2.cvtColor(frames[0], cv2.COLOR_RGB2BGR))358 359 # Save segmentation metadata (Stage 2 expects this)360 seg_info = {361 "total_frames": len(frames),362 "frame_width": frames[0].shape[1],363 "frame_height": frames[0].shape[0],364 "fps": fps,365 "video_path": str(input_video_path),366 "instruction": instruction,367 "primary_points_by_frame": {"0": points},368 "first_appears_frame": 0,369 }370 with open(output_dir / "segmentation_info.json", 'w') as f:371 json.dump(seg_info, f, indent=2)372 373 progress(0.3, desc="Stage 1 complete.")374 375 # ── Stage 2: VLM Analysis (repo code) ───────────────────────────────────376 analysis = None377 if api_key:378 progress(0.35, desc="Stage 2: VLM analysis (calling Gemini)...")379 try:380 video_info = {381 "video_path": str(input_video_path),382 "instruction": instruction,383 "output_dir": str(output_dir),384 "multi_frame_grids": True,385 }386 client = openai.OpenAI(387 api_key=api_key,388 base_url="https://generativelanguage.googleapis.com/v1beta/openai/",389 )390 analysis = vlm_process_video(video_info, client, DEFAULT_VLM_MODEL)391 progress(0.55, desc="Stage 2 complete.")392 except Exception as e:393 gr.Warning(f"VLM analysis failed: {e}. Generating binary mask only.")394 analysis = None395 else:396 gr.Warning("No GEMINI_API_KEY set. Generating binary mask only (no VLM analysis).")397 398 # ── Stage 3: Grey Mask Generation (repo logic + transformers SAM3) ──────399 grey_mask_path = output_dir / "grey_mask.mp4"400 vlm_analysis_path = output_dir / "vlm_analysis.json"401 402 if analysis and vlm_analysis_path.exists():403 progress(0.6, desc="Stage 3: Generating grey masks (SAM3 segmentation)...")404 try:405 video_info_3 = {406 "video_path": str(input_video_path),407 "output_dir": str(output_dir),408 "min_grid": 8,409 }410 # Uses the repo's process_video_grey_masks with our TransformersSam3Segmenter411 process_video_grey_masks(video_info_3, seg_model)412 progress(0.8, desc="Stage 3 complete.")413 except Exception as e:414 gr.Warning(f"Stage 3 failed: {e}. Generating binary mask only.")415 416 # ── Stage 4: Combine into Quadmask (repo code) ─────────────────────────417 quadmask_path = output_dir / "quadmask_0.mp4"418 if grey_mask_path.exists():419 progress(0.85, desc="Stage 4: Combining into quadmask...")420 combine_process_video(black_mask_path, grey_mask_path, quadmask_path)421 else:422 shutil.copy2(black_mask_path, quadmask_path)423 424 progress(0.9, desc="Creating visualizations...")425 426 # ── Visualization outputs ───────────────────────────────────────────────427 overlay_path = create_quadmask_visualization(str(input_video_path), str(quadmask_path))428 color_path = create_quadmask_color_video(str(quadmask_path))429 430 analysis_text = ""431 if vlm_analysis_path.exists():432 with open(vlm_analysis_path) as f:433 analysis_text = f.read()434 else:435 analysis_text = "No VLM analysis available."436 437 progress(1.0, desc="Done!")438 return str(quadmask_path), overlay_path, color_path, analysis_text439 440 441# ══════════════════════════════════════════════════════════════════════════════442# GRADIO UI443# ══════════════════════════════════════════════════════════════════════════════444 445def on_video_upload(video_path):446 if not video_path:447 return None, None, "[]", gr.update(interactive=False)448 frames, _ = extract_frames(video_path, max_frames=1)449 if not frames:450 return None, None, "[]", gr.update(interactive=False)451 return frames[0], frames[0], "[]", gr.update(interactive=True)452 453 454def on_frame_select(clean_frame, points_json, evt: gr.SelectData):455 if clean_frame is None:456 return None, points_json457 points = json.loads(points_json) if points_json else []458 x, y = evt.index459 points.append([int(x), int(y)])460 annotated = draw_points_on_image(clean_frame, points)461 return annotated, json.dumps(points)462 463 464def on_clear_points(clean_frame):465 if clean_frame is not None:466 return clean_frame, "[]"467 return None, "[]"468 469 470DESCRIPTION = """471# VOID VLM-Mask-Reasoner — Quadmask Generation472 473Generate **4-level semantic masks** (quadmasks) for interaction-aware video inpainting with [VOID](https://github.com/Netflix/void-model).474 475**Pipeline:** Click points on object → SAM2 segments it → Gemini VLM reasons about interactions → SAM3 segments affected objects → Quadmask generated476 477Use the generated quadmask with the [VOID inpainting demo](https://huggingface.co/spaces/sam-motamed/VOID).478"""479 480QUADMASK_EXPLAINER = """481### Quadmask format482 483| Pixel Value | Color | Meaning |484|-------------|-------|---------|485| **0** (black) | Red overlay | Primary object to remove |486| **63** (dark grey) | Yellow overlay | Overlap of primary + affected zone |487| **127** (mid grey) | Green overlay | Affected region (shadows, reflections, physics) |488| **255** (white) | Original | Background — keep as-is |489"""490 491with gr.Blocks(title="VOID VLM-Mask-Reasoner", theme=gr.themes.Default()) as demo:492 gr.Markdown(DESCRIPTION)493 494 points_state = gr.Textbox(value="[]", visible=False, label="points_json_api", elem_id="points_json_api")495 clean_frame_state = gr.State(None)496 497 with gr.Row():498 with gr.Column(scale=1):499 video_input = gr.Video(label="Upload Video", sources=["upload"])500 frame_display = gr.Image(501 label="Click to select primary object points (click multiple spots on the object)",502 interactive=True, type="numpy",503 )504 with gr.Row():505 clear_btn = gr.Button("Clear Points", size="sm")506 points_display = gr.Textbox(label="Selected Points", value="[]",507 interactive=False, max_lines=2)508 instruction_input = gr.Textbox(509 label="Edit instruction — describe what to remove",510 placeholder="e.g., remove the person", lines=1,511 )512 generate_btn = gr.Button("Generate Quadmask", variant="primary", size="lg")513 514 with gr.Column(scale=1):515 output_quadmask_file = gr.File(label="Download lossless quadmask_0.mp4 (use this with VOID)")516 with gr.Tabs():517 with gr.TabItem("Quadmask Overlay"):518 output_overlay = gr.Video(label="Quadmask overlay on original video")519 with gr.TabItem("Raw Quadmask"):520 output_color = gr.Video(label="Color-coded quadmask")521 with gr.TabItem("VLM Analysis"):522 output_analysis = gr.Code(label="VLM Analysis JSON", language="json")523 524 video_input.change(525 fn=on_video_upload, inputs=[video_input],526 outputs=[frame_display, clean_frame_state, points_state, generate_btn],527 )528 points_state.change(lambda p: p, inputs=points_state, outputs=points_display)529 frame_display.select(530 fn=on_frame_select, inputs=[clean_frame_state, points_state],531 outputs=[frame_display, points_state],532 )533 clear_btn.click(534 fn=on_clear_points, inputs=[clean_frame_state],535 outputs=[frame_display, points_state],536 )537 generate_btn.click(538 fn=run_pipeline, inputs=[video_input, points_state, instruction_input],539 outputs=[output_quadmask_file, output_overlay, output_color, output_analysis],540 )541 542 gr.Markdown(QUADMASK_EXPLAINER)543 544if __name__ == "__main__":545 demo.launch()546