Update app.py

app.py

@@ -1,225 +1,222 @@
-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
-from torchvision.transforms.functional import rgb_to_grayscale
-import spaces
-
-last_file = None
-img_mode = "RGBA"
-
-@spaces.GPU(duration=120)
-def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
-    """Real-ESRGAN function to restore (and upscale) images.
-    """
-    if not img:
-        return
-
-    # Define model parameters
-    if model_name == 'RealESRGAN_x4plus':  # x4 RRDBNet model
-        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':  # x4 RRDBNet model
-        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':  # x4 RRDBNet model with 6 blocks
-        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':  # x2 RRDBNet model
-        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-general-x4v3':  # x4 VGG-style model (S size)
-        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'
-        ]
-
-    # Determine model paths
-    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 will be updated
-            model_path = load_file_from_url(
-                url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
-
-    # Use dni to control the 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]
-
-    # Restorer Class
-    upsampler = RealESRGANer(
-        scale=netscale,
-        model_path=model_path,
-        dni_weight=dni_weight,
-        model=model,
-        tile=0,
-        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.3.pth',
-            upscale=outscale,
-            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
+import cv2
+import torch
+from flask import Flask, request, jsonify, send_file
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from gfpgan.utils import GFPGANer
+from realesrgan.utils import RealESRGANer
+import uuid
+import tempfile
+
+app = Flask(__name__)
+
+# モデルの初期化
+model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+model_path = 'realesr-general-x4v3.pth'
+half = True if torch.cuda.is_available() else False
+upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+
+os.makedirs('output', exist_ok=True)
+
+@app.route('/api/restore', methods=['POST'])
+def restore_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:
-        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
-        if img_mode == 'RGBA':  # RGBA images should be saved in png format
-            extension = 'png'
+        # リクエストからパラメータを取得
+        if 'file' not in request.files:
+            return jsonify({'error': 'No file uploaded'}), 400
+        
+        file = request.files['file']
+        version = request.form.get('version', 'v1.4')
+        scale = float(request.form.get('scale', 2))
+        # weight = float(request.form.get('weight', 50)) / 100  # CodeFormer用のweightパラメータが必要な場合
+        
+        # 一時ファイルに保存
+        temp_dir = tempfile.mkdtemp()
+        input_path = os.path.join(temp_dir, file.filename)
+        file.save(input_path)
+        
+        # 画像処理
+        extension = os.path.splitext(os.path.basename(str(input_path)))[1]
+        img = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)
+        
+        if len(img.shape) == 3 and img.shape[2] == 4:
+            img_mode = 'RGBA'
+        elif len(img.shape) == 2:
+            img_mode = None
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
         else:
-            extension = 'jpg'
-
-        out_filename = f"output_{rnd_string(8)}.{extension}"
-        cv2.imwrite(out_filename, output)
-        global last_file
-        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"
+            img_mode = None
+
+        h, w = img.shape[0:2]
+        if h < 300:
+            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
+
+        # バージョンに応じてモデルを選択
+        if version == 'v1.2':
+            face_enhancer = GFPGANer(
+                model_path='GFPGANv1.2.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
+        elif version == 'v1.3':
+            face_enhancer = GFPGANer(
+                model_path='GFPGANv1.3.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
+        elif version == 'v1.4':
+            face_enhancer = GFPGANer(
+                model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
+        elif version == 'RestoreFormer':
+            face_enhancer = GFPGANer(
+                model_path='RestoreFormer.pth', upscale=2, arch='RestoreFormer', channel_multiplier=2, bg_upsampler=upsampler)
+        elif version == 'CodeFormer':
+            face_enhancer = GFPGANer(
+                model_path='CodeFormer.pth', upscale=2, arch='CodeFormer', channel_multiplier=2, bg_upsampler=upsampler)
+        elif version == 'RealESR-General-x4v3':
+            face_enhancer = GFPGANer(
+                model_path='realesr-general-x4v3.pth', upscale=2, arch='realesr-general', channel_multiplier=2, bg_upsampler=upsampler)
+
+        # 画像を拡張
+        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+        
+        # スケール調整
+        if scale != 2:
+            interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
+            h, w = img.shape[0:2]
+            output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
+        
+        # 出力ファイルを保存
+        output_filename = f'output_{uuid.uuid4().hex}'
+        if img_mode == 'RGBA':
+            output_path = os.path.join('output', f'{output_filename}.png')
+            cv2.imwrite(output_path, output)
+            mimetype = 'image/png'
         else:
-            img_mode = "RGB"
-        properties = f"Resolution: Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
-        return properties
-
-
-def main():
-    # Gradio Interface
-    with gr.Blocks(title="Real-ESRGAN Gradio Demo", theme="dark") as demo:
-
-        gr.Markdown(
-            """# <div align="center"> Ilaria Upscaler 💖 </div>  
-
-        Do not use images over 750x750 especially with 4x the resolution upscaling, it will give you an error.
+            output_path = os.path.join('output', f'{output_filename}.jpg')
+            cv2.imwrite(output_path, output)
+            mimetype = 'image/jpeg'
         
-        Hugginface port of [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN).
-        """
-        )
-
-        with gr.Accordion("Upscaling option"):
-            with gr.Row():
-                model_name = gr.Dropdown(label="Upscaler model",
-                                         choices=["RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B",
-                                                  "RealESRGAN_x2plus", "realesr-general-x4v3"],
-                                         value="RealESRGAN_x4plus_anime_6B", show_label=True)
-                denoise_strength = gr.Slider(label="Denoise Strength",
-                                             minimum=0, maximum=1, step=0.1, value=0.5)
-                outscale = gr.Slider(label="Resolution upscale",
-                                     minimum=1, maximum=6, step=1, value=4, show_label=True)
-                face_enhance = gr.Checkbox(label="Face Enhancement (GFPGAN)",
-                )
+        # 結果を返す
+        return send_file(output_path, mimetype=mimetype, as_attachment=True, download_name=os.path.basename(output_path))
+    
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/')
+def index():
+    return """
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>Image Upscaling & Restoration API</title>
+        <style>
+            body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
+            .container { border: 1px solid #ddd; padding: 20px; border-radius: 5px; }
+            .form-group { margin-bottom: 15px; }
+            label { display: block; margin-bottom: 5px; }
+            input, select { width: 100%; padding: 8px; box-sizing: border-box; }
+            button { background-color: #4CAF50; color: white; padding: 10px 15px; border: none; border-radius: 4px; cursor: pointer; }
+            button:hover { background-color: #45a049; }
+            #result { margin-top: 20px; }
+            #preview { max-width: 100%; margin-top: 10px; }
+        </style>
+    </head>
+    <body>
+        <h1>Image Upscaling & Restoration API</h1>
+        <div class="container">
+            <form id="uploadForm" enctype="multipart/form-data">
+                <div class="form-group">
+                    <label for="file">Upload Image:</label>
+                    <input type="file" id="file" name="file" required>
+                </div>
+                <div class="form-group">
+                    <label for="version">Version:</label>
+                    <select id="version" name="version">
+                        <option value="v1.2">v1.2</option>
+                        <option value="v1.3">v1.3</option>
+                        <option value="v1.4" selected>v1.4</option>
+                        <option value="RestoreFormer">RestoreFormer</option>
+                        <option value="CodeFormer">CodeFormer</option>
+                        <option value="RealESR-General-x4v3">RealESR-General-x4v3</option>
+                    </select>
+                </div>
+                <div class="form-group">
+                    <label for="scale">Rescaling factor:</label>
+                    <input type="number" id="scale" name="scale" value="2" step="0.1" min="1" max="4" required>
+                </div>
+                <!-- CodeFormer用のweightパラメータが必要な場合 -->
+                <!--
+                <div class="form-group">
+                    <label for="weight">Weight (only for CodeFormer):</label>
+                    <input type="range" id="weight" name="weight" min="0" max="100" value="50">
+                    <span id="weightValue">50</span>
+                </div>
+                -->
+                <button type="submit">Process Image</button>
+            </form>
+            
+            <div id="result">
+                <h3>Result:</h3>
+                <div id="outputContainer" style="display: none;">
+                    <img id="preview" src="" alt="Processed Image">
+                    <a id="downloadLink" href="#" download>Download Image</a>
+                </div>
+            </div>
+        </div>
+        
+        <script>
+            document.getElementById('uploadForm').addEventListener('submit', function(e) {
+                e.preventDefault();
                 
-        with gr.Row():
-            with gr.Group():
-                input_image = gr.Image(label="Input Image", type="pil", image_mode="RGBA")
-                input_image_properties = gr.Textbox(label="Image Properties", max_lines=1)
-            output_image = gr.Image(label="Output Image", image_mode="RGBA")
-        with gr.Row():
-            reset_btn = gr.Button("Remove images")
-            restore_btn = gr.Button("Upscale")
-
-        # Event listeners:
-        input_image.change(fn=image_properties, inputs=input_image, outputs=input_image_properties)
-        restore_btn.click(fn=realesrgan,
-                          inputs=[input_image, model_name, denoise_strength, face_enhance, outscale],
-                          outputs=output_image)
-        reset_btn.click(fn=reset, inputs=[], outputs=[output_image, input_image])
-        # reset_btn.click(None, inputs=[], outputs=[input_image], _js="() => (null)\n")
-        # Undocumented method to clear a component's value using Javascript
-
-        gr.Markdown(
-            """Made with love by Ilaria 💖 | Support me on [Ko-Fi](https://ko-fi.com/ilariaowo) | Join [AI Hub](https://discord.gg/aihub)
-        """
-        )
-
-    demo.launch()
-
+                const formData = new FormData();
+                formData.append('file', document.getElementById('file').files[0]);
+                formData.append('version', document.getElementById('version').value);
+                formData.append('scale', document.getElementById('scale').value);
+                // formData.append('weight', document.getElementById('weight').value); // CodeFormer用
+                
+                fetch('/api/restore', {
+                    method: 'POST',
+                    body: formData
+                })
+                .then(response => {
+                    if (!response.ok) {
+                        return response.json().then(err => { throw new Error(err.error || 'Unknown error'); });
+                    }
+                    return response.blob();
+                })
+                .then(blob => {
+                    const url = URL.createObjectURL(blob);
+                    const preview = document.getElementById('preview');
+                    const downloadLink = document.getElementById('downloadLink');
+                    const outputContainer = document.getElementById('outputContainer');
+                    
+                    preview.src = url;
+                    downloadLink.href = url;
+                    downloadLink.download = 'restored_' + document.getElementById('file').files[0].name;
+                    outputContainer.style.display = 'block';
+                })
+                .catch(error => {
+                    alert('Error: ' + error.message);
+                });
+            });
+            
+            // CodeFormer用のweightパラメータが必要な場合
+            // document.getElementById('weight').addEventListener('input', function() {
+            //     document.getElementById('weightValue').textContent = this.value;
+            // });
+        </script>
+    </body>
+    </html>
+    """
 
-if __name__ == "__main__":
-    main()
+if __name__ == '__main__':
+    # ウェイトファイルをダウンロード（存在しない場合）
+    if not os.path.exists('realesr-general-x4v3.pth'):
+        os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+    if not os.path.exists('GFPGANv1.2.pth'):
+        os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.2.pth -P .")
+    if not os.path.exists('GFPGANv1.3.pth'):
+        os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P .")
+    if not os.path.exists('GFPGANv1.4.pth'):
+        os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
+    if not os.path.exists('RestoreFormer.pth'):
+        os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth -P .")
+    if not os.path.exists('CodeFormer.pth'):
+        os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/CodeFormer.pth -P .")
+    
+    app.run(host='0.0.0.0', port=5000, debug=True)