Update app.py

app.py

@@ -4,7 +4,8 @@ import numpy as np
 from transformers import DPTForDepthEstimation, DPTImageProcessor
 import gradio as gr
 import torch.nn.utils.prune as prune
-from DepthVisualizer import DepthVisualizer
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -42,6 +43,15 @@ def preprocess_image(image):
     image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float().to(device)
     return image / 255.0
 
+def plot_depth_map(depth_map):
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection='3d')
+    x, y = np.meshgrid(range(depth_map.shape[1]), range(depth_map.shape[0]))
+    ax.plot_surface(x, y, depth_map, cmap='viridis')
+    ax.set_zlim(0, 1)
+    plt.close(fig)
+    return fig
+
 @torch.inference_mode()
 def process_frame(image):
     if image is None:
@@ -53,13 +63,15 @@ def process_frame(image):
     # Normalize depth map
     depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
     
-    # Convert depth map to point cloud
-    point_cloud = visualizer.depth_map_to_point_cloud(depth_map)
+    # Create 3D plot
+    fig = plot_depth_map(depth_map)
     
-    # Render point cloud
-    rendered_image = visualizer.render_frame(point_cloud)
+    # Convert plot to image
+    fig.canvas.draw()
+    img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
     
-    return rendered_image
+    return img
 
 interface = gr.Interface(
     fn=process_frame,