Spaces:
Sleeping
Sleeping
Daniel Meyer
commited on
Commit
•
271d34f
1
Parent(s):
e2b322d
Add application file
Browse files- app.py +27 -0
- requirements.txt +3 -0
app.py
ADDED
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import numpy as np
|
3 |
+
from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
|
4 |
+
|
5 |
+
extractor = AutoFeatureExtractor.from_pretrained("d4niel92/my-segmentation-model")
|
6 |
+
model = SegformerForSemanticSegmentation.from_pretrained("d4niel92/my-segmentation-model")
|
7 |
+
|
8 |
+
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']
|
9 |
+
|
10 |
+
|
11 |
+
def classify(im):
|
12 |
+
inputs = extractor(images=im, return_tensors="pt").to("cuda")
|
13 |
+
outputs = model(**inputs)
|
14 |
+
logits = outputs.logits
|
15 |
+
classes = logits[0].detach().cpu().numpy().argmax(axis=0)
|
16 |
+
annotations = []
|
17 |
+
for c, class_name in enumerate(class_labels):
|
18 |
+
mask = np.array(classes == c, dtype=int)
|
19 |
+
mask = np.repeat(np.repeat(mask, 5, axis=0), 5, axis=1) # scaling up the masks
|
20 |
+
annotations.append((mask, class_name))
|
21 |
+
im = np.repeat(np.repeat(im, 5, axis=0), 5, axis=1) # scaling up the images
|
22 |
+
return im, annotations
|
23 |
+
|
24 |
+
|
25 |
+
interface = gr.Interface(classify, gr.Image(type="pil", shape=(128, 128)), gr.AnnotatedImage())
|
26 |
+
|
27 |
+
interface.launch(share=True)
|
requirements.txt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
torch
|
2 |
+
transformers
|
3 |
+
Pillow
|