SegFormer / app.py
Karol Majek
nonsense 9dcc578
1 from matplotlib.pyplot import axis
2 import gradio as gr
3 import requests
4 import numpy as np
5 from torch import nn
6 from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
7 import requests
8
9 url1 = 'https://cdn.pixabay.com/photo/2014/09/07/21/52/city-438393_1280.jpg'
10 r = requests.get(url1, allow_redirects=True)
11 open("city1.jpg", 'wb').write(r.content)
12 url2 = 'https://cdn.pixabay.com/photo/2016/02/19/11/36/canal-1209808_1280.jpg'
13 r = requests.get(url2, allow_redirects=True)
14 open("city2.jpg", 'wb').write(r.content)
15
16 def cityscapes_palette():
17 return [[128, 64, 128],[244, 35, 232],[70, 70, 70],[102, 102, 156],[190, 153, 153],
18 [153, 153, 153],[250, 170, 30],[220, 220, 0],[107, 142, 35],[152, 251, 152],
19 [70, 130, 180], [220, 20, 60], [255, 0, 0], [0, 0, 142], [0, 0, 70],
20 [0, 60, 100], [0, 80, 100], [0, 0, 230], [119, 11, 32]]
21
22 model_name = "nvidia/segformer-b5-finetuned-cityscapes-1024-1024"
23
24 feature_extractor = SegformerFeatureExtractor.from_pretrained(model_name)
25 model = SegformerForSemanticSegmentation.from_pretrained(model_name)
26
27 def inference(image):
28 image = image.resize((1024,1024))
29 inputs = feature_extractor(images=image, return_tensors="pt")
30 outputs = model(**inputs)
31
32 # First, rescale logits to original image size
33 logits = nn.functional.interpolate(outputs.logits.detach().cpu(),
34 size=image.size[::-1], # (height, width)
35 mode='bilinear',
36 align_corners=False)
37
38 # Second, apply argmax on the class dimension
39 seg = logits.argmax(dim=1)[0]
40 color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
41 palette = np.array(cityscapes_palette())
42 for label, color in enumerate(palette):
43 color_seg[seg == label, :] = color
44
45 # Show image + mask
46 img = np.array(image) * 0.5 + color_seg * 0.5
47 img = img.astype(np.uint8)
48
49 merged = np.concatenate((np.concatenate((np.array(image), color_seg), axis=1), np.concatenate((np.zeros_like(image), img), axis=1)), axis=0)
50 return merged
51
52
53
54 title = "Transformers - SegFormer B5 @ 1024px"
55 description = "demo for SegFormer. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below.\nModel: nvidia/segformer-b5-finetuned-cityscapes-1024-1024"
56 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2105.15203'>SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers</a> | <a href='https://huggingface.co/transformers/model_doc/segformer.html#segformerforsemanticsegmentation'>Segformer page</a></p>"
57
58 gr.Interface(
59 inference,
60 [gr.inputs.Image(type="pil", label="Input")],
61 gr.outputs.Image(type="numpy", label="Output"),
62 title=title,
63 description=description,
64 article=article,
65 examples=[
66 ["city1.jpg"],
67 ["city2.jpg"]
68 ]).launch()
69