soham07ml/controlnet-interior-design
0
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 st.write("For the gathering and inference of the metadata we used the Fondant framework (https://github.com/ml6team/fondant) provided by ML6 (https://www.ml6.eu/), which is an open source "299 "data centric framework for data preparation. The pipeline used for training this controlnet will soon be available as an "300 "example pipeline within Fondant and can be easily adapted for building your own dataset."301 )302 st.write("### About the model")303 st.write(304 "These were then used to train the controlnet model to generate quality interior design images by using "305 "the segmentation maps and prompts as conditioning information for the model. "306 "By training on segmentation maps, the enduser has a very finegrained control over which objects they "307 "want to place in their room. "308 "The resulting model is then used in a community pipeline that supports image2image and inpainting, "309 "so the user can keep elements of their room and change specific parts of the image."310 ""311 )312 313 st.write("### Trivia")314 st.write("The first time someone uses the demo after startup, the models still need to be loaded into memory. "315 "After this initial load, the model is cached as a resource and can be used for all the users. "316 "To avoid simultaneous requests, we have implemented a queueing mechanism that ensures that only one "317 "user accesses the model at a time (similar to the Gradio framework).\n"318 )319 st.write("To enable the features in the demo, we calculate the underlying segmentation maps and categories that "320 "are present in the image. This allows us to hide some of the manual work for the user, and "321 "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."322 )323 324 # st.write("### News: Fondant - an open source data-centric framework for Foundation model finetuning")325 # st.write("The ML6 team is proud to announce that we are open sourcing our Fondant framework, which is a "326 # "data-centric framework that allows you to prepare large scale multimodal datasets with ease. We have implemented the components "327 # "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."328 # )329 # 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.")330 331 st.write("### Testing images")332 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")333 334 st.session_state['example_image_0'] = Image.open("content/example_0.png")335 st.session_state['example_image_1'] = Image.open("content/example_1.jpg")336 st.session_state['example_image_2'] = Image.open("content/example_2.jpg")337 st.session_state['example_image_3'] = Image.open("content/example_3.jpg")338 339 col_im_0, col_im_1 = st.columns(2)340 341 with col_im_0:342 st.image(st.session_state['example_image_0'], caption="Example image 1", use_column_width=True)343 if st.button("Use example 1"):344 move_image('example_image_0', 'initial_image', remove_state=True, rerun=True)345 346 st.image(st.session_state['example_image_2'], caption="Example image 3", use_column_width=True)347 if st.button("Use example 3"):348 move_image('example_image_2', 'initial_image', remove_state=True, rerun=True)349 with col_im_1:350 st.image(st.session_state['example_image_1'], caption="Example image 2", use_column_width=True)351 if st.button("Use example 2"):352 move_image('example_image_1', 'initial_image', remove_state=True, rerun=True)353 354 st.image(st.session_state['example_image_3'], caption="Example image 4", use_column_width=True)355 if st.button("Use example 4"):356 move_image('example_image_3', 'initial_image', remove_state=True, rerun=True)357 358 st.write("## Generated examples")359 make_image_row(Image.open("content/output_1.png"), Image.open("content/regen_example.png"))360 make_image_row(Image.open("content/keep background 2.png"), Image.open("content/output_0.png"))361 make_image_row(Image.open("content/segmentation window.png"), Image.open("content/output_3.png"))362 363 st.write("## Example video")364 st.write("### Video 1")365 st.video(open('content/controlnet_sprint_demo.mp4', 'rb').read())366 st.write("### Video 2")367 st.video(open('content/controlnet_demo_video_2.mp4', 'rb').read())368 369 else:370 make_prompt_row()371 372 _reset_state = check_reset_state()373 374 if generation_mode == "Inpainting":375 make_inpainting_explanation()376 elif generation_mode == "Segmentation":377 make_segmentation_explanation()378 elif generation_mode == "Regenerate":379 make_regeneration_explanation()380 381 col1, col2 = st.columns(2)382 with col1:383 make_editing_canvas(canvas_color=color_chooser,384 brush=brush,385 _reset_state=_reset_state,386 generation_mode=generation_mode,387 paint_mode=paint_mode388 )389 390 with col2:391 make_output_image()392 393if __name__ == "__main__":394 main()395 396