from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
import gradio as gr
import numpy as np

extractor = AutoFeatureExtractor.from_pretrained("marcelarosalesj/my-segmentation-model")
model = SegformerForSemanticSegmentation.from_pretrained("marcelarosalesj/my-segmentation-model")

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)
  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.launch(debug=True)