Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,853 Bytes
b42e4a2 511a9f4 b42e4a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
import gradio as gr
import spaces
import torch
import numpy as np
from PIL import Image
from transformers import pipeline
import matplotlib.pyplot as plt
import io
model_pipeline = pipeline("image-segmentation", model="sergiopaniego/test-segformer-b0-segments-sidewalk-finetuned")
id2label = {0: 'unlabeled', 1: 'flat-road', 2: 'flat-sidewalk', 3: 'flat-crosswalk', 4: 'flat-cyclinglane', 5: 'flat-parkingdriveway', 6: 'flat-railtrack', 7: 'flat-curb', 8: 'human-person', 9: 'human-rider', 10: 'vehicle-car', 11: 'vehicle-truck', 12: 'vehicle-bus', 13: 'vehicle-tramtrain', 14: 'vehicle-motorcycle', 15: 'vehicle-bicycle', 16: 'vehicle-caravan', 17: 'vehicle-cartrailer', 18: 'construction-building', 19: 'construction-door', 20: 'construction-wall', 21: 'construction-fenceguardrail', 22: 'construction-bridge', 23: 'construction-tunnel', 24: 'construction-stairs', 25: 'object-pole', 26: 'object-trafficsign', 27: 'object-trafficlight', 28: 'nature-vegetation', 29: 'nature-terrain', 30: 'sky', 31: 'void-ground', 32: 'void-dynamic', 33: 'void-static', 34: 'void-unclear'}
sidewalk_palette = [
[0, 0, 0], # unlabeled
[216, 82, 24], # flat-road
[255, 255, 0], # flat-sidewalk
[125, 46, 141], # flat-crosswalk
[118, 171, 47], # flat-cyclinglane
[161, 19, 46], # flat-parkingdriveway
[255, 0, 0], # flat-railtrack
[0, 128, 128], # flat-curb
[190, 190, 0], # human-person
[0, 255, 0], # human-rider
[0, 0, 255], # vehicle-car
[170, 0, 255], # vehicle-truck
[84, 84, 0], # vehicle-bus
[84, 170, 0], # vehicle-tramtrain
[84, 255, 0], # vehicle-motorcycle
[170, 84, 0], # vehicle-bicycle
[170, 170, 0], # vehicle-caravan
[170, 255, 0], # vehicle-cartrailer
[255, 84, 0], # construction-building
[255, 170, 0], # construction-door
[255, 255, 0], # construction-wall
[33, 138, 200], # construction-fenceguardrail
[0, 170, 127], # construction-bridge
[0, 255, 127], # construction-tunnel
[84, 0, 127], # construction-stairs
[84, 84, 127], # object-pole
[84, 170, 127], # object-trafficsign
[84, 255, 127], # object-trafficlight
[170, 0, 127], # nature-vegetation
[170, 84, 127], # nature-terrain
[170, 170, 127], # sky
[170, 255, 127], # void-ground
[255, 0, 127], # void-dynamic
[255, 84, 127], # void-static
[255, 170, 127], # void-unclear
]
def get_output_figure(pil_img, results):
plt.figure(figsize=(16, 10))
plt.imshow(pil_img)
image_array = np.array(pil_img)
segmentation_map = np.zeros_like(image_array)
for result in results:
mask = np.array(result['mask'])
label = result['label']
label_index = list(id2label.values()).index(label)
color = sidewalk_palette[label_index]
for c in range(3):
segmentation_map[:, :, c] = np.where(mask, color[c], segmentation_map[:, :, c])
plt.imshow(segmentation_map, alpha=0.5)
plt.axis('off')
return plt.gcf()
@spaces.GPU
def detect(image):
results = model_pipeline(image)
print(results)
output_figure = get_output_figure(image, results)
buf = io.BytesIO()
output_figure.savefig(buf, bbox_inches='tight')
buf.seek(0)
output_pil_img = Image.open(buf)
return output_pil_img
with gr.Blocks() as demo:
gr.Markdown("# Semantic segmentation with SegFormer fine tuned on segments/sidewalk")
gr.Markdown(
"""
This application uses a fine tuned SegFormer for sematic segmenation over an input image.
This version was trained using segments/sidewalk dataset.
You can load an image and see the predicted segmentation.
"""
)
gr.Interface(
fn=detect,
inputs=gr.Image(label="Input image", type="pil"),
outputs=[
gr.Image(label="Output prediction", type="pil")
]
)
demo.launch(show_error=True) |