initial

app.py

@@ -0,0 +1,67 @@
+from matplotlib.pyplot import axis
+import gradio as gr
+import requests
+import numpy as np
+from torch import nn
+from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
+import requests
+
+url1 = 'https://cdn.pixabay.com/photo/2014/09/07/21/52/city-438393_1280.jpg'
+r = requests.get(url1, allow_redirects=True)
+open("city1.jpg", 'wb').write(r.content)
+url2 = 'https://cdn.pixabay.com/photo/2016/02/19/11/36/canal-1209808_1280.jpg'
+r = requests.get(url2, allow_redirects=True)
+open("city2.jpg", 'wb').write(r.content)
+
+def cityscapes_palette():
+    return [[128, 64, 128],[244, 35, 232],[70, 70, 70],[102, 102, 156],[190, 153, 153],
+            [153, 153, 153],[250, 170, 30],[220, 220, 0],[107, 142, 35],[152, 251, 152],
+            [70, 130, 180], [220, 20, 60], [255, 0, 0], [0, 0, 142], [0, 0, 70],
+            [0, 60, 100], [0, 80, 100], [0, 0, 230], [119, 11, 32]]
+
+model_name = "nvidia/segformer-b5-finetuned-cityscapes-1024-1024"
+
+feature_extractor = SegformerFeatureExtractor.from_pretrained(model_name)
+model = SegformerForSemanticSegmentation.from_pretrained(model_name)
+
+def inference(image):
+    inputs = feature_extractor(images=image.resize((1024,1024)), return_tensors="pt")
+    outputs = model(**inputs)
+
+    # First, rescale logits to original image size
+    logits = nn.functional.interpolate(outputs.logits.detach().cpu(),
+                    size=image.size[::-1], # (height, width)
+                    mode='bilinear',
+                    align_corners=False)
+
+    # Second, apply argmax on the class dimension
+    seg = logits.argmax(dim=1)[0]
+    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
+    palette = np.array(cityscapes_palette())
+    for label, color in enumerate(palette):
+        color_seg[seg == label, :] = color
+
+    # Show image + mask
+    img = np.array(image) * 0.5 + color_seg * 0.5
+    img = img.astype(np.uint8)
+
+    merged = np.concatenate((np.concatenate((np.array(image), color_seg), axis=1), np.concatenate((np.zeros_like(image), img), axis=1)), axis=0)
+    return merged
+
+
+
+title = "Transformers - SegFormer B5 @ 1024px"
+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"
+article = "<p style='text-align: center'><a href='https://huggingface.co/transformers/model_doc/segformer.html#segformerforsemanticsegmentation'>Segformer page</a></p>"
+
+gr.Interface(
+    inference, 
+    [gr.inputs.Image(type="pil", label="Input")], 
+    gr.outputs.Image(type="numpy", label="Output"),
+    title=title,
+    description=description,
+    article=article,
+    examples=[
+            ["city1.jpg"],
+            ["city2.jpg"]
+        ]).launch()

requirements.txt

@@ -0,0 +1,2 @@
+opencv-python-headless
+transformers