pytorch/DeepLabV3

1import torch2model = torch.hub.load('pytorch/vision:v0.9.0', 'deeplabv3_resnet101', pretrained=True)3model.eval()4# Download an example image from the pytorch website5import urllib6url, filename = ("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")7try: urllib.URLopener().retrieve(url, filename)8except: urllib.request.urlretrieve(url, filename)9# sample execution (requires torchvision)10from PIL import Image11from torchvision import transforms12import gradio as gr13import matplotlib.pyplot as plt14 15 16def inference(input_image):17    preprocess = transforms.Compose([18        transforms.ToTensor(),19        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),20    ])21 22    input_tensor = preprocess(input_image)23    input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model24 25    # move the input and model to GPU for speed if available26    if torch.cuda.is_available():27        input_batch = input_batch.to('cuda')28        model.to('cuda')29 30    with torch.no_grad():31        output = model(input_batch)['out'][0]32    output_predictions = output.argmax(0)33    # create a color pallette, selecting a color for each class34    palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])35    colors = torch.as_tensor([i for i in range(21)])[:, None] * palette36    colors = (colors % 255).numpy().astype("uint8")37 38    # plot the semantic segmentation predictions of 21 classes in each color39    r = Image.fromarray(output_predictions.byte().cpu().numpy()).resize(input_image.size)40    r.putpalette(colors)41    plt.imshow(r)42    return plt43 44title = "DEEPLABV3-RESNET101"45description = "demo for DEEPLABV3-RESNET101, DeepLabV3 model with a ResNet-101 backbone. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."46article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1706.05587'>Rethinking Atrous Convolution for Semantic Image Segmentation</a> | <a href='https://github.com/pytorch/vision/blob/master/torchvision/models/segmentation/deeplabv3.py'>Github Repo</a></p>"47 48gr.Interface(49    inference, 50    gr.inputs.Image(type="pil", label="Input"), 51    gr.outputs.Image(type="plot", label="Output"),52    title=title,53    description=description,54    article=article,55    examples=[56        ["dog.jpg"]57    ]).launch()