import sys import asyncio import tempfile import uuid import os import cv2 import numpy as np import torch import torch.nn.functional as F import imageio import gradio as gr # 1. ZeroGPU & Local CPU Fallback Logic try: import spaces USING_SPACES = True except ImportError: USING_SPACES = False def adaptive_gpu_decorator(func): """Applies ZeroGPU decorator if on Hugging Face, otherwise runs normally.""" if USING_SPACES: return spaces.GPU()(func) return func from megaflow.model import MegaFlow from megaflow.utils.basic import gridcloud2d from megaflow.utils.visualizer import Visualizer from megaflow.utils.flow_viz import flow_to_image # 2. Global Setup and Model Loading device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Loading pre-trained MegaFlow models onto {device}...") model_track = MegaFlow.from_pretrained("megaflow-track", device=device) model_track.eval() model_flow = MegaFlow.from_pretrained("megaflow-flow", device=device) model_flow.eval() # 3. Shared Utilities def calculate_dynamic_size(orig_h, orig_w, target_fix_width, patch_size=14, mode="track"): if mode == "track": new_w = target_fix_width new_h = round(orig_h * (new_w / orig_w) / patch_size) * patch_size else: if orig_w >= orig_h: new_w = target_fix_width new_h = round(orig_h * (new_w / orig_w) / patch_size) * patch_size else: new_h = target_fix_width new_w = round(orig_w * (new_h / orig_h) / patch_size) * patch_size return int(new_h), int(new_w) def get_video_frames(input_path, fix_width, mode="track"): cap = cv2.VideoCapture(input_path) if not cap.isOpened(): raise gr.Error(f"Cannot open video: {input_path}") native_fps = cap.get(cv2.CAP_PROP_FPS) if native_fps <= 0 or np.isnan(native_fps): native_fps = 24.0 frames, orig_shape = [], None while True: ret, frame = cap.read() if not ret: break frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) if orig_shape is None: orig_shape = frame.shape[:2] new_h, new_w = calculate_dynamic_size(orig_shape[0], orig_shape[1], fix_width, mode=mode) frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_LINEAR) frames.append(frame) cap.release() return frames, orig_shape, native_fps # 4. Tracking Wrapper @adaptive_gpu_decorator @torch.inference_mode() def run_tracking(video_in, grid_size, iters, restore_size): if not video_in: raise gr.Error("Please upload or select a video first.") frames_np, native_size, fps = get_video_frames(video_in, fix_width=518, mode="track") if len(frames_np) < 2: raise gr.Error("Video requires at least 2 frames.") input_image = [torch.from_numpy(f).permute(2, 0, 1).float() for f in frames_np] frames = torch.stack(input_image, dim=0)[None].to(device) B, T, _, H, W = frames.shape grid_xy = gridcloud2d(1, H, W, norm=False, device=device).float() grid_xy = grid_xy.permute(0, 2, 1).reshape(1, 1, 2, H, W) # Use float32 on CPU to avoid autocast errors, bfloat16/float16 on GPU compute_dtype = torch.bfloat16 if device == "cuda" and torch.cuda.is_bf16_supported() else (torch.float16 if device == "cuda" else torch.float32) with torch.autocast(device_type=device, dtype=compute_dtype, enabled=(device == "cuda")): results = model_track.forward_track(frames, num_reg_refine=iters) flows_e = results["flow_final"] traj_maps = flows_e.to(device) + grid_xy traj_sub = traj_maps[..., ::grid_size, ::grid_size] pred_tracks = traj_sub.flatten(3).permute(0, 1, 3, 2) if restore_size and native_size is not None: orig_H, orig_W = native_size pred_tracks[..., 0] *= (orig_W / W) pred_tracks[..., 1] *= (orig_H / H) frames = F.interpolate(frames[0], size=(orig_H, orig_W), mode='bilinear', align_corners=True).unsqueeze(0) output_dir = tempfile.mkdtemp() filename = f"track_{uuid.uuid4().hex}" vis = Visualizer(save_dir=output_dir, pad_value=0, linewidth=1, tracks_leave_trace=0, fps=fps) vis.visualize(frames, pred_tracks, filename=filename, opacity=0.5) return os.path.join(output_dir, f"{filename}.mp4") # 5. Flow Wrapper @adaptive_gpu_decorator @torch.inference_mode() def run_flow(video_in, window_size, iters, restore_size): if not video_in: raise gr.Error("Please upload or select a video first.") frames_np, native_size, fps = get_video_frames(video_in, fix_width=952, mode="flow") if len(frames_np) < 2: raise gr.Error("Video requires at least 2 frames.") input_image = [torch.from_numpy(f).permute(2, 0, 1).float() for f in frames_np] input_scene = torch.stack(input_image, dim=0)[None] B, T, C, H, W = input_scene.shape output_dir = tempfile.mkdtemp() output_path = os.path.join(output_dir, f"flow_{uuid.uuid4().hex}.mp4") video_writer = imageio.get_writer(output_path, fps=fps, codec='libx264', macro_block_size=None) infer_window = window_size for start in range(0, T - 1, infer_window - 1): end = min(start + infer_window, T) chunk = input_scene[:, start:end].to(device) compute_dtype = torch.bfloat16 if device == "cuda" and torch.cuda.is_bf16_supported() else (torch.float16 if device == "cuda" else torch.float32) with torch.autocast(device_type=device, dtype=compute_dtype, enabled=(device == "cuda")): results_dict = model_flow(chunk, num_reg_refine=iters) flow_pr = results_dict['flow_preds'][-1] if native_size is not None and restore_size: scaled_flow = F.interpolate(flow_pr.view(-1, 2, H, W), size=native_size, mode='bilinear', align_corners=True) scaled_flow[:, 0, :, :] *= (native_size[1] / W) scaled_flow[:, 1, :, :] *= (native_size[0] / H) flow_pr = scaled_flow.view(*flow_pr.shape[:2], 2, *native_size) flow_all = flow_pr[0].permute(0, 2, 3, 1).cpu().numpy() for t, flow in enumerate(flow_all): flow_vis_rgb = flow_to_image(flow, convert_to_bgr=False) video_writer.append_data(flow_vis_rgb) video_writer.close() return output_path # 6. UI Design with gr.Blocks(theme=gr.themes.Soft(), title="MegaFlow Demo") as demo: gr.HTML( """

