import json

import gradio as gr
import tensorflow as tf
import tensorflow.keras as keras
from gradio import inputs, outputs

SIZE = 256
DEVICE = "/cpu:0"


with open("./tags.json", "rt", encoding="utf-8") as f:
    tags = json.load(f)


with tf.device(DEVICE):
    base_model = keras.applications.resnet.ResNet50(
        include_top=False, weights=None, input_shape=(SIZE, SIZE, 3)
    )
    model = keras.Sequential(
        [
            base_model,
            keras.layers.Conv2D(filters=len(tags), kernel_size=(1, 1), padding="same"),
            keras.layers.BatchNormalization(epsilon=1.001e-5),
            keras.layers.GlobalAveragePooling2D(name="avg_pool"),
            keras.layers.Activation("sigmoid"),
        ]
    )
    model.load_weights("tf_model.h5")


def predict(img, hide: float):
    with tf.device(DEVICE):
        img = tf.image.resize_with_pad(img, SIZE, SIZE)
        img = tf.image.per_image_standardization(img)
        data = tf.expand_dims(img, 0)
        out, *_ = model(data)
    return {
        tag: confidence
        for i, tag in enumerate(tags)
        if (confidence := float(out[i].numpy())) >= hide
    }


image = inputs.Image(label="Upload your image here!")
hide_threshold = inputs.Slider(
    label="Hide confidence lower than", default=0.5, maximum=1, minimum=0
)

labels = outputs.Label(label="Tags", type="confidences")

interface = gr.Interface(predict, inputs=[image, hide_threshold], outputs=[labels])
interface.launch()