Ai-Swaper/core/models.py

import sys
import types
import urllib.request
import importlib
import importlib.util
from pathlib import Path
import numpy as np
import cv2

from core.config import MODELS_DIR, get_providers


def dl_file(url, dest, log=print):
    """Download a file with progress reporting."""
    def hook(count, block, total):
        if total > 0 and count % 100 == 0:
            pct = min(100, count * block * 100 // total)
            log(f"     ... {pct}%  ({count*block/1_048_576:.0f} MB)")
    urllib.request.urlretrieve(url, dest, reporthook=hook)


def ensure_torchvision_patch(log_fn=None):
    """Alias torchvision.transforms.functional_tensor for GFPGAN compatibility."""
    if "torchvision.transforms.functional_tensor" in sys.modules:
        return
    try:
        from torchvision.transforms import functional as tvf
        alias = types.ModuleType("torchvision.transforms.functional_tensor")
        alias.rgb_to_grayscale = tvf.rgb_to_grayscale
        sys.modules["torchvision.transforms.functional_tensor"] = alias
    except Exception as e:
        if log_fn:
            log_fn(f"  Hinweis: torchvision-Patch fehlgeschlagen ({e}).")


class FaceRestorer:
    """GFPGANv1.4 face restoration wrapper."""
    MODEL_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth"
    MODEL_PATH = MODELS_DIR / "GFPGANv1.4.pth"
    _FACEXLIB_FILES = {
        "detection_Resnet50_Final.pth": "https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth",
        "parsing_parsenet.pth": "https://github.com/xinntao/facexlib/releases/download/v0.2.2/parsing_parsenet.pth",
    }

    def __init__(self, log_fn=print):
        self.log = log_fn
        self.restorer = None
        self.upscale = 1
        self.only_center_face = False
        self.min_face_px = 72
        self.blend_alpha = 0.75
        self.pad_ratio = 0.22
        self.load_attempted = False
        self.load_failed = False
        self.last_error = None
        self._error_logged = False

    @staticmethod
    def is_available(log_fn=None):
        if importlib.util.find_spec("gfpgan") is None:
            return False
        ensure_torchvision_patch(log_fn)
        try:
            importlib.import_module("gfpgan")
            return True
        except Exception:
            return False

    def ensure_model(self):
        MODELS_DIR.mkdir(parents=True, exist_ok=True)
        if not self.MODEL_PATH.exists() or self.MODEL_PATH.stat().st_size < 100_000:
            self.log("  Lade GFPGANv1.4.pth (~330 MB) ...")
            dl_file(self.MODEL_URL, self.MODEL_PATH, self.log)
        self._ensure_facexlib_weights()

    def _ensure_facexlib_weights(self):
        if not self.is_available():
            return
        try:
            import facexlib
            weights_dir = Path(facexlib.__file__).resolve().parent / "weights"
            weights_dir.mkdir(parents=True, exist_ok=True)
            for filename, url in self._FACEXLIB_FILES.items():
                dst = weights_dir / filename
                if dst.exists() and dst.stat().st_size > 100_000:
                    continue
                self.log(f"  Lade {filename} ...")
                dl_file(url, dst, self.log)
        except Exception as e:
            self.log(f"  Hinweis: facexlib-Modelle konnten nicht vorab geladen werden ({e}).")

    def load(self):
        if self.restorer is not None:
            return
        if self.load_failed:
            raise RuntimeError(self.last_error or "GFPGAN konnte nicht geladen werden.")
        self.load_attempted = True
        try:
            if not self.is_available(log_fn=self.log):
                raise RuntimeError("GFPGAN ist nicht installiert.")
            self.ensure_model()
            ensure_torchvision_patch(self.log)
            from gfpgan import GFPGANer
            providers = get_providers(self.log)
            device = "cuda" if "CUDAExecutionProvider" in providers else "cpu"
            self.restorer = GFPGANer(
                model_path=str(self.MODEL_PATH),
                upscale=self.upscale,
                arch="clean",
                channel_multiplier=2,
                bg_upsampler=None,
                device=device,
            )
            self.log("  GFPGAN Gesichtswiederherstellung bereit.")
        except Exception as e:
            self.load_failed = True
            self.last_error = str(e)
            raise

    def _bbox_to_rect(self, bbox, img_w, img_h):
        try:
            x1, y1, x2, y2 = [float(v) for v in bbox]
        except Exception:
            return None
        bw = max(1.0, x2 - x1)
        bh = max(1.0, y2 - y1)
        pad_x = max(8, int(round(bw * self.pad_ratio)))
        pad_y = max(8, int(round(bh * self.pad_ratio)))
        rx1 = max(0, int(round(x1)) - pad_x)
        ry1 = max(0, int(round(y1)) - pad_y)
        rx2 = min(int(img_w), int(round(x2)) + pad_x)
        ry2 = min(int(img_h), int(round(y2)) + pad_y)
        if rx2 <= rx1 + 2 or ry2 <= ry1 + 2:
            return None
        return rx1, ry1, rx2, ry2

    def restore_faces(self, bgr_img, bboxes, blend_alpha=None):
        """Restore facial features using GFPGAN within detected bounding boxes."""
        if bgr_img is None or not bboxes:
            return bgr_img
        if self.load_failed:
            return bgr_img
        try:
            self.load()
        except Exception as e:
            if not self._error_logged:
                self.log(f"  Hinweis: GFPGAN nicht bereit ({e}).")
                self._error_logged = True
            return bgr_img

        alpha = float(self.blend_alpha if blend_alpha is None else blend_alpha)
        alpha = max(0.0, min(1.0, alpha))
        out = bgr_img.copy()
        h, w = out.shape[:2]
        restored_any = False

        for bbox in bboxes:
            rect = self._bbox_to_rect(bbox, w, h)
            if rect is None:
                continue
            x1, y1, x2, y2 = rect
            rw, rh = x2 - x1, y2 - y1
            if rw < self.min_face_px or rh < self.min_face_px:
                continue
            crop = out[y1:y2, x1:x2]
            if crop.size == 0:
                continue
            try:
                _, _, restored = self.restorer.enhance(
                    crop,
                    has_aligned=False,
                    only_center_face=True,
                    paste_back=True,
                )
            except Exception:
                continue
            if restored is None:
                continue
            if restored.shape[:2] != crop.shape[:2]:
                restored = cv2.resize(restored, (crop.shape[1], crop.shape[0]), interpolation=cv2.INTER_AREA)

            mixed = cv2.addWeighted(restored.astype(np.float32), alpha, crop.astype(np.float32), 1.0 - alpha, 0.0)
            out[y1:y2, x1:x2] = np.clip(mixed, 0, 255).astype(np.uint8)
            restored_any = True

        return out if restored_any else bgr_img