import gradio as gr

from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
from PIL import Image
import requests
from matplotlib import gridspec
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf


def sepia():
    feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
    model = SegformerForSemanticSegmentation.from_pretrained("segments-tobias/segformer-b0-finetuned-segments-sidewalk")
    url = "https://segmentsai-prod.s3.eu-west-2.amazonaws.com/assets/admin-tobias/439f6843-80c5-47ce-9b17-0b2a1d54dbeb.jpg"
    image = Image.open(requests.get(url, stream=True).raw)
    
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    
    logits = outputs.logits

demo = gr.Interface(fn=sepia,
                    inputs=gr.Image(shape=(400, 600)),
                    outputs=['plot'],
                    examples=[],
                    allow_flagging='never')


demo.launch()