Update app.py

app.py

@@ -44,7 +44,6 @@ def install_cuda_toolkit():
 
 install_cuda_toolkit()
 
-
 # Utility to select first image from a folder
 def first_image_from_dir(directory):
     patterns = ["*.jpg", "*.png", "*.jpeg"]
@@ -55,14 +54,27 @@ def first_image_from_dir(directory):
         return None
     return sorted(files)[0]
 
+# Function to reset the UI and state
+def reset_all():
+    return (
+        None,  # crop_img
+        None,  # normals_img
+        None,  # uv_img
+        None,  # track_img
+        "Awaiting new image upload...",  # status
+        {},    # state
+        gr.update(interactive=True),   # preprocess_btn
+        gr.update(interactive=False),  # normals_btn
+        gr.update(interactive=False),  # uv_map_btn
+        gr.update(interactive=False)   # track_btn
+    )
+
 # 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
+        return "❌ Please upload an image first.", None, state, gr.update(interactive=True), gr.update(interactive=False)
 
-    # 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)
@@ -73,96 +85,87 @@ def preprocess_image(image_array, state):
     img.save(saved_image_path)
     state["image_path"] = saved_image_path
     
-    # Step 1b: Run the preprocessing script
     try:
         p = subprocess.run([
-            "python", "scripts/run_preprocessing.py",
-            "--video_or_images_path", saved_image_path
+            "python", "scripts/run_preprocessing.py", "--video_or_images_path", saved_image_path
         ], check=True, capture_output=True, text=True)
     except subprocess.CalledProcessError as e:
         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
+        return err, None, {}, gr.update(interactive=True), gr.update(interactive=False)
 
     crop_dir = os.path.join(base_dir, "cropped")
     image = first_image_from_dir(crop_dir)
-    return "✅ Preprocessing complete", image, state
-
+    return "✅ Step 1 complete. Ready for Normals.", image, state, gr.update(interactive=False), gr.update(interactive=True)
 
 # Step 2: Normals inference → normals image
 @spaces.GPU()
 def step2_normals(state):
     session_id = state.get("session_id")
     if not session_id:
-        return "❌ Please preprocess an image first.", None, state
+        return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=False), gr.update(interactive=False)
 
     try:
-        # Execute the network inference for normals
         p = subprocess.run([
-            "python", "scripts/network_inference.py",
-            "model.prediction_type=normals", f"video_name={session_id}"
+            "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:
         err = f"❌ Normal map failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
-        return err, None, state
+        return err, None, state, gr.update(interactive=True), gr.update(interactive=False)
 
     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
+    return "✅ Step 2 complete. Ready for UV Map.", image, state, gr.update(interactive=False), gr.update(interactive=True)
 
 # Step 3: UV map inference → uv map image
 @spaces.GPU()
 def step3_uv_map(state):
     session_id = state.get("session_id")
     if not session_id:
-        return "❌ Please preprocess an image first.", None, state
+        return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=False), gr.update(interactive=False)
 
     try:
-        # 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}"
+            "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:
         err = f"❌ UV map failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
-        return err, None, state
+        return err, None, state, gr.update(interactive=True), gr.update(interactive=False)
 
     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
+    return "✅ Step 3 complete. Ready for Tracking.", image, state, gr.update(interactive=False), gr.update(interactive=True)
 
 # Step 4: Tracking → final tracking image
 @spaces.GPU()
 def step4_track(state):
     session_id = state.get("session_id")
     if not session_id:
-        return "❌ Please preprocess an image first.", None, state
+        return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=False)
 
     script = os.path.join(os.environ["PIXEL3DMM_CODE_BASE"], "scripts", "track.py")
     try:
-        # Execute the tracking script
         p = subprocess.run([
-            "python", script,
-            f"video_name={session_id}"
+            "python", script, f"video_name={session_id}"
         ], check=True, capture_output=True, text=True)
     except subprocess.CalledProcessError as e:
         err = f"❌ Tracking failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
-        return err, None, state
+        return err, None, state, gr.update(interactive=True)
 
     tracking_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], session_id, "frames")
     image = first_image_from_dir(tracking_dir)
-    return "✅ Step 4: Tracking complete", image, state
+    return "✅ Pipeline complete!", image, state, gr.update(interactive=False)
 
 # Build Gradio UI
 demo = gr.Blocks()
 
 with demo:
     gr.Markdown("## Image Processing Pipeline")
+    gr.Markdown("Upload an image, then click the buttons in order. Uploading a new image will reset the process.")
     with gr.Row():
         with gr.Column():
             image_in = gr.Image(label="Upload Image", type="numpy", height=512)
-            status = gr.Textbox(label="Status", lines=2, interactive=False)
+            status = gr.Textbox(label="Status", lines=2, interactive=False, value="Upload an image to start.")
             state = gr.State({})
         with gr.Column():
             with gr.Row():
@@ -173,21 +176,41 @@ with demo:
                 track_img = gr.Image(label="Tracking", height=256)
 
     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
-    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])
+        preprocess_btn = gr.Button("Step 1: Preprocess", interactive=True)
+        normals_btn = gr.Button("Step 2: Normals", interactive=False)
+        uv_map_btn = gr.Button("Step 3: UV Map", interactive=False)
+        track_btn = gr.Button("Step 4: Track", interactive=False)
+
+    # Define component list for reset
+    outputs_for_reset = [crop_img, normals_img, uv_img, track_img, status, state, preprocess_btn, normals_btn, uv_map_btn, track_btn]
+
+    # Pipeline execution logic
+    preprocess_btn.click(
+        fn=preprocess_image,
+        inputs=[image_in, state],
+        outputs=[status, crop_img, state, preprocess_btn, normals_btn]
+    )
+    normals_btn.click(
+        fn=step2_normals,
+        inputs=[state],
+        outputs=[status, normals_img, state, normals_btn, uv_map_btn]
+    )
+    uv_map_btn.click(
+        fn=step3_uv_map,
+        inputs=[state],
+        outputs=[status, uv_img, state, uv_map_btn, track_btn]
+    )
+    track_btn.click(
+        fn=step4_track,
+        inputs=[state],
+        outputs=[status, track_img, state, track_btn]
+    )
 
+    # Event to reset everything when a new image is uploaded
+    image_in.upload(fn=reset_all, inputs=None, outputs=outputs_for_reset)
 
 # ------------------------------------------------------------------
 # START THE GRADIO SERVER
 # ------------------------------------------------------------------
 demo.queue()
-demo.launch(share=True, ssr_mode=False)
-
+demo.launch(share=True, ssr_mode=False)