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