haritsahm commited on
Commit
4e097eb
·
1 Parent(s): ef9685e

Fix when no mask result found

Browse files
Files changed (2) hide show
  1. app.py +5 -0
  2. utils/utils.py +2 -3
app.py CHANGED
@@ -208,6 +208,11 @@ def process_everything(automask_model, show_mask, radius_width):
208
  bg_image = np.asarray(bg_image)
209
  masks = utils.model_predict_masks_everything(automask_model, bg_image)
210
  im_masked = utils.show_everything(masks)
 
 
 
 
 
211
  im_masked = Image.fromarray(im_masked).convert('RGBA')
212
  result_image = Image.alpha_composite(Image.fromarray(bg_image).convert('RGBA'),im_masked).convert("RGB")
213
  result_image = result_image.resize(scaled_wh)
 
208
  bg_image = np.asarray(bg_image)
209
  masks = utils.model_predict_masks_everything(automask_model, bg_image)
210
  im_masked = utils.show_everything(masks)
211
+
212
+ if len(im_masked) == 0:
213
+ st.warning("No Masks Found", icon="❗")
214
+ return np.asarray(bg_image)
215
+
216
  im_masked = Image.fromarray(im_masked).convert('RGBA')
217
  result_image = Image.alpha_composite(Image.fromarray(bg_image).convert('RGBA'),im_masked).convert("RGB")
218
  result_image = result_image.resize(scaled_wh)
utils/utils.py CHANGED
@@ -47,12 +47,11 @@ def get_model(checkpoint='checkpoint/sam_vit_b_01ec64.pth'):
47
 
48
 
49
  def show_everything(sorted_anns):
 
 
50
  #sorted_anns = sorted(anns, key=(lambda x: x['stability_score']), reverse=True)
51
  h, w = sorted_anns[0]['segmentation'].shape[-2:]
52
  #sorted_anns = sorted_anns[:int(len(sorted_anns) * stability_score/100)]
53
- if sorted_anns == []:
54
- st.warning("No Masks Found", icon="❗")
55
- return np.zeros((h,w,4)).astype(np.uint8)
56
  mask = np.zeros((h,w,4))
57
  for ann in sorted_anns:
58
  m = ann['segmentation']
 
47
 
48
 
49
  def show_everything(sorted_anns):
50
+ if len(sorted_anns) == 0:
51
+ return np.array([])
52
  #sorted_anns = sorted(anns, key=(lambda x: x['stability_score']), reverse=True)
53
  h, w = sorted_anns[0]['segmentation'].shape[-2:]
54
  #sorted_anns = sorted_anns[:int(len(sorted_anns) * stability_score/100)]
 
 
 
55
  mask = np.zeros((h,w,4))
56
  for ann in sorted_anns:
57
  m = ann['segmentation']