app.py

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


@tf.function
def process_data(content):
    img = tf.io.decode_jpeg(content, channels=3)
    img = tf.image.resize_with_pad(img, SIZE, SIZE)
    img = tf.image.per_image_standardization(img)
    return img


def predict(img, size):
    with tf.device(DEVICE):
        img = tf.image.resize_with_pad(img, size, size)
        img = tf.image.per_image_standardization(img)
        data = process_data(image)
        data = tf.expand_dims(data, 0)
        out = model(data)[0]
    return dict((tags[i], out[i].numpy()) for i in range(len(tags)))


image = inputs.Image(label="Upload your image here!")
size = inputs.Number(label="Image resize", default=SIZE)

labels = outputs.Label(label="Tags")

gr.Interface(predict, inputs=[image, size], outputs=[labels])