Update app.py

app.py

@@ -25,12 +25,18 @@ transform = transforms.Compose([
 ])
 
 def predict(image):
-    img = transform(image).unsqueeze(0)  # shape: (1,3,128,128)
+    orig_w, orig_h = image.size   # original size of uploaded image
+    img = transform(image).unsqueeze(0)  # (1,3,128,128)
     with torch.no_grad():
         pred = model(img)
+
     mask = pred.squeeze(0).squeeze(0).cpu().numpy()
-    mask = (mask * 255).astype(np.uint8)   # keep gray levels
-    return Image.fromarray(mask)
+    mask = (mask * 255).astype(np.uint8)  # grayscale mask
+
+    # Resize back to original size
+    mask_img = Image.fromarray(mask).resize((orig_w, orig_h), Image.NEAREST)
+    return mask_img
+
 
 # Gradio interface
 demo = gr.Interface(