Update app.py

app.py

@@ -38,6 +38,8 @@ def inference_mask1(
     for i in range(len(prompt)):
         mask = torch.from_numpy(np.stack(prompt[i]['layers'], axis=0).any(0).any(-1)).cpu()
         mask[mask > 0] = 1
+        return (mask.detach().cpu().numpy() * 255).astype(np.uint8)
+
         if mask.sum() == 0:
             break
         mask = F.interpolate(mask.unsqueeze(0).unsqueeze(0).float(), query_img.size()[-2:], mode='nearest').squeeze(0).squeeze(0)
@@ -45,7 +47,6 @@ def inference_mask1(
         support_img = Image.fromarray(prompt[i]['background'][..., :3])
         support_img = transformation(support_img)
         support_imgs.append(support_img)
-    return (support_imgs[0].detach().cpu().numpy() * 255).astype(np.uint8)
     model = DCAMA('resnet50', 'resnet50_a1h-35c100f8.pth', True)
     model.eval()
     model.cpu()