fix: sharp output, source skin tone preservation, 4K upscaling pipeline

- swapper.py: sharpen ORIGINAL crop (not bilateral-blurred) — unsharp mask
strength 2.3x on raw crop; mild bilateral d=5/30 only removes compression
artefacts without killing texture detail
- skin_tone.py: match_skin_tone now accepts face_mask; LAB transfer confined
to face region via feathered alpha — hair/neck/background no longer shifted
- core/super_res.py (NEW): GFPGAN face restoration + RealESRGAN x4 upscale;
Lanczos fallback when models absent — always delivers 4K download
- web_app.py: complete pipeline — add match_skin_tone (source tone, face mask),
laplacian_blend + poisson_blend, 4K hi-res PNG saved for download, preview
JPEG at 92 quality; input cap raised 1024→2048 for more InsightFace detail
- quality_checker.py: _alignment_score handles None/empty arrays — returns
50.0 instead of propagating nan (was showing 0.0/100 in UI)
- full_pipeline.py: enable_super_res=True by default; pass face_mask to
match_skin_tone; Lanczos fallback on super_res failure; remove unused imports
- scripts/download_models.py: add RealESRGAN_x4plus.pth + GFPGANv1.4.pth
download entries so `python scripts/download_models.py` fetches all models

core/quality_checker.py

@@ -5,8 +5,8 @@ import numpy as np
 def compute_quality_score(
     swapped: np.ndarray,
     target: np.ndarray,
-    src_landmarks: np.ndarray,
-    tgt_landmarks: np.ndarray
+    src_landmarks,   # type: np.ndarray | None
+    tgt_landmarks,   # type: np.ndarray | None
 ) -> dict:
     """
     Compute quality metrics for the swap result.
@@ -25,16 +25,19 @@ def compute_quality_score(
     }
 
 
-def _alignment_score(src_lm: np.ndarray, tgt_lm: np.ndarray) -> float:
+def _alignment_score(src_lm, tgt_lm) -> float:
     """Mean pixel deviation between landmark sets, converted to 0-100 score."""
     if src_lm is None or tgt_lm is None:
         return 50.0
     n = min(len(src_lm), len(tgt_lm))
+    if n == 0:
+        return 50.0
     diff = np.linalg.norm(src_lm[:n] - tgt_lm[:n], axis=1)
-    mean_err = diff.mean()
-    # Map: 0px -> 100, 10px -> 0
+    mean_err = float(diff.mean())
+    if not np.isfinite(mean_err):
+        return 50.0
     score = max(0.0, 100.0 - mean_err * 10.0)
-    return round(float(score), 1)
+    return round(score, 1)
 
 
 def _blend_quality(swapped: np.ndarray, target: np.ndarray) -> float:
@@ -50,7 +53,7 @@ def _blend_quality(swapped: np.ndarray, target: np.ndarray) -> float:
     discontinuity = np.abs(edges).mean()
     # Map: 0 -> 100, 30 -> 0
     score = max(0.0, 100.0 - discontinuity * (100.0 / 30.0))
-    return round(float(score), 1)
+    return round(score, 1)
 
 
 def _compute_delta_e(swapped: np.ndarray, target: np.ndarray) -> float:
@@ -94,4 +97,4 @@ def _naturalness_score(swapped: np.ndarray) -> float:
     noise_penalty = min(noise / 1000.0, 20.0)
 
     score = max(0.0, 100.0 - clip_penalty - sat_penalty - noise_penalty)
-    return round(float(score), 1)
+    return round(score, 1)

core/skin_tone.py

@@ -75,29 +75,50 @@ def analyze_skin_tone(image: np.ndarray, face_bbox) -> dict:
 def match_skin_tone(
     swapped_img: np.ndarray,
     target_img: np.ndarray,
-    src_tone: dict,
-    tgt_tone: dict,
-    strength: float = 0.9
+    src_tone: dict,  # noqa: ARG001 — kept for API compatibility
+    tgt_tone: dict,  # noqa: ARG001 — kept for API compatibility
+    strength: float = 0.9,
+    face_mask=None,  # type: np.ndarray | None
 ) -> np.ndarray:
     """
