#!/usr/bin/env python3
"""Detect faces in images and write annotated copies.

  detect_faces.py IMAGE [IMAGE ...] [--out DIR] [--min-size PX]

Uses OpenCV's YuNet detector when its model loads (~/models/face_detection_yunet_2023mar.onnx),
otherwise the Haar cascades shipped with opencv-data (frontal, plus profile in both directions).
Prints one line per image: the number of faces, then x,y,w,h boxes in pixels.
"""
import argparse
import os

import cv2
import numpy as np

HAAR = '/usr/share/opencv4/haarcascades'
YUNET = os.path.expanduser('~/models/face_detection_yunet_2023mar.onnx')


def yunet_detector(w, h):
    if not os.path.exists(YUNET):
        return None
    try:
        cv2.FaceDetectorYN.create(YUNET, '', (640, 480), 0.6, 0.3, 5000)
    except cv2.error:
        return None
    by_size = {}

    def run(img):
        # This OpenCV's YuNet needs input sides that are multiples of 32: work at about 640 wide.
        h0, w0 = img.shape[:2]
        scale = 640 / max(w0, h0) if max(w0, h0) > 640 else 1.0
        tw, th = max(32, round(w0 * scale / 32) * 32), max(32, round(h0 * scale / 32) * 32)
        small = cv2.resize(img, (tw, th), interpolation=cv2.INTER_AREA)
        if (tw, th) not in by_size:   # OpenCV 4.6 breaks if the input size changes after create()
            by_size[(tw, th)] = cv2.FaceDetectorYN.create(YUNET, '', (tw, th), 0.6, 0.3, 5000)
        _, faces = by_size[(tw, th)].detect(small)
        if faces is None:
            return []
        fx, fy = w0 / tw, h0 / th
        return [(int(f[0] * fx), int(f[1] * fy), int(f[2] * fx), int(f[3] * fy)) for f in faces]
    return run


def haar_detector(min_size):
    frontal = cv2.CascadeClassifier(os.path.join(HAAR, 'haarcascade_frontalface_alt2.xml'))
    profile = cv2.CascadeClassifier(os.path.join(HAAR, 'haarcascade_profileface.xml'))

    def run(img):
        gray = cv2.equalizeHist(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY))
        ms = (min_size, min_size)
        boxes = [tuple(b) for b in frontal.detectMultiScale(gray, 1.1, 5, minSize=ms)]
        for flip in (False, True):          # the profile cascade only sees one side
            g = cv2.flip(gray, 1) if flip else gray
            for x, y, w, h in profile.detectMultiScale(g, 1.1, 5, minSize=ms):
                boxes.append((gray.shape[1] - x - w if flip else x, y, w, h))
        return merge(boxes)
    return run


def merge(boxes, overlap=0.4):
    """Drop boxes that mostly overlap an earlier one (frontal and profile can both fire)."""
    kept = []
    for x, y, w, h in boxes:
        for kx, ky, kw, kh in kept:
            iw = min(x + w, kx + kw) - max(x, kx)
            ih = min(y + h, ky + kh) - max(y, ky)
            if iw > 0 and ih > 0 and iw * ih > overlap * min(w * h, kw * kh):
                break
        else:
            kept.append((int(x), int(y), int(w), int(h)))
    return kept


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('images', nargs='+')
    p.add_argument('--out', default='.')
    p.add_argument('--min-size', type=int, default=40)
    args = p.parse_args()
    os.makedirs(args.out, exist_ok=True)
    detector, name = None, None
    for path in args.images:
        img = cv2.imread(path)
        if img is None:
            print(f'{path}: unreadable')
            continue
        if detector is None:
            detector = yunet_detector(img.shape[1], img.shape[0])
            name = 'yunet' if detector else 'haar'
            detector = detector or haar_detector(args.min_size)
        boxes = detector(img)
        for x, y, w, h in boxes:
            cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
        dest = os.path.join(args.out, 'faces-' + os.path.basename(path).rsplit('.', 1)[0] + '.jpg')
        cv2.imwrite(dest, img)
        print(f'{path}: {len(boxes)} face(s) [{name}] {boxes} -> {dest}')


if __name__ == '__main__':
    main()