Update app.py

app.py

@@ -37,6 +37,7 @@ 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()
+        mask[mask > 0] = 1
         if mask.sum() == 0:
             break
         mask = F.interpolate(mask.unsqueeze(0).unsqueeze(0).float(), query_img.size()[-2:], mode='nearest').squeeze(0).squeeze(0)
@@ -44,6 +45,7 @@ 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()