app.py

# Code taken (and slightly adopted) from https://huggingface.co/spaces/havas79/Real-ESRGAN_Demo/blob/main/app.py - credit where credit is due. I am not showcasing code here, but demoing my own trained models ;)

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

last_file = None
img_mode = "RGBA"
firstrun = True

def realesrgan(img, model_name, face_enhance):

    # remove last upscale when doing this new upscale to prevent memory being full
    # added firstrun variable since I was running into global declaration problems?
    if not firstrun:
        global last_file
        
        if last_file:
            print(f"Deleting {last_file} ...")
            os.remove(last_file)
            last_file = None
            
    if not img:
        return

    imgwidth, imgheight = img.size

    if imgwidth > 1000 or imgheight > 1000:
        return error("Input Image too big")

    # Define model parameters
    if model_name == '4xNomos8kSC':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
    elif model_name == '4xHFA2k':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
    elif model_name == '4xLSDIR':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
    elif model_name == '4xLSDIRplusN':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
    elif model_name == '4xLSDIRplusC':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
    elif model_name == '4xLSDIRplusR':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
    elif model_name == '2xParimgCompact':
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')
        netscale = 2
    elif model_name == '2xHFA2kCompact':
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')
        netscale = 2
    elif model_name == '4xLSDIRCompactN':
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
        netscale = 4
    elif model_name == '4xLSDIRCompactC3':
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
        netscale = 4
    elif model_name == '4xLSDIRCompactR3':
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
        netscale = 4

    # Determine model paths
    model_path = os.path.join('weights', model_name + '.pth')

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

    # Use GFPGAN for face enhancement
    if face_enhance:
        from gfpgan import GFPGANer
        face_enhancer = GFPGANer(
            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth',
            upscale=netscale,
            arch='clean',
            channel_multiplier=2,
            bg_upsampler=upsampler)

    # 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:
        if face_enhance:
            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
        else:
            output, _ = upsampler.enhance(img, netscale)
    except RuntimeError as error:
        print('Error', error)
        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
    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)
        global last_file
        last_file = out_filename
        if firstrun:
            global firstrun
            firstrun = False
        return out_filename
        

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


def reset():
    """Resets the Image components of the Gradio interface and deletes
    the last processed image
    """
    global last_file
    if last_file:
        print(f"Deleting {last_file} ...")
        os.remove(last_file)
        last_file = None
    return gr.update(value=None), gr.update(value=None)


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


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


