import gradio as gr

import tensorflow as tf
import tensorflow_hub as hub

ckpt_type = '1k'
tf_hub_url = 'gs://cloud-tpu-checkpoints/efficientnet/v2/hub/efficientnetv2-s/classification'

m = hub.KerasLayer(tf_hub_url, trainable=False)
m.build([None, 224, 224, 3])  # Batch input shape.

def get_imagenet_labels(filename):
  labels = []
  with open(filename, 'r') as f:
    for line in f:
      labels.append(line.split('	')[1][:-1])  # split and remove line break.
  return labels
  
classes = get_imagenet_labels("imagenet1k_labels.txt")

def classify(image):
  image = tf.keras.preprocessing.image.img_to_array(image)
  image = (image - 128.) / 128.
  logits = m(tf.expand_dims(image, 0), False)
  pred = tf.keras.layers.Softmax()(logits)
  idx = tf.argsort(logits[0])[::-1][0].numpy()
  return classes[idx]
  
title = "Interactive demo: EfficientNetV2"
description = "Demo for Google's EfficientNetV2. EfficientNetV2 (accepted at ICML 2021) consists of convolutional neural networks that aim for fast training speed for relatively small-scale datasets, such as ImageNet1k."
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2104.00298'>EfficientNetV2: Smaller Models and Faster Training</a> | <a href='https://github.com/google/automl/tree/master/efficientnetv2'>Github Repo</a> | <a href='https://ai.googleblog.com/2021/09/toward-fast-and-accurate-neural.html'>Blog Post</a></p>"

iface = gr.Interface(fn=classify, 
  inputs=gr.inputs.Image(label="image", shape=(224,224)),
  outputs='text',
  title=title,
  description=description,
  enable_queue=True,
  examples=[['panda.jpeg'], ["llamas.jpeg"], ["hot_dog.png"]],
  article=article)

iface.launch()