SnapOtter/packages/ai/python/enhance_faces.py

"""Face enhancement using GFPGAN or CodeFormer with MediaPipe detection."""
import sys
import json
import os
import types

# basicsr imports torchvision.transforms.functional_tensor which was removed
# in torchvision >= 0.17. This shim must exist before basicsr is imported.
try:
    import torchvision.transforms.functional_tensor  # noqa: F401
except (ImportError, ModuleNotFoundError):
    try:
        import torchvision.transforms.functional as _F
        import torchvision.transforms

        _shim = types.ModuleType("torchvision.transforms.functional_tensor")
        for _attr in dir(_F):
            if not _attr.startswith("_"):
                setattr(_shim, _attr, getattr(_F, _attr))
        sys.modules["torchvision.transforms.functional_tensor"] = _shim
        torchvision.transforms.functional_tensor = _shim
    except ImportError:
        pass


def emit_progress(percent, stage):
    """Emit structured progress to stderr for bridge.ts to capture."""
    print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)


_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models")

GFPGAN_MODEL_PATH = os.environ.get(
    "GFPGAN_MODEL_PATH",
    os.path.join(_MODELS_BASE, "gfpgan", "GFPGANv1.3.pth"),
)

CODEFORMER_MODEL_PATH = os.environ.get(
    "CODEFORMER_MODEL_PATH",
    os.path.join(_MODELS_BASE, "codeformer", "codeformer.pth"),
)


# ── Model path for new mp.tasks API ─────────────────────────────────

_FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.tflite"
_DOCKER_MODEL_PATH = os.path.join(_MODELS_BASE, "mediapipe", "blaze_face_short_range.tflite")
_LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models")
_LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "blaze_face_short_range.tflite")


def _ensure_face_detect_model():
    """Resolve face detector model. Docker path first, then local dev."""
    if os.path.exists(_DOCKER_MODEL_PATH):
        return _DOCKER_MODEL_PATH
    if os.path.exists(_LOCAL_MODEL_PATH):
        return _LOCAL_MODEL_PATH
    from offline_guard import ensure_download_allowed
    ensure_download_allowed("Face detection model (blaze_face_short_range.tflite)")
    os.makedirs(_LOCAL_MODEL_DIR, exist_ok=True)
    import urllib.request
    emit_progress(15, "Downloading face detection model")
    urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _LOCAL_MODEL_PATH)
    return _LOCAL_MODEL_PATH


_MAX_DETECT_DIM = 1920


def _downscale_for_detection(img_array):
    """Downscale image if needed so MediaPipe can detect faces reliably."""
    h, w = img_array.shape[:2]
    longest = max(h, w)
    if longest <= _MAX_DETECT_DIM:
        return img_array, 1.0

    import cv2
    scale = _MAX_DETECT_DIM / longest
    new_w = int(w * scale)
    new_h = int(h * scale)
    resized = cv2.resize(img_array, (new_w, new_h), interpolation=cv2.INTER_AREA)
    return resized, 1.0 / scale


def detect_faces_mediapipe(img_array, sensitivity):
    """Detect faces using MediaPipe with dual-model approach.

    Returns a list of {x, y, w, h} dicts for each detected face.
    Tries legacy mp.solutions API first, falls back to mp.tasks.
    Large images are downscaled before detection for reliability.
    """
    import mediapipe as mp

    min_confidence = max(0.1, 1.0 - sensitivity)
    scaled, inv_scale = _downscale_for_detection(img_array)

    try:
        mp_face = mp.solutions.face_detection

        all_detections = []
        for model_sel in [0, 1]:
            detector = mp_face.FaceDetection(
                model_selection=model_sel,
                min_detection_confidence=min_confidence,
            )
            results = detector.process(scaled)
            detector.close()
            if results.detections:
                all_detections.extend(results.detections)

        if not all_detections:
            return []

        ih, iw = scaled.shape[:2]
        faces = []
        for detection in all_detections:
            bbox = detection.location_data.relative_bounding_box
            faces.append({
                "x": int(bbox.xmin * iw * inv_scale),
                "y": int(bbox.ymin * ih * inv_scale),
                "w": int(bbox.width * iw * inv_scale),
                "h": int(bbox.height * ih * inv_scale),
            })
        return faces

    except AttributeError:
        model_path = _ensure_face_detect_model()
        options = mp.tasks.vision.FaceDetectorOptions(
            base_options=mp.tasks.BaseOptions(model_asset_path=model_path),
            running_mode=mp.tasks.vision.RunningMode.IMAGE,
            min_detection_confidence=min_confidence,
        )
        detector = mp.tasks.vision.FaceDetector.create_from_options(options)
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=scaled)
        result = detector.detect(mp_image)
        detector.close()

        faces = []
        for detection in result.detections:
            bbox = detection.bounding_box
            faces.append({
                "x": int(bbox.origin_x * inv_scale),
                "y": int(bbox.origin_y * inv_scale),
                "w": int(bbox.width * inv_scale),
                "h": int(bbox.height * inv_scale),
            })
        return faces


