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
# ウェイトファイルをダウンロード（存在しない場合）
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 = 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 '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; }
        #apiUsage { background-color: #f5f5f5; padding: 15px; border-radius: 5px; margin-top: 20px; font-family: monospace; white-space: pre-wrap; }
        #apiUsage h3 { margin-top: 0; }
        #formDataPreview { max-height: 200px; overflow-y: auto; margin-bottom: 10px; }
        .code-block { background-color: #f8f8f8; padding: 10px; border-radius: 4px; border-left: 3px solid #4CAF50; }
        .comment { color: #666; font-style: italic; }
    </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>
            <button type="submit">Process Image</button>
        </form>
        
        <div id="apiUsage">
            <h3>API Usage:</h3>
            <p>Here's how to call the API with the current form data:</p>
            <div id="formDataPreview">Fill out the form to see the API call code...</div>
            <div id="fetchCode"></div>
        </div>
        
        <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>
        // フォームの変更を監視してAPI使用例を更新
        function updateApiUsage() {
            const form = document.getElementById('uploadForm');
            const fileInput = document.getElementById('file');
            const version = document.getElementById('version').value;
            const scale = document.getElementById('scale').value;
            
            // フォームデータのJSONを作成
            const formJson = {
                version: version,
                scale: parseFloat(scale)
            };
            
            if (fileInput.files.length > 0) {
                const file = fileInput.files[0];
                formJson.file = {
                    name: file.name,
                    type: file.type,
                    size: file.size + ' bytes'
                };
            } else {
                formJson.file = 'No file selected';
            }
            
            updateFormDataPreview(formJson);
            updateFetchCode();
        }
        
        function updateFormDataPreview(formJson) {
            const previewDiv = document.getElementById('formDataPreview');
            previewDiv.innerHTML = '<strong>Form data (JSON):</strong><br>' + 
                JSON.stringify(formJson, null, 2) + 
                ' <span class="comment">// この値がフォーム部分です。</span>';
        }
        
        function updateFetchCode() {
            // 現在のURLからベースURLを取得（パス、パラメータ、ハッシュを含めない）
            const baseUrl = window.location.origin;
            const apiUrl = baseUrl + '/api/restore';
            
            const fetchCodeDiv = document.getElementById('fetchCode');
            fetchCodeDiv.innerHTML = `
<div class="code-block">
// JavaScript fetch code:
fetch('${apiUrl}', {
    method: 'POST',
    headers: {
        'Content-Type': 'multipart/form-data'
    },
    body: formData
})
.then(response => {
    if (!response.ok) {
        return response.json().then(err => { throw new Error(err.error || 'Unknown error'); });
    }
    return response.blob();
})
.then(blob => {
    // Process the returned image blob
    const url = URL.createObjectURL(blob);
    console.log('Image processed successfully', url);
})
.catch(error => {
    console.error('Error:', error.message);
});
</div>`;
        }
        
        // フォーム要素の変更を監視
        document.getElementById('file').addEventListener('change', updateApiUsage);
        document.getElementById('version').addEventListener('change', updateApiUsage);
        document.getElementById('scale').addEventListener('input', updateApiUsage);
        
        // 初期表示
        updateApiUsage();
        
        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);
            
            // 現在のURLからベースURLを取得（パス、パラメータ、ハッシュを含めない）
            const baseUrl = window.location.origin;
            const apiUrl = baseUrl + '/api/restore';
            
            fetch(apiUrl, {
                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);
            });
        });
    </script>
</body>
</html>


    """

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