Update app.py

app.py

@@ -1,41 +1,16 @@
 import os
-
 import cv2
-import gradio as gr
 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
 
-os.system("pip freeze")
-# download weights
-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 .")
-
-torch.hub.download_url_to_file(
-    'https://thumbs.dreamstime.com/b/tower-bridge-traditional-red-bus-black-white-colors-view-to-tower-bridge-london-black-white-colors-108478942.jpg',
-    'a1.jpg')
-torch.hub.download_url_to_file(
-    'https://media.istockphoto.com/id/523514029/photo/london-skyline-b-w.jpg?s=612x612&w=0&k=20&c=kJS1BAtfqYeUDaORupj0sBPc1hpzJhBUUqEFfRnHzZ0=',
-    'a2.jpg')
-torch.hub.download_url_to_file(
-    'https://i.guim.co.uk/img/media/06f614065ed82ca0e917b149a32493c791619854/0_0_3648_2789/master/3648.jpg?width=700&quality=85&auto=format&fit=max&s=05764b507c18a38590090d987c8b6202',
-    'a3.jpg')
-torch.hub.download_url_to_file(
-    'https://i.pinimg.com/736x/46/96/9e/46969eb94aec2437323464804d27706d--victorian-london-victorian-era.jpg',
-    'a4.jpg')
+app = Flask(__name__)
 
-# background enhancer with RealESRGAN
+# モデルの初期化
 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
@@ -43,17 +18,30 @@ upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, ti
 
 os.makedirs('output', exist_ok=True)
 
-
-# def inference(img, version, scale, weight):
-def inference(img, version, scale):
-    # weight /= 100
-    print(img, version, scale)
+@app.route('/api/restore', methods=['POST'])
+def restore_image():
     try:
-        extension = os.path.splitext(os.path.basename(str(img)))[1]
-        img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
+        # リクエストからパラメータを取得
+        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:  # for gray inputs
+        elif len(img.shape) == 2:
             img_mode = None
             img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
         else:
@@ -63,80 +51,172 @@ def inference(img, version, scale):
         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)
+                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)
+                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)
+                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)
+                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)
+            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)
-
-        try:
-            # _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
-            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
-        except RuntimeError as error:
-            print('Error', error)
+            face_enhancer = GFPGANer(
+                model_path='realesr-general-x4v3.pth', upscale=2, arch='realesr-general', channel_multiplier=2, bg_upsampler=upsampler)
 
-        try:
-            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)
-        except Exception as error:
-            print('wrong scale input.', error)
-        if img_mode == 'RGBA':  # RGBA images should be saved in png format
-            extension = 'png'
+        # 画像を拡張
+        _, _, 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:
-            extension = 'jpg'
-        save_path = f'output/out.{extension}'
-        cv2.imwrite(save_path, output)
-
-        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
-        return output, save_path
-    except Exception as error:
-        print('global exception', error)
-        return None, None
+            output_path = os.path.join('output', f'{output_filename}.jpg')
+            cv2.imwrite(output_path, output)
+            mimetype = 'image/jpeg'
+        
+        # 結果を返す
+        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();
+                
+                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>
+    """
 
-title = "Image Upscaling & Restoration(esp. Face) using GFPGAN Algorithm"
-description = r"""Gradio demo for <a href='https://github.com/TencentARC/GFPGAN' target='_blank'><b>GFPGAN: Towards Real-World Blind Face Restoration and Upscalling of the image with a Generative Facial Prior</b></a>.<br>
-Practically the algorithm is used to restore your **old photos** or improve **AI-generated faces**.<br>
-To use it, simply just upload the concerned image.<br>
-"""
-article = r"""
-[![download](https://img.shields.io/github/downloads/TencentARC/GFPGAN/total.svg)](https://github.com/TencentARC/GFPGAN/releases)
-[![GitHub Stars](https://img.shields.io/github/stars/TencentARC/GFPGAN?style=social)](https://github.com/TencentARC/GFPGAN)
-[![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2101.04061)
-<center><img src='https://visitor-badge.glitch.me/badge?page_id=dj_face_restoration_GFPGAN' alt='visitor badge'></center>
-"""
-demo = gr.Interface(
-    inference, [
-        gr.inputs.Image(type="filepath", label="Input"),
-        # gr.inputs.Radio(['v1.2', 'v1.3', 'v1.4', 'RestoreFormer', 'CodeFormer'], type="value", default='v1.4', label='version'),
-        gr.inputs.Radio(['v1.2', 'v1.3', 'v1.4', 'RestoreFormer','CodeFormer','RealESR-General-x4v3'], type="value", default='v1.4', label='version'),
-        gr.inputs.Number(label="Rescaling factor", default=2),
-        # gr.Slider(0, 100, label='Weight, only for CodeFormer. 0 for better quality, 100 for better identity', default=50)
-    ], [
-        gr.outputs.Image(type="numpy", label="Output (The whole image)"),
-        gr.outputs.File(label="Download the output image")
-    ],
-    title=title,
-    description=description,
-    article=article,
-    # examples=[['AI-generate.jpg', 'v1.4', 2, 50], ['lincoln.jpg', 'v1.4', 2, 50], ['Blake_Lively.jpg', 'v1.4', 2, 50],
-    #           ['10045.png', 'v1.4', 2, 50]]).launch()
-    examples=[['a1.jpg', 'v1.4', 2], ['a2.jpg', 'v1.4', 2], ['a3.jpg', 'v1.4', 2],['a4.jpg', 'v1.4', 2]])
+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 .")
     
-demo.queue(concurrency_count=4)
-demo.launch()
+    app.run(host='0.0.0.0', port=5000, debug=True)