Update app.py

app.py

@@ -8,13 +8,14 @@ import shutil
 import time
 import gradio as gr
 import sys
-
+from PIL import Image
 
 # Set environment variables
 os.environ["PIXEL3DMM_CODE_BASE"] = f"{os.getcwd()}"
 os.environ["PIXEL3DMM_PREPROCESSED_DATA"] = f"{os.getcwd()}/proprocess_results"
 os.environ["PIXEL3DMM_TRACKING_OUTPUT"] = f"{os.getcwd()}/tracking_results"
 
+
 def sh(cmd): subprocess.check_call(cmd, shell=True)
 
 # only do this once per VM restart
@@ -44,224 +45,149 @@ def install_cuda_toolkit():
 install_cuda_toolkit()
 
 
-# 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"))
+# Utility to select first image from a folder
+def first_image_from_dir(directory):
+    patterns = ["*.jpg", "*.png", "*.jpeg"]
+    files = []
+    for p in patterns:
+        files.extend(glob.glob(os.path.join(directory, p)))
     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)
+    return sorted(files)[0]
 
-    # 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: Preprocess the input image (Save and Crop)
+@spaces.GPU()
+def preprocess_image(image_array, state):
+    # Check if an image was uploaded
+    if image_array is None:
+        return "❌ Please upload an image first.", None, state
 
-# Step 1: Trim video based on user-defined duration and fps based on user-defined duration and fps
-def step1_trim(video_path, seconds, fps, state):
+    # Step 1a: Save the uploaded image
     session_id = str(uuid.uuid4())
     base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id)
+    os.makedirs(base_dir, exist_ok=True)
     state.update({"session_id": session_id, "base_dir": base_dir})
 
-    tmp = tempfile.mkdtemp()
-    trimmed = os.path.join(tmp, f"{session_id}.mp4")
+    img = Image.fromarray(image_array)
+    saved_image_path = os.path.join(base_dir, f"{session_id}.png")
+    img.save(saved_image_path)
+    state["image_path"] = saved_image_path
     
+    # Step 1b: Run the preprocessing script
     try:
-        # capture both stdout & stderr
         p = 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)
-
-        all_output = []
-        
-        for line in p.stdout:
-            print(line, end="")   # real-time echo
-            all_output.append(line)
-        
+            "python", "scripts/run_preprocessing.py",
+            "--video_or_images_path", saved_image_path
+        ], check=True, capture_output=True, text=True)
     except subprocess.CalledProcessError as e:
-        # e.stdout contains everything
-        err = f"❌ Preprocess failed (exit {e.returncode}).\n\n{e.stdout}"
-        return err, None, state
-
-
-    state["trimmed_path"] = trimmed
-    return f"✅ Step 1: Trimmed to {seconds}s @{fps}fps", state
-
-# Step 2: Preprocessing → cropped video
-@spaces.GPU()
-def step2_preprocess(state):
-    session_id = state["session_id"]
-    base_dir = state["base_dir"]
-    trimmed = state["trimmed_path"]
-
-    try:
-        # capture both stdout & stderr
-        p =     subprocess.run([
-        "python", "scripts/run_preprocessing.py",
-        "--video_or_images_path", trimmed
-    ], check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-    except subprocess.CalledProcessError as e:
-        # e.stdout contains everything
-        err = f"❌ Preprocess failed (exit {e.returncode}).\n\n{e.stdout}"
+        err = f"❌ Preprocess failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
+        # Clean up created directory on failure
+        shutil.rmtree(base_dir)
         return err, None, state
 
     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
+    image = first_image_from_dir(crop_dir)
+    return "✅ Preprocessing complete", image, state
 
-# Step 3: Normals inference → normals video
+
+# Step 2: Normals inference → normals image
 @spaces.GPU()
-def step3_normals(state):
-    session_id = state["session_id"]
-    base_dir = state["base_dir"]
+def step2_normals(state):
+    session_id = state.get("session_id")
+    if not session_id:
+        return "❌ Please preprocess an image first.", None, state
 
     try:
-        # capture both stdout & stderr
+        # Execute the network inference for normals
         p = subprocess.run([
-        "python", "scripts/network_inference.py",
-        "model.prediction_type=normals", f"video_name={session_id}"
-    ], check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+            "python", "scripts/network_inference.py",
+            "model.prediction_type=normals", f"video_name={session_id}"
+        ], check=True, capture_output=True, text=True)
     except subprocess.CalledProcessError as e:
-        # e.stdout contains everything
-        err = f"❌ Normal map failed (exit {e.returncode}).\n\n{e.stdout}"
+        err = f"❌ Normal map failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
         return err, None, state
