Python×OpenCVで「ブレ・半目写真」を自動で除外！顔認識付き画像選別ツールを作ってみました

python select_good_images.py --input ./photos --output ./selected --verbose

from __future__ import annotations

import argparse
import logging
import shutil
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Tuple, Dict, Any, Optional
import os
from concurrent.futures import ProcessPoolExecutor

import cv2
import numpy as np

# --- グローバル変数 (ワーカープロセスごと) ---
# これらは init_worker で初期化されます
worker_face_cascade: Optional[cv2.CascadeClassifier] = None
worker_profile_cascade: Optional[cv2.CascadeClassifier] = None
worker_eye_cascade: Optional[cv2.CascadeClassifier] = None
worker_cfg: Optional[Config] = None
# ------------------------------------------

Rect = Tuple[int, int, int, int]


@dataclass
class Config:
    # IO
    input_dir: Path = Path("E:/fuyuppi/screenshot/15")
    output_dir: Path = Path("./output_good_images15")

    # Cascade paths
    cascade_face: Path = Path("haarcascade_frontalface_default.xml")
    cascade_eye: Path = Path("haarcascade_eye.xml")
    cascade_profile: Path = Path("haarcascade_profileface.xml")  # ← 横顔追加

    # Thresholds
    blur_face_min: float = 150.0
    blur_full_min: float = 75.0
    min_open_eye_height_px: int = 10
    min_open_eye_height_ratio: float = 0.05

    # Selection policy
    require_all_faces: bool = False

    # Detection params
    face_scale_factor: float = 1.1
    face_min_neighbors: int = 5
    face_min_size: Tuple[int, int] = (30, 30)

    eye_scale_factor: float = 1.1
    eye_min_neighbors: int = 5
    eye_min_size: Tuple[int, int] = (15, 15)

    eye_region_upper_ratio: float = 0.6

    image_patterns: List[str] = field(
        default_factory=lambda: ["*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG"]
    )

    check_open_eyes: bool = False
    use_full_blur_gate: bool = False
    verbose: bool = False


def setup_logger(verbose: bool) -> None:
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%H:%M:%S",
    )


def resolve_cascade_path(path: Path, fallback_name: str) -> Path:
    if path.exists():
        return path
    try:
        candidate = Path(cv2.data.haarcascades) / fallback_name
        if candidate.exists():
            return candidate
    except Exception:
        pass
    return path


def load_cascade(path: Path) -> cv2.CascadeClassifier:
    clf = cv2.CascadeClassifier(str(path))
    if clf.empty():
        raise FileNotFoundError(f"Failed to load cascade: {path}")
    return clf


def list_images(input_dir: Path, patterns: List[str]) -> List[Path]:
    files: List[Path] = []
    for pat in patterns:
        files.extend(input_dir.glob(pat))
    return sorted(set(files))


def laplacian_variance(gray: np.ndarray) -> float:
    return float(cv2.Laplacian(gray, cv2.CV_64F).var())


def detect_faces(
    gray: np.ndarray,
    face_cascade: cv2.CascadeClassifier,
    profile_cascade: cv2.CascadeClassifier,
    cfg: Config,
) -> List[Rect]:
    # 正面顔
    faces = face_cascade.detectMultiScale(
        gray,
        scaleFactor=cfg.face_scale_factor,
        minNeighbors=cfg.face_min_neighbors,
        minSize=cfg.face_min_size,
    )

    # 横顔（右向き）
    profiles = profile_cascade.detectMultiScale(
        gray,
        scaleFactor=cfg.face_scale_factor,
        minNeighbors=cfg.face_min_neighbors,
        minSize=cfg.face_min_size,
    )

    # 横顔（左向き）→画像を反転して検出
    flipped = cv2.flip(gray, 1)
    profiles_flipped = profile_cascade.detectMultiScale(
        flipped,
        scaleFactor=cfg.face_scale_factor,
        minNeighbors=cfg.face_min_neighbors,
        minSize=cfg.face_min_size,
    )
    w = gray.shape[1]
    profiles_flipped_corrected = [
        (w - x - w_, y, w_, h_) for (x, y, w_, h_) in profiles_flipped
    ]

    all_faces = list(faces) + list(profiles) + profiles_flipped_corrected
    return all_faces


def detect_open_eyes_count(
    face_gray: np.ndarray,
    face_h: int,
    eye_cascade: cv2.CascadeClassifier,
    cfg: Config,
) -> int:
    h, w = face_gray.shape[:2]
    upper_h = max(1, int(cfg.eye_region_upper_ratio * h))
    eye_region = face_gray[:upper_h, :]
    eye_region_eq = cv2.equalizeHist(eye_region)

    eyes = eye_cascade.detectMultiScale(
        eye_region_eq,
        scaleFactor=cfg.eye_scale_factor,
        minNeighbors=cfg.eye_min_neighbors,
        minSize=cfg.eye_min_size,
    )
    if len(eyes) == 0:
        return 0

    th_px = max(cfg.min_open_eye_height_px, int(cfg.min_open_eye_height_ratio * face_h))
    filtered = [(ex, ey, ew, eh) for (ex, ey, ew, eh) in eyes if ew > eh]
    if not filtered:
        filtered = list(eyes)
    filtered.sort(key=lambda r: r[3], reverse=True)
    top = filtered[:2]
    open_count = sum(1 for (_, _, _, eh) in top if eh >= th_px)
    return open_count


