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")