import matplotlib.pyplot as plt
import gradio as gr
from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
import torch
import numpy as np

extractor = AutoFeatureExtractor.from_pretrained("hshetty/segmentation-model-finetuned-on-semantic-sidewalk-3e-4-e5")
model = SegformerForSemanticSegmentation.from_pretrained("hshetty/segmentation-model-finetuned-on-semantic-sidewalk-3e-4-e5")

def classify(im):
  inputs = extractor(images=im, return_tensors="pt")
  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",
             title="Self Driving Car App- Semantic Segmentation",
             description="This is a self driving car app using Semantic Segmentation as part of week 2 end to end vision application project on CoRise by Abubakar Abid!",
             examples=["/static-proxy?url=https%3A%2F%2Fdatasets-server.huggingface.co%2Fassets%2Fsegments%2Fsidewalk-semantic%2F--%2Fsegments--sidewalk-semantic-2%2Ftrain%2F3%2Fpixel_values%2Fimage.jpg",
                       "/static-proxy?url=https%3A%2F%2Fdatasets-server.huggingface.co%2Fassets%2Fsegments%2Fsidewalk-semantic%2F--%2Fsegments--sidewalk-semantic-2%2Ftrain%2F5%2Fpixel_values%2Fimage.jpg",
                       "/static-proxy?url=https%3A%2F%2Fdatasets-server.huggingface.co%2Fassets%2Fsegments%2Fsidewalk-semantic%2F--%2Fsegments--sidewalk-semantic-2%2Ftrain%2F20%2Fpixel_values%2Fimage.jpg"])
# FILL HERE

interface.launch()