pytorch/MobileNet_v2

1import torch2from PIL import Image3from torchvision import transforms4import gradio as gr5import os6 7 8os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")9 10model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=True)11model.eval()12 13torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")14 15 16def inference(input_image):17    preprocess = transforms.Compose([18        transforms.Resize(256),19        transforms.CenterCrop(224),20        transforms.ToTensor(),21        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),22    ])23    input_tensor = preprocess(input_image)24    input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model25 26    # move the input and model to GPU for speed if available27    if torch.cuda.is_available():28        input_batch = input_batch.to('cuda')29        model.to('cuda')30 31    with torch.no_grad():32        output = model(input_batch)33    # The output has unnormalized scores. To get probabilities, you can run a softmax on it.34    probabilities = torch.nn.functional.softmax(output[0], dim=0)35 36    # Read the categories37    with open("imagenet_classes.txt", "r") as f:38        categories = [s.strip() for s in f.readlines()]39    # Show top categories per image40    top5_prob, top5_catid = torch.topk(probabilities, 5)41    result = {}42    for i in range(top5_prob.size(0)):43        result[categories[top5_catid[i]]] = top5_prob[i].item()44    return result45 46inputs = gr.inputs.Image(type='pil')47outputs = gr.outputs.Label(type="confidences",num_top_classes=5)48 49title = "MOBILENET V2"50description = "Gradio demo for MOBILENET V2, Efficient networks optimized for speed and memory, with residual blocks. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."51article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1801.04381'>MobileNetV2: Inverted Residuals and Linear Bottlenecks</a> | <a href='https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenet.py'>Github Repo</a></p>"52 53examples = [54            ['dog.jpg']55]56gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch()