File size: 1,736 Bytes
6beeea2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
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))
|