app.py

import numpy as np
import gradio as gr
import torch
from torch import nn
from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor

#extractor = AutoFeatureExtractor.from_pretrained("andresgtn/segformer-b0-finetuned-ade-64-64-finetuned-semantic-sidewalk")
extractor = SegformerFeatureExtractor() 
model = SegformerForSemanticSegmentation.from_pretrained("andresgtn/segformer-b0-finetuned-ade-64-64-finetuned-semantic-sidewalk")

def rescale_output_image(logits, image):

  upsampled_logits = nn.functional.interpolate(
    logits,
    size=image.shape[::-1][1:][::-1], # (height, width)
    mode='bilinear',
    align_corners=False
    )
  pred_seg = upsampled_logits.argmax(dim=1)[0]
  return pred_seg
  
# classify function
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)
  #classes = rescale_output_image(logits, im).detach().cpu().numpy()
  classes = rescale_output_image(logits, im).detach().numpy()
  colors = np.array([[128,0,0], [128,128,0], [0, 0, 128], 	[128,0,128], [0, 0, 0]])
  return colors[classes]

# define gradio interface 
interface = gr.Interface(classify, gr.Image(), 'image')# FILL HERE

interface.launch(debug=True)