Update app_recon.py

app_recon.py

@@ -35,7 +35,7 @@ from rembg import remove
 from segment_anything import sam_model_registry, SamPredictor
 from datetime import datetime
 import time
-
+import trimesh
 
 import sys
 sys.path.append("../")
@@ -181,7 +181,16 @@ def reconstruction(image, files):
     _, surface, _,  _, _ = bilateral_normal_integration_function(normal_np, mask_output_temp, k=2, K=None, max_iter=100, tol=1e-4, cg_max_iter=5000, cg_tol=1e-3)
     ply_path = os.path.join(dir_name, f"{name_base}_recon.ply")
     surface.save(ply_path, binary=False)
-    return ply_path, masked_image
+
+    obj_path = ply_path.replace('ply', 'obj')
+    mesh = trimesh.load(ply_path)
+    T2 = np.array([[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
+    mesh.apply_transform(T2)
+    mesh.export(obj_path)
+
+    torch.cuda.empty_cache()
+    
+    return obj_path, masked_image, [ply_path]
 
 def run_demo():
 
@@ -279,7 +288,7 @@ def run_demo():
 
         with gr.Row():
             with gr.Column():
-                masked_image = gr.Image(interactive=False, label="Masked foreground.")
+                masked_image = gr.Image(interactive=False, height=320, label="Masked foreground.")
             with gr.Column():
                 reconstructed_3d = gr.Model3D(
                     label = 'Bini post-processed 3D model', height=320, interactive=False,
@@ -290,7 +299,14 @@ def run_demo():
             #     interactive=False,
             # )
 
+        with gr.Row():
+            reconstructed_file = gr.Files(
+                label = "3D Mesh (plyfile)",
+                elem_id = "download",
+                interactive=False,
+            )
 
+        
         run_btn.click(fn=depth_normal, 
                         inputs=[input_image, denoising_steps,
                                 ensemble_size,
@@ -301,7 +317,7 @@ def run_demo():
                         )
         recon_btn.click(fn=reconstruction,
                         inputs=[input_image, files],
-                        outputs=[reconstructed_3d, masked_image]
+                        outputs=[reconstructed_3d, masked_image, reconstructed_file]
                         )
         demo.queue().launch(share=True, max_threads=80)