pytorch/GhostNet

1import os2import torch3from PIL import Image4from torchvision import transforms5import gradio as gr6 7os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")8 9model = torch.hub.load('huawei-noah/ghostnet', 'ghostnet_1x', pretrained=True)10model.eval()11# Download an example image from the pytorch website12torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")13 14def inference(input_image):15    preprocess = transforms.Compose([16        transforms.Resize(256),17        transforms.CenterCrop(224),18        transforms.ToTensor(),19        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),20    ])21    input_tensor = preprocess(input_image)22    input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model23 24    # move the input and model to GPU for speed if available25    if torch.cuda.is_available():26        input_batch = input_batch.to('cuda')27        model.to('cuda')28 29    with torch.no_grad():30        output = model(input_batch)31    # The output has unnormalized scores. To get probabilities, you can run a softmax on it.32    probabilities = torch.nn.functional.softmax(output[0], dim=0)33 34    # Read the categories35    with open("imagenet_classes.txt", "r") as f:36        categories = [s.strip() for s in f.readlines()]37    # Show top categories per image38    top5_prob, top5_catid = torch.topk(probabilities, 5)39    result = {}40    for i in range(top5_prob.size(0)):41        result[categories[top5_catid[i]]] = top5_prob[i].item()42    return result43 44inputs = gr.inputs.Image(type='pil')45outputs = gr.outputs.Label(type="confidences",num_top_classes=5)46 47title = "GHOSTNET"48description = "Gradio demo for GHOSTNET, Efficient networks by generating more features from cheap operations. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."49article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1911.11907'>GhostNet: More Features from Cheap Operations</a> | <a href='https://github.com/huawei-noah/CV-Backbones'>Github Repo</a></p>"50 51examples = [52            ['dog.jpg']53]54gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch()