add masking

app.py

@@ -168,8 +168,14 @@ def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, pain
 
         if st.button("generate image", key='generate_button'):
             image = get_image()
-            real_seg = segment_image(Image.fromarray(image))
-            mask, seg = preprocess_seg_mask(canvas, real_seg)
+            print(chosen_colors)
+
+            segmentation = st.session_state['seg']
+            mask = np.zeros_like(segmentation)
+            for color in chosen_colors:
+                # if the color is in the segmentation, set mask to 1
+                mask[np.where((segmentation == color).all(axis=2))] = 1
+            print(mask)
 
             with st.spinner(text="Generating image"):
                 result_image = make_image_controlnet(image=image,