ai_complete / app.py
Essa20001's picture
Update app.py
9d318b8 verified
import streamlit as st
import numpy as np
from PIL import Image
from io import BytesIO
import keras
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 = 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",
)