Update app.py

app.py

@@ -36,7 +36,7 @@ def inference_mask1(
         mask = torch.from_numpy(np.stack(prompt[i]['layers'], axis=0).any(0).any(-1)).cpu()
         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'])
+        support_img = Image.fromarray(prompt[i]['background'][..., :3])
         support_img = transformation(support_img)
         support_imgs.append(support_img)
     model = DCAMA('resnet50', 'resnet50_a1h-35c100f8.pth', True)