-    Adjust swapped image skin tone to match target.
-    Uses LAB color space transfer with strength control.
+    Adjust swapped face skin tone to match target, confined to face_mask region.
+    Passing face_mask prevents hair/neck/background from being colour-shifted.
+    When face_mask is None the transfer falls back to the whole image (legacy).
     """
     src_lab = cv2.cvtColor(swapped_img, cv2.COLOR_BGR2LAB).astype(np.float32)
     tgt_lab = cv2.cvtColor(target_img,  cv2.COLOR_BGR2LAB).astype(np.float32)
 
-    # Per-channel mean/std transfer
-    for ch in range(3):
-        src_mean, src_std = src_lab[:, :, ch].mean(), src_lab[:, :, ch].std()
-        tgt_mean, tgt_std = tgt_lab[:, :, ch].mean(), tgt_lab[:, :, ch].std()
+    # Determine which pixels to compute statistics from
+    if face_mask is not None and face_mask.size > 0:
+        stat_mask = face_mask > 128
+    else:
+        stat_mask = np.ones(src_lab.shape[:2], dtype=bool)
 
-        if src_std < 1e-6:
+    corrected_lab = src_lab.copy()
+    for ch in range(3):
+        src_vals = src_lab[:, :, ch][stat_mask]
+        tgt_vals = tgt_lab[:, :, ch][stat_mask]
+        if src_vals.std() < 1e-6 or tgt_vals.std() < 1e-6:
             continue
-        corrected = (src_lab[:, :, ch] - src_mean) * (tgt_std / src_std) + tgt_mean
-        src_lab[:, :, ch] = src_lab[:, :, ch] * (1 - strength) + corrected * strength
+        corrected_ch = (
+            (src_lab[:, :, ch] - src_vals.mean()) *
+            (tgt_vals.std() / src_vals.std()) +
+            tgt_vals.mean()
+        )
+        corrected_lab[:, :, ch] = (
+            src_lab[:, :, ch] * (1 - strength) + corrected_ch * strength
+        )
+
+    # Apply correction only inside the face mask (feathered at boundary)
+    if face_mask is not None and face_mask.size > 0:
+        alpha = cv2.GaussianBlur(face_mask.astype(np.float32) / 255.0, (21, 21), 7)
+        alpha = np.stack([alpha] * 3, axis=-1)
+        result_lab = src_lab * (1.0 - alpha) + corrected_lab * alpha
+    else:
+        result_lab = corrected_lab
 
     result = cv2.cvtColor(
-        np.clip(src_lab, 0, 255).astype(np.uint8), cv2.COLOR_LAB2BGR
+        np.clip(result_lab, 0, 255).astype(np.uint8), cv2.COLOR_LAB2BGR
     )
     return result
 

core/super_res.py

@@ -0,0 +1,111 @@
+import cv2
+import numpy as np
+import os
+
+_realesrgan_instance = None
+_gfpgan_instance = None
+
+MODELS_DIR = "models"
+REALESRGAN_MODEL = os.path.join(MODELS_DIR, "RealESRGAN_x4plus.pth")
+GFPGAN_MODEL     = os.path.join(MODELS_DIR, "GFPGANv1.4.pth")
+
+
+def _load_realesrgan():
+    global _realesrgan_instance
+    if _realesrgan_instance is not None:
+        return _realesrgan_instance
+    if not os.path.exists(REALESRGAN_MODEL):
+        print(f"[super_res] RealESRGAN model not found: {REALESRGAN_MODEL}")
+        return None
+    try:
+        from basicsr.archs.rrdbnet_arch import RRDBNet
+        from realesrgan import RealESRGANer
+
+        model = RRDBNet(
+            num_in_ch=3, num_out_ch=3, num_feat=64,
+            num_block=23, num_grow_ch=32, scale=4
+        )
+        _realesrgan_instance = RealESRGANer(
+            scale=4,
+            model_path=REALESRGAN_MODEL,
+            model=model,
+            tile=512,
+            tile_pad=10,
+            pre_pad=0,
+            half=False,  # CPU-safe — no float16
+        )
+        print("[super_res] RealESRGAN x4 loaded OK")
+        return _realesrgan_instance
+    except Exception as e:
+        print(f"[super_res] RealESRGAN load failed: {e}")
+        return None
+
+
+def _load_gfpgan():
+    global _gfpgan_instance
+    if _gfpgan_instance is not None:
+        return _gfpgan_instance
+    if not os.path.exists(GFPGAN_MODEL):
+        print(f"[super_res] GFPGAN model not found: {GFPGAN_MODEL}")
+        return None
+    try:
+        from gfpgan import GFPGANer
+
+        _gfpgan_instance = GFPGANer(
+            model_path=GFPGAN_MODEL,
+            upscale=1,          # restore only — upscaling done by RealESRGAN
+            arch="clean",
+            channel_multiplier=2,
+        )
+        print("[super_res] GFPGAN loaded OK")
+        return _gfpgan_instance
+    except Exception as e:
+        print(f"[super_res] GFPGAN load failed: {e}")
+        return None
+
+
+def enhance_resolution(image: np.ndarray, scale: int = 4) -> np.ndarray:
+    """
+    Two-stage quality enhancement:
+      1. GFPGAN — restores face detail lost in InsightFace's 128x128 resize
+         (no resolution change, just sharpness/texture restoration on faces).
+      2. RealESRGAN x4 — upscales the whole image to ~4K.
+    Falls back to Lanczos resize when models are not downloaded.
+
+    scale: 2 or 4 (4 = ~4K from a 1024px input)
+    """
+    result = image.copy()
+
+    # --- Stage 1: face restoration ----------------------------------------
+    gfpgan = _load_gfpgan()
+    if gfpgan is not None:
+        try:
+            _, _, restored = gfpgan.enhance(
+                result,
+                has_aligned=False,
+                only_center_face=False,
+                paste_back=True,
+            )
+            if restored is not None and restored.shape == result.shape:
+                result = restored
+                print("[super_res] GFPGAN face restoration applied")
+        except Exception as e:
+            print(f"[super_res] GFPGAN enhance failed: {e}")
+
+    # --- Stage 2: full-image upscaling ------------------------------------
+    upsampler = _load_realesrgan()
+    if upsampler is not None:
+        try:
+            output, _ = upsampler.enhance(result, outscale=scale)
+            if output is not None:
+                h, w = output.shape[:2]
+                print(f"[super_res] RealESRGAN {scale}x done → {w}×{h}")
+                return output
+        except Exception as e:
+            print(f"[super_res] RealESRGAN enhance failed: {e}")
+
+    # --- Fallback: Lanczos resize -----------------------------------------
+    h, w = result.shape[:2]
+    tw, th = w * scale, h * scale
+    print(f"[super_res] Lanczos fallback {scale}x: {w}×{h} → {tw}×{th}")
+    return cv2.resize(result, (tw, th), interpolation=cv2.INTER_LANCZOS4)

core/swapper.py

@@ -114,10 +114,12 @@ def _restore_glasses_region(swapped: np.ndarray, original_target: np.ndarray, fa
 
 def _sharpen_face_region(image: np.ndarray, tgt_faces: list) -> np.ndarray:
     """
