Update app_recon.py

app_recon.py

@@ -105,7 +105,6 @@ def sam_init():
 
 sam_predictor = sam_init()
 
-@spaces.GPU
 def sam_segment(predictor, input_image, *bbox_coords):
     bbox = np.array(bbox_coords)
     image = np.asarray(input_image)
@@ -194,10 +193,12 @@ def seg_foreground(image_file):
     return masked_image, mask
 
 @spaces.GPU
-def reconstruction(mask, files):
+def reconstruction(image_file, files):
 
     torch.cuda.empty_cache()   
 
+    masked_image, mask = seg_foreground(image_file)
+
     mask = mask[:, :, 0] > 0.5
     depth_np = np.load(files[0])
     normal_np = np.load(files[1])
@@ -221,7 +222,7 @@ def reconstruction(mask, files):
 
     torch.cuda.empty_cache()
     
-    return obj_path, [ply_path]
+    return obj_path, [ply_path], masked_image
 
 def run_demo():
 
@@ -340,10 +341,9 @@ def run_demo():
                                 domain],
                         outputs=[depth, normal, files]
                         )
-        recon_btn.click(fn=seg_foreground, inputs=[input_image], outputs=[masked_image, mask]
-                        ).success(fn=reconstruction,
-                        inputs=[mask, files],
-                        outputs=[reconstructed_3d, reconstructed_file]
+        recon_btn.click(fn=reconstruction,
+                        inputs=[input_image, mask, files],
+                        outputs=[reconstructed_3d, reconstructed_file, masked_image]
                         )
         demo.queue().launch(share=True, max_threads=80)