Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 from transformers import DPTForDepthEstimation, DPTImageProcessor
 import gradio as gr
 import torch.nn.utils.prune as prune
+from DepthVisualizer import DepthVisualizer
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -31,6 +32,11 @@ model = model.to(device)
 
 processor = DPTImageProcessor.from_pretrained("Intel/dpt-swinv2-tiny-256")
 
+color_map = cv2.applyColorMap(np.arange(256, dtype=np.uint8), cv2.COLORMAP_INFERNO)
+color_map = torch.from_numpy(color_map).to(device)
+
+visualizer = DepthVisualizer()
+
 def preprocess_image(image):
     image = cv2.resize(image, (128, 128))
     image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float().to(device)
@@ -47,14 +53,13 @@ def process_frame(image):
     # Normalize depth map
     depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
     
-    # Create a more visually informative depth map
-    depth_color = cv2.applyColorMap((depth_map * 255).astype(np.uint8), cv2.COLORMAP_INFERNO)
+    # Convert depth map to point cloud
+    point_cloud = visualizer.depth_map_to_point_cloud(depth_map)
     
-    # Blend original image with depth map for context
-    original_resized = cv2.resize(image, (128, 128))
-    blended = cv2.addWeighted(original_resized, 0.6, depth_color, 0.4, 0)
+    # Render point cloud
+    rendered_image = visualizer.render_frame(point_cloud)
     
-    return blended
+    return rendered_image
 
 interface = gr.Interface(
     fn=process_frame,
@@ -63,4 +68,4 @@ interface = gr.Interface(
     live=True
 )
 
-interface.launch()
+interface.launch()