def enhance_with_gfpgan(img_array, only_center_face):
    """Enhance faces using GFPGAN. Returns the enhanced image array."""
    import torch
    from gfpgan import GFPGANer
    from gpu import torch_gpu_available

    if not os.path.exists(GFPGAN_MODEL_PATH):
        raise FileNotFoundError(f"GFPGAN model not found: {GFPGAN_MODEL_PATH}")

    # GFPGANer resolves its facexlib helper weights relative to the cwd and
    # downloads them from GitHub when missing; resolve them from the bundle
    # first so no download is needed (strict offline mode errors instead).
    from offline_guard import prepare_gfpgan_helper_weights
    prepare_gfpgan_helper_weights(_MODELS_BASE)

    use_gpu = torch_gpu_available()
    device = torch.device("cuda" if use_gpu else "cpu")

    enhancer = GFPGANer(
        model_path=GFPGAN_MODEL_PATH,
        upscale=1,
        arch="clean",
        channel_multiplier=2,
        bg_upsampler=None,
        device=device,
    )
    _, _, output = enhancer.enhance(
        img_array,
        has_aligned=False,
        only_center_face=only_center_face,
        paste_back=True,
    )
    return output


def enhance_with_codeformer(img_array, fidelity_weight):
    """Enhance faces using CodeFormer via codeformer-pip.

    The codeformer-pip package provides inference_app() which handles
    face detection, alignment, restoration, and paste-back internally.
    fidelity_weight controls quality vs fidelity (0 = quality, 1 = fidelity).

    NOTE: inference_app() expects a file path, not a numpy array. We save
    to a temp file and pass the path. The function returns a file path to
    the result which we read back.
    """
    import tempfile
    import cv2
    import numpy as np
    import torch
    from gpu import torch_gpu_available

    use_gpu = torch_gpu_available()

    # codeformer-pip downloads four weights into a cwd-relative tree at import
    # time when they are missing; resolve the bundled ones first so only a
    # genuinely unbundled weight can trigger the download fallback.
    from offline_guard import prepare_codeformer_weights
    prepare_codeformer_weights(_MODELS_BASE)

    _orig_cuda_check = torch.cuda.is_available
    if not use_gpu:
        torch.cuda.is_available = lambda: False
    try:
        from codeformer.app import inference_app

        img_bgr = img_array[:, :, ::-1].copy()

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_in:
            cv2.imwrite(tmp_in.name, img_bgr)
            tmp_in_path = tmp_in.name

        try:
            result_path = inference_app(
                image=tmp_in_path,
                background_enhance=False,
                face_upsample=False,
                upscale=1,
                codeformer_fidelity=fidelity_weight,
            )
        finally:
            os.unlink(tmp_in_path)

        if result_path is None:
            raise RuntimeError("CodeFormer returned no result (face detection may have failed)")

        restored_bgr = cv2.imread(str(result_path), cv2.IMREAD_COLOR)
        if restored_bgr is None:
            raise RuntimeError("CodeFormer output file could not be read")
    finally:
        torch.cuda.is_available = _orig_cuda_check

    restored_rgb = restored_bgr[:, :, ::-1].copy()
    return restored_rgb


