Spaces:
Runtime error
Runtime error
| # 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" | |
| def realesrgan(img, model_name, face_enhance): | |
| global last_file | |
| # remove last upscale when doing this new upscale to prevent memory being full | |
| 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) | |
| last_file = out_filename | |
| 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() | |