from transformers import SegformerFeatureExtractor, SegformerForImageClassification from PIL import Image import requests import gradio as gr def seg(image): feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/mit-b0") model = SegformerForImageClassification.from_pretrained("nvidia/mit-b0") print(model) inputs = feature_extractor(images=image, return_tensors="pt") outputs = model(**inputs) logits = outputs.logits # model predicts one of the 1000 ImageNet classes predicted_class_idx = logits.argmax(-1).item() return model.config.id2label[predicted_class_idx] iface = gr.Interface(fn=seg, inputs=gr.inputs.Image(type='pil'), outputs='label') iface.launch()