Update app.py

app.py

@@ -13,7 +13,6 @@ from torchvision import transforms
 from PIL import Image
 import torch.nn.functional as F
 import ipdb
-import ipdb
 
 img_mean = [0.485, 0.456, 0.406]
 img_std = [0.229, 0.224, 0.225]
@@ -86,7 +85,7 @@ def inference_mask1(
     pred_mask, simi, simi_map = model.predict_mask_nshot(batch, nshot=nshot)
     pred_mask = pred_mask.detach().cpu().numpy()[0]
     output_img = query_img_np.copy()
-    output_img[pred_mask] = np.array([255, 0, 0])
+    output_img[pred_mask > 0] = np.array([255, 0, 0])
     output_img = (output_img).astype(np.uint8)
     return output_img