Add new features to segment everything

app.py

@@ -10,10 +10,10 @@ from PIL import Image
 from streamlit_drawable_canvas import st_canvas
 from utils import utils
 
-SAM_MODEL = utils.get_model('vit_b')
+PREDICTOR_MODEL, AUTOMASK_MODEL = utils.get_model('vit_b')
 
 
-def box_process(model, show_mask, radius_width):
+def process_box(predictor_model, show_mask, radius_width):
     bg_image = st.session_state['image']
     width, height   = bg_image.size[:2]
     container_width = 700
@@ -44,6 +44,7 @@ def box_process(model, show_mask, radius_width):
             st.session_state.rerun_once = True
 
         st.session_state.display_result = True
+        st.warning("Mask view is disabled", icon="❗")
         if st.session_state.rerun_once:
             st.experimental_rerun()
         else:
@@ -64,11 +65,12 @@ def box_process(model, show_mask, radius_width):
             input_box.append([x,y,x+w,y+h])
 
         masks = []
-        if model:
-            masks = utils.model_predict_masks_box(model, center_point, center_label, input_box)
+        if predictor_model:
+            masks = utils.model_predict_masks_box(predictor_model, center_point, center_label, input_box)
 
             if len(masks) == 0:
-                return bg_image
+                st.warning("No Masks Found", icon="❗")
+                return np.asarray(bg_image)
 
             bg_image = np.asarray(bg_image)
             color = np.concatenate([random.choice(utils.get_color()), np.array([0.6])], axis=0)
@@ -84,7 +86,7 @@ def box_process(model, show_mask, radius_width):
     return np.asarray(bg_image)
 
 
-def click_process(model, show_mask, radius_width):
+def process_click(predictor_model, show_mask, radius_width):
 
     bg_image = st.session_state['image']
     width, height   = bg_image.size[:2]
@@ -114,6 +116,7 @@ def click_process(model, show_mask, radius_width):
             st.session_state.rerun_once = True
 
         st.session_state.display_result = True
+        st.warning("Mask view is disabled", icon="❗")
         if st.session_state.rerun_once:
             st.experimental_rerun()
         else:
@@ -135,11 +138,12 @@ def click_process(model, show_mask, radius_width):
                 input_labels.append(0)
 
         masks = []
-        if model:
-            masks = utils.model_predict_masks_click(model, input_points, input_labels)
+        if predictor_model:
+            masks = utils.model_predict_masks_click(predictor_model, input_points, input_labels)
 
             if len(masks) == 0:
-                return bg_image
+                st.warning("No Masks Found", icon="❗")
+                return np.asarray(bg_image)
 
             bg_image = np.asarray(bg_image)
             color = np.concatenate([random.choice(utils.get_color()), np.array([0.6])], axis=0)
@@ -147,6 +151,7 @@ def click_process(model, show_mask, radius_width):
             im_masked = Image.fromarray(im_masked).convert('RGBA')
             result_image = Image.alpha_composite(Image.fromarray(bg_image).convert('RGBA'),im_masked).convert("RGB")
             result_image = result_image.resize(scaled_hw)
+            st.session_state.display_result = True
             return result_image
         else:
             return np.asarray(bg_image)
@@ -154,16 +159,65 @@ def click_process(model, show_mask, radius_width):
     return np.asarray(bg_image)
 
 
-def image_preprocess_callback(model):
+def process_everything(automask_model, show_mask, radius_width):
+    bg_image = st.session_state['image']
+    width, height   = bg_image.size[:2]
+    container_width = 700
+    scale           = container_width/width
+    scaled_hw = (container_width, int(height * scale))
+
+    if 'result_image' not in st.session_state:
+        st.session_state.result_image = bg_image.resize(scaled_hw)
+
+    dummy_canvas = st_canvas(
+        fill_color="rgba(255, 255, 0, 0.8)",
+        background_image = bg_image,
+        drawing_mode='freedraw',
+        width = container_width,
+        height = height * scale,
+        point_display_radius = radius_width,
+        stroke_width=2,
+        update_streamlit=False,
+        key="everything",)
+
+    if not show_mask:
+        if 'rerun_once' in st.session_state:
+            if st.session_state.rerun_once:
+                st.session_state.rerun_once = False
+        else:
+            st.session_state.rerun_once = True
+
+        st.session_state.display_result = True
+        st.warning("Mask view is disabled", icon="❗")
+        if st.session_state.rerun_once:
+            st.experimental_rerun()
+        else:
+            return np.asarray(bg_image)
+
+    if automask_model:
+        bg_image = np.asarray(bg_image)
+        masks = utils.model_predict_masks_everything(automask_model, bg_image)
+        im_masked = utils.show_everything(masks)
+        im_masked = Image.fromarray(im_masked).convert('RGBA')
+        result_image = Image.alpha_composite(Image.fromarray(bg_image).convert('RGBA'),im_masked).convert("RGB")
+        result_image = result_image.resize(scaled_hw)
+        st.session_state.display_result = True
+        return result_image
+
+    else:
+        return np.asarray(bg_image)
+
+
+def image_preprocess_callback(predictor_model, option):
     if 'uploaded_image' not in st.session_state:
         return
     if st.session_state.uploaded_image is not None:
         with st.spinner(text="Uploading image..."):
             image = Image.open(st.session_state.uploaded_image).convert("RGB")
