Update app.py

app.py

@@ -1,101 +1,45 @@
-import sys
 import os
-
-# 依存関係のインストール
-os.system("git clone https://github.com/sczhou/CodeFormer.git")
-os.system("cd CodeFormer && pip install -r requirements.txt")
-os.system("cd CodeFormer && python basicsr/setup.py develop")
-sys.path.append(os.path.abspath('CodeFormer'))
-sys.path.append(os.path.abspath('CodeFormer/CodeFormer'))
-# ウェイトファイルをダウンロード（毎回消えるので毎回必ず実行。）
-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 .")
-
 import cv2
 import torch
-from flask import Flask, request, jsonify, send_file
+from flask import Flask, request, jsonify, send_file, render_template_string
 from basicsr.archs.srvgg_arch import SRVGGNetCompact
 from gfpgan.utils import GFPGANer
 from realesrgan.utils import RealESRGANer
-import uuid
 import tempfile
-from torchvision.transforms.functional import normalize
-from torchvision import transforms
-from PIL import Image
-from basicsr.utils import img2tensor, tensor2img
-from facexlib.utils.face_restoration_helper import FaceRestoreHelper
-from codeformer.archs.codeformer_arch import CodeFormer
+import uuid
 
 app = Flask(__name__)
 
-# モデルの初期化
+# Initialize models
 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)
 
+# Ensure output directory exists
 os.makedirs('output', exist_ok=True)
 
-def restore_with_codeformer(img, scale=2, weight=0.5):
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    net = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(device)
-    net.load_state_dict(torch.load('CodeFormer.pth')['params_ema'])
-    net.eval()
-
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    img = Image.fromarray(img)
-    
-    face_helper = FaceRestoreHelper(
-        upscale_factor=scale, face_size=512, crop_ratio=(1, 1), use_parse=True,
-        device=device)
+# Download weights if not exists
+def download_weights():
+    weights = {
+        'realesr-general-x4v3.pth': 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth',
+        'GFPGANv1.2.pth': 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.2.pth',
+        'GFPGANv1.3.pth': 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
+        'GFPGANv1.4.pth': 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth',
+        'RestoreFormer.pth': 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth',
+        'CodeFormer.pth': 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/CodeFormer.pth'
+    }
     
-    face_helper.clean_all()
-    face_helper.read_image(img)
-    face_helper.get_face_landmarks_5(only_center_face=False, resize=640)
-    face_helper.align_warp_face()
+    for weight_file, url in weights.items():
+        if not os.path.exists(weight_file):
+            os.system(f"wget {url} -O {weight_file}")
 
-    for idx, cropped_face in enumerate(face_helper.cropped_faces):
-        cropped_face_t = img2tensor(cropped_face / 255.0, bgr2rgb=False, float32=True)
-        normalize(cropped_face_t, [0.5], [0.5], inplace=True)
-        cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
-        with torch.no_grad():
-            output = net(cropped_face_t, w=weight, adain=True)[0]
-        restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
-        face_helper.add_restored_face(restored_face)
-    
-    restored_img = face_helper.paste_faces_to_input_image()
-    return cv2.cvtColor(restored_img, cv2.COLOR_RGB2BGR)
+download_weights()
 
-@app.route('/api/restore', methods=['POST'])
-def restore_image():
+def process_image(img_path, version, scale, weight=0.5):
     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', 0.5))  # 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)
+        extension = os.path.splitext(os.path.basename(str(img_path)))[1]
+        img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
         
         if len(img.shape) == 3 and img.shape[2] == 4:
             img_mode = 'RGBA'
@@ -109,56 +53,59 @@ def restore_image():
         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)
-            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
         elif version == 'v1.3':
             face_enhancer = GFPGANer(
                 model_path='GFPGANv1.3.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
-            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
         elif version == 'v1.4':
             face_enhancer = GFPGANer(
                 model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
-            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
         elif version == 'RestoreFormer':
             face_enhancer = GFPGANer(
                 model_path='RestoreFormer.pth', upscale=2, arch='RestoreFormer', channel_multiplier=2, bg_upsampler=upsampler)
-            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
         elif version == 'CodeFormer':
-            output = restore_with_codeformer(img, scale=scale, weight=weight)
+            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,  map_location=torch.device('cpu'))
-            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+                model_path='realesr-general-x4v3.pth', upscale=2, arch='realesr-general', channel_multiplier=2, bg_upsampler=upsampler)
+
+        try:
+            if version == 'CodeFormer':
+                _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
+            else:
+                _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+        except RuntimeError as error:
+            print('Error', error)
+            raise Exception(f"Enhancement error: {str(error)}")
 
-        # スケール調整
-        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)
+        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)
+
+        # Save to temporary file
+        output_filename = f"output_{uuid.uuid4().hex}.jpg"
+        output_path = os.path.join('output', output_filename)
         
-        # 出力ファイルを保存
-        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'
+            cv2.imwrite(output_path, output, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
         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
+            cv2.imwrite(output_path, output, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
+
+        return output_path
+    except Exception as error:
+        print('Global exception', error)
+        raise Exception(f"Processing error: {str(error)}")
 
 @app.route('/')
 def index():
-    return """
+    return render_template_string('''
 <!DOCTYPE html>
 <html>
 <head>
@@ -278,9 +225,9 @@ def index():
                 reader.onload = function(e) {
                     const dataURL = e.target.result;
                     if (dataURL.length > 40) {
-                        filePreview = "${dataURL.substring(0, 20)}...${dataURL.substring(dataURL.length - 20)}";
+                        filePreview = `"${dataURL.substring(0, 20)}...${dataURL.substring(dataURL.length - 20)}"`;
                     } else {
-                        filePreview = "${dataURL}";
+                        filePreview = `"${dataURL}"`;
                     }
                     updateFetchCode(apiUrl, version, scale, weight, filePreview);
                 };
@@ -393,7 +340,42 @@ fetch('${apiUrl}', {
     </script>
 </body>
 </html>
-    """
+    ''')
+
+@app.route('/api/restore', methods=['POST'])
+def api_restore():
+    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', 0.5)) if version == 'CodeFormer' else None
+    
+    if file.filename == '':
+        return jsonify({'error': 'No selected file'}), 400
+    
+    try:
+        # Save uploaded file to temp location
+        temp_dir = tempfile.mkdtemp()
+        input_path = os.path.join(temp_dir, file.filename)
+        file.save(input_path)
+        
+        # Process image
+        output_path = process_image(input_path, version, scale, weight)
+        
+        # Return the processed image
+        return send_file(output_path, mimetype='image/jpeg')
+    
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+    
+    finally:
+        # Clean up temp files
+        if 'input_path' in locals() and os.path.exists(input_path):
+            os.remove(input_path)
+        if 'temp_dir' in locals() and os.path.exists(temp_dir):
+            os.rmdir(temp_dir)
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860, debug=True)
+    app.run(host='0.0.0.0', port=5000, debug=True)