BertChristiaens commited on
Commit
5a67d9b
1 Parent(s): f8c7d9d
Files changed (1) hide show
  1. app.py +9 -15
app.py CHANGED
@@ -154,30 +154,24 @@ def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, pain
154
  if 'seg' not in st.session_state:
155
  with st.spinner(text="Preparing image segmentation"):
156
  image = get_image()
157
- print("Preparing image segmentation")
158
  real_seg = np.array(segment_image(Image.fromarray(image)))
159
  st.session_state['seg'] = real_seg
160
- # get unique RGB colors from this segmentation map
161
- st.image(Image.fromarray(real_seg), width=512)
162
- print(real_seg.shape)
163
- unique_colors = np.unique(real_seg.reshape(-1, real_seg.shape[2]), axis=0)
164
- unique_colors = [tuple(color) for color in unique_colors]
165
- print(unique_colors)
166
- st.session_state['unique_colors'] = unique_colors
167
- # multiselect to choose multiple colors
168
- chosen_colors = st.multiselect(
169
- label="Choose colors", options=unique_colors, default=None, key='unique_colors', format_func=map_colors_rgb
170
- )
171
- print(chosen_colors)
172
 
173
  if st.button("generate image", key='generate_button'):
174
  image = get_image()
175
- print("Preparing image segmentation")
176
  real_seg = segment_image(Image.fromarray(image))
177
  mask, seg = preprocess_seg_mask(canvas, real_seg)
178
 
179
  with st.spinner(text="Generating image"):
180
- print("Making image")
181
  result_image = make_image_controlnet(image=image,
182
  mask_image=mask,
183
  controlnet_conditioning_image=seg,
 
154
  if 'seg' not in st.session_state:
155
  with st.spinner(text="Preparing image segmentation"):
156
  image = get_image()
 
157
  real_seg = np.array(segment_image(Image.fromarray(image)))
158
  st.session_state['seg'] = real_seg
159
+
160
+ if 'unique_colors' not in st.session_state:
161
+ unique_colors = np.unique(real_seg.reshape(-1, real_seg.shape[2]), axis=0)
162
+ unique_colors = [tuple(color) for color in unique_colors]
163
+ st.session_state['unique_colors'] = unique_colors
164
+
165
+ chosen_colors = st.multiselect(
166
+ label="Choose colors", options=unique_colors, default=None, key='unique_colors', format_func=map_colors_rgb
167
+ )
 
 
 
168
 
169
  if st.button("generate image", key='generate_button'):
170
  image = get_image()
 
171
  real_seg = segment_image(Image.fromarray(image))
172
  mask, seg = preprocess_seg_mask(canvas, real_seg)
173
 
174
  with st.spinner(text="Generating image"):
 
175
  result_image = make_image_controlnet(image=image,
176
  mask_image=mask,
177
  controlnet_conditioning_image=seg,