marcelarosalesj commited on
Commit
0988bc0
1 Parent(s): 2ab8068

Add application file

Browse files

Signed-off-by: Marcela Rosales <marcelarosalesj@gmail.com>

Files changed (1) hide show
  1. app.py +21 -0
app.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
2
+ import gradio as gr
3
+ import numpy as np
4
+
5
+ extractor = AutoFeatureExtractor.from_pretrained("marcelarosalesj/my-segmentation-model")
6
+ model = SegformerForSemanticSegmentation.from_pretrained("marcelarosalesj/my-segmentation-model")
7
+
8
+ def classify(im):
9
+ inputs = extractor(images=im, return_tensors="pt").to("cuda")
10
+ outputs = model(**inputs)
11
+ logits = outputs.logits
12
+ classes = logits[0].detach().cpu().numpy().argmax(axis=0)
13
+ colors = np.array([[128,0,0], [128,128,0], [0, 0, 128], [128,0,128], [0, 0, 0]])
14
+ return colors[classes]
15
+
16
+
17
+ interface = gr.Interface(fn=classify,
18
+ inputs="image",
19
+ outputs="image")
20
+
21
+ interface.launch(debug=True)