import os
os.system("hub install deoldify==1.0.1")
import gradio as gr
import paddlehub as hub
from pathlib import Path
from datetime import datetime
from typing import Optional

model = hub.Module(name='deoldify')
# NOTE:  Max is 45 with 11GB video cards. 35 is a good default
render_factor=40

def colorize_image(image):
    # now = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
    if not os.path.exists("./output"):
        os.makedirs("./output")
    # if image is not None:
    #     image.save(f"./output/{now}-input.jpg")    
    model.predict(image.name)
    return f'./output/DeOldify/'+Path(image.name).stem+".png"

# def inference(img, version, scale, weight):
def inferenceColorize(img, scale):
    # weight /= 100
    print(img, scale)
    '''
    try:
        extension = os.path.splitext(os.path.basename(str(img)))[1]
        img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
        if len(img.shape) == 3 and img.shape[2] == 4:
            img_mode = 'RGBA'
        elif len(img.shape) == 2:  # for gray inputs
            img_mode = None
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        else:
            img_mode = None
        h, w = img.shape[0:2]
        if h < 300:
            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)

        try:
            # _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
        except RuntimeError as error:
            print('Error', error)
        
        try:
            if scale != 2:
                interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
                h, w = img.shape[0:2]
                output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
        except Exception as error:
            print('wrong scale input.', error)
            
        if img_mode == 'RGBA':  # RGBA images should be saved in png format
            extension = 'png'
        else:
            extension = 'jpg'
    '''
    os.makedirs('output', exist_ok=True)
    model.predict(img.name)
    save_path = f'output/DeOldify/'+Path(img.name).stem+'.{extension}'
    
    '''
        cv2.imwrite(save_path, output)
        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
    '''
    return img, save_path
    '''
    except Exception as error:
        print('global exception', error)
        return None, None
    '''

colorize = gr.Interface(
    inferenceColorize, [
        gr.Markdown("Colorize old black & white photos"),
        gr.Image(type="filepath", label="Input"),
        
        gr.inputs.Number(label="Rescaling factor", default=2),
        gr.Button("Upscale & Colorize!")
        
        # gr.Slider(0, 100, label='Weight, only for CodeFormer. 0 for better quality, 100 for better identity', default=50)
    ], [
        gr.outputs.Image(type="numpy", label="Output (The colorized image)"),
        gr.outputs.File(label="Download the colorized image")
    ])
    
colorize.queue(concurrency_count=4)
colorize.launch()
'''
def create_interface():
    with gr.Blocks() as enhancer:
        gr.Markdown("Colorize old black & white photos")
        with gr.Column(scale=1, label = "Colorize photo", visible=True) as colorize_column:
            colorize_input = gr.Image(type="file")
            colorize_button = gr.Button("Colorize!")
            colorize_output = gr.outputs.Image(type="numpy", label="Output (The Colorized image)")
            download_colorize_button = gr.outputs.File(label="Download the colorized image!")
            colorize_button.click(colorize_image, inputs=colorize_input, outputs=colorize_output)
    enhancer.launch()

def run_code():
    create_interface()

# The main function
run_code()
'''