Update get_depth_map

lucataco · Sep 12, 2023 · 6d137da · 6d137da
1 parent 6c93b5b
commit 6d137da
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 1 deletion.
diff --git a/output.png b/output.png
diff --git a/predict.py b/predict.py
@@ -59,9 +59,11 @@ def get_depth_map(self, image):
         with torch.no_grad(), torch.autocast("cuda"):
             depth_map = self.depth_estimator(image).predicted_depth
 
+        height, width = image.shape[2], image.shape[3]
+
         depth_map = torch.nn.functional.interpolate(
             depth_map.unsqueeze(1),
-            size=(1024, 1024),
+            size=(height, width),
             mode="bicubic",
             align_corners=False,
         )

diff --git a/salon.jpeg b/salon.jpeg