imkaushalpatel/GoogleNet

1import os2import torch3from PIL import Image4from torchvision import transforms5import gradio as gr6 7model = torch.hub.load('pytorch/vision:v0.9.0', 'googlenet', pretrained=True)8model.eval()9 10torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")11 12# sample execution (requires torchvision)13def inference(input_image):14    preprocess = transforms.Compose([15        transforms.Resize(256),16        transforms.CenterCrop(224),17        transforms.ToTensor(),18        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),19    ])20    input_tensor = preprocess(input_image)21    input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model22 23    # move the input and model to GPU for speed if available24    if torch.cuda.is_available():25        input_batch = input_batch.to('cuda')26        model.to('cuda')27 28    with torch.no_grad():29        output = model(input_batch)30    # The output has unnormalized scores. To get probabilities, you can run a softmax on it.31    probabilities = torch.nn.functional.softmax(output[0], dim=0)32    # Download ImageNet labels33    os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")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', label="Original Image")45outputs = gr.outputs.Image(type="pil", label="Output Image")46 47title = "GOOGLENET"48description = "Gradio demo for GOOGLENET, GoogLeNet was based on a deep convolutional neural network architecture codenamed Inception which won ImageNet 2014. 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/1409.4842'>Going Deeper with Convolutions</a> | <a href='https://github.com/pytorch/vision/blob/master/torchvision/models/googlenet.py'>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()