ai_complete / app.py
Essa20001's picture
Update app.py
16d1a1a verified
raw
history blame
1.44 kB
import streamlit as st
import numpy as np
from PIL import Image
from io import BytesIO
import tensorflow as tf
st.title("Complete image using an AI model trained on Flickr images with the Pix2pix architecture.")
image = st.file_uploader("Upload an image", type=["jpg", "png","jpeg"])
model = tf.keras.models.load_model('complete_.keras')
if image :
button = st.button("Complete")
image = Image.open(image)
image = image.convert("RGB")
image = image.resize((128,128-80))
padded_width = image.width + 80 # Add the padding width to the original width
padded_height =image.height # Keep the original height
padded_image = Image.new("RGB", (padded_width, padded_height), color="white")
padded_image.paste(image, (0, 0))
image = padded_image
image = np.array(image)
if button:
image = image - 127.5
image = image / 127.5
image.shape = (1,128,128,1)
result = model(image,training = True)
result = (result * 127.5) + 127.5
numpy_array = np.array(result.numpy()[0] , dtype=np.uint8)
pillow_image = Image.fromarray(numpy_array)
output_path = "output_image.jpg"
pillow_image.save(output_path)
st.image([output_path], caption='Colored Image', use_column_width=True)
st.download_button(
label="Download Result Image",
data=BytesIO(numpy_array.tobytes()),
file_name="output_image.jpg",
key="download_button",
)