import os
import tensorflow as tf
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tensorflow import keras

import requests
import PIL
import io
import matplotlib.pyplot as plt


from keras_cv_attention_models import convnext

import gradio as gr

mm = convnext.ConvNeXtBase()

def inference(img):
    img = img
    imm = keras.applications.imagenet_utils.preprocess_input(img, mode='torch')
    image_input = tf.expand_dims(tf.image.resize(imm, mm.input_shape[1:3]), 0)
    
    pred = mm(image_input)
    pred_np = pred.numpy()
    
    pred_names = keras.applications.imagenet_utils.decode_predictions(pred.numpy())[0]
    
    result = {}
    
    for i in range(5):
        result[pred_names[i][1]] = pred_names[i][2].item()
    
    return result

inputs = gr.inputs.Image(type='numpy')
outputs = gr.outputs.Label(type="confidences",num_top_classes=5)

title = "ConvNeXt"
description = "Gradio demo for ConvNeXt. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."

article = "<p style='text-align: center'><a href='A ConvNet for the 2020s'>MobileNetV2: Inverted Residuals and Linear Bottlenecks</a> | <a href='https://github.com/facebookresearch/ConvNeXt'>Github Repo</a> | <a href='https://github.com/leondgarse/keras_cv_attention_models'>pretrained ConvNeXt model from keras_cv_attention_models</a></p>"


gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, analytics_enabled=False).launch()