Update app.py

app.py

@@ -78,11 +78,11 @@ transforms = GOSNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 def predict(image):
     H,W = image.shape[:2]
     depth = DAMV2.infer_image(image)
-    image = cv2.resize(image, (1024,1024))
-    depth = cv2.resize(depth, (1024,1024))
-    image = torch.from_numpy(image).permute(2,0,1).float()/255
+    image = torch.nn.functional.interpolate(torch.from_numpy(image).permute(2,0,1)[None,...],size=[1024,1024],mode='bilinear',align_corners=True)[0]
+    depth = torch.nn.functional.interpolate(torch.from_numpy(depth)[None,None,...],size=[1024,1024],mode='bilinear',align_corners=True)
+    image = torch.divide(image,255.0)
+    depth = torch.divide(depth,255.0)
     image = transforms(image).unsqueeze(0)
-    depth = torch.from_numpy(depth).unsqueeze(0).unsqueeze(0)/255
     DIS_map = model.inference(image.to(device),depth.to(device))[0][0][0].cpu()
     DIS_map = cv2.resize(np.array(DIS_map), (W,H))
     # return cv2.resize(np.array(depth[0][0]), (W,H))