def evaluate_face(
    full_gray: np.ndarray,
    face_rect: Rect,
    eye_cascade: cv2.CascadeClassifier,
    cfg: Config,
) -> Tuple[bool, Dict[str, Any]]:
    x, y, w, h = face_rect
    roi_gray = full_gray[y : y + h, x : x + w]
    lap_face = laplacian_variance(roi_gray)

    if lap_face < cfg.blur_face_min:
        return (
            False,
            {
                "rect": (int(x), int(y), int(w), int(h)),
                "lap_face": lap_face,
                "open_eye_count": 0,
                "reason": f"face_blur {lap_face:.2f} < {cfg.blur_face_min}",
            },
        )

    if not cfg.check_open_eyes:
        return (
            True,
            {
                "rect": (int(x), int(y), int(w), int(h)),
                "lap_face": lap_face,
                "open_eye_count": None,
                "reason": "ok_blur_only",
            },
        )

    open_eye_count = detect_open_eyes_count(roi_gray, h, eye_cascade, cfg)
    if open_eye_count < 2:
        return (
            False,
            {
                "rect": (int(x), int(y), int(w), int(h)),
                "lap_face": lap_face,
                "open_eye_count": open_eye_count,
                "reason": f"half-eye (open={open_eye_count})",
            },
        )

    return (
        True,
        {
            "rect": (int(x), int(y), int(w), int(h)),
            "lap_face": lap_face,
            "open_eye_count": open_eye_count,
            "reason": "ok",
        },
    )


def evaluate_image(
    img_path: Path,
    face_cascade: cv2.CascadeClassifier,
    profile_cascade: cv2.CascadeClassifier,
    eye_cascade: cv2.CascadeClassifier,
    cfg: Config,
) -> Tuple[bool, Dict[str, Any]]:
    img = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
    if img is None:
        return False, {"reason": "imread_failed"}

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    lap_full = laplacian_variance(gray)

    if cfg.use_full_blur_gate and lap_full < cfg.blur_full_min:
        return False, {"reason": f"full_blur {lap_full:.2f} < {cfg.blur_full_min}", "lap_full": lap_full}

    faces = detect_faces(gray, face_cascade, profile_cascade, cfg)
    if len(faces) == 0:
        return False, {"reason": "no_face", "lap_full": lap_full}

    face_results = []
    any_pass = False
    all_pass = True

    for rect in faces:
        ok, detail = evaluate_face(gray, rect, eye_cascade, cfg)
        face_results.append({"ok": ok, **detail})
        if ok:
            any_pass = True
        else:
            all_pass = False

    selected = all_pass if cfg.require_all_faces else any_pass

    return selected, {
        "lap_full": lap_full,
        "faces": face_results,
        "face_count": len(faces),
        "policy": "all_faces" if cfg.require_all_faces else "any_face",
    }


# --- ここから並列処理用の新関数 ---

def init_worker(cfg: Config) -> None:
    """
    ProcessPoolExecutor の各ワーカープロセスを初期化する。
    カスケードファイルはプロセスごとに1回だけロードされる。
    """
    global worker_face_cascade, worker_profile_cascade, worker_eye_cascade, worker_cfg
    
    # logging.info(f"Worker {os.getpid()} initializing...") # デバッグ時以外はコメントアウト
    worker_cfg = cfg
    
    try:
        face_path = resolve_cascade_path(cfg.cascade_face, "haarcascade_frontalface_default.xml")
        eye_path = resolve_cascade_path(cfg.cascade_eye, "haarcascade_eye.xml")
        profile_path = resolve_cascade_path(cfg.cascade_profile, "haarcascade_profileface.xml")

        worker_face_cascade = load_cascade(face_path)
        worker_eye_cascade = load_cascade(eye_path)
        worker_profile_cascade = load_cascade(profile_path)
        
    except Exception as e:
        logging.error(f"Worker {os.getpid()} failed to initialize: {e}")

def process_single_image(img_path: Path) -> Tuple[Path, bool, Dict[str, Any]]:
    """
    単一の画像を評価するワーカー関数。
    ワーカープロセスのグローバルカスケードを使用する。
    """
    global worker_face_cascade, worker_profile_cascade, worker_eye_cascade, worker_cfg

    if not all([worker_face_cascade, worker_profile_cascade, worker_eye_cascade, worker_cfg]):
        msg = f"Worker {os.getpid()} not initialized."
        logging.error(msg)
        return (img_path, False, {"reason": msg})

    try:
        selected, info = evaluate_image(
            img_path,
            worker_face_cascade,
            worker_profile_cascade,
            worker_eye_cascade,
            worker_cfg
        )
        return (img_path, selected, info)
    except Exception as e:
        logging.exception(f"評価中にエラーが発生しました: {img_path.name}: {e}")
        return (img_path, False, {"reason": f"Exception: {e}"})

