import streamlit as st
from streamlit_drawable_canvas import st_canvas
from PIL import Image
from typing import Union
import random
import numpy as np
import os
import time
from models import make_image_controlnet, make_inpainting, segment_image
from config import HEIGHT, WIDTH, POS_PROMPT, NEG_PROMPT, COLOR_MAPPING, map_colors, map_colors_rgb
from palette import COLOR_MAPPING_CATEGORY
from preprocessing import preprocess_seg_mask, get_image, get_mask
# wide layout
st.set_page_config(layout="wide")
def on_upload() -> None:
"""Upload image to the canvas."""
if 'input_image' in st.session_state and st.session_state['input_image'] is not None:
image = Image.open(st.session_state['input_image']).convert('RGB')
st.session_state['initial_image'] = image
if 'seg' in st.session_state:
del st.session_state['seg']
if 'unique_colors' in st.session_state:
del st.session_state['unique_colors']
if 'output_image' in st.session_state:
del st.session_state['output_image']
def check_reset_state() -> bool:
"""Check whether the UI elements need to be reset
Returns:
bool: True if the UI elements need to be reset, False otherwise
"""
if ('reset_canvas' in st.session_state and st.session_state['reset_canvas']):
st.session_state['reset_canvas'] = False
return True
st.session_state['reset_canvas'] = False
return False
def move_image(source: Union[str, Image.Image],
dest: str,
rerun: bool = True,
remove_state: bool = True) -> None:
"""Move image from source to destination.
Args:
source (Union[str, Image.Image]): source image
dest (str): destination image location
rerun (bool, optional): rerun streamlit. Defaults to True.
remove_state (bool, optional): remove the canvas state. Defaults to True.
"""
source_image = source if isinstance(source, Image.Image) else st.session_state[source]
if remove_state:
st.session_state['reset_canvas'] = True
if 'seg' in st.session_state:
del st.session_state['seg']
if 'unique_colors' in st.session_state:
del st.session_state['unique_colors']
st.session_state[dest] = source_image
if rerun:
st.experimental_rerun()
def on_change_radio() -> None:
"""Reset the UI elements when the radio button is changed."""
st.session_state['reset_canvas'] = True
def make_canvas_dict(canvas_color, brush, paint_mode, _reset_state):
canvas_dict = dict(
fill_color=canvas_color,
stroke_color=canvas_color,
background_color="#FFFFFF",
background_image=st.session_state['initial_image'] if 'initial_image' in st.session_state else None,
stroke_width=brush,
initial_drawing={'version': '4.4.0', 'objects': []} if _reset_state else None,
update_streamlit=True,
height=512,
width=512,
drawing_mode=paint_mode,
key="canvas",
)
return canvas_dict
def make_prompt_row():
col_0_0, col_0_1 = st.columns(2)
with col_0_0:
st.text_input(label="Positive prompt", value="a photograph of a room, interior design, 4k, high resolution", key='positive_prompt')
with col_0_1:
st.text_input(label="Negative prompt", value="", key='negative_prompt')
def make_sidebar():
with st.sidebar:
input_image = st.file_uploader("", type=["png", "jpg"], key='input_image', on_change=on_upload)
generation_mode = st.selectbox("Generation mode", ["Re-generate objects",
"Segmentation conditioning",
"Inpainting"], on_change=on_change_radio)
if generation_mode == "Segmentation conditioning":
paint_mode = st.sidebar.selectbox("Painting mode", ("freedraw", "polygon"))
if paint_mode == "freedraw":
brush = st.slider("Stroke width", 5, 140, 100, key='slider_seg')
else:
brush = 5
category_chooser = st.sidebar.selectbox("Filter on category", list(
COLOR_MAPPING_CATEGORY.keys()), index=0, key='category_chooser')
chosen_colors = list(COLOR_MAPPING_CATEGORY[category_chooser].keys())
color_chooser = st.sidebar.selectbox(
"Choose a color", chosen_colors, index=0, format_func=map_colors, key='color_chooser'
)
elif generation_mode == "Re-generate objects":
color_chooser = "rgba(0, 0, 0, 0.0)"
paint_mode = 'freedraw'
brush = 0
else:
paint_mode = st.sidebar.selectbox("Painting mode", ("freedraw", "polygon"))
if paint_mode == "freedraw":
brush = st.slider("Stroke width", 5, 140, 100, key='slider_seg')
else:
brush = 5
color_chooser = "#000000"
return input_image, generation_mode, brush, color_chooser, paint_mode
def make_output_image():
if 'output_image' in st.session_state:
output_image = st.session_state['output_image']
if isinstance(output_image, np.ndarray):
output_image = Image.fromarray(output_image)
if isinstance(output_image, Image.Image):
output_image = output_image.resize((512, 512))
else:
output_image = Image.new('RGB', (512, 512), (255, 255, 255))
st.write("#### Output image")
st.image(output_image, width=512)
if st.button("Move to input image"):
move_image('output_image', 'initial_image', remove_state=True, rerun=True)
def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, paint_mode):
st.write("#### Input image")
canvas_dict = make_canvas_dict(
canvas_color=canvas_color,
paint_mode=paint_mode,
brush=brush,
_reset_state=_reset_state
)
if generation_mode == "Segmentation conditioning":
canvas = st_canvas(
**canvas_dict,
)
if st.button("generate image", key='generate_button'):
image = get_image()
print("Preparing image segmentation")
real_seg = segment_image(Image.fromarray(image))
mask, seg = preprocess_seg_mask(canvas, real_seg)
with st.spinner(text="Generating image"):
print("Making image")
result_image = make_image_controlnet(image=image,
mask_image=mask,
controlnet_conditioning_image=seg,
positive_prompt=st.session_state['positive_prompt'],
negative_prompt=st.session_state['negative_prompt'],
seed=random.randint(0, 100000) # nosec
)[0]
if isinstance(result_image, np.ndarray):
result_image = Image.fromarray(result_image)
st.session_state['output_image'] = result_image
elif generation_mode == "Re-generate objects":
canvas = st_canvas(
**canvas_dict,
)
if 'seg' not in st.session_state:
with st.spinner(text="Preparing image segmentation"):
image = get_image()
real_seg = np.array(segment_image(Image.fromarray(image)))
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 which concepts you want to regenerate in the image",
options=st.session_state['unique_colors'],
key='chosen_colors',
default=st.session_state['unique_colors'],
format_func=map_colors_rgb,
)
with st.expander("Explanation", expanded=False):
st.write("This mode allows you to choose which objects you want to re-generate in the image. "
"Use the selection dropdown to add or remove objects. If you are ready, press the generate button"
" to generate the image, which can take up to 30 seconds. If you want to improve the generated image, click"
" the 'move image to input' button."
)
if st.button("generate image", key='generate_button'):
image = get_image()
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
with st.spinner(text="Generating image"):
result_image = make_image_controlnet(image=image,
mask_image=mask,
controlnet_conditioning_image=segmentation,
positive_prompt=st.session_state['positive_prompt'],
negative_prompt=st.session_state['negative_prompt'],
seed=random.randint(0, 100000) # nosec
)[0]
if isinstance(result_image, np.ndarray):
result_image = Image.fromarray(result_image)
st.session_state['output_image'] = result_image
elif generation_mode == "Inpainting":
image = get_image()
canvas = st_canvas(
**canvas_dict,
)
if st.button("generate images", key='generate_button'):
canvas_mask = canvas.image_data
if not isinstance(canvas_mask, np.ndarray):
canvas_mask = np.array(canvas_mask)
mask = get_mask(canvas_mask)
with st.spinner(text="Generating new images"):
print("Making image")
result_image = make_inpainting(positive_prompt=st.session_state['positive_prompt'],
image=image,
mask_image=mask,
negative_prompt=st.session_state['negative_prompt'],
)[0]
if isinstance(result_image, np.ndarray):
result_image = Image.fromarray(result_image)
st.session_state['output_image'] = result_image
def main():
# center text
st.write("## Controlnet sprint - interior design", unsafe_allow_html=True)
input_image, generation_mode, brush, color_chooser, paint_mode = make_sidebar()
# check if there is an input_image
if not ('input_image' in st.session_state and st.session_state['input_image'] is not None):
print("Image not present")
st.success("Upload an image to start")
else:
make_prompt_row()
_reset_state = check_reset_state()
col1, col2 = st.columns(2)
with col1:
make_editing_canvas(canvas_color=color_chooser,
brush=brush,
_reset_state=_reset_state,
generation_mode=generation_mode,
paint_mode=paint_mode
)
with col2:
make_output_image()
if __name__ == "__main__":
main()