Spaces:
Sleeping
Sleeping
File size: 1,234 Bytes
271d34f 960819e 271d34f 591dbdc |
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 |
import gradio as gr
import numpy as np
from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
extractor = AutoFeatureExtractor.from_pretrained("d4niel92/my-segmentation-model")
model = SegformerForSemanticSegmentation.from_pretrained("d4niel92/my-segmentation-model")
class_labels = ['unlabeled', 'paved-area', 'dirt', 'grass', 'gravel', 'water', 'rocks', 'pool', 'vegetation', 'roof', 'wall', 'window', 'door', 'fence', 'fence-pole', 'person', 'dog', 'car', 'bicycle', 'tree', 'bald-tree', 'ar-marker', 'obstacle', 'conflicting']
def classify(im):
inputs = extractor(images=im, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
classes = logits[0].detach().cpu().numpy().argmax(axis=0)
annotations = []
for c, class_name in enumerate(class_labels):
mask = np.array(classes == c, dtype=int)
mask = np.repeat(np.repeat(mask, 5, axis=0), 5, axis=1) # scaling up the masks
annotations.append((mask, class_name))
im = np.repeat(np.repeat(im, 5, axis=0), 5, axis=1) # scaling up the images
return im, annotations
interface = gr.Interface(classify, gr.Image(type="pil", shape=(128, 128)), gr.AnnotatedImage())
interface.launch()
|