Update app.py

app.py

@@ -37,6 +37,8 @@ def inference_mask1(
     support_imgs = []
     for i in range(len(prompt)):
         mask = torch.from_numpy(np.stack(prompt[i]['layers'], axis=0).any(0).any(-1)).cpu()
+        if mask.sum() == 0:
+            break
         mask = F.interpolate(mask.unsqueeze(0).unsqueeze(0).float(), query_img.size()[-2:], mode='nearest').squeeze(0).squeeze(0)
         support_masks.append(mask)
         support_img = Image.fromarray(prompt[i]['background'][..., :3])
@@ -81,7 +83,8 @@ def inference_mask1(
     nshot = support_masks.size(1)
     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 * 0.5 + 0.5 * np.array([[[1, 0, 0]]]) * np.expand_dims(pred_mask, axis=-1)
+    output_img = query_img_np.copy()
+    output_img[pred_mask] = np.array([1, 0, 0])
     output_img = (output_img * 255).astype(np.uint8)
     return output_img