from _typeshed import OpenBinaryModeUpdating
from types import resolve_bases
import requests
import gradio as gr
import torch
from timm import create_model
from tim.data import reslove_data_config
from timm.data.transformer import create_transform

IMAGENET_1k_URL = "https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt"
LABELS = requests.get(IMAGENET_1k_URL).text.strip().split('\n')

model = create_model('restnet50',pretrained=True)

transofrm = create_transform(**resolve_data_config{}, model=model)

model.eval()

def predict_fn(img):
    img = img.convert('RGB')
    img = transofrm(img).unsqueez(0)

    with torch.no_grad():
        out = model(img)
    probabilites = torch.nn.functional.softmax(out[0], dim=0)

    values , indices = torch.topk(probabilites, k=5)

    # return