| import datasets |
| import gradio as gr |
| import numpy as np |
| import torch |
| from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation |
|
|
| extractor = AutoFeatureExtractor.from_pretrained('saved_model_files') |
|
|
| labels = {0: 'road/sidewalk/path', 1: 'human', 2: 'vehicles', 3:'other objects', 4:'nature and background'} |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| model = SegformerForSemanticSegmentation.from_pretrained('saved_model_files', |
| num_labels = len(labels), |
| id2label={str(i) : c for i, c in enumerate(labels)}, |
| label2id={c: str(i) for i, c in enumerate(labels)}, |
| ignore_mismatched_sizes=True) |
| model.eval() |
| model.to(device) |
|
|
| def classify(im): |
| inputs = extractor(images=im, return_tensors='pt').to(device) |
| 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 = gr.Interface(fn=classify, inputs='image', outputs='image', examples=['1.png', '2.png', '3.png'], title='Image Segmentation App', description='Perform segmentation on pictures of outdoor scenes', flagging_dir='flagged_examples/') |
|
|
| interface.launch(debug=True) |
