Spaces:
Sleeping
Sleeping
import streamlit as st | |
import numpy as np | |
import cv2 | |
from PIL import Image | |
import matplotlib.pyplot as plt | |
from tensorflow.keras.models import Model | |
from tensorflow.keras.layers import Input, Conv2D, LeakyReLU, Concatenate, Conv2DTranspose, Activation, BatchNormalization, Dropout | |
from tensorflow.keras.initializers import RandomNormal | |
# Define the generator model architecture | |
def define_generator(image_shape=(256, 256, 3)): | |
init = RandomNormal(stddev=0.02) | |
in_image = Input(shape=image_shape) | |
# Encoder model | |
e1 = define_encoder_block(in_image, 64, batchnorm=False) | |
e2 = define_encoder_block(e1, 128) | |
e3 = define_encoder_block(e2, 256) | |
e4 = define_encoder_block(e3, 512) | |
e5 = define_encoder_block(e4, 512) | |
e6 = define_encoder_block(e5, 512) | |
e7 = define_encoder_block(e6, 512) | |
# Bottleneck | |
b = Conv2D(512, (4, 4), strides=(2, 2), padding='same', kernel_initializer=init)(e7) | |
b = Activation('relu')(b) | |
# Decoder model | |
d1 = decoder_block(b, e7, 512) | |
d2 = decoder_block(d1, e6, 512) | |
d3 = decoder_block(d2, e5, 512) | |
d4 = decoder_block(d3, e4, 512, dropout=False) | |
d5 = decoder_block(d4, e3, 256, dropout=False) | |
d6 = decoder_block(d5, e2, 128, dropout=False) | |
d7 = decoder_block(d6, e1, 64, dropout=False) | |
# Output | |
g = Conv2DTranspose(image_shape[2], (4, 4), strides=(2, 2), padding='same', kernel_initializer=init)(d7) | |
out_image = Activation('tanh')(g) | |
# Define model | |
model = Model(in_image, out_image) | |
return model | |
def define_encoder_block(layer_in, n_filters, batchnorm=True): | |
init = RandomNormal(stddev=0.02) | |
g = Conv2D(n_filters, (4, 4), strides=(2, 2), padding='same', kernel_initializer=init)(layer_in) | |
if batchnorm: | |
g = BatchNormalization()(g, training=True) | |
g = LeakyReLU(alpha=0.2)(g) | |
return g | |
def decoder_block(layer_in, skip_in, n_filters, dropout=True): | |
init = RandomNormal(stddev=0.02) | |
g = Conv2DTranspose(n_filters, (4, 4), strides=(2, 2), padding='same', kernel_initializer=init)(layer_in) | |
g = BatchNormalization()(g, training=True) | |
if dropout: | |
g = Dropout(0.5)(g, training=True) | |
g = Concatenate()([g, skip_in]) | |
g = Activation('relu')(g) | |
return g | |
def load_and_preprocess_image(image): | |
custom_image = np.array(image) | |
custom_image = cv2.resize(custom_image, (256, 256)) | |
#custom_image = cv2.cvtColor(custom_image, cv2.COLOR_BGR2RGB) | |
custom_image = (custom_image.astype(np.float32) - 127.5) / 127.5 | |
custom_image = np.expand_dims(custom_image, axis=0) | |
return custom_image | |
# Streamlit UI | |
st.title("Image Processing with Keras Model") | |
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
if uploaded_file is not None: | |
# Load and preprocess image | |
image = Image.open(uploaded_file).convert('RGB') | |
st.image(image, caption='Original Image', use_column_width=True) | |
custom_image = load_and_preprocess_image(image) | |
# Define and load model | |
model = define_generator() | |
model.load_weights('model_weights.h5') | |
# Generate image | |
gen_image = model.predict(custom_image) | |
gen_image = np.squeeze(gen_image, axis=0) | |
gen_image = ((gen_image + 1) * 127.5).astype(np.uint8) | |
gen_image_pil = Image.fromarray(gen_image) | |
st.image(gen_image_pil, caption='Generated Image', use_column_width=True) | |
# Plotting results | |
fig, ax = plt.subplots(1, 2, figsize=(10, 5)) | |
ax[0].imshow(image) | |
ax[0].set_title("Original Image") | |
ax[0].axis('off') | |
ax[1].imshow(gen_image_pil) | |
ax[1].set_title("Generated Image") | |
ax[1].axis('off') | |
st.pyplot(fig) | |