import os import torch import torch.nn.functional as F import torchvision.transforms as T from uniformer import uniformer_small from imagenet_class_index import imagenet_classnames import gradio as gr from huggingface_hub import hf_hub_download def inference(img): image = img image_transform = T.Compose( [ T.Resize(224), T.CenterCrop(224), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] ) image = image_transform(image) # The model expects inputs of shape: B x C x H x W image = image.unsqueeze(0) prediction = model(image) prediction = F.softmax(prediction, dim=1).flatten() return {imagenet_id_to_classname[str(i)]: float(prediction[i]) for i in range(1000)} # Device on which to run the model # Set to cuda to load on GPU device = "cpu" model_path = hf_hub_download(repo_id="Sense-X/uniformer_image", filename="uniformer_small_in1k.pth") # Pick a pretrained model model = uniformer_small() state_dict = torch.load(model_path, map_location='cpu') model.load_state_dict(state_dict['model']) # Set to eval mode and move to desired device model = model.to(device) model = model.eval() # Create an id to label name mapping imagenet_id_to_classname = {} for k, v in imagenet_classnames.items(): imagenet_id_to_classname[k] = v[1] inputs = gr.inputs.Image(type='pil') label = gr.outputs.Label(num_top_classes=5) title = "UniFormer-S" description = "Gradio demo for UniFormer: To use it, simply upload your image, or click one of the examples to load them. Read more at the links below." article = "

UniFormer: Unifying Convolution and Self-attention for Visual Recognition | Github Repo

" gr.Interface( inference, inputs, outputs=label, title=title, description=description, article=article, examples=[['library.jpeg'], ['cat.png'], ['dog.png'], ['panda.png']] ).launch(enable_queue=True, cache_examples=True)