app.py

import streamlit as st
# wide layout
st.set_page_config(layout="wide")

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
from preprocessing import preprocess_seg_mask, get_image, get_mask
from explanation import make_inpainting_explanation, make_regeneration_explanation, make_segmentation_explanation



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 make_image_row(image_0, image_1):
    col_0, col_1 = st.columns(2)
    with col_0:
        st.image(image_0, use_column_width=True)
    with col_1:
        st.image(image_1, use_column_width=True)

def check_reset_state() -> bool:
    """Check whether the UI elements need to be reset"""
    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."""
    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 = {
        '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 realistic photograph of a room, high resolution",
                     key='positive_prompt')
    with col_0_1:
        st.text_input(label="Negative prompt", 
                     value="watermark, banner, logo, watermark, contactinfo, text, deformed, blurry, blur, lowres", 
                     key='negative_prompt')

def make_sidebar():
    with st.sidebar:
        input_image = st.file_uploader("Upload Image", 
                                     type=["png", "jpg"], 
                                     key='input_image', 
                                     on_change=on_upload)
        generation_mode = 'Inpainting'
        paint_mode = 'freedraw'
        brush = 30 if paint_mode == "freedraw" else 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":
        canvas = st_canvas(**canvas_dict)
    elif generation_mode == "Inpainting":
        image = get_image()  # Assuming this function exists in your preprocessing module
        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)  # Assuming this function exists in your preprocessing module

            with st.spinner(text="Generating new images"):
                print("Making image")
                result_image = make_inpainting(
                    positive_prompt=st.session_state['positive_prompt'],
                    image=Image.fromarray(image),
                    mask_image=mask,
                    negative_prompt=st.session_state['negative_prompt'],
                )
                if isinstance(result_image, np.ndarray):
                    result_image = Image.fromarray(result_image)
                st.session_state['output_image'] = result_image

def main():
    st.write("## Virtual Staging")

    input_image, generation_mode, brush, color_chooser, paint_mode = make_sidebar()

    if not ('initial_image' in st.session_state and st.session_state['initial_image'] is not None):
        st.write("Please upload an image to begin")
    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()