from fastai.vision.all import *

import gradio as gr

from PIL import Image, ImageEnhance
import torchvision.transforms as T
import pathlib
plt = platform.system()
print(str(plt))
if plt == 'linux': pathlib.WindowsPath = pathlib.PosixPath

def name_to_hrs (r): return float(round(float(os.path.basename(r)[0:-4].split("_")[1][1:])*(minutes/60)+5,2))
def validation_split (r):  return os.path.basename(r)[0:-4].split("_")[3] == "R0003" or os.path.basename(r)[0:-4].split("_")[3] == "R0006"
def get_label_filename(name): return path/'labels'/f'{name.stem}_annotationLabels.tif'
zebrafish_age_predictor = load_learner('FishAge.pkl')
zebrafish_classifier = load_learner('FishSegmentation.pkl')
examples = ["5hr.tif", "12hr.tif", "24hr.tif"]
def process_zebrafish_image(img):

    # out_pl.clear_output()
    # enhancer = ImageEnhance.Brightness(img)
    # factor = 10 
    # im_output = enhancer.enhance(factor)
    # with out_pl: display(im_output.to_thumb(256,256))
    age,tensor, tensor=zebrafish_age_predictor.predict(img)
    
    pred,pred_idx,probs=zebrafish_classifier.predict(img)

    # with out_pl_mask: pred.show(alpha=1, vmin=0, vmax=3, title='mask')
    # lbl_pred.value = f'Predicted age: {age[0]};'
    #return dict(zip(pred, map(float,age)))
    image_out = T.ToPILImage(pred)
    text_out = dict(zip("Age prediction", map(float,age)))
    return (image_out, text_out )

intf = gr.Interface(fn=process_zebrafish_image, inputs=gr.inputs.Image(shape=(512, 512)), outputs=[gr.outputs.Image(), gr.outputs.Label()]).launch(share=True)
intf.launch(inline=False)