Update app_recon.py

app_recon.py

@@ -75,6 +75,9 @@ pipe = DepthNormalEstimationPipeline(vae=vae,
                             unet=unet,
                             scheduler=scheduler)
 
+
+outputs_dir = "./outs"
+
 try:
     import xformers
     pipe.enable_xformers_memory_efficient_attention()
@@ -159,6 +162,7 @@ def depth_normal(img_path,
     normal_np = pipe_out.normal_np
 
     path_output_dir = os.path.splitext(os.path.basename(img_path))[0] + datetime.now().strftime('%Y%m%d-%H%M%S')
+    path_output_dir = os.path.join(path_output_dir, outputs_dir)
     os.makedirs(path_output_dir, exist_ok=True)
 
     name_base = os.path.splitext(os.path.basename(img_path))[0]
@@ -199,7 +203,10 @@ def reconstruction(image_file, files):
 
     masked_image, mask = seg_foreground(image_file)
 
-    mask = mask[:, :, 0] > 0.5
+    # import pdb
+    # pdb.set_trace()
+
+    mask = np.array(mask) > 0.5
     depth_np = np.load(files[0])
     normal_np = np.load(files[1])
 
@@ -211,7 +218,7 @@ def reconstruction(image_file, files):
 
     normal_np[:, :, 0] *= -1
     _, 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")
+    ply_path = os.path.join(outputs_dir, dir_name, f"{name_base}_recon.ply")
     surface.save(ply_path, binary=False)
 
     obj_path = ply_path.replace('ply', 'obj')
@@ -224,6 +231,7 @@ def reconstruction(image_file, files):
     
     return obj_path, [ply_path], masked_image
 
+# @spaces.GPU
 def run_demo():
 
 
@@ -342,7 +350,7 @@ def run_demo():
                         outputs=[depth, normal, files]
                         )
         recon_btn.click(fn=reconstruction,
-                        inputs=[input_image, mask, files],
+                        inputs=[input_image, files],
                         outputs=[reconstructed_3d, reconstructed_file, masked_image]
                         )
         demo.queue().launch(share=True, max_threads=80)