Create app.py

app.py

@@ -0,0 +1,195 @@
+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)
+