import gradio as gr
from fastai.vision.all import *
from PIL import Image, ImageDraw, ImageFont
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as mcolors

# import os
# Load a pre-trained image classification model
import pathlib
plt = platform.system()
if plt == 'Windows': pathlib.PosixPath = pathlib.WindowsPath
if plt == 'Linux': pathlib.WindowsPath = pathlib.PosixPath

root = os.path.dirname(__file__)

def acc_camvid():
    pass
    
model = load_learner("./models/model.pkl")

def process(imagep):
    # Predict and create the image as before
    pred = model.predict(imagep)
    a = pred[0]
    i = np.stack([(a**17) % 255, (a**11) % 255, (a**9) % 255], axis=2)
    img = Image.fromarray(i.astype('uint8'), mode='RGB')

    imagep = Image.open(imagep)
    # imagep = Image.fromarray(imagep)

    imagep = imagep.convert("RGBA")
    img = img.convert("RGBA")
    img = img.resize(imagep.size)
    alpha = Image.new('L', img.size, int(0.6 * 255))
    img.putalpha(alpha)
    combined = Image.alpha_composite(imagep, img)
    
    return combined.convert("RGB")



# Sample images for user to choose from
sample_images = ["./sample_images/street.jpg", "./sample_images/market.jpg","./sample_images/day.jpg"]

iface = gr.Interface(
    fn=process,
    inputs=gr.Image(label="Select an image", type="filepath"),
    outputs='image',
    live=False,
    title="Traffic image - semantic segmentation",
    description="Upload a road traffic image or select one of the examples below",
    examples=sample_images
)


iface.launch()