import gradio as gr
import subprocess
import os
import numpy as np
from PIL import Image

os.environ['data_raw'] = 'data_raw/'
os.environ['nnUNet_raw_data_base'] = 'nnUNet_raw_data_base/'
os.environ['nnUNet_preprocessed'] = 'nnUNet_preprocessed/'
os.environ['RESULTS_FOLDER'] = 'calvingfronts/'

model_path = 'Task500_Glacier_zonefronts'


def run_front_detection(input_img):
    image_gray = input_img.convert("L")
    image_gray.save('data_raw/test.png')
    subprocess.run(
        ['python3', 'nnunet/dataset_conversion/Task500_Glacier_inference.py', '-data_percentage', '100', '-base',
         os.environ['data_raw']])
    cmd = [
        'python3', 'nnunet/inference/predict_simple.py',
        '-i', os.path.join(os.environ['nnUNet_raw_data_base'], 'nnUNet_raw_data/Task500_Glacier_zonefronts/imagesTs/'),
        '-o', os.path.join(os.environ['RESULTS_FOLDER'], 'fold_0'),
        '-t', '500','-m','2d','-f','0','-p', 'nnUNetPlansv2.1', '-tr','nnUNetTrainerV2', '-model_folder_name',
        model_path
    ]
    subprocess.run(cmd)
    subprocess.run(['python3', 'nnunet/dataset_conversion/Task500_Glacier_reverse.py', '-i',
                    os.path.join(os.environ['RESULTS_FOLDER'], 'fold_0')])
    os.remove('data_raw/test.png')
    os.remove('nnUNet_raw_data_base/nnUNet_raw_data/Task500_Glacier_zonefronts/imagesTs/test_0000.nii.gz')


    front = Image.open(os.path.join(os.environ['RESULTS_FOLDER'], 'tifs/test_front.png'))
    os.remove(os.path.join(os.environ['RESULTS_FOLDER'],'fold_0/test.nii.gz'))
    os.remove(os.path.join(os.environ['RESULTS_FOLDER'], 'tifs/test_front.png'))

    return front

demo = gr.Interface(run_front_detection, gr.Image(type='pil'), "image",interface_size=(1000, 1000))
demo.launch()