from turtle import title
import gradio as gr
from huggingface_hub import from_pretrained_keras
import tensorflow as tf
import numpy as np
from PIL import Image
import io
import base64


model = tf.keras.models.load_model("./tf_model.h5")


def predict(image):
    img = np.array(image)

    im = tf.image.resize(img, (128, 128))
    im = tf.cast(im, tf.float32) / 255.0
    pred_mask = model.predict(im[tf.newaxis, ...])
    
    return pred_mask[0]
    

title = '<h1 style="text-align: center;">Segment Pets</h1>'

description = """
## About
This space demonstrates the use of a semantic segmentation model to segment pets and classify them 
according to the pixels.


## 🚀 To run
Upload a pet image and hit submit or select one from the given examples
"""

inputs = gr.inputs.Image(label="Upload a pet image", type = 'pil', optional=False)
outputs = [
    gr.outputs.Image(label="Segmentation")
    # , gr.outputs.Textbox(type="auto",label="Pet Prediction")
]

examples = [
    "./examples/cat_1.jpg",
    "./examples/cat_2.jpg",
    "./examples/dog_1.jpg",
    "./examples/dog_2.jpg",
]



interface = gr.Interface(fn=predict, 
    inputs=inputs,
    outputs=outputs,
    title = title, 
    description=description,
    examples=examples
    )
interface.launch()