Update app.py

app.py

@@ -41,23 +41,24 @@ css = """
 
 set_seed(666)
 
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+default_steps = 10
+model = PixelPerfectDepth(sampling_steps=default_steps)
+ckpt_path = hf_hub_download(
+    repo_id="gangweix/Pixel-Perfect-Depth",
+    filename="ppd.pth",
+    repo_type="model"
+)
+state_dict = torch.load(ckpt_path, map_location="cpu")
+model.load_state_dict(state_dict, strict=False)
+model = model.to(DEVICE).eval()
+
+moge_model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal").to(DEVICE).eval()
+
+
 def main(share=True):
     print("Initializing Pixel-Perfect Depth Demo...")
-    
-    # DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-    default_steps = 10
-    model = PixelPerfectDepth(sampling_steps=default_steps)
-    ckpt_path = hf_hub_download(
-        repo_id="gangweix/Pixel-Perfect-Depth",
-        filename="ppd.pth",
-        repo_type="model"
-    )
-    state_dict = torch.load(ckpt_path, map_location="cpu")
-    model.load_state_dict(state_dict, strict=False)
-    model = model.to(DEVICE).eval()
-
-    moge_model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal").to(DEVICE).eval()
 
     cmap = matplotlib.colormaps.get_cmap('Spectral')
 
@@ -90,7 +91,7 @@ def main(share=True):
         moge_depth, mask, intrinsics = predict_moge_depth(resize_image)
 
         # relative depth -> metric depth
-        metric_depth = recover_metric_depth_ransac(ppd_depth.squeeze().cpu().numpy(), moge_depth, mask)
+        metric_depth = recover_metric_depth_ransac(ppd_depth, moge_depth, mask)
         intrinsics[0, 0] *= resize_W 
         intrinsics[1, 1] *= resize_H
         intrinsics[0, 2] *= resize_W
@@ -112,7 +113,7 @@ def main(share=True):
         tmp_raw_depth = tempfile.NamedTemporaryFile(suffix='.npy', delete=False)
         np.save(tmp_raw_depth.name, depth)
 
-        depth_vis = (depth - depth.min()) / (depth.max() - depth.min()+1e-5) * 255.0
+        depth_vis = (depth - depth.min()) / (depth.max() - depth.min() + 1e-5) * 255.0
         depth_vis = depth_vis.astype(np.uint8)
         colored_depth = (cmap(depth_vis)[:, :, :3] * 255).astype(np.uint8)