Update app.py

app.py

@@ -33,7 +33,7 @@ processor = DPTImageProcessor.from_pretrained("Intel/dpt-swinv2-tiny-256")
 
 color_map = torch.from_numpy(cv2.applyColorMap(np.arange(256, dtype=np.uint8), cv2.COLORMAP_INFERNO)).to(device)
 
-input_tensor = torch.zeros((1, 3, 128, 128), dtype=torch.float32, device=device)
+input_tensor = torch.zeros((1, 3, 72, 128), dtype=torch.float32, device=device)
 
 def preprocess_image(image):
     return cv2.resize(image, (128, 72), interpolation=cv2.INTER_AREA).transpose(2, 0, 1).astype(np.float32) / 255.0
@@ -49,13 +49,10 @@ def process_frame(image):
     depth_map = predicted_depth.squeeze()
     
     # Discretization on GPU
-    num_bins = 1000
-    depth_min, depth_max = depth_map.min(), depth_map.max()
-    bins = torch.linspace(depth_min, depth_max, num_bins, device=device)
-    depth_map = torch.bucketize(depth_map, bins)
+    num_bins = 256  # Match the color_map size
+    depth_map = torch.quantile(depth_map, torch.linspace(0, 1, num_bins, device=device))
+    depth_map = torch.bucketize(depth_map.squeeze(), depth_map)
     
-    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
-    depth_map = (depth_map * 255).byte()
     depth_map_colored = color_map[depth_map]
     
     return cv2.cvtColor(depth_map_colored.cpu().numpy(), cv2.COLOR_BGR2RGB)