salaz055 commited on
Commit
4030c59
1 Parent(s): f26c601

Create app.py

Browse files
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 torch
4
+
5
+ extractor = AutoFeatureExtractor.from_pretrained("salaz055/my_extractor_segmentation_model")
6
+ model = SegformerForSemanticSegmentation.from_pretrained("salaz055/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
+ interface = gr.Interface(fn = classify,
17
+ inputs = gr.Image(type = 'pil'),
18
+ outputs = 'image',
19
+ title = "Image Segmentation",
20
+ description = "Use for semantic image segmentation. Finetuned on sidewalk-semantic")
21
+ interface.launch()