#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.10"
# dependencies = ["opencv-python-headless>=4.10,<5"]
# ///
"""Put solid black bars over the eyes of faces in image files.

Usage: uv run anon_eyes.py IMG_1116.JPG [other_image.jpg ...]
"""

from __future__ import annotations

import argparse
import hashlib
import os
from pathlib import Path
import tempfile
from urllib.request import urlopen

import cv2


MODEL_URL = (
    "https://github.com/opencv/opencv_zoo/raw/refs/heads/main/"
    "models/face_detection_yunet/face_detection_yunet_2023mar.onnx"
)
MODEL_SHA256 = "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4"
MODEL_PATH = Path.home() / ".cache" / "anon_eyes" / "face_detection_yunet_2023mar.onnx"


def get_model() -> Path:
    if MODEL_PATH.exists() and hashlib.sha256(MODEL_PATH.read_bytes()).hexdigest() == MODEL_SHA256:
        return MODEL_PATH

    MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    print("Downloading the tiny face detector (first run only)...")
    with urlopen(MODEL_URL, timeout=30) as response:
        data = response.read()
    if hashlib.sha256(data).hexdigest() != MODEL_SHA256:
        raise RuntimeError("Downloaded face detector failed its checksum")
    with tempfile.NamedTemporaryFile(dir=MODEL_PATH.parent, delete=False) as tmp:
        tmp.write(data)
        temp_path = Path(tmp.name)
    os.replace(temp_path, MODEL_PATH)
    return MODEL_PATH


def detect_faces(image):
    """Return face boxes and eye landmarks in original image coordinates."""
    height, width = image.shape[:2]
    model = str(get_model())
    candidates = []
    # Two sizes help with both small group faces and large close-up portraits.
    for max_side in sorted({min(max(width, height), size) for size in (640, 1600)}):
        scale = max_side / max(width, height)
        small = cv2.resize(image, None, fx=scale, fy=scale) if scale < 1 else image
        detector = cv2.FaceDetectorYN.create(
            model, "", (small.shape[1], small.shape[0]), score_threshold=0.7
        )
        _, faces = detector.detect(small)
        if faces is not None:
            candidates.extend((float(face[-1]), face[:8] / scale) for face in faces)

    def overlap(a, b):
        ax, ay, aw, ah = a[:4]
        bx, by, bw, bh = b[:4]
        intersection = max(0, min(ax + aw, bx + bw) - max(ax, bx)) * max(
            0, min(ay + ah, by + bh) - max(ay, by)
        )
        return intersection / (aw * ah + bw * bh - intersection)

    found = []
    for _, face in sorted(candidates, key=lambda item: item[0], reverse=True):
        if all(overlap(face, other) < 0.4 for other in found):
            found.append(face)
    return found


def cover_eyes(image, face):
    _, _, width, height, eye1_x, eye1_y, eye2_x, eye2_y = face
    image_height, image_width = image.shape[:2]
    # Cover both landmarks generously, including tilted eyes and eyelids.
    left = max(0, round(min(eye1_x, eye2_x) - width * 0.15))
    right = min(image_width, round(max(eye1_x, eye2_x) + width * 0.15))
    top = max(0, round(min(eye1_y, eye2_y) - height * 0.09))
    bottom = min(image_height, round(max(eye1_y, eye2_y) + height * 0.09))
    cv2.rectangle(image, (left, top), (right, bottom), (0, 0, 0), -1)


def process(path: Path) -> None:
    if path.stem.endswith("_out"):
        print(f"Skipping output file: {path}")
        return
    image = cv2.imread(str(path), cv2.IMREAD_COLOR)
    if image is None:
        raise ValueError(f"Could not read image: {path}")

    faces = detect_faces(image)
    for face in faces:
        cover_eyes(image, face)

    output = path.with_name(f"{path.stem}_out{path.suffix}")
    options = [cv2.IMWRITE_JPEG_QUALITY, 95] if path.suffix.lower() in (".jpg", ".jpeg") else []
    if not cv2.imwrite(str(output), image, options):
        raise OSError(f"Could not write image: {output}")
    print(f"{path} -> {output} ({len(faces)} face(s))")
    if not faces:
        print("  Warning: no faces detected; output is unchanged.")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("images", nargs="+", type=Path, help="image files to process")
    args = parser.parse_args()
    for path in args.images:
        process(path)


if __name__ == "__main__":
    main()