import torch from torchvision import transforms import gradio as gr import requests from PIL import Image #load models from pytorch hub model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5m-cls.pt').eval() # load from PyTorch Hub model.classify = True # load imagenet 1000 labels response = requests.get("https://git.io/JJkYN") labels = response.text.split("\n") def preprocess_image(inp): # Define the preprocessing steps preprocess = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # Apply the preprocessing steps to the image image = preprocess(inp) # Convert the image to a PyTorch tensor image = torch.tensor(image).unsqueeze(0) return image def predict(inp): with torch.no_grad(): prediction = torch.nn.functional.softmax(model(preprocess_image(inp))[0], dim=0) print(prediction) confidences = {labels[i]: float(prediction[i]) for i in range(1000)} return confidences gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs="label",labels=labels).launch() #outputs=gr.Label(num_top_classes=5))