initial commit

.gitignore

@@ -0,0 +1,141 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Virtual environments
+venv/
+env/
+.venv/
+.env/
+ENV/
+.env.bak/
+venv.bak/
+env.bak/
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+cover/
+coverage/
+nosetests.xml
+*.log
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# dotenv
+.env
+
+# Virtualenv
+.venv/
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+.idea/
+
+# VS Code
+.vscode/
+
+# Local environment
+weights/

app.py

@@ -0,0 +1,132 @@
+import torch
+import cv2
+import numpy as np
+import os
+from PIL import Image
+import gradio as gr
+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
+
+def enhance_image(image, model_name, denoise_strength, outscale, tile, face_enhance, ext):
+    # Convert PIL image to OpenCV format (BGR)
+    img = np.array(image.convert('RGB'))[:, :, ::-1]  # Convert RGB to BGR
+    
+    # Model configuration
+    if model_name == 'RealESRGAN_x4plus':
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
+    elif model_name == 'RealESRNet_x4plus':
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
+    elif model_name == 'RealESRGAN_x4plus_anime_6B':
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
+    elif model_name == 'RealESRGAN_x2plus':
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        netscale = 2
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
+    elif model_name == 'realesr-animevideov3':
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth']
+    elif model_name == 'realesr-general-x4v3':
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+        netscale = 4
+        file_url = [
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
+        ]
+    else:
+        return "Error: Invalid model name."
+
+    # Download model weights if not available locally
+    model_path = os.path.join('weights', model_name + '.pth')
+    if not os.path.isfile(model_path):
+        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+        for url in file_url:
+            model_path = load_file_from_url(url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
+
+    # Handle denoise strength
+    dni_weight = None
+    if model_name == 'realesr-general-x4v3' and denoise_strength != 1:
+        wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
+        model_path = [model_path, wdn_model_path]
+        dni_weight = [denoise_strength, 1 - denoise_strength]
+
+    # Create upsampler
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=tile,
+        tile_pad=10,
+        pre_pad=0,
+        half=False,
+        gpu_id=None
+    )
+
+    # Handle 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.3.pth',
+            upscale=outscale,
+            arch='clean',
+            channel_multiplier=2,
+            bg_upsampler=upsampler
+        )
+
+    # Process the image
+    try:
+        if face_enhance:
+            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+        else:
+            output, _ = upsampler.enhance(img, outscale=outscale)
+    except RuntimeError as error:
+        return f'Error: {error}'
+    else:
+        # Convert BGR back to RGB
+        output = output[:, :, ::-1]  # Convert BGR to RGB
+        output = np.clip(output, 0, 255).astype(np.uint8)
+        output_image = Image.fromarray(output)
+        return output_image
+
+def create_gradio_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("## Real-ESRGAN Image Enhancement")
+        
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(type='pil', label="Input Image")
+                model_name = gr.Dropdown(["RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B", 
+                                        "RealESRGAN_x2plus", "realesr-animevideov3", "realesr-general-x4v3"], 
+                                        label="Model Name")
+                denoise_strength = gr.Slider(0, 1, value=0.5, step=0.1, label="Denoise Strength")
+                outscale = gr.Slider(1, 4, value=4, step=1, label="Output Scale")
+                tile = gr.Slider(128, 512, value=256, step=64, label="Tile Size")
+                face_enhance = gr.Checkbox(False, label="Enable Face Enhancement")
+                ext = gr.Dropdown(['auto', 'jpg', 'png'], value='auto', label="Output Extension")
+                
+                generate_button = gr.Button("Generate")
+
+            with gr.Column():
+                output_image = gr.Image(type='pil', label="Output Image")
+        
+        generate_button.click(
+            lambda image, model_name, denoise_strength, outscale, tile, face_enhance, ext: enhance_image(
+                image, model_name, denoise_strength, outscale, tile, face_enhance, ext
+            ), 
+            inputs=[image_input, model_name, denoise_strength, outscale, tile, face_enhance, ext], 
+            outputs=[output_image]
+        )
+        
+    demo.launch()
+
+if __name__ == '__main__':
+    create_gradio_interface()

requirements.txt

@@ -0,0 +1,9 @@
+torch
+opencv-python-headless
+numpy
+pillow
+gradio
+basicsr
+realesrgan
+gfpgan
+requests