Update app.py

app.py

@@ -32,30 +32,30 @@ model = torch.quantization.quantize_dynamic(
 processor = DPTImageProcessor.from_pretrained("Intel/dpt-swinv2-tiny-256")
 
 color_map = cv2.applyColorMap(np.arange(256, dtype=np.uint8), cv2.COLORMAP_INFERNO)
-
-input_tensor = torch.zeros((1, 3, 128, 128), dtype=torch.float32, device=device)
+color_map_gpu = torch.from_numpy(color_map).to(device)
 
 def preprocess_image(image):
-    return cv2.resize(image, (128, 72), interpolation=cv2.INTER_AREA).transpose(2, 0, 1).astype(np.float32) / 255.0
+    resized = cv2.resize(image, (128, 72), interpolation=cv2.INTER_AREA)
+    return torch.from_numpy(resized.transpose(2, 0, 1)).float().div(255).unsqueeze(0).to(device)
 
 @torch.inference_mode()
 def process_frame(image):
     if image is None:
         return None
-    preprocessed = preprocess_image(image)
-    input_tensor = torch.from_numpy(preprocessed).unsqueeze(0).to(device)
+    input_tensor = preprocess_image(image)
     
     predicted_depth = model(input_tensor).predicted_depth
-    depth_map = predicted_depth.squeeze().cpu().numpy()
-    num_bins = 1000
-    depth_map = np.digitize(depth_map, bins=np.linspace(depth_map.min(), depth_map.max(), num_bins)) - 1
-    unique_values = np.unique(depth_map)
-    print(f"Number of unique depth values: {len(unique_values)}")
+    depth_map = predicted_depth.squeeze()
+    
+    # Rounding instead of discretization
+    depth_map = torch.round(depth_map * 100) / 100
+    
     depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
-    depth_map = (depth_map * 255).astype(np.uint8)
-    depth_map_colored = cv2.applyColorMap(depth_map, cv2.COLORMAP_INFERNO)
+    depth_map = (depth_map * 255).byte()
+    
+    depth_map_colored = color_map_gpu[depth_map]
     
-    return cv2.cvtColor(depth_map_colored, cv2.COLOR_BGR2RGB)
+    return cv2.cvtColor(depth_map_colored.cpu().numpy(), cv2.COLOR_BGR2RGB)
 
 interface = gr.Interface(
     fn=process_frame,