Update app.py

app.py

@@ -51,44 +51,26 @@ model.eval()
 processor = DPTImageProcessor.from_pretrained("Intel/dpt-swinv2-tiny-256")
 
 def preprocess_image(image):
-    image = cv2.resize(image, (128, 72))
+    image = cv2.resize(image, (255, 143))
     image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float().to(device)
     return image / 255.0
 
 def plot_depth_map(depth_map, original_image):
-    fig = plt.figure(figsize=(32, 9))
+    fig = plt.figure(figsize=(16, 9))
+    ax = fig.add_subplot(111, projection='3d')
+    x, y = np.meshgrid(range(depth_map.shape[1]), range(depth_map.shape[0]))
     
-    # Increase resolution of the meshgrid
-    x, y = np.meshgrid(np.linspace(0, depth_map.shape[1]-1, 255), np.linspace(0, depth_map.shape[0]-1, 255))
+    # Normalize depth map to [0, 1] range
+    norm = plt.Normalize(depth_map.min(), depth_map.max())
+    colors = plt.cm.viridis(norm(depth_map))
     
-    # Interpolate depth map
-    depth_interp = interp2d(np.arange(depth_map.shape[1]), np.arange(depth_map.shape[0]), depth_map)
-    z = depth_interp(np.linspace(0, depth_map.shape[1]-1, 255), np.linspace(0, depth_map.shape[0]-1, 255))
+    ax.plot_surface(x, y, depth_map, facecolors=colors, shade=False)
+    ax.set_zlim(0, 1)
     
-    # Interpolate colors
-    original_image_resized = cv2.resize(original_image, (depth_map.shape[1], depth_map.shape[0]))
-    colors = original_image_resized.reshape(-1, original_image_resized.shape[1], 3) / 255.0
-    colors_interp = interp2d(np.arange(colors.shape[1]), np.arange(colors.shape[0]), colors.reshape(-1, colors.shape[1]), kind='linear')
-    new_colors = colors_interp(np.linspace(0, colors.shape[1]-1, 255), np.linspace(0, colors.shape[0]-1, 255))
-    
-    # Plot with depth map color
-    ax1 = fig.add_subplot(121, projection='3d')
-    surf1 = ax1.plot_surface(x, y, z, facecolors=plt.cm.viridis(z), shade=False)
-    ax1.set_zlim(0, 1)
-    ax1.view_init(elev=150, azim=90)
-    ax1.set_title("Depth Map Color")
-    plt.axis('off')
-    
-    # Plot with RGB color
-    ax2 = fig.add_subplot(122, projection='3d')
-    surf2 = ax2.plot_surface(x, y, z, facecolors=new_colors, shade=False)
-    ax2.set_zlim(0, 1)
-    ax2.view_init(elev=150, azim=90)
-    ax2.set_title("RGB Color")
+    # Adjust the view to look down at an angle from a higher position
+    ax.view_init(elev=70, azim=90) 
     plt.axis('off')
-    
-    plt.tight_layout()
-    plt.show()
+    plt.close(fig)
     
     fig.canvas.draw()
     img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)