def main():
    input_path = sys.argv[1]
    output_path = sys.argv[2]
    settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}

    model_choice = settings.get("model", "auto")
    strength = float(settings.get("strength", 0.8))
    only_center_face = settings.get("onlyCenterFace", False)
    sensitivity = float(settings.get("sensitivity", 0.5))

    try:
        emit_progress(10, "Preparing")
        from PIL import Image
        import numpy as np

        img = Image.open(input_path).convert("RGB")
        img_array = np.array(img)

        # Detect faces with MediaPipe
        try:
            emit_progress(20, "Scanning for faces")
            faces = detect_faces_mediapipe(img_array, sensitivity)
        except ImportError:
            print(
                json.dumps(
                    {
                        "success": False,
                        "error": "Face detection requires MediaPipe. Install with: pip install mediapipe",
                    }
                )
            )
            sys.exit(1)

        num_faces = len(faces)
        emit_progress(30, f"Found {num_faces} face{'s' if num_faces != 1 else ''}")

        # No faces found - save original unchanged
        if num_faces == 0:
            img.save(output_path)
            print(
                json.dumps(
                    {
                        "success": True,
                        "facesDetected": 0,
                        "faces": [],
                        "model": "none",
                    }
                )
            )
            return

        emit_progress(40, "Loading AI model")

        # Redirect stdout to stderr for the ENTIRE AI pipeline.
        # Libraries like basicsr, gfpgan, and torch print download
        # progress and init messages to stdout which would corrupt
        # our JSON result.
        stdout_fd = None
        try:
            stdout_fd = os.dup(1)
            sys.stdout.flush()  # Flush before redirect to avoid mixing buffers
            os.dup2(2, 1)
        except OSError:
            # os.dup may fail on Windows when launched via child_process.spawn
            # with piped stdio — fall back to just suppressing sys.stdout
            stdout_fd = None
        sys.stdout = sys.stderr  # Python-level redirect regardless of OS

        enhanced = None
        model_used = None

        try:
            if model_choice == "gfpgan":
                enhanced = enhance_with_gfpgan(img_array, only_center_face)
                model_used = "gfpgan"

            elif model_choice == "codeformer":
                fidelity_weight = 1.0 - strength
                enhanced = enhance_with_codeformer(img_array, fidelity_weight)
                model_used = "codeformer"

            elif model_choice == "auto":
                try:
                    fidelity_weight = 1.0 - strength
                    enhanced = enhance_with_codeformer(img_array, fidelity_weight)
                    model_used = "codeformer"
                except Exception as e:
                    import traceback
                    print(f"[enhance-faces] CodeFormer failed, falling back to GFPGAN: {e}", file=sys.stderr, flush=True)
                    traceback.print_exc(file=sys.stderr)
                    emit_progress(50, "Falling back to GFPGAN")
                    enhanced = enhance_with_gfpgan(img_array, only_center_face)
                    model_used = "gfpgan"

        finally:
            # Restore stdout after ALL AI processing
            sys.stdout.flush()
            if stdout_fd is not None:
                os.dup2(stdout_fd, 1)
                os.close(stdout_fd)
            sys.stdout = sys.__stdout__  # Restore Python-level stdout

        if enhanced is None:
            raise RuntimeError("Face enhancement failed: no model available")

        emit_progress(85, "Enhancement complete")

        # Alpha blend result with original based on strength.
        # For CodeFormer, strength is already applied via fidelity_weight,
        # so skip the blend to avoid double-applying.
        # For GFPGAN (which has no fidelity knob), blend with original.
        if strength < 1.0 and model_used != "codeformer":
            blended = (
                img_array.astype(np.float32) * (1.0 - strength)
                + enhanced.astype(np.float32) * strength
            )
            enhanced = np.clip(blended, 0, 255).astype(np.uint8)

        emit_progress(95, "Saving result")
        Image.fromarray(enhanced).save(output_path)

        print(
            json.dumps(
                {
                    "success": True,
                    "facesDetected": num_faces,
                    "faces": faces,
                    "model": model_used,
                }
            )
        )

    except ImportError:
        print(
            json.dumps(
                {
                    "success": False,
                    "error": "Pillow is not installed. Install with: pip install Pillow",
                }
            )
        )
        sys.exit(1)
    except Exception as e:
        print(json.dumps({"success": False, "error": str(e)}))
        sys.exit(1)


if __name__ == "__main__":
    main()