-    Sharpen + even skin tone in each swapped face region.
-    - Unsharp mask restores detail lost in InsightFace's 128x128 internal resize.
-    - Bilateral filter smooths out skin blotchiness while preserving edges.
-    - Colour correction at the face boundary blends the two skin tones.
+    Recover detail lost in InsightFace's 128x128 internal resize.
+    Strategy:
+      - Mild bilateral filter to remove compression artefacts (not texture).
+      - Unsharp mask on the ORIGINAL crop (not the blurred version) for true
+        high-frequency recovery; 2.3x strength gives clean edges without halos.
+      - Blend sharp + smooth so skin stays natural while edges are crisp.
     """
     result = image.copy()
     h, w   = image.shape[:2]
@@ -132,14 +134,17 @@ def _sharpen_face_region(image: np.ndarray, tgt_faces: list) -> np.ndarray:
         if crop.size == 0:
             continue
 
-        # 1. Even skin tone — bilateral filter (smooth blotches, keep edges)
-        smooth = cv2.bilateralFilter(crop, d=9, sigmaColor=60, sigmaSpace=60)
+        # 1. Mild bilateral — removes compression blotches while keeping edges
+        smooth = cv2.bilateralFilter(crop, d=5, sigmaColor=30, sigmaSpace=30)
 
-        # 2. Sharpen — unsharp mask on top of smoothed
-        blur   = cv2.GaussianBlur(smooth, (0, 0), sigmaX=2.5)
-        sharp  = cv2.addWeighted(smooth, 1.6, blur, -0.6, 0)
+        # 2. Unsharp mask on ORIGINAL crop — true detail recovery
+        blur  = cv2.GaussianBlur(crop, (0, 0), sigmaX=1.5)
+        sharp = cv2.addWeighted(crop, 2.3, blur, -1.3, 0)
 
-        # 3. Feathered paste — fade near the crop edges so no hard border
+        # 3. Composite: 55% sharp detail + 45% smooth skin base
+        enhanced = cv2.addWeighted(sharp, 0.55, smooth, 0.45, 0)
+
+        # 4. Feathered paste — no hard border at crop edges
         fh, fw = crop.shape[:2]
         feather = np.ones((fh, fw), dtype=np.float32)
         border  = max(4, pad // 2)
@@ -152,7 +157,7 @@ def _sharpen_face_region(image: np.ndarray, tgt_faces: list) -> np.ndarray:
         feather = np.stack([feather] * 3, axis=-1)
 
         result[y1p:y2p, x1p:x2p] = (
-            sharp.astype(np.float32) * feather +
+            enhanced.astype(np.float32) * feather +
             result[y1p:y2p, x1p:x2p].astype(np.float32) * (1 - feather)
         ).astype(np.uint8)
 

pipeline/full_pipeline.py

@@ -1,4 +1,3 @@
-import cv2
 import numpy as np
 
 from core.detector import detect_faces
@@ -10,17 +9,16 @@ from core.skin_tone import analyze_skin_tone, match_skin_tone
 from core.neck_integrator import seamless_hair_to_neck_blend
 from core.color_corrector import harmonize_colors
 from core.quality_checker import compute_quality_score
-from utils.image_io import resize_keep_aspect
 
 
 def run_full_pipeline(
     source: np.ndarray,
     target: np.ndarray,
     blend_strength: float = 0.85,
-    tone_match_strength: float = 0.9,
+    tone_match_strength: float = 0.6,
     hair_preserve: float = 0.8,
     neck_blend_strength: float = 0.75,
-    enable_super_res: bool = False,
+    enable_super_res: bool = True,
     progress_callback=None,
 ) -> dict:
     """
