# Code taken and  adopted from https://huggingface.co/spaces/havas79/Real-ESRGAN_Demo/blob/main/app.py - credit where credit is due.
# These models are exclusively models I trained myself.

# When using GPU option set gpu_id=0 in restorer class
# Running times with 1024 tile:
# Upscaling a 256x256 image to 512x512: 7sec with free cpu, 3sec with cpu upgrade, ~1sec with T4 Small GPU
# Upscaling a 1024x1024 image to 2048x2048: ~50sec with free cpu, ~45sec with cpu upgrade, ~3.5sec with T4 Small GPU!
# More powerful GPU's will make this even faster, especially if people are upscaling even larger images, also increase tilesize if more VRAM available so most upscales would be un-tiled and no possible seams problem and also faster
# Maybe add my weak, medium and strong otf handling models as a choice if they have images with strong noise, compression and deblur going on. Maybe also otf jpg one. I have trained multiple compact models, thought id keep it simple here though, these are my latest models.

import torch
print(f"Is CUDA available: {torch.cuda.is_available()}")
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")

import gradio as gr
import cv2
import numpy
import os
import random
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url

from realesrgan import RealESRGANer
from realesrgan.archs.srvgg_arch import SRVGGNetCompact

# global variables
last_file = None
img_mode = "RGB"
choice= "2x Fast Upscale"

# Upscale function
def upscale(img, choice):
    
    # variables
    global last_file
    model_path = ""

    # remove last upscale when doing this new upscale to prevent memory being full   
    if last_file:
        os.remove(last_file)
        last_file = None
        
    # There is no input image to upscale
    if not img:
        return error("Input Image not detected")
    
    # Get image dimenstions
    imgwidth, imgheight = img.size

    if choice == '2x Fast Upscale':
        model_path = os.path.join('weights', '2xNomosUni_compact_multijpg_ldl.pth')
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')

    elif choice == '2x Fast Real Upscale':
        model_path = os.path.join('weights', '2xNomosUni_compact_otf_medium.pth')
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')

    else:
        model_path = os.path.join('weights', '2xNomosUni_compact_multijpg_ldl.pth')
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')

    # Restorer Class
    upsampler = RealESRGANer(
        scale=2,
        model_path=model_path,
        dni_weight=None,
        model=model,
        tile=1024,
        tile_pad=10,
        pre_pad=10,
        half=False,
        gpu_id=0,
        #gpu_id=None, #when running on cpu
    )

    # Convert the input PIL image to cv2 image, so that it can be processed by realesrgan
    cv_img = numpy.array(img)
    img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)

    # Apply restoration
    try:
        output, _ = upsampler.enhance(img, 2)
    except RuntimeError as error:
        print('Error when upscaling', error)
    else:
        # Save restored image and return it to the output Image component
        extension = 'jpg'
        out_filename = f"output_{rnd_string(16)}.{extension}"
        cv2.imwrite(out_filename, output)
        last_file = out_filename
        return out_filename
        
# Get random image file name for newly created image
def rnd_string(x):
    """Returns a string of 'x' random characters
    """
    characters = "abcdefghijklmnopqrstuvwxyz_0123456789"
    result = "".join((random.choice(characters)) for i in range(x))
    return result

# Reset function to reset inputs and also delete last file
def reset():
    """
    Resets the Image components of the Gradio interface and deletes
    the last processed image
    """
    global last_file
    if last_file:
        os.remove(last_file)
        last_file = None
    return gr.update(value=None), gr.update(value=None)

# Check for transparency function
def has_transparency(img):
    """This function works by first checking to see if a "transparency" property is defined
    in the image's info -- if so, we return "True". Then, if the image is using indexed colors
    (such as in GIFs), it gets the index of the transparent color in the palette
    (img.info.get("transparency", -1)) and checks if it's used anywhere in the canvas
    (img.getcolors()). If the image is in RGBA mode, then presumably it has transparency in
    it, but it double-checks by getting the minimum and maximum values of every color channel
    (img.getextrema()), and checks if the alpha channel's smallest value falls below 255.
    https://stackoverflow.com/questions/43864101/python-pil-check-if-image-is-transparent
    """
    if img.info.get("transparency", None) is not None:
        return True
    if img.mode == "P":
        transparent = img.info.get("transparency", -1)
        for _, index in img.getcolors():
            if index == transparent:
                return True
    elif img.mode == "RGBA":
        extrema = img.getextrema()
        if extrema[3][0] < 255:
            return True
    return False

# Get image properties function
def image_properties(img):
    """
    Returns the dimensions (width and height) and color mode of the input image and
    also sets the global img_mode variable to be used by the realesrgan function
    """
    global img_mode
    if img:
        if has_transparency(img):
            img_mode = "RGBA"
        else:
            img_mode = "RGB"
        properties = f"Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
        return properties

# Gradio Interface, Event Listeners, launch command
def main():
    # Gradio Interface
    with gr.Blocks(title="Fast 2x Image Upscaler") as demo:
        
        gr.Markdown(
            """# <div align="center">Fast 2x Image Upscaler</div> 
            My self-trained general fast 2x upscaling model which handles some jpg compression and dof. If you have some noise and/or blur in your image, your can also try the Real model instead.
            """
        )

        with gr.Row():
            
            with gr.Group():
                input_image = gr.Image(label="Source Image", type="pil", image_mode="RGB")
                input_image_properties = gr.Textbox(label="Image Properties:", max_lines=1)
                choice = gr.Radio(choices=["2x Fast Upscale", "2x Fast Real Upscale"], label="Model Selection", info="Use Real if you have noise and/or blur in your image, otherwise leave default", value="2x Fast Upscale")
            
            with gr.Group():
                output_image = gr.Image(label="2x Upscaled Image", type="pil", image_mode="RGB", interactive=False, height=512)
                output_image_properties = gr.Textbox(label="Image Properties", max_lines=1)
       
        with gr.Row():
            upscale_btn = gr.Button("Upscale")
            reset_btn = gr.Button("Reset")
        
        with gr.Group():
            gr.Examples(examples="examples", inputs=[input_image, choice], outputs=output_image, fn=upscale, cache_examples=True)
            gr.Markdown(
                """ 

                My self trained 2x universal dof preserving upscaling model, which is a SRVGGNetCompact model. This handles jpg compression and preserves depth of field.     
                The Real model denoises, deblurs and handles jpg compression, and can be used for such a case like shown in the examples.   

                Upscaling a 256x256 image to 512x512: ~1sec if this space is running on T4 Small GPU (~7sec with free cpu, ~3 sec with cpu upgrade)   
                Upscaling a 1024x1024 image to 2048x2048: ~3.5sec if this space is running on T4 Small GPU (~50sec with free cpu, ~45sec with cpu upgrade)   

                These have been trained with [musl's neosr](https://github.com/muslll/neosr), [kim's datasetdestroyer](https://github.com/Kim2091/helpful-scripts/tree/main/Dataset%20Destroyer) and the real-esrgan otf pipeline.
                
                Find more of my models on [openmodeldb](https://openmodeldb.info/?q=Helaman)   
                    
                Author: Philip Hofmann   
                License: CC BY 4.0
                """)
            
        # Event listeners:
        input_image.change(fn=image_properties, inputs=input_image, outputs=input_image_properties)
        output_image.change(fn=image_properties, inputs=output_image, outputs=output_image_properties)
        upscale_btn.click(fn=upscale, inputs=[input_image, choice], outputs=output_image)
        reset_btn.click(fn=reset, inputs=[], outputs=[input_image, output_image])
            
    demo.launch()

if __name__ == "__main__":
    main()