Update app.py

app.py

@@ -6,6 +6,7 @@ import gradio as gr
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
 import torch.nn as nn
+from scipy.interpolate import interp2d
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -55,19 +56,38 @@ def preprocess_image(image):
     return image / 255.0
 
 def plot_depth_map(depth_map, original_image):
-    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]))
+    fig = plt.figure(figsize=(32, 9))
     
+    # 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))
+    
+    # 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))
+    
+    # Interpolate colors
     original_image_resized = cv2.resize(original_image, (depth_map.shape[1], depth_map.shape[0]))
-    colors = original_image_resized.reshape(depth_map.shape[0], depth_map.shape[1], 3) / 255.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))
     
-    ax.plot_surface(x, y, depth_map, facecolors=colors, shade=False)
-    ax.set_zlim(0, 1)
+    # 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')
     
-    ax.view_init(elev=150, azim=90)
+    # 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")
     plt.axis('off')
     
+    plt.tight_layout()
     plt.show()
     
     fig.canvas.draw()