# --- ここまで ---


def process_images(cfg: Config) -> int:
    """
    並列処理で画像フォルダ全体を処理する (修正版)
    """
    cfg.output_dir.mkdir(parents=True, exist_ok=True)

    images = list_images(cfg.input_dir, cfg.image_patterns)
    if not images:
        logging.warning(f"入力フォルダに画像が見つかりません: {cfg.input_dir}")
        return 0

    logging.info(f"合計 {len(images)} 枚の画像を処理します (並列処理)...")
    selected_count = 0
    
    # CPUコア数をワーカー数として設定 (Noneで自動設定)
    # initializer で各ワーカーの初期化 (カスケードロード) を行う
    worker_count = max(1, os.cpu_count() - 1)
    logging.info(f"並列処理ワーカー数: {worker_count} (全コア数: {os.cpu_count()})")
    with ProcessPoolExecutor(max_workers=worker_count, initializer=init_worker, initargs=(cfg,)) as executor:
    
    
        # executor.map が画像のリストを各ワーカーに分配して処理
        results = executor.map(process_single_image, images)

        # 結果をメインスレッドで受け取り、ログ出力とファイルコピーを行う
        for img_path, selected, info in results:
            logging.info(f"--- 評価完了: {img_path.name} ---")
            
            if "lap_full" in info:
                logging.debug(f"全体シャープネス(LaplacianVar): {info['lap_full']:.2f}")
            if "faces" in info:
                for i, fr in enumerate(info["faces"], start=1):
                    rect = fr.get("rect", (-1, -1, -1, -1))
                    logging.debug(
                        f"顔#{i} rect={rect} lap_face={fr.get('lap_face', -1):.2f} "
                        f"open_eyes={fr.get('open_eye_count', 'NA')} ok={fr.get('ok', False)} reason={fr.get('reason', '')}"
                    )
            
            if not selected:
                logging.info(f"選別外: {info.get('reason', 'criteria_not_met')}")
                continue

            dst = cfg.output_dir / img_path.name
            try:
                shutil.copy2(img_path, dst)
                selected_count += 1
                logging.info(f"選別通過: {img_path.name} -> {dst}")
            except Exception as e:
                logging.exception(f"コピーに失敗しました: {img_path.name}: {e}")

    logging.info("=" * 32)
    logging.info(f"処理完了。選別通過枚数: {selected_count} 枚")
    logging.info("=" * 32)
    return selected_count


def build_arg_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="ブレ/半目を除外する画像選別ツール（横顔対応）")
    p.add_argument("--input", type=Path, default=Path(r"E:\fuyuppi\screenshot\15"), help="入力ディレクトリ")
    p.add_argument("--output", type=Path, default=Path("./output_good_images15"), help="出力ディレクトリ")
    p.add_argument("--face-cascade", type=Path, default=Path("haarcascade_frontalface_default.xml"))
    p.add_argument("--eye-cascade", type=Path, default=Path("haarcascade_eye.xml"))
    p.add_argument("--profile-cascade", type=Path, default=Path("haarcascade_profileface.xml"), help="横顔カスケードのパス")
    p.add_argument("--face-blur-th", type=float, default=150.0)
    p.add_argument("--full-blur-th", type=float, default=75.0)
    p.add_argument("--min-eye-px", type=int, default=10)
    p.add_argument("--min-eye-ratio", type=float, default=0.05)

    grp = p.add_mutually_exclusive_group()
    grp.add_argument("--any-face", dest="any_face", action="store_true")
    grp.add_argument("--all-faces", dest="any_face", action="store_false")
    p.set_defaults(any_face=True)

    p.add_argument("--check-eyes", action="store_true")
    p.add_argument("--use-full-blur-gate", action="store_true")
    p.add_argument("--verbose", action="store_true")
    return p


def main(argv: List[str]) -> int:
    parser = build_arg_parser()
    args = parser.parse_args(argv)

    cfg = Config(
        input_dir=args.input,
        output_dir=args.output,
        cascade_face=args.face_cascade,
        cascade_eye=args.eye_cascade,
        cascade_profile=args.profile_cascade,
        blur_face_min=args.face_blur_th,
        blur_full_min=args.full_blur_th,
        min_open_eye_height_px=args.min_eye_px,
        min_open_eye_height_ratio=args.min_eye_ratio,
        require_all_faces=not args.any_face,
        check_open_eyes=args.check_eyes,
        use_full_blur_gate=args.use_full_blur_gate,
        verbose=args.verbose,
    )

    setup_logger(cfg.verbose)

    if not cfg.input_dir.exists() or not cfg.input_dir.is_dir():
        logging.error(f"入力ディレクトリが存在しません: {cfg.input_dir}")
        return 2

    try:
        count = process_images(cfg)
    except FileNotFoundError as e:
        logging.error(str(e))
        return 3
    except Exception as e:
        logging.exception(f"予期せぬエラー: {e}")
        return 1

    return 0 if count >= 0 else 1


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))