-        if model:
+        if predictor_model and option != 'Everything':
             np_image = np.asanyarray(image)
             with st.spinner(text="Extracing embeddings.."):
-                model.set_image(np_image)
+                predictor_model.set_image(np_image)
         st.session_state.image = image
     else:
         with st.spinner(text="Cleaning up!"):
@@ -173,8 +227,8 @@ def image_preprocess_callback(model):
                 st.session_state.image = None
             if 'result_image' in st.session_state:
                 del st.session_state['result_image']
-            if model:
-                model.reset_image()
+            if predictor_model:
+                predictor_model.reset_image()
 
 def main():
     with open('index.html', encoding='utf-8') as f:
@@ -202,21 +256,29 @@ def main():
         st.write("Upload Image")
         st.file_uploader(label='Upload image',type=['png','jpg','tif'], key='uploaded_image', on_change=image_preprocess_callback, args=(SAM_MODEL,), label_visibility="hidden")
 
+    result_image = None
     canvas_input, canvas_output = st.columns(2)
+
     if 'image' in st.session_state:
-        result_image = None
         with canvas_input:
             st.write("Select Interest Area/Objects")
             if st.session_state.image is not None:
-                if option == 'Click':
-                    with st.spinner(text="Computing masks"):
-                        result_image = click_process(SAM_MODEL, show_mask, radius_width)
-                elif option == 'Box':
-                        result_image = box_process(SAM_MODEL, show_mask, radius_width)
-        with canvas_output:
-            if result_image is not None:
-                st.write("Result")
-                st.image(result_image)
+                with st.spinner(text="Computing masks"):
+                    if option == 'Click':
+                        result_image = process_click(PREDICTOR_MODEL, show_mask, radius_width)
+                    elif option == 'Box':
+                        result_image = process_box(PREDICTOR_MODEL, show_mask, radius_width)
+                    else:
+                        result_image = process_everything(AUTOMASK_MODEL, show_mask, radius_width)
+        if 'display_result' in st.session_state:
+            if st.session_state.display_result:
+                with canvas_output:
+                    if result_image is not None:
+                        st.write("Result")
+                        st.image(result_image)
+        else:
+            st.warning("No result found, please set input prompt", icon="⚠️")
+        st.success('Process completed！', icon="✅")
 
     else:
         st.cache_data.clear()

utils/utils.py

@@ -1,5 +1,5 @@
 
-from segment_anything import SamPredictor, sam_model_registry
+from segment_anything import SamPredictor, SamAutomaticMaskGenerator, sam_model_registry
 import torch
 import numpy as np
 from distinctipy import distinctipy
@@ -19,20 +19,19 @@ def get_model(model):
     device         = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     build_sam      = sam_model_registry[model]
     model          = build_sam(checkpoint=get_checkpoint_path(model)).to(device)
-    predictor      = SamPredictor(model)
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-    return predictor
+    predictor      = SamPredictor(model)
+    mask_generator = SamAutomaticMaskGenerator(model)
+    return predictor, mask_generator
 
 
-@st.cache_data
 def show_everything(sorted_anns):
-    if len(sorted_anns) == 0:
-        return
     #sorted_anns = sorted(anns, key=(lambda x: x['stability_score']), reverse=True)
     h, w        = sorted_anns[0]['segmentation'].shape[-2:]
     #sorted_anns = sorted_anns[:int(len(sorted_anns) * stability_score/100)]
     if sorted_anns == []:
+        st.warning("No Masks Found", icon="❗")
         return np.zeros((h,w,4)).astype(np.uint8)
     mask = np.zeros((h,w,4))
     for ann in sorted_anns:
@@ -40,13 +39,13 @@ def show_everything(sorted_anns):
         color  = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
         mask  += m.reshape(h,w,1) * color.reshape(1, 1, -1)
     mask = mask * 255
-    st.success('Process completed！', icon="✅")
     return mask.astype(np.uint8)
 
 
 def show_click(masks, colors):
     h, w         = masks[0].shape[-2:]
     masks_total = np.zeros((h,w,4)).astype(np.uint8)
+
     for mask, color in zip(masks, colors):
         if np.array_equal(mask,np.array([])):continue
         masks = np.zeros((h,w,4)).astype(np.uint8)
@@ -54,7 +53,7 @@ def show_click(masks, colors):
         masks = masks.astype(bool).astype(np.uint8)
         masks = masks * 255 * color.reshape(1, 1, -1)
         masks_total += masks.astype(np.uint8)
-    st.success('Process completed!', icon="✅")
+
     return masks_total
 
 def model_predict_masks_click(model,input_points,input_labels):
@@ -66,6 +65,7 @@ def model_predict_masks_click(model,input_points,input_labels):
         point_labels=input_labels,
         multimask_output=False,
     )
+
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
@@ -93,3 +93,12 @@ def model_predict_masks_box(model,center_point,center_label,input_box):
         torch.cuda.empty_cache()
 
     return masks
+
+
+def model_predict_masks_everything(mask_generator, image):
+    masks = mask_generator.generate(image)
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    return masks