AideepImage/interior-design
1
1import streamlit as st2# wide layout3st.set_page_config(layout="wide")4 5from streamlit_drawable_canvas import st_canvas6from PIL import Image7from typing import Union8import random9import numpy as np10import os11import time12 13from models import make_image_controlnet, make_inpainting14from segmentation import segment_image15from config import HEIGHT, WIDTH, POS_PROMPT, NEG_PROMPT, COLOR_MAPPING, map_colors, map_colors_rgb16from palette import COLOR_MAPPING_CATEGORY17from preprocessing import preprocess_seg_mask, get_image, get_mask18from explanation import make_inpainting_explanation, make_regeneration_explanation, make_segmentation_explanation19 20 21def on_upload() -> None:22 """Upload image to the canvas."""23 if 'input_image' in st.session_state and st.session_state['input_image'] is not None:24 image = Image.open(st.session_state['input_image']).convert('RGB')25 st.session_state['initial_image'] = image26 if 'seg' in st.session_state:27 del st.session_state['seg']28 if 'unique_colors' in st.session_state:29 del st.session_state['unique_colors']30 if 'output_image' in st.session_state:31 del st.session_state['output_image']32 33def make_image_row(image_0, image_1):34 col_0, col_1 = st.columns(2)35 with col_0:36 st.image(image_0, use_column_width=True)37 with col_1:38 st.image(image_1, use_column_width=True)39 40 41def check_reset_state() -> bool:42 """Check whether the UI elements need to be reset43 Returns:44 bool: True if the UI elements need to be reset, False otherwise45 """46 if ('reset_canvas' in st.session_state and st.session_state['reset_canvas']):47 st.session_state['reset_canvas'] = False48 return True49 st.session_state['reset_canvas'] = False50 return False51 52 53def move_image(source: Union[str, Image.Image],54 dest: str,55 rerun: bool = True,56 remove_state: bool = True) -> None:57 """Move image from source to destination.58 Args:59 source (Union[str, Image.Image]): source image60 dest (str): destination image location61 rerun (bool, optional): rerun streamlit. Defaults to True.62 remove_state (bool, optional): remove the canvas state. Defaults to True.63 """64 source_image = source if isinstance(source, Image.Image) else st.session_state[source]65 66 if remove_state:67 st.session_state['reset_canvas'] = True68 if 'seg' in st.session_state:69 del st.session_state['seg']70 if 'unique_colors' in st.session_state:71 del st.session_state['unique_colors']72 73 st.session_state[dest] = source_image74 st.session_state['dest'] = source_image75 if rerun:76 st.experimental_rerun()77 78 79def on_change_radio() -> None:80 """Reset the UI elements when the radio button is changed."""81 st.session_state['reset_canvas'] = True82 83 84def make_canvas_dict(canvas_color, brush, paint_mode, _reset_state):85 canvas_dict = dict(86 fill_color=canvas_color,87 stroke_color=canvas_color,88 background_color="#FFFFFF",89 background_image=st.session_state['initial_image'] if 'initial_image' in st.session_state else None,90 stroke_width=brush,91 initial_drawing={'version': '4.4.0', 'objects': []} if _reset_state else None,92 update_streamlit=True,93 height=512,94 width=512,95 drawing_mode=paint_mode,96 key="canvas",97 )98 return canvas_dict 99 100def make_prompt_row():101 col_0_0, col_0_1 = st.columns(2)102 with col_0_0:103 st.text_input(label="Positive prompt", value="a photograph of a room, interior design, 4k, high resolution", key='positive_prompt')104 with col_0_1:105 st.text_input(label="Negative prompt", value="lowres, watermark, banner, logo, watermark, contactinfo, text, deformed, blurry, blur, out of focus, out of frame, surreal, ugly", key='negative_prompt')106 107def make_sidebar():108 with st.sidebar:109 input_image = st.file_uploader("", type=["png", "jpg"], key='input_image', on_change=on_upload)110 generation_mode = st.selectbox("Generation mode", ["Regenerate",111 "Segmentation",112 "Inpainting"], on_change=on_change_radio)113 114 115 if generation_mode == "Segmentation":116 paint_mode = st.sidebar.selectbox("Painting mode", ("freedraw", "polygon"))117 if paint_mode == "freedraw":118 brush = st.slider("Stroke width", 5, 140, 100, key='slider_seg')119 else:120 brush = 5121 122 category_chooser = st.sidebar.selectbox("Filter on category", list(123 COLOR_MAPPING_CATEGORY.keys()), index=0, key='category_chooser')124 125 chosen_colors = list(COLOR_MAPPING_CATEGORY[category_chooser].keys())126 127 color_chooser = st.sidebar.selectbox(128 "Choose a color", chosen_colors, index=0, format_func=map_colors, key='color_chooser'129 )130 131 elif generation_mode == "Regenerate":132 color_chooser = "rgba(0, 0, 0, 0.0)"133 paint_mode = 'freedraw'134 brush = 0135 136 else:137 paint_mode = st.sidebar.selectbox("Painting mode", ("freedraw", "polygon"))138 if paint_mode == "freedraw":139 brush = st.slider("Stroke width", 5, 140, 100, key='slider_seg')140 else:141 brush = 5142 143 color_chooser = "#000000"144 return input_image, generation_mode, brush, color_chooser, paint_mode145 146 147def make_output_image():148 if 'output_image' in st.session_state:149 output_image = st.session_state['output_image']150 if isinstance(output_image, np.ndarray):151 output_image = Image.fromarray(output_image)152 153 if isinstance(output_image, Image.Image):154 output_image = output_image.resize((512, 512))155 else:156 output_image = Image.new('RGB', (512, 512), (255, 255, 255))157 158 st.write("#### Output image")159 st.image(output_image, width=512)160 if st.button("Move to input image"):161 move_image('output_image', 'initial_image', remove_state=True, rerun=True)162 163def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, paint_mode):164 st.write("#### Input image")165 canvas_dict = make_canvas_dict(166 canvas_color=canvas_color,167 paint_mode=paint_mode,168 brush=brush,169 _reset_state=_reset_state170 )171 if generation_mode == "Segmentation":172 canvas = st_canvas(173 **canvas_dict,174 )175 176 if st.button("generate image", key='generate_button'):177 image = get_image()178 print("Preparing image segmentation")179 real_seg = segment_image(Image.fromarray(image))180 mask, seg = preprocess_seg_mask(canvas, real_seg)181 182 with st.spinner(text="Generating image"):183 print("Making image")184 result_image = make_image_controlnet(image=image,185 mask_image=mask,186 controlnet_conditioning_image=seg,187 positive_prompt=st.session_state['positive_prompt'],188 negative_prompt=st.session_state['negative_prompt'],189 seed=random.randint(0, 100000) # nosec190 )191 if isinstance(result_image, np.ndarray):192 result_image = Image.fromarray(result_image)193 st.session_state['output_image'] = result_image194 195 196 elif generation_mode == "Regenerate":197 canvas = st_canvas(198 **canvas_dict,199 )200 if 'seg' not in st.session_state:201 with st.spinner(text="Preparing image segmentation"):202 image = get_image()203 real_seg = np.array(segment_image(Image.fromarray(image)))204 st.session_state['seg'] = real_seg205 206 if 'unique_colors' not in st.session_state:207 real_seg = st.session_state['seg']208 unique_colors = np.unique(real_seg.reshape(-1, real_seg.shape[2]), axis=0)209 unique_colors = [tuple(color) for color in unique_colors]210 st.session_state['unique_colors'] = unique_colors211 212 with st.expander("Explanation", expanded=True):213 st.write("This mode allows you to choose which objects you want to re-generate in the image. "214 "Use the selection dropdown to add or remove objects. If you are ready, press the generate button"215 " to generate the image, which can take up to 30 seconds. If you want to improve the generated image, click"216 " the 'move image to input' button."217 )218 219 chosen_colors = st.multiselect(220 label="Choose which concepts you want to regenerate in the image",221 options=st.session_state['unique_colors'],222 key='chosen_colors',223 default=st.session_state['unique_colors'],224 format_func=map_colors_rgb,225 )226 227 if st.button("generate image", key='generate_button'):228 image = get_image()229 print(chosen_colors)230 231 segmentation = st.session_state['seg']232 mask = np.zeros_like(segmentation)233 for color in chosen_colors:234 # if the color is in the segmentation, set mask to 1235 mask[np.where((segmentation == color).all(axis=2))] = 1236 237 with st.spinner(text="Generating image"):238 result_image = make_image_controlnet(image=image,239 mask_image=mask,240 controlnet_conditioning_image=segmentation,241 positive_prompt=st.session_state['positive_prompt'],242 negative_prompt=st.session_state['negative_prompt'],243 seed=random.randint(0, 100000) # nosec244 )245 if isinstance(result_image, np.ndarray):246 result_image = Image.fromarray(result_image)247 st.session_state['output_image'] = result_image248 249 elif generation_mode == "Inpainting":250 image = get_image()251 252 canvas = st_canvas(253 **canvas_dict,254 )255 256 if st.button("generate images", key='generate_button'):257 canvas_mask = canvas.image_data258 if not isinstance(canvas_mask, np.ndarray):259 canvas_mask = np.array(canvas_mask)260 mask = get_mask(canvas_mask)261 262 with st.spinner(text="Generating new images"):263 print("Making image")264 result_image = make_inpainting(positive_prompt=st.session_state['positive_prompt'],265 image=Image.fromarray(image),266 mask_image=mask,267 negative_prompt=st.session_state['negative_prompt'],268 )269 if isinstance(result_image, np.ndarray):270 result_image = Image.fromarray(result_image)271 st.session_state['output_image'] = result_image272 273def main():274 # center text275 st.write("## Controlnet sprint - interior design", unsafe_allow_html=True)276 277 input_image, generation_mode, brush, color_chooser, paint_mode = make_sidebar()278 279 # check if there is an input_image280 if not ('initial_image' in st.session_state and st.session_state['initial_image'] is not None):281 st.success("Upload an image to start")282 st.write("Welcome to the interior design controlnet demo! "283 "You can start by uploading a picture of your room, after which you will see "284 "a good variety of options to edit your current room to generate the room of your dreams! "285 "You can choose between inpainting, Segmentation and re-generating objects, which "286 "use our custom trained controlnet model. The main idea is that you can iterate over the "287 "generated images, because you will rarely get something perfect in one step (although it's possible). "288 "We added functionality to load in the generated image into the input, so you can keep "289 "changing the image until you are satisfied."290 )291 with st.expander("Useful information", expanded=True):292 st.write("### About the dataset")293 st.write("To make this demo as good as possible, our team spend a lot of time training a custom model. "294 "We used the LAION5B dataset to build our custom dataset, which contains 130k images of 15 types of rooms "295 "in almost 30 design styles. After fetching all these images, we started adding metadata such as "296 "captions (from the BLIP captioning model) and segmentation maps (from the HuggingFace UperNetForSemanticSegmentation model). "297 )298 299 st.write("### About the model")300 st.write(301 "These were then used to train the controlnet model to generate quality interior design images by using "302 "the segmentation maps and prompts as conditioning information for the model. "303 "By training on segmentation maps, the enduser has a very finegrained control over which objects they "304 "want to place in their room. "305 "The resulting model is then used in a community pipeline that supports image2image and inpainting, "306 "so the user can keep elements of their room and change specific parts of the image."307 ""308 )309 310 st.write("### Trivia")311 312 313 st.write("To enable the features in the demo, we calculate the underlying segmentation maps and categories that "314 "are present in the image. This allows us to hide some of the manual work for the user, and "315 "by doing this, the users don't need to make a segmentation map in an external tool. Everything needed can be done within this demo."316 )317 318 # st.write("### News: Fondant - an open source data-centric framework for Foundation model finetuning")319 # st.write("The ML6 team is proud to announce that we are open sourcing our Fondant framework, which is a "320 # "data-centric framework that allows you to prepare large scale multimodal datasets with ease. We have implemented the components "321 # "that we used to train this controlnet model in Fondant as an example pipeline, and we are excited to see what you can do with it! In the future we will add a whole library of plug-and-play data preparation components, such as different ML models and filtering steps, in addition to dataset scraping components that connect to LAION5B."322 # )323 # st.write("The framework is built on top of kubeflow pipelines and abstracts all the complexity of efficient storing and moving of large datasets, so you can focus on implemented just that piece of code that you need without worrying about the rest. We also build it to run on each Cloud provider or VM. You can find the code on our github page: https://github.com/ml6team/fondant.")324 325 st.write("### Testing images")326 st.write("If you don't have any pictures close, you can use one of these images to test the model by clicking on the 'use example X' buttons")327 328 st.session_state['example_image_0'] = Image.open("content/example_0.png")329 st.session_state['example_image_1'] = Image.open("content/example_1.jpg")330 st.session_state['example_image_2'] = Image.open("content/example_2.jpg")331 st.session_state['example_image_3'] = Image.open("content/example_3.jpg")332 333 col_im_0, col_im_1 = st.columns(2)334 335 with col_im_0:336 st.image(st.session_state['example_image_0'], caption="Example image 1", use_column_width=True)337 if st.button("Use example 1"):338 move_image('example_image_0', 'initial_image', remove_state=True, rerun=True)339 340 st.image(st.session_state['example_image_2'], caption="Example image 3", use_column_width=True)341 if st.button("Use example 3"):342 move_image('example_image_2', 'initial_image', remove_state=True, rerun=True)343 with col_im_1:344 st.image(st.session_state['example_image_1'], caption="Example image 2", use_column_width=True)345 if st.button("Use example 2"):346 move_image('example_image_1', 'initial_image', remove_state=True, rerun=True)347 348 st.image(st.session_state['example_image_3'], caption="Example image 4", use_column_width=True)349 if st.button("Use example 4"):350 move_image('example_image_3', 'initial_image', remove_state=True, rerun=True)351 352 st.write("## Generated examples")353 make_image_row(Image.open("content/output_1.png"), Image.open("content/regen_example.png"))354 make_image_row(Image.open("content/keep background 2.png"), Image.open("content/output_0.png"))355 make_image_row(Image.open("content/segmentation window.png"), Image.open("content/output_3.png"))356 357 st.write("## Example video")358 st.write("### Video 1")359 st.video(open('content/controlnet_sprint_demo.mp4', 'rb').read())360 st.write("### Video 2")361 st.video(open('content/controlnet_demo_video_2.mp4', 'rb').read())362 363 else:364 make_prompt_row()365 366 _reset_state = check_reset_state()367 368 if generation_mode == "Inpainting":369 make_inpainting_explanation()370 elif generation_mode == "Segmentation":371 make_segmentation_explanation()372 elif generation_mode == "Regenerate":373 make_regeneration_explanation()374 375 col1, col2 = st.columns(2)376 with col1:377 make_editing_canvas(canvas_color=color_chooser,378 brush=brush,379 _reset_state=_reset_state,380 generation_mode=generation_mode,381 paint_mode=paint_mode382 )383 384 with col2:385 make_output_image()386 387if __name__ == "__main__":388 main()389 390