Update app.py

app.py

@@ -26,6 +26,7 @@ def inference_mask1(
             query_img,
             *prompt,
               ):
+    org_qry_imsize = query_img.size
     query_img = Image.fromarray(query_img)
     query_img_np = np.asarray(query_img)
     query_img = transformation(query_img)
@@ -73,10 +74,10 @@ def inference_mask1(
         "support_masks": support_masks,
         "support_imgs": support_img,
         "query_img": query_img,
+        "org_query_imsize": [torch.tensor([org_qry_imsize[0]]), torch.tensor([org_qry_imsize[1]])],
     }
     nshot = support_masks.size(1)
     pred_mask, simi, simi_map = model.predict_mask_nshot(batch, nshot=nshot)
-    pred_mask = F.interpolate(pred_mask, shape)
     pred_mask = pred.detach().cpu().numpy()
     output_img = query_img_np * 0.5 + 0.5 * np.array([1, 0, 0]) * np.expand_dims(pred_mask, axis=0)
     output_img = (output_img * 255).astype(np.uint8)