add segmentation

app.py

@@ -29,11 +29,11 @@ def on_upload() -> None:
     if 'input_image' in st.session_state and st.session_state['input_image'] is not None:
         image = Image.open(st.session_state['input_image']).convert('RGB')
         st.session_state['initial_image'] = image
-        # st.session_state['history'] = [{'image': image.resize((512, 512)),
-        #                                 'message': "initial image",
-        #                                 "positive_prompt": "",
-        #                                 "negative_prompt": "",
-        #                                 "index": 0}]
+        if 'seg' in st.session_state:
+            del st.session_state['seg']
+        if 'unique_colors' in st.session_state:
+            del st.session_state['unique_colors']
+
 
 
 def check_reset_state() -> bool:
@@ -63,6 +63,10 @@ def move_image(source: Union[str, Image.Image],
 
     if remove_state:
         st.session_state['reset_canvas'] = True
+        if 'seg' in st.session_state:
+            del st.session_state['seg']
+        if 'unique_colors' in st.session_state:
+            del st.session_state['unique_colors']
 
     st.session_state[dest] = source_image
     if rerun:
@@ -161,8 +165,32 @@ def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, pain
         brush=brush,
         _reset_state=_reset_state
     )
+    if generation_mode == "Segmentation conditioning":
+        canvas = st_canvas(
+            **canvas_dict,
+        )
+
+        if st.button("generate image", key='generate_button'):
+            image = get_image()
+            print("Preparing image segmentation")
+            real_seg = segment_image(Image.fromarray(image))
+            mask, seg = preprocess_seg_mask(canvas, real_seg)
+
+            with st.spinner(text="Generating image"):
+                print("Making image")
+                result_image = make_image_controlnet(image=image,
+                                                        mask_image=mask,
+                                                        controlnet_conditioning_image=seg,
+                                                        positive_prompt=st.session_state['positive_prompt'],
+                                                        negative_prompt=st.session_state['negative_prompt'],
+                                                        seed=random.randint(0, 100000) # nosec
+                                                        )[0]
+                if isinstance(result_image, np.ndarray):
+                    result_image = Image.fromarray(result_image)
+                st.session_state['output_image'] = result_image
+
 
-    if generation_mode == "Re-generate objects":
+    elif generation_mode == "Re-generate objects":
         canvas = st_canvas(
             **canvas_dict,
         )