add segmentation maps

app.py

@@ -151,6 +151,19 @@ def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, pain
         canvas = st_canvas(
             **canvas_dict,
         )
+        if 'seg' not in st.session_state:
+            image = get_image()
+            print("Preparing image segmentation")
+            real_seg = segment_image(Image.fromarray(image))
+            st.session_state['seg'] = real_seg
+            # get unique RGB colors from this segmentation map
+            unique_colors = np.unique(real_seg.reshape(-1, real_seg.shape[2]), axis=0)
+            st.session_state['unique_colors'] = unique_colors
+            # multiselect to choose multiple colors
+            chosen_colors = st.multiselect(
+                "Choose colors", list(unique_colors), format_func=map_colors, key='chosen_colors'
+            )
+            print(chosen_colors)
 
         if st.button("generate image", key='generate_button'):
             image = get_image()