import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
from io import BytesIO

st.title("Colorize black and white 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('generator_color.keras')

if image :
  button = st.button("Colore")
  image = Image.open(image)
  image = image.convert("L")
  image = image.resize((128,128))
  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=False)
    st.download_button(
            label="Download Colored Image",
            data=BytesIO(numpy_array.tobytes()),
            file_name="output_image.jpg",
            key="download_button",
            help="Click to download the colored image.",
        )