app.py

import os
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__)
# ウェイトファイルをダウンロード（存在しない場合）
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 .")
# モデルの初期化
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 '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:
            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:
            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>
    """

if __name__ == '__main__':
    
    app.run(host='0.0.0.0', port=5000, debug=True)