File size: 1,601 Bytes
d7afb92 4f47f02 d7afb92 9cf9edd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
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/') # FILL HERE
interface.launch(debug=True) |