Image reshape

app.py

@@ -11,7 +11,8 @@ from misc import load_config
 from torchvision import transforms as T
 
 import gradio as gr
-    
+
+MAX_IM_SIZE = 512
 NORMALIZE = T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
 CACHE = True
 
@@ -48,7 +49,14 @@ def predict(img_input):
     img_pil = Image.open(img_input)
     img = img_pil.convert("RGB")
 
-    t = T.Compose([T.ToTensor(), NORMALIZE])
+    # Image transformations
+    transforms = [T.ToTensor()]
+    # Resize image if needed
+    if img.size[0] > MAX_IM_SIZE or img.size[1] > MAX_IM_SIZE:
+        transforms.append(T.Resize(max_size=MAX_IM_SIZE))
+    transforms.append(NORMALIZE)
+    t = T.Compose(transforms)
+
     img_t = t(img)[None,:,:,:]
     inputs = img_t