haritsahm
commited on
Commit
·
4e097eb
1
Parent(s):
ef9685e
Fix when no mask result found
Browse files- app.py +5 -0
- 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']
|