Add box selection feature

app.py

@@ -13,6 +13,77 @@ from utils import utils
 SAM_MODEL = utils.get_model('vit_b')
 
 
+def box_process(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)
+
+    box_canvas = st_canvas(
+        fill_color="rgba(255, 255, 0, 0)",
+        background_image = bg_image,
+        drawing_mode='rect',
+        stroke_color = "rgba(0, 255, 0, 0.6)",
+        stroke_width = radius_width,
+        width = container_width,
+        height = height * scale,
+        point_display_radius = 12,
+        update_streamlit=True,
+        key="box"
+    )
+
+    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
+        if st.session_state.rerun_once:
+            st.experimental_rerun()
+        else:
+            return np.asarray(bg_image)
+
+    elif box_canvas.json_data is not None:
+        df = pd.json_normalize(box_canvas.json_data["objects"])
+        center_point,center_label,input_box = [],[],[]
+        center_point, center_label, input_box = [], [], []
+        for _, row in df.iterrows():
+            x, y, w,h = row["left"], row["top"], row["width"], row["height"]
+            x = int(x/scale)
+            y = int(y/scale)
+            w = int(w/scale)
+            h = int(h/scale)
+            center_point.append([x+w/2,y+h/2])
+            center_label.append([1])
+            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 len(masks) == 0:
+                return bg_image
+
+            bg_image = np.asarray(bg_image)
+            color = np.concatenate([random.choice(utils.get_color()), np.array([0.6])], axis=0)
+            im_masked = utils.show_click(masks,color)
+            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)
+
+    return np.asarray(bg_image)
+
+
 def click_process(model, show_mask, radius_width):
 
     bg_image = st.session_state['image']
@@ -24,7 +95,7 @@ def click_process(model, show_mask, radius_width):
     if 'result_image' not in st.session_state:
         st.session_state.result_image = bg_image.resize(scaled_hw)
 
-    canvas_result = st_canvas(
+    click_canvas = st_canvas(
         fill_color="rgba(255, 255, 0, 0.8)",
         background_image = bg_image,
         drawing_mode='point',
@@ -35,13 +106,21 @@ def click_process(model, show_mask, radius_width):
         update_streamlit=True,
         key="point",)
 
-    # ! Warn: Can cause infinite loop or high cpu usage
     if not show_mask:
-        print("rerun no mask")
-        st.experimental_rerun()
+        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
+        if st.session_state.rerun_once:
+            st.experimental_rerun()
+        else:
+            return np.asarray(bg_image)
 
-    elif canvas_result.json_data is not None:
-        df = pd.json_normalize(canvas_result.json_data["objects"])
+    elif click_canvas.json_data is not None:
+        df = pd.json_normalize(click_canvas.json_data["objects"])
         input_points = []
         input_labels = []
 
@@ -88,6 +167,8 @@ def image_preprocess_callback(model):
         st.session_state.image = image
     else:
         with st.spinner(text="Cleaning up!"):
+            if 'display_result' in st.session_state:
+                st.session_state.display_result = False
             if 'image' in st.session_state:
                 st.session_state.image = None
             if 'result_image' in st.session_state:
@@ -130,16 +211,15 @@ def main():
                 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)
 
-    # else:
-        # print(f'embedding is empty - {option} - {show_mask} - {radius_width}')
-        # if 'image' in st.session_state:
-            # if st.session_state.image is None:
-                # st.session_state.clear()
+    else:
+        st.cache_data.clear()
 
 
 if __name__ == '__main__':

utils/utils.py

@@ -70,3 +70,26 @@ def model_predict_masks_click(model,input_points,input_labels):
         torch.cuda.empty_cache()
 
     return masks
+
+def model_predict_masks_box(model,center_point,center_label,input_box):
+    masks = np.array([])
+    for i in range(len(center_label)):
+        if center_point[i] == []:continue
+        center_point_1 = np.array([center_point[i]])
+        center_label_1 = np.array(center_label[i])
+        input_box_1 = np.array(input_box[i])
+        mask, _, _ = model.predict(
+            point_coords=center_point_1,
+            point_labels=center_label_1,
+            box=input_box_1,
+            multimask_output=False,
+        )
+        try:
+            masks = masks + mask
+        except:
+            masks = mask
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    return masks