def main():
    # Gradio Interface
    with gr.Blocks(title="Self-trained ESRGAN models demo", theme="dark") as demo:
        
        gr.Markdown(
            """# <div align="center"> Upscale image </div> 
        Here I demo some of my self-trained models (only those trained on the SRVGGNet or RRDBNet archs). All my self-trained models can be found on the [openmodeldb](https://openmodeldb.info/?q=Helaman&sort=date-desc) or on [my github repo](https://github.com/phhofm/models).
        """
        )

        with gr.Group():
            with gr.Group():
                model_name = gr.Dropdown(label="Model to be used",
                                         choices=["2xHFA2kCompact", "2xParimgCompact", "4xLSDIRCompactN", "4xLSDIRCompactC3", "4xLSDIRCompactR3", "4xNomos8kSC", "4xHFA2k", "4xLSDIR", "4xLSDIRplusN", "4xLSDIRplusC", "4xLSDIRplusR"], value="4xLSDIRCompactC3", 
                                         info="See model infos at the bottom of this page")
                face_enhance = gr.Checkbox(label="Face Enhancement using GFPGAN (Doesn't work for anime images)",value=False, show_label=True)

        with gr.Group():
                input_image = gr.Image(label="Source Image", type="pil", image_mode="RGB")
                input_image_properties = gr.Textbox(label="Image Properties - Demo will throw error if input image has either width or height > 1000. Output download is jpg for smaller size. Use models locally to circument these limits.", max_lines=1)
        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.Markdown(""" **Examples are not pre-cached. You need to press the Upscale Button after selecting one**""")
            gr.Examples(examples="examples",inputs=[input_image, model_name, face_enhance],outputs=output_image,fn=realesrgan, cache_examples=False)
            gr.Markdown(
    """
    **Model infos**  
    *SRVGGNetCompact models - in general faster, but less powerful, than RRDBNet*  
    2xHFA2kCompact - use for upscaling anime images 2x, faster than 4xHFA2k but less powerful (SRVGGNetCompact)    
    2xParimgCompact - upscaling photos 2x, fast (SRVGGNetCompact)  
    4xLSDIRCompactN - upscale a good quality photo (no degradations) 4x, faster than 4xLSDIRN but less powerful (SRVGGNetCompact)  
    4xLSDIRCompactC3 - upscale a jpg compressed photo 4x, fast (SRVGGNetCompact)  
    4xLSDIRCompactR3 - upscale a degraded photo 4x, fast (SRVGGNetCompact) (too strong, best used for interpolation like 4xLSDIRCompactN (or C) 75% 4xLSDIRCompactR3 25% to add little degradation handling to the previous one)  



    *RRDBNet models - in general more powerful than SRVGGNetCompact, but very slow in this demo*  
    4xNomos8kSC - use for upscaling photos 4x or can also be tried out on anime   
    4xHFA2k - use for upscaling anime images 4x   
    4xLSDIR - upscale a good quality photo (no degradation) 4x  
    4xLSDIRplusN - upscale a good quality photo (no degradation) 4x  
    4xLSDIRplusC - upscale a jpg compressed photo 4x   
    4xLSDIRplusR - upscale a degraded photo 4x (too strong, best used for interpolation like 4xLSDIRplusN (or C) 75% 4xLSDIRplusR 25% to add little degradation handling to the previous one)    



    *Models that I trained that are not featured here, but available on [openmodeldb](https://openmodeldb.info/?q=Helaman&sort=date-desc) or on [github](https://github.com/phhofm/models):*   
    4xNomos8kSCHAT-L - Photo upscaler (handles little bit of jpg compression and blur), [HAT-L](https://github.com/XPixelGroup/HAT) model (good output but very slow since huge)     
    4xNomos8kSCHAT-S - Photo upscaler (handles little bit of jpg compression and blur), [HAT-S](https://github.com/XPixelGroup/HAT) model   
    4xNomos8kSCSRFormer - Photo upscaler (handles little bit of jpg compression and blur), [SRFormer](https://github.com/HVision-NKU/SRFormer) base model (also good and slow since also big model)    
    2xHFA2kAVCOmniSR - Anime frame upscaler that handles AVC (h264) video compression, [OmniSR](https://github.com/Francis0625/Omni-SR) model  
    2xHFA2kAVCOmniSR_Sharp - Anime frame upscaler that handles AVC (h264) video compression with sharper outputs, [OmniSR](https://github.com/Francis0625/Omni-SR) model  
    4xHFA2kAVCSRFormer_light - Anime frame upscaler that handles AVC (h264) video compression, [SRFormer](https://github.com/HVision-NKU/SRFormer) lightweight model  
    2xHFA2kAVCEDSR_M - Anime frame upscaler that handles AVC (h264) video compression, [EDSR-M](https://github.com/LimBee/NTIRE2017) model  
    2xHFA2kAVCCompact - Anime frame upscaler that handles AVC (h264) video compression, [SRVGGNet](https://github.com/xinntao/Real-ESRGAN) (also called Real-ESRGAN Compact) model  
    4xHFA2kLUDVAESwinIR_light - Anime image upscaler that handles various realistic degradations, [SwinIR](https://github.com/JingyunLiang/SwinIR) light model  
    4xHFA2kLUDVAEGRL_small - Anime image upscaler that handles various realistic degradations, [GRL](https://github.com/ofsoundof/GRL-Image-Restoration) small model  
    4xHFA2kLUDVAESRFormer_light - Anime image upscaler that handles various realistic degradations, [SRFormer](https://github.com/HVision-NKU/SRFormer) light model  
    4xLexicaHAT - An AI generated image upscaler, does not handle any degradations, [HAT](https://github.com/XPixelGroup/HAT) base model  
    2xLexicaSwinIR - An AI generated image upscaler, does not handle any degradations, [SwinIR](https://github.com/JingyunLiang/SwinIR) base model  
    2xLexicaRRDBNet - An AI generated image upscaler, does not handle any degradations, RRDBNet base model  
    2xLexicaRRDBNet_Sharp - An AI generated image upscaler with sharper outputs, does not handle any degradations, RRDBNet base model  
    4xHFA2kLUDVAESAFMN - dropped model since there were artifacts on the outputs when training with [SAFMN](https://github.com/sunny2109/SAFMN) arch  



    *The following are not models I had trained, but rather interpolations I had created, they are available on my [repo](https://github.com/phhofm/models) and can be tried out locally with chaiNNer:*  
    4xLSDIRplus (4xLSDIRplusC + 4xLSDIRplusR)
    4xLSDIRCompact3 (4xLSDIRCompactC3 + 4xLSDIRCompactR3)  
    4xLSDIRCompact2 (4xLSDIRCompactC2 + 4xLSDIRCompactR2)  
    4xInt-Ultracri (UltraSharp + Remacri)  
    4xInt-Superscri (Superscale + Remacri)  
    4xInt-Siacri(Siax + Remacri)  
    4xInt-RemDF2K (Remacri + RealSR_DF2K_JPEG)  
    4xInt-RemArt (Remacri + VolArt)  
    4xInt-RemAnime (Remacri + AnimeSharp)  
    4xInt-RemacRestore (Remacri + UltraMix_Restore)  
    4xInt-AnimeArt (AnimeSharp + VolArt)  
    2xInt-LD-AnimeJaNai (LD-Anime + AnimeJaNai)  
    """)

        # 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=realesrgan, inputs=[input_image, model_name, face_enhance], outputs=output_image)
        reset_btn.click(fn=reset, inputs=[], outputs=[output_image, input_image])

    demo.launch()


if __name__ == "__main__":
    main()