Add matplotlib.pyplot import and modify depth visualization code

app.py

@@ -9,6 +9,7 @@ import torch.nn.functional as F
 from torchvision.transforms import Compose, Normalize
 import tempfile
 from gradio_imageslider import ImageSlider
+import matplotlib.pyplot as plt
 
 from iebins.networks.NewCRFDepth import NewCRFDepth
 from iebins.util.transfrom import Resize, NormalizeImage, PrepareForNet
@@ -94,22 +95,18 @@ with gr.Blocks(css=css) as demo:
         # image = torch.from_numpy(image).unsqueeze(0)
         image = torch.autograd.Variable(image.unsqueeze(0))
 
-        pred_depths_r_list, _, _ = predict_depth(model, image)
+        pred_depths_r_list, _, _ = model(image)
         image_flipped = flip_lr(image)
         pred_depths_r_list_flipped, _, _ = model(image_flipped)
         pred_depth = post_process_depth(
             pred_depths_r_list[-1], pred_depths_r_list_flipped[-1])
 
         pred_depth = pred_depth.cpu().numpy().squeeze()
+        colored_depth = cv2.applyColorMap(
+            pred_depth, cv2.COLORMAP_INFERNO)[:, :, ::-1]
 
-        raw_depth = Image.fromarray(pred_depth.cpu().numpy().astype('uint16'))
         tmp = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
-        raw_depth.save(tmp.name)
-
-        depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
-        depth = depth.cpu().numpy().astype(np.uint8)
-        colored_depth = cv2.applyColorMap(
-            depth, cv2.COLORMAP_INFERNO)[:, :, ::-1]
+        pred_depth.save(tmp.name)
 
         return [(original_image, colored_depth), tmp.name]