Spaces:
Sleeping
Sleeping
# pretrained Resnet-18 mode | |
import torch | |
model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).eval() | |
# define a function that takes in the user input, which in this case is an image, and returns the prediction. | |
'''The prediction should be returned as a dictionary whose keys are class name and values are confidence probabilities. | |
We will load the class names from this text file. | |
''' | |
import requests | |
from PIL import Image | |
from torchvision import transforms | |
# Download human-readable labels for ImageNet. | |
response = requests.get("https://git.io/JJkYN") | |
labels = response.text.split("\n") | |
def predict(inp): | |
inp = transforms.ToTensor()(inp).unsqueeze(0) | |
with torch.no_grad(): | |
prediction = torch.nn.functional.softmax(model(inp)[0], dim=0) | |
confidences = {labels[i]: float(prediction[i]) for i in range(1000)} | |
return confidences | |
'''The function converts the input image into a PIL Image and subsequently into a PyTorch tensor. | |
After processing the tensor through the model, it returns the predictions in the form of a dictionary named confidences. | |
The dictionary's keys are the class labels, and its values are the corresponding confidence probabilities. | |
In this section, we define a predict function that processes an input image to return prediction probabilities. | |
The function first converts the image into a PyTorch tensor and then forwards it through the pretrained model. | |
We use the softmax function in the final step to calculate the probabilities of each class. | |
The softmax function is crucial because it converts the raw output logits from the model, which can be any real number, into probabilities that sum up to 1. | |
This makes it easier to interpret the model’s outputs as confidence levels for each class.''' | |
# Creating a Gradio interface | |
import gradio as gr | |
gr.Interface(fn=predict, | |
inputs=gr.Image(type="pil"), # creates the component and handles the preprocessing to convert that to a PIL image | |
outputs=gr.Label(num_top_classes=3), #a Label, which displays the top labels in a nice form. Since we don't want to show all 1,000 class labels, we will customize it to show only the top 3 images by constructing it as | |
examples=["lion.jpg", "cheetah.jpg"]).launch() |