from fastai.vision.all import *
import gradio as gr
import numpy as np
from skimage import color
import pathlib

# temp = pathlib.PosixPath
# pathlib.PosixPath = pathlib.WindowsPath

def get_label(filepath):
    return None

learn = load_learner("resnet18_5epochs.pkl")

def segment(img):
	mask, _, _ = learn.predict(img)
	result = color.label2rgb(np.array(mask),img, colors=['red', 'green', 'blue', 'chocolate', 'cyan', 'purple', 'yellow', 'darkorange', 'turquoise', 'lime'] , bg_label=0)
	return result

TITLE = "Pixel-level Face Segmentation"
DESCRIPTION = "Resnet18 fine-tuned for face features segmentation via fastai and [Mut1ny's Face Segmentation Dataset](https://store.mut1ny.com/product/face-head-segmentation-dataset-community-edition?v=a04a1cb60875).\
	\n I am now trying to reproduce this fine-tuning in Pytorch, you can check my progress on this [Kaggle notebook](https://www.kaggle.com/eliotlacroix/pixel-level-face-segmentation)"
INPUT = gr.Image(shape=(224, 224))
OUTPUT = gr.ImageMask(shape=(224, 224), label='Resnet18')
EXAMPLES = [f for f in pathlib.Path('examples').iterdir()]

intf = gr.Interface(fn=segment, inputs=INPUT, outputs=OUTPUT, examples=EXAMPLES, title=TITLE, description=DESCRIPTION)
intf.launch(inline=False)