import spaces import os import subprocess import tempfile import uuid import glob import shutil import time import gradio as gr # Set environment variables os.environ["PIXEL3DMM_CODE_BASE"] = "./" os.environ["PIXEL3DMM_PREPROCESSED_DATA"] = "./proprocess_results" os.environ["PIXEL3DMM_TRACKING_OUTPUT"] = "./tracking_results" # Utility to stitch frames into a video def make_video_from_frames(frames_dir, out_path, fps=15): if not os.path.isdir(frames_dir): return None files = glob.glob(os.path.join(frames_dir, "*.jpg")) + glob.glob(os.path.join(frames_dir, "*.png")) if not files: return None ext = files[0].split('.')[-1] pattern = os.path.join(frames_dir, f"%05d.{ext}") subprocess.run([ "ffmpeg", "-y", "-i", pattern, "-r", str(fps), out_path ], check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) return out_path # Function to probe video for duration and frame rate def get_video_info(video_path): """ Probes the uploaded video and returns updated slider configs: - seconds slider: max = int(duration) - fps slider: max = int(orig_fps) """ if not video_path: # Return default slider updates when no video is uploaded return gr.update(maximum=10, value=3, step=1), gr.update(maximum=30, value=15, step=1) # Use ffprobe to get JSON metadata cmd = [ "ffprobe", "-v", "quiet", "-print_format", "json", "-show_streams", video_path ] res = subprocess.run(cmd, capture_output=True, text=True) try: import json data = json.loads(res.stdout) stream = next(s for s in data.get('streams', []) if s.get('codec_type') == 'video') duration = float(stream.get('duration') or data.get('format', {}).get('duration', 0)) fr = stream.get('r_frame_rate', '0/1') num, den = fr.split('/') orig_fps = float(num) / float(den) if float(den) else 30 except Exception: duration, orig_fps = 10, 30 # Configure sliders based on actual video properties seconds_cfg = gr.update(maximum=int(duration), value=min(int(duration), 3), step=1) fps_cfg = gr.update(maximum=int(orig_fps), value=min(int(orig_fps), 15), step=1) return seconds_cfg, fps_cfg # Step 1: Trim video based on user-defined duration and fps based on user-defined duration and fps @space.GPU() def step1_trim(video_path, seconds, fps, state): session_id = str(uuid.uuid4()) base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id) state.update({"session_id": session_id, "base_dir": base_dir}) tmp = tempfile.mkdtemp() trimmed = os.path.join(tmp, f"{session_id}.mp4") subprocess.run([ "ffmpeg", "-y", "-i", video_path, "-t", str(seconds), # user-specified duration "-r", str(fps), # user-specified fps trimmed ], check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) state["trimmed_path"] = trimmed return f"✅ Step 1: Trimmed to {seconds}s @{fps}fps", state # Step 2: Preprocessing → cropped video @space.GPU() def step2_preprocess(state): session_id = state["session_id"] base_dir = state["base_dir"] trimmed = state["trimmed_path"] subprocess.run([ "python", "scripts/run_preprocessing.py", "--video_or_images_path", trimmed ], check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) crop_dir = os.path.join(base_dir, "cropped") out = os.path.join(os.path.dirname(trimmed), f"crop_{session_id}.mp4") video = make_video_from_frames(crop_dir, out) return "✅ Step 2: Preprocessing complete", video, state # Step 3: Normals inference → normals video @space.GPU() def step3_normals(state): session_id = state["session_id"] base_dir = state["base_dir"] subprocess.run([ "python", "scripts/network_inference.py", "model.prediction_type=normals", f"video_name={session_id}" ], check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) normals_dir = os.path.join(base_dir, "p3dmm", "normals") out = os.path.join(os.path.dirname(state["trimmed_path"]), f"normals_{session_id}.mp4") video = make_video_from_frames(normals_dir, out) return "✅ Step 3: Normals inference complete", video, state # Step 4: UV map inference → uv map video @space.GPU() def step4_uv_map(state): session_id = state["session_id"] base_dir = state["base_dir"] subprocess.run([ "python", "scripts/network_inference.py", "model.prediction_type=uv_map", f"video_name={session_id}" ], check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) uv_dir = os.path.join(base_dir, "p3dmm", "uv_map") out = os.path.join(os.path.dirname(state["trimmed_path"]), f"uv_map_{session_id}.mp4") video = make_video_from_frames(uv_dir, out) return "✅ Step 4: UV map inference complete", video, state # Step 5: Tracking → final tracking video @space.GPU() def step5_track(state): session_id = state["session_id"] script = os.path.join(os.environ["PIXEL3DMM_CODE_BASE"], "scripts", "track.py") cmd = [ "python", script, f"video_name={session_id}" ] try: # capture both stdout & stderr p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, check=True) except subprocess.CalledProcessError as e: # e.stdout contains everything err = f"❌ Tracking failed (exit {e.returncode}).\n\n{e.stdout}" return err, None, state # if we get here, it succeeded: tracking_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], session_id, "frames") out = os.path.join(os.path.dirname(state["trimmed_path"]), f"result_{session_id}.mp4") video = make_video_from_frames(tracking_dir, out) return "✅ Step 5: Tracking complete", video, state # Build Gradio UI demo = gr.Blocks() with demo: gr.Markdown("## Video Processing Pipeline") with gr.Row(): with gr.Column(): video_in = gr.Video(label="Upload video", height=512) # Sliders for duration and fps seconds_slider = gr.Slider(label="Duration (seconds)", minimum=2, maximum=10, step=1, value=3) fps_slider = gr.Slider(label="Frame Rate (fps)", minimum=15, maximum=30, step=1, value=15) status = gr.Textbox(label="Status", lines=2, interactive=False) state = gr.State({}) with gr.Column(): with gr.Row(): crop_vid = gr.Video(label="Preprocessed", height=256) normals_vid = gr.Video(label="Normals", height=256) with gr.Row(): uv_vid = gr.Video(label="UV Map", height=256) track_vid = gr.Video(label="Tracking", height=256) run_btn = gr.Button("Run Pipeline") # Update sliders after video upload video_in.change(fn=get_video_info, inputs=video_in, outputs=[seconds_slider, fps_slider]) # Pipeline execution (run_btn.click(fn=step1_trim, inputs=[video_in, seconds_slider, fps_slider, state], outputs=[status, state]) .then(fn=step2_preprocess, inputs=[state], outputs=[status, crop_vid, state]) .then(fn=step3_normals, inputs=[state], outputs=[status, normals_vid, state]) .then(fn=step4_uv_map, inputs=[state], outputs=[status, uv_vid, state]) .then(fn=step5_track, inputs=[state], outputs=[status, track_vid, state]) ) # ------------------------------------------------------------------ # START THE GRADIO SERVER # ------------------------------------------------------------------ demo.queue() demo.launch(share=True)