# Code taken and  adopted from https://huggingface.co/spaces/havas79/Real-ESRGAN_Demo/blob/main/app.py - credit where credit is due.

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"

# 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 imgwidth > 512 or imgheight > 512:
        return error("Input Image too big")

    # Define model parameters
    if choice == '2x Fast Upscale':
        model_path = os.path.join('weights', '2xNomosUni_compact_multijpg.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 Upscale':
        model_path = os.path.join('weights', '2xNomosUni_esrgan_multijpg.pth')
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)

    # Restorer Class
    upsampler = RealESRGANer(
        scale=2,
        model_path=model_path,
        dni_weight=None,
        model=model,
        tile=128,
        tile_pad=10,
        pre_pad=10,
        half=False,
        gpu_id=None,
    )

    # 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="Self-trained 2x general upscaler models") as demo:
        
        gr.Markdown(
            """# <div align="center"> Upscale Image </div> 
            Here I demo two of my self-trained general 2x upscaler models which handle some jpg compression and dof. You can try more models on my older [huggingface space](https://huggingface.co/spaces/Phips/upscale_demo) or download models from [openmodeldb](https://openmodeldb.info/?sort=date-desc) and run them locally with [chaiNNer](https://github.com/chaiNNer-org/chaiNNer).
            """
        )

        with gr.Row():
            
            with gr.Group():
                input_image = gr.Image(label="Source Image", type="pil", image_mode="RGB")
                input_image_properties = gr.Textbox(label="Input Image needs to have width and hight smaller than 512. Use models unrestriced locally. Image Properties:", max_lines=1)
                choice = gr.Radio(choices=["2x Fast Upscale", "2x Upscale"], label="Model Selection", info="See infos at the bottom of this page", value="2x Fast Upscale")
            
            with gr.Group():
                output_image = gr.Image(label="Upscaled Image", type="pil", image_mode="RGB", interactive=False)
                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(
                """
                **Details**  
                
                These two 2x models are a Compact(SRVGGNet) for the '2x Fast Upscale' and an ESRGAN(RRDBNet) for the '2x Upscale' upscaling model which I recently trained and released (december 23)    
                2x Fast Upscale: [2xNomosUni_compact_multijpg](https://openmodeldb.info/models/2x-NomosUni-compact-multijpg)  
                2x Upscale: 2xNomosUni_esrgan_multijpg (not on openmodeldb yet, but on my [google drive](https://drive.google.com/drive/folders/12zKVS74mz0NtBKGlkx_r0ytUVoeDiTX3?usp=drive_link))  
                
                These two models are general upscalers with the goal to handle jpg compression and preserve depth of field for the most part.  
                
                I trained these using musl's [neosr](https://github.com/muslll/neosr) and Kim's [Dataset Destroyer](https://github.com/Kim2091/helpful-scripts/tree/main/Dataset%20Destroyer) on musl's universal nomos_uni dataset.  
                
                You can find more information about upscaling model training on the [training info repo](https://github.com/Upscale-Community/training-info).  
                If you have questions or simply be up to date on new community upscaling models released can of course also join our upscaling discord community [Enhance Everything](https://discord.gg/enhance-everything-547949405949657098) and watch the model-releases channel.  
                
                You can also run these two and way more models locally on your own GPU (so waay faster than on this cpu space) with [chaiNNer](https://github.com/chaiNNer-org/chaiNNer).  
                Here my [google drive folder](https://drive.google.com/drive/folders/1coYgan0VDo578MVO1LUpjpsxdY3LMyJW?usp=drive_link) with my self trained models.  
                Find a lot of models on our [Open Model Database](https://openmodeldb.info/?sort=date-desc).  
                
                I published my models under 'Helaman', but my real name is [Philip Hofmann](https://github.com/Phhofm). I got into upscaling in Summer 22 when Midjourney entered open beta.    
                After discovering and using [chaiNNer](https://github.com/chaiNNer-org/chaiNNer) with the [upscale wiki model database](https://upscale.wiki/w/index.php?title=Model_Database&oldid=1571), I thought that having visual outputs instead of only textual model descriptions would be nice, to not just read about but visually see what these models do.  
                So I gathered all of the upscaling models on there. Created a [youtube vid](https://youtu.be/0TYRDmQ5LZk) to compare ESRGAN models, made a [reddit post](https://www.reddit.com/r/StableDiffusion/comments/yev37i/comparison_of_upscaling_models_for_ai_generated/), and created a whole [Interactive Visual Comparison of Upscaling Models website](https://phhofm.github.io/upscale/) built with [vitepress](https://vitepress.dev/) (which had reached 1.0.0-alpha.26 at that time) to compare the visual outputs of over 300 different upsaling models.  
                Instead of only using and comparing upscaling models, I started learning about and training models myself, and released my very first upscaling model in march 23 called [4xLSDIRCompact](https://openmodeldb.info/models/4x-LSDIRCompact), a Compact model based on the [LSDIR](https://data.vision.ee.ethz.ch/yawli/) dataset.  
                Since then I have trained and released over 50 models of different networks/architectures like SRVGGNet, [RDDBNet](https://github.com/xinntao/Real-ESRGAN), [SwinIR](https://github.com/JingyunLiang/SwinIR), [SRFormer](https://github.com/HVision-NKU/SRFormer) (my model got mentioned on the readme), [GRL](https://github.com/ofsoundof/GRL-Image-Restoration), [OmniSR](https://github.com/Francis0625/Omni-SR), [EDSR](https://github.com/sanghyun-son/EDSR-PyTorch), [HAT](https://github.com/XPixelGroup/HAT), [DAT](https://github.com/zhengchen1999/DAT) (my model got mentioned on the readme), and [SPAN](https://github.com/hongyuanyu/SPAN).   
                Helped with testing and bug reporting of [neosr](https://github.com/muslll/neosr). Released the datasets FaceUp and SSDIR and made a [youtube video](https://www.youtube.com/watch?v=TBiVIzQkptI) about it.  
                It has been fun and fascinating so far :D  
                
                To keep up on latest sisr (single image super resolution) networks (architectures) one can follow the [papers with code image super resolution task](https://paperswithcode.com/task/image-super-resolution/latest) where interesting papers with code bases get published frequently. Also the [Awesome Image Super Resolution Repo](https://github.com/ChaofWang/Awesome-Super-Resolution). And of course be in our discord discussing about them in the image-networks channel. Its a fascinating and huge field with a lot of stuff to learn about.    
                 
                """)
            
        # 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()