#!/usr/bin/env python3
"""
オンデバイス顔ぼかしカメラ（Raspberry Pi / Jetson / PC 共通）
- カメラ映像の顔を「端末内」で自動検出し、ぼかし/モザイク/塗りつぶし。クラウド送信なし。
- 顔検出は OpenCV の FaceDetectorYN（YuNet, ONNX）。mediapipe不要でPi/Jetson/PCで動作。

使い方:
    python3 face_blur_cam.py                  # USB/標準カメラ・ぼかし
    python3 face_blur_cam.py --mode mosaic    # モザイク
    python3 face_blur_cam.py --picamera       # Raspberry Pi 内蔵カメラ(CSI)
    python3 face_blur_cam.py --save out.mp4   # 録画も保存
    python3 face_blur_cam.py --no-display     # 画面表示なし(ヘッドレス)
    q キー または Ctrl+C で終了。

提供: エッジAIラボ https://ai-edge-lab.com/demos/face-mosaic/
"""
import argparse
import os
import sys
import time
import urllib.request

import cv2

MODEL_FILE = "face_detection_yunet.onnx"
MODEL_URL = "https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx"


def ensure_model() -> str:
    if not os.path.exists(MODEL_FILE):
        print("顔検出モデル(YuNet)をダウンロード中...")
        urllib.request.urlretrieve(MODEL_URL, MODEL_FILE)
    return MODEL_FILE


def open_camera(use_picamera: bool, width: int, height: int):
    """Piカメラ(CSI) or USB/標準カメラ。返り値: (種別, ハンドル)"""
    if use_picamera:
        try:
            from picamera2 import Picamera2  # Raspberry Pi OS に同梱
        except ImportError:
            print("picamera2 が見つかりません。Raspberry Pi OS で実行するか、"
                  "USBカメラ（--picamera を外す）をご利用ください。", file=sys.stderr)
            sys.exit(1)
        try:
            picam = Picamera2()
            picam.configure(picam.create_preview_configuration(
                main={"format": "RGB888", "size": (width, height)}))
            picam.start()
            time.sleep(0.5)
        except Exception as e:  # カメラ未接続など
            print("Piカメラを認識できませんでした（カメラモジュール未接続の可能性）。\n"
                  "  ・CSIカメラの接続を確認してください\n"
                  "  ・USBカメラを使う場合は --picamera を外してください\n"
                  f"  詳細: {e}", file=sys.stderr)
            sys.exit(1)
        return ("picamera", picam)
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
    if not cap.isOpened():
        print("カメラを開けませんでした。USBカメラの接続を確認するか、"
              "Raspberry Pi の内蔵(CSI)カメラなら --picamera を付けて実行してください。", file=sys.stderr)
        sys.exit(1)
    return ("cv2", cap)


def read_frame(kind, src):
    if kind == "picamera":
        frame = src.capture_array()
        return True, cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    return src.read()


def mask_face(frame, x, y, w, h, mode, strength):
    """検出boxを顔全体(額/髪/あご)まで広げて隠す。"""
    H, W = frame.shape[:2]
    mx, my = int(w * 0.18), int(h * 0.22)
    x0, y0 = max(0, x - mx), max(0, y - my)
    x1, y1 = min(W, x + w + mx), min(H, y + h + my)
    if x1 <= x0 or y1 <= y0:
        return
    roi = frame[y0:y1, x0:x1]
    if mode == "mosaic":
        bw = max(1, (x1 - x0) // max(4, strength))
        bh = max(1, (y1 - y0) // max(4, strength))
        small = cv2.resize(roi, (bw, bh), interpolation=cv2.INTER_LINEAR)
        frame[y0:y1, x0:x1] = cv2.resize(small, (x1 - x0, y1 - y0), interpolation=cv2.INTER_NEAREST)
    elif mode == "fill":
        frame[y0:y1, x0:x1] = (0, 0, 0)
    else:  # blur
        frame[y0:y1, x0:x1] = cv2.GaussianBlur(roi, (0, 0), max(1, strength))


def main():
    ap = argparse.ArgumentParser(description="オンデバイス顔ぼかしカメラ")
    ap.add_argument("--mode", choices=["blur", "mosaic", "fill"], default="blur", help="隠し方")
    ap.add_argument("--strength", type=int, default=25, help="強さ(大きいほど強い)")
    ap.add_argument("--picamera", action="store_true", help="Raspberry Pi 内蔵カメラ(CSI)を使う")
    ap.add_argument("--width", type=int, default=640)
    ap.add_argument("--height", type=int, default=480)
    ap.add_argument("--conf", type=float, default=0.6, help="検出しきい値(下げると拾いやすい)")
    ap.add_argument("--save", default=None, help="録画ファイル名(例 out.mp4)")
    ap.add_argument("--no-display", action="store_true", help="画面表示なし(ヘッドレス)")
    args = ap.parse_args()

    model = ensure_model()
    # create(model, config, input_size, score_threshold, nms_threshold, top_k)
    detector = cv2.FaceDetectorYN.create(model, "", (args.width, args.height), args.conf, 0.3, 5000)

    kind, src = open_camera(args.picamera, args.width, args.height)
    writer = None
    t0, frames, fps = time.time(), 0, 0.0
    print("実行中。'q' または Ctrl+C で終了します。（すべて端末内処理・送信なし）")
    try:
        while True:
            ok, frame = read_frame(kind, src)
            if not ok:
                break
            h, w = frame.shape[:2]
            detector.setInputSize((w, h))
            _, faces = detector.detect(frame)
            n = 0
            if faces is not None:
                for f in faces:
                    fx, fy, fw, fh = f[:4].astype(int)
                    mask_face(frame, fx, fy, fw, fh, args.mode, args.strength)
                    n += 1
            frames += 1
            if time.time() - t0 >= 1.0:
                fps = frames / (time.time() - t0)
                t0, frames = time.time(), 0
            cv2.putText(frame, f"faces:{n}  {fps:.1f}fps  local/no-cloud", (8, 24),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            if args.save:
                if writer is None:
                    writer = cv2.VideoWriter(args.save, cv2.VideoWriter_fourcc(*"mp4v"), 15, (w, h))
                writer.write(frame)
            if not args.no_display:
                cv2.imshow("face-blur  (q=quit)", frame)
                if cv2.waitKey(1) & 0xFF == ord("q"):
                    break
    except KeyboardInterrupt:
        pass
    finally:
        if writer:
            writer.release()
        if kind == "cv2":
            src.release()
        else:
            src.stop()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()