-        
-    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
+    normals_dir = os.path.join(state["base_dir"], "p3dmm", "normals")
+    image = first_image_from_dir(normals_dir)
+    return "✅ Step 2: Normals inference complete", image, state
+
+# Step 3: UV map inference → uv map image
 @spaces.GPU()
-def step4_uv_map(state):
-    session_id = state["session_id"]
-    base_dir = state["base_dir"]
+def step3_uv_map(state):
+    session_id = state.get("session_id")
+    if not session_id:
+        return "❌ Please preprocess an image first.", None, state
 
     try:
-        # capture both stdout & stderr
+        # Execute the network inference for UV map
         p = 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)
+            "python", "scripts/network_inference.py",
+            "model.prediction_type=uv_map", f"video_name={session_id}"
+        ], check=True, capture_output=True, text=True)
     except subprocess.CalledProcessError as e:
-        # e.stdout contains everything
-        err = f"❌ UV map failed (exit {e.returncode}).\n\n{e.stdout}"
+        err = f"❌ UV map failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
         return err, None, state
 
-    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
+    uv_dir = os.path.join(state["base_dir"], "p3dmm", "uv_map")
+    image = first_image_from_dir(uv_dir)
+    return "✅ Step 3: UV map inference complete", image, state
 
-# Step 5: Tracking → final tracking video
+# Step 4: Tracking → final tracking image
 @spaces.GPU()
-def step5_track(state):
-    session_id = state["session_id"]
+def step4_track(state):
+    session_id = state.get("session_id")
+    if not session_id:
+        return "❌ Please preprocess an image first.", None, state
+
     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)
+        # Execute the tracking script
+        p = subprocess.run([
+            "python", script,
+            f"video_name={session_id}"
+        ], check=True, capture_output=True, text=True)
     except subprocess.CalledProcessError as e:
-        # e.stdout contains everything
-        err = f"❌ Tracking failed (exit {e.returncode}).\n\n{e.stdout}"
+        err = f"❌ Tracking failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
         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
+    image = first_image_from_dir(tracking_dir)
+    return "✅ Step 4: Tracking complete", image, state
 
 # Build Gradio UI
 demo = gr.Blocks()
 
 with demo:
-    gr.Markdown("## Video Processing Pipeline")
+    gr.Markdown("## Image 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({})
+            image_in = gr.Image(label="Upload Image", type="numpy", height=512)
+            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)
+                crop_img = gr.Image(label="Preprocessed", height=256)
+                normals_img = gr.Image(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_1 = gr.Button("Run Pipeline 1")
-    run_btn_2 = gr.Button("Run Pipeline 2")
-    run_btn_3 = gr.Button("Run Pipeline 3")
-    run_btn_4 = gr.Button("Run Pipeline 4")
-    run_btn_5 = gr.Button("Run Pipeline 5")
+                uv_img = gr.Image(label="UV Map", height=256)
+                track_img = gr.Image(label="Tracking", height=256)
 
-    # Update sliders after video upload
-    video_in.change(fn=get_video_info, inputs=video_in, outputs=[seconds_slider, fps_slider])
+    with gr.Row():
+        preprocess_btn = gr.Button("Step 1: Preprocess")
+        normals_btn = gr.Button("Step 2: Normals")
+        uv_map_btn = gr.Button("Step 3: UV Map")
+        track_btn = gr.Button("Step 4: Track")
 
     # Pipeline execution
-    run_btn_1.click(fn=step1_trim, inputs=[video_in, seconds_slider, fps_slider, state], outputs=[status, state])
-    run_btn_2.click(fn=step2_preprocess, inputs=[state], outputs=[status, crop_vid, state])
-    run_btn_3.click(fn=step3_normals, inputs=[state], outputs=[status, normals_vid, state])
-    run_btn_4.click(fn=step4_uv_map, inputs=[state], outputs=[status, uv_vid, state])
-    run_btn_5.click(fn=step5_track, inputs=[state], outputs=[status, track_vid, 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])
-    
+    preprocess_btn.click(fn=preprocess_image, inputs=[image_in, state], outputs=[status, crop_img, state])
+    normals_btn.click(fn=step2_normals, inputs=[state], outputs=[status, normals_img, state])
+    uv_map_btn.click(fn=step3_uv_map, inputs=[state], outputs=[status, uv_img, state])
+    track_btn.click(fn=step4_track, inputs=[state], outputs=[status, track_img, state])
+
 
 # ------------------------------------------------------------------
 # START THE GRADIO SERVER
 # ------------------------------------------------------------------
 demo.queue()
-
 demo.launch(share=True, ssr_mode=False)