import cv2 as cv
from keras.utils import get_file
import numpy as np
import streamlit as st
import tensorflow as tf
from keras.models import load_model
from streamlit_drawable_canvas import st_canvas

st.title('Facades pix2pix 🏢')

drawing_mode = st.sidebar.selectbox(
    "Drawing tool:", ("freedraw", "line", "rect", "circle", "transform")
)

stroke_color = st.sidebar.color_picker("Stroke color : ", "#000")

fill_color = st.sidebar.color_picker("Fill color : ", "#eee")

background_color = st.sidebar.color_picker("Background_color: ", "#fff")

stroke_width = st.sidebar.slider("Stroke width : ", 0, 25, 0)

canvas_result = st_canvas(
    fill_color=fill_color,
    stroke_width=stroke_width,
    stroke_color=stroke_color,
    background_color=background_color,
    height=256,
    width=256,
    drawing_mode=drawing_mode,
    point_display_radius=1 if drawing_mode == 'point' else 0,
    key="canvas",
)

reshaped_data = np.reshape(canvas_result.image_data, (256, 256, 4))
red_channel = reshaped_data[:, :, 0]
green_channel = reshaped_data[:, :, 1]
blue_channel = reshaped_data[:, :, 2]
image = np.dstack((red_channel, green_channel, blue_channel))

resize = cv.resize(image, (256, 256))

if st.button("Generate image"):
    pix2pix_model_path = get_file(
        origin="https://huggingface.co/CineAI/Pix2Pix/resolve/main/pix2pix.h5",
        file_hash="555576fa2ea8be3192c547a3355f4ac419ff9c2b94abcce5a2252b873055ce6b",
    )

    model = load_model(pix2pix_model_path)
    extra_dim = np.expand_dims(resize, axis=0)
    image_generated = model(extra_dim)
    image_generated = (image_generated + 1) / 2.0
    image_generated = tf.cast(image_generated * 255, tf.uint8)
    st.image(np.squeeze(image_generated, axis=0))