@@ -58,8 +56,11 @@ def run_full_pipeline(
     swapped = swap_face_insightface(source, target)
 
     _progress(55, "Matching skin tones...")
-    swapped = match_skin_tone(swapped, target, src_tone, tgt_tone,
-                              strength=tone_match_strength)
+    swapped = match_skin_tone(
+        swapped, source, src_tone, src_tone,
+        strength=tone_match_strength,
+        face_mask=tgt_masks.get("face_mask"),
+    )
 
     _progress(65, "Blending hair, face, neck seamlessly...")
     swapped = seamless_hair_to_neck_blend(
@@ -83,12 +84,16 @@ def run_full_pipeline(
     swapped = harmonize_colors(swapped, target, tgt_masks)
 
     if enable_super_res:
-        _progress(92, "Enhancing resolution...")
+        _progress(92, "Enhancing resolution to 4K...")
         try:
             from core.super_res import enhance_resolution
-            swapped = enhance_resolution(swapped)
-        except Exception:
-            pass
+            swapped = enhance_resolution(swapped, scale=4)
+        except Exception as e:
+            import cv2
+            print(f"[pipeline] super_res skipped: {e}")
+            h, w = swapped.shape[:2]
+            swapped = cv2.resize(swapped, (w * 4, h * 4),
+                                 interpolation=cv2.INTER_LANCZOS4)
 
     _progress(97, "Computing quality metrics...")
     quality = compute_quality_score(swapped, target, src_lm, tgt_lm)

scripts/download_models.py

@@ -2,7 +2,6 @@
 Download pretrained model weights required by the pipeline.
 Run: python scripts/download_models.py
 """
-import os
 import sys
 import urllib.request
 from pathlib import Path
@@ -14,13 +13,21 @@ MODELS_DIR.mkdir(exist_ok=True)
 MODELS = {
     "inswapper_128.onnx": (
         "https://huggingface.co/deepinsight/inswapper/resolve/main/inswapper_128.onnx",
-        "InsightFace face-swap model (required)"
+        "InsightFace face-swap model (required)",
+    ),
+    "RealESRGAN_x4plus.pth": (
+        "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
+        "RealESRGAN x4 upscaler — enables 4K output (optional but recommended)",
+    ),
+    "GFPGANv1.4.pth": (
+        "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth",
+        "GFPGAN face restoration — fixes InsightFace blur (optional but recommended)",
     ),
 }
 
 # Note: BiSeNet and RetinaFace weights are downloaded automatically by
-# their respective Python packages on first use. Only the InsightFace
-# inswapper model requires a manual download.
+# their respective Python packages on first use. Only the models listed
+# above require a manual download.
 
 
 def _progress_hook(block_num, block_size, total_size):

web_app.py

@@ -30,8 +30,9 @@ from PIL import Image
 from core.detector import detect_faces
 from core.swapper import swap_face_insightface
 from core.segmentor import segment_hair_neck_skin
-from core.skin_tone import analyze_skin_tone
+from core.skin_tone import analyze_skin_tone, match_skin_tone
 from core.neck_integrator import seamless_hair_to_neck_blend
+from core.blender import laplacian_blend, poisson_blend
 from core.quality_checker import compute_quality_score
 from utils.image_io import save_image, resize_keep_aspect
 
@@ -178,9 +179,9 @@ def api_swap():
         neck_blend     = int(request.form.get("neck_blend",    75)) / 100.0
         blend_strength = int(request.form.get("blend_strength",85)) / 100.0
 
-        # -- resize -----------------------------------------------------------
-        source = resize_keep_aspect(source, 1024)
-        target = resize_keep_aspect(target, 1024)
+        # -- resize (2048 gives InsightFace more texture to work with) --------
+        source = resize_keep_aspect(source, 2048)
+        target = resize_keep_aspect(target, 2048)
 
         # -- face detection ---------------------------------------------------
         faces_src = _safe_detect(source)
@@ -190,20 +191,31 @@ def api_swap():
         if not faces_tgt:
             return jsonify({"ok": False, "error": "No face detected in target image"}), 400
 
-        # -- pipeline ---------------------------------------------------------
-        src_tone  = analyze_skin_tone(source, faces_src[0])
-        tgt_tone  = analyze_skin_tone(target, faces_tgt[0])
-        delta_e   = (
+        # -- skin tone analysis -----------------------------------------------
+        src_tone = analyze_skin_tone(source, faces_src[0])
+        tgt_tone = analyze_skin_tone(target, faces_tgt[0])
+        delta_e  = (
             (src_tone["L"] - tgt_tone["L"]) ** 2 +
             (src_tone["a"] - tgt_tone["a"]) ** 2 +
             (src_tone["b"] - tgt_tone["b"]) ** 2
         ) ** 0.5
 
+        # -- segmentation (needed for masking) --------------------------------
+        tgt_masks = segment_hair_neck_skin(target)
+        face_mask = tgt_masks.get("face_mask")
+
         # 1. Core face swap (InsightFace inswapper_128)
         swapped = swap_face_insightface(source, target)
 
-        # 2. Extend blend from face → hair boundary → neck
-        tgt_masks = segment_hair_neck_skin(target)
+        # 2. Match face skin tone to source — only inside face region so hair
+        #    and neck are not colour-shifted.
+        swapped = match_skin_tone(
+            swapped, source, src_tone, src_tone,
+            strength=0.6,
+            face_mask=face_mask,
+        )
+
+        # 3. Seamless hair → face → neck blend
         swapped = seamless_hair_to_neck_blend(
             source_img=swapped, target_img=target,
             src_masks=tgt_masks, tgt_masks=tgt_masks,
@@ -213,22 +225,35 @@ def api_swap():
             blend_strength=blend_strength,
         )
 
-        # Skin tone is preserved naturally by InsightFace — no override needed
+        # 4. Laplacian pyramid + Poisson seamless-clone to remove paste edges
+        if face_mask is not None:
+            swapped = laplacian_blend(target, swapped, face_mask, levels=4)
+            swapped = poisson_blend(swapped, target, face_mask)
 
+        # -- quality metrics --------------------------------------------------
         quality = compute_quality_score(
             swapped, target,
-            np.zeros((0, 2), dtype=np.float32),
-            np.zeros((0, 2), dtype=np.float32),
+            None,   # landmarks not extracted in web pipeline — returns 50.0
+            None,
         )
 
-        # -- save output -------------------------------------------------------
+        # -- 4K upscale for download (Lanczos fallback when SR models absent) -
+        try:
+            from core.super_res import enhance_resolution
+            hi_res = enhance_resolution(swapped, scale=4)
+        except Exception:
+            h, w = swapped.shape[:2]
+            hi_res = cv2.resize(swapped, (w * 4, h * 4),
+                                interpolation=cv2.INTER_LANCZOS4)
+
+        # -- save 4K PNG output -----------------------------------------------
         out_name = f"swap_{uuid.uuid4().hex[:8]}.png"
         out_path = os.path.join(OUTPUT_DIR, out_name)
-        save_image(swapped, out_path)
+        save_image(hi_res, out_path)
 
         return jsonify({
             "ok": True,
-            "result_image": _encode_image(swapped, fmt="JPEG", quality=90),
+            "result_image": _encode_image(swapped, fmt="JPEG", quality=92),
             "quality": quality,
             "delta_e": round(delta_e, 2),
             "src_tone": src_tone,