Update app.py

app.py

@@ -29,6 +29,8 @@ from PIL import Image
 import matplotlib.pyplot as plt
 
 from mono.utils.unproj_pcd import reconstruct_pcd, save_point_cloud
+from mono.utils.transform import gray_to_colormap
+from mono.utils.visualization import vis_surface_normal
 import gradio as gr
 
 cfg = Config.fromfile('./mono/configs/HourglassDecoder/vit.raft5.small.py')
@@ -64,10 +66,19 @@ def depth_normal(img):
         pred_normal = output['normal_out_list'][0][:, :3, :, :] 
         H, W = pred_normal.shape[2:]
         pred_normal[:, :, pad[0]:H-pad[1], pad[2]:W-pad[3]]
-        
-    output = pred_depth.cpu().numpy()
-    formatted = (output * 255 / np.max(output)).astype('uint8')
-    img = Image.fromarray(formatted)
+
+    pred_depth = pred_depth.squeeze().cpu().numpy()
+    pred_depth[pred_depth<0] = 0
+    pred_color = gray_to_colormap(pred_depth)
+
+    pred_normal = pred_normal.squeeze()
+    if pred_normal.size(0) == 3
+        pred_normal == pred_normal.permute(1,2,0)
+    pred_color_normal = vis_surface_normal(pred_normal)
+    
+    ##formatted = (output * 255 / np.max(output)).astype('uint8')
+    img = Image.fromarray(pred_color)
+    img_normal = Image.fromarray(pred_color_normal)
     return img
         
 #inputs =  gr.inputs.Image(type='pil', label="Original Image")