|
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)) |
|
|
|
|