app.py

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()