Add image normalization and variable conversion in app.py

app.py

@@ -6,7 +6,7 @@ from PIL import Image
 import spaces
 import torch
 import torch.nn.functional as F
-from torchvision.transforms import Compose
+from torchvision.transforms import Compose, Normalize
 import tempfile
 from gradio_imageslider import ImageSlider
 
@@ -74,9 +74,15 @@ with gr.Blocks(css=css) as demo:
 
         h, w = image.shape[:2]
 
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
-        image = transform({'image': image})['image']
-        image = torch.from_numpy(image).unsqueeze(0).to(DEVICE)
+        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
+        # image = transform({'image': image})['image']
+        # image = torch.from_numpy(image).unsqueeze(0).to(DEVICE)
+
+        image = torch.from_numpy(image.transpose((2, 0, 1))).to(DEVICE)
+        image = torch.from_numpy(image.transpose((2, 0, 1)))
+        image = Normalize(mean=[0.485, 0.456, 0.406], std=[
+            0.229, 0.224, 0.225])(image)
+        image = torch.autograd.Variable(image.unsqueeze(0).cuda())
 
         depth = predict_depth(model, image)
         depth = F.interpolate(depth[None], (h, w),