Update app.py

app.py

@@ -7,7 +7,19 @@ 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 facelib.utils.face_restoration_helper import FaceRestoreHelper
+from codeformer.archs.codeformer_arch import CodeFormer
+
+# 依存関係のインストール
+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")
+
+# ウェイトファイルをダウンロード（毎回消えるので毎回必ず実行。）
 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'):
@@ -31,6 +43,36 @@ upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, ti
 
 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)
+    
+    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 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)
+
 @app.route('/api/restore', methods=['POST'])
 def restore_image():
     try:
@@ -41,7 +83,7 @@ def restore_image():
         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パラメータが必要な場合
+        weight = float(request.form.get('weight', 0.5))  # CodeFormer用のweightパラメータ
         
         # 一時ファイルに保存
         temp_dir = tempfile.mkdtemp()
@@ -49,7 +91,7 @@ def restore_image():
         file.save(input_path)
         
         # 画像処理
-        extension = os.path.splitext(os.path.basename(str(input_path)))[1]
+        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:
@@ -68,25 +110,26 @@ def restore_image():
         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':
-            face_enhancer = GFPGANer(
-                model_path='CodeFormer.pth', upscale=2, arch='CodeFormer', channel_multiplier=2, bg_upsampler=upsampler)
+            output = restore_with_codeformer(img, scale=scale, weight=weight)
         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)
+                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)
 
-        # 画像を拡張
-        _, _, 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
@@ -179,6 +222,10 @@ def index():
                 <label for="scale">Rescaling factor:</label>
                 <input type="number" id="scale" name="scale" value="2" step="0.1" min="1" max="4" required>
             </div>
+            <div class="form-group" id="weightGroup" style="display: none;">
+                <label for="weight">CodeFormer Weight (0-1):</label>
+                <input type="number" id="weight" name="weight" value="0.5" step="0.1" min="0" max="1">
+            </div>
             <button type="submit" id="submitButton">Process Image</button>
         </form>
         <div id="loading" class="loader"></div>
@@ -198,11 +245,23 @@ def index():
     </div>
     
     <script>
+        // CodeFormerが選択された時にweightパラメータを表示
+        document.getElementById('version').addEventListener('change', function() {
+            const weightGroup = document.getElementById('weightGroup');
+            if (this.value === 'CodeFormer') {
+                weightGroup.style.display = 'block';
+            } else {
+                weightGroup.style.display = 'none';
+            }
+            updateApiUsage();
+        });
+
         // フォームの変更を監視してAPI使用例を更新
         function updateApiUsage() {
             const fileInput = document.getElementById('file');
             const version = document.getElementById('version').value;
             const scale = document.getElementById('scale').value;
+            const weight = document.getElementById('weight').value;
             
             // 現在のURLからベースURLを取得（パス、パラメータ、ハッシュを含めない）
             const baseUrl = window.location.origin;
@@ -216,27 +275,32 @@ 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, filePreview);
+                    updateFetchCode(apiUrl, version, scale, weight, filePreview);
                 };
                 reader.readAsDataURL(file);
             } else {
-                updateFetchCode(apiUrl, version, scale, filePreview);
+                updateFetchCode(apiUrl, version, scale, weight, filePreview);
             }
         }
         
-        function updateFetchCode(apiUrl, version, scale, filePreview) {
+        function updateFetchCode(apiUrl, version, scale, weight, filePreview) {
             const fetchCodeDiv = document.getElementById('fetchCode');
-            fetchCodeDiv.innerHTML = `
-// JavaScript fetch example:
+            let code = `// JavaScript fetch example:
 const formData = new FormData();
 formData.append('file', ${filePreview});
 formData.append('version', '${version}');
-formData.append('scale', ${scale});
+formData.append('scale', ${scale});`;
+
+            if (version === 'CodeFormer') {
+                code += `
+formData.append('weight', ${weight});`;
+            }
 
+            code += `
 fetch('${apiUrl}', {
     method: 'POST',
     body: formData
@@ -256,12 +320,15 @@ fetch('${apiUrl}', {
 .catch(error => {
     console.error('Error:', error.message);
 });`;
+
+            fetchCodeDiv.innerHTML = code;
         }
         
         // フォーム要素の変更を監視
         document.getElementById('file').addEventListener('change', updateApiUsage);
         document.getElementById('version').addEventListener('change', updateApiUsage);
         document.getElementById('scale').addEventListener('input', updateApiUsage);
+        document.getElementById('weight').addEventListener('input', updateApiUsage);
         
         // 初期表示
         updateApiUsage();
@@ -281,6 +348,11 @@ fetch('${apiUrl}', {
             formData.append('version', document.getElementById('version').value);
             formData.append('scale', document.getElementById('scale').value);
             
+            // CodeFormerが選択されている場合はweightも追加
+            if (document.getElementById('version').value === 'CodeFormer') {
+                formData.append('weight', document.getElementById('weight').value);
+            }
+            
             // 現在のURLからベースURLを取得（パス、パラメータ、ハッシュを含めない）
             const baseUrl = window.location.origin;
             const apiUrl = baseUrl + '/api/restore';
@@ -321,5 +393,4 @@ fetch('${apiUrl}', {
     """
 
 if __name__ == '__main__':
-    
     app.run(host='0.0.0.0', port=7860, debug=True)