MegaFlow: Zero-Shot Large Displacement Optical Flow

Dingxi Zhang1     Fangjinhua Wang1     Marc Pollefeys1,2     Haofei Xu1,3
1 ETH Zurich       2 Microsoft       3 University of Tübingen, Tübingen AI Center
Project Page arXiv Models
""" ) with gr.Row(): vid_input = gr.Video(label="Input Video", height=400) vid_output = gr.Video(label="Output Result", height=400, interactive=False) with gr.Tabs(): # Tab 1: Optical Flow with gr.TabItem("🌈 1. Optical Flow"): gr.Markdown("### Compute dense optical flow between consecutive frames\nConfigure your parameters below, then click **Run Optical Flow** to generate the sequence.") with gr.Row(): flow_window_size = gr.Slider(minimum=2, maximum=16, step=1, value=4, label="Sliding Window Size") flow_iters = gr.Slider(minimum=1, maximum=16, step=1, value=8, label="Refinement Iterations") flow_restore = gr.Checkbox(value=True, label="Restore Original Resolution") flow_btn = gr.Button("🚀 Run Optical Flow", variant="primary", size="lg") # Tab 2: Point Tracking with gr.TabItem("🎯 2. Point Tracking"): gr.Markdown("### Track specific points across the video sequence\nConfigure your parameters below, then click **Run Point Tracking** to track the grid.") gr.Markdown("*Note: As a zero-shot tracking application of our flow model, point visibility is not explicitly predicted, resulting in tracking through occlusions.*") with gr.Row(): track_grid_size = gr.Slider(minimum=2, maximum=16, step=2, value=8, label="Grid Size (Subsampling)") track_iters = gr.Slider(minimum=1, maximum=16, step=1, value=8, label="Refinement Iterations") track_restore = gr.Checkbox(value=True, label="Restore Original Resolution") track_btn = gr.Button("🚀 Run Point Tracking", variant="primary", size="lg") gr.Markdown("### Try an Example") example_videos = [ ["assets/chamaleon.mp4"], ["assets/apple.mp4"], ["assets/longboard.mp4"], ] gr.Examples( examples=example_videos, inputs=vid_input, label="Click any video to load it instantly into the Input Video player", examples_per_page=5 ) track_btn.click( fn=run_tracking, inputs=[vid_input, track_grid_size, track_iters, track_restore], outputs=vid_output ) flow_btn.click( fn=run_flow, inputs=[vid_input, flow_window_size, flow_iters, flow_restore], outputs=vid_output ) if __name__ == "__main__": demo.launch()