Spaces:
Runtime error
Runtime error
| from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation | |
| import gradio as gr | |
| import numpy as np | |
| extractor = AutoFeatureExtractor.from_pretrained("marcelarosalesj/my-segmentation-model") | |
| model = SegformerForSemanticSegmentation.from_pretrained("marcelarosalesj/my-segmentation-model") | |
| def classify(im): | |
| inputs = extractor(images=im, return_tensors="pt").to("cuda") | |
| outputs = model(**inputs) | |
| logits = outputs.logits | |
| classes = logits[0].detach().cpu().numpy().argmax(axis=0) | |
| colors = np.array([[128,0,0], [128,128,0], [0, 0, 128], [128,0,128], [0, 0, 0]]) | |
| return colors[classes] | |
| interface = gr.Interface(fn=classify, | |
| inputs="image", | |
| outputs="image") | |
| interface.launch(debug=True) |