download the pie

app.py

@@ -218,28 +218,32 @@ def step4_track(state):
 # New: run all steps sequentially
 @spaces.GPU()
 def run_pipeline(image_array, state):
-    # Step 1: Preprocess
+    # Step 1
     status1, crop_img, state, _, _ = preprocess_image(image_array, state)
     if "❌" in status1:
-        return status1, None, None, None, None, {}
-    # Step 2: Normals
+        return status1, None, None, None, None, None, {}
+    # Step 2
     status2, normals_img, state, _, _ = step2_normals(state)
-    # Step 3: UV Map
+    # Step 3
     status3, uv_img, state, _, _ = step3_uv_map(state)
-    # Step 4: Tracking
+    # Step 4
     status4, track_img, state, _ = step4_track(state)
+    # Locate mesh (.ply)
+    mesh_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], state.get("session_id"), "mesh")
+    mesh_file = first_file_from_dir(mesh_dir, "ply")
+
     final_status = "\n".join([status1, status2, status3, status4])
-    return final_status, crop_img, normals_img, uv_img, track_img, state
+    return final_status, crop_img, normals_img, uv_img, track_img, mesh_file, state
 
 # Build Gradio UI
 demo = gr.Blocks()
 with demo:
     gr.Markdown("## Image Processing Pipeline (Single Button)")
-    gr.Markdown("Upload an image and click 'Run Pipeline' to execute all steps.")
+    gr.Markdown("Upload an image and click 'Run Pipeline' to execute all steps and download the mesh.")
     with gr.Row():
         with gr.Column():
             image_in = gr.Image(label="Upload Image", type="numpy", height=512)
-            status = gr.Textbox(label="Status", lines=4, interactive=True, value="Upload an image to start.")
+            status = gr.Textbox(label="Status", lines=6, interactive=True, value="Upload an image to start.")
             state = gr.State({})
         with gr.Column():
             with gr.Row():
@@ -248,16 +252,17 @@ with demo:
             with gr.Row():
                 uv_img = gr.Image(label="UV Map", height=256)
                 track_img = gr.Image(label="Tracking", height=256)
+            mesh_file = gr.File(label="Download Mesh (.ply)")
     run_btn = gr.Button("Run Pipeline")
 
-    # Single button click
     run_btn.click(
         fn=run_pipeline,
         inputs=[image_in, state],
-        outputs=[status, crop_img, normals_img, uv_img, track_img, state]
+        outputs=[status, crop_img, normals_img, uv_img, track_img, mesh_file, state]
     )
 
-    image_in.upload(fn=reset_all, inputs=None, outputs=[crop_img, normals_img, uv_img, track_img, status, state])
+    image_in.upload(fn=reset_all, inputs=None, outputs=[crop_img, normals_img, uv_img, track_img, mesh_file, status, state, run_btn])
+
 
 demo.queue()
-demo.launch(share=True, ssr_mode=False)
+demo.launch(share=True, ssr_mode=False)