add colors

app.py

@@ -158,12 +158,17 @@ def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, pain
                 st.session_state['seg'] = real_seg
 
         if 'unique_colors' not in st.session_state:
+            real_seg = st.session_state['seg']
             unique_colors = np.unique(real_seg.reshape(-1, real_seg.shape[2]), axis=0)
             unique_colors = [tuple(color) for color in unique_colors]
             st.session_state['unique_colors'] = unique_colors
 
         chosen_colors = st.multiselect(
-            label="Choose colors", options=unique_colors, default=None, key='unique_colors', format_func=map_colors_rgb
+            label="Choose colors",
+            options=st.session_state['unique_colors'],
+            default=None,
+            key='unique_colors',
+            format_func=map_colors_rgb
         )
 
         if st.button("generate image", key='generate_button'):