import streamlit as st
import numpy as np
from PIL import Image
import os
import tensorflow_hub as hub
import tensorflow as tf

os.environ['TFHUB_MODEL_LOAD_FORMAT'] = 'COMPRESSED'

def tensor_to_image(tensor):
    tensor = tensor * 255
    tensor = np.array(tensor, dtype=np.uint8)
    if np.ndim(tensor) > 3:
        assert tensor.shape[0] == 1
        tensor = tensor[0]
    return Image.fromarray(tensor)

def load(file_uploader,fst=True):
    max_dim = 512
    image = file_uploader
    if fst:
     image = tf.image.decode_image(image.read(), channels=3)
    else: 
        image = tf.image.decode_image(image, channels=3)
    image = tf.image.convert_image_dtype(image, tf.float32)

    shape = tf.cast(tf.shape(image)[:-1], tf.float32)
    long_dim = max(shape)
    scale = max_dim / long_dim

    new_shape = tf.cast(shape * scale, tf.int32)

    image = tf.image.resize(image, (256, 256))  # Resize to a consistent shape (256, 256)
    image = image[tf.newaxis, :]
    return image
    
st.title("Sketch to Real Image Pix2Pix")
inp = st.file_uploader("Sketch input")
style = load('style.jpg',False)

if inp is not None:
    img = load(inp)
    hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
    
    stylized_image = hub_model(img, img)[0]
    pred = tensor_to_image(stylized_image)
    
    toshow = np.array(img[0].numpy() * 255, dtype=np.uint8)  # Convert to uint8
    st.image(toshow, caption="Uploaded Image")
    st.image(pred, caption="Generated Real Face")