Fork facesaver -> anime-facesaver: real anime face detector instead of COCO 'person' proxy

Original saved any frame with a big COCO person box (heads, hands,
shoes). Switches to yolov8x6_animeface (HF: Fuyucchi), keeps the
scene-change gate, adds --gpu pinning via CUDA_VISIBLE_DEVICES before
torch init, --download-model, --preview-dir QC mode, optional edge
margin and min face-area fraction.

.gitignore

@@ -0,0 +1,8 @@
+.venv/
+__pycache__/
+*.pyc
+models/
+output/
+preview/
+*.pt
+*.part

README.md

@@ -0,0 +1,83 @@
+# anime-facesaver
+
+Scene-gated still extractor for building anime training datasets from
+video. Fork of `facesaver`, created Oct 2026.
+
+## Why this fork exists
+
+The original facesaver ran stock COCO `yolov11l.pt` and accepted
+class 0 — **`person`** — as a "proxy for face" (their comment, not ours).
+COCO has no face class, so every frame containing a large enough person
+passed: full-body shots, headless crops, a hand at 75% confidence. That
+is why the old Gabriel Dropout output contained stills of shoes.
+
+anime-facesaver uses a dedicated anime-face detector instead:
+[`yolov8x6_animeface`](https://huggingface.co/Fuyucchi/yolov8_animeface)
+— yolov8x6 fine-tuned at 1280px on 10k manually annotated anime images
+(P=0.957 R=0.924 mAP50=0.955 on its test set). It detects faces. Only
+faces. The scene-change gate from the original is preserved: at most one
+still is saved per detected scene, which keeps near-duplicates out of
+the dataset.
+
+## Setup
+
+```bash
+uv venv --python 3.12 .venv           # or: python3.12 -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt       # add --index-url for your CUDA ver
+python anime_facesaver.py --download-model ...   # fetches weights once
+```
+
+Weights land in `models/yolov8x6_animeface.pt` (~195 MB, AGPL-3.0).
+
+## Usage
+
+```bash
+python anime_facesaver.py \
+    -I /home/alex/Documents/image-training-datasets/gabriel-dropout-show \
+    -O /home/alex/Documents/image-training-datasets/gabriel-dropout-v2 \
+    --gpu 1 \
+    --preview-dir /home/alex/Documents/image-training-datasets/gabriel-dropout-v2-preview
+```
+
+### GPU pinning (important on the shared .80 box)
+
+`--gpu` sets `CUDA_VISIBLE_DEVICES` **before** torch initializes, so no
+context leaks onto the wrong card:
+
+| flag      | effect                                        |
+|-----------|-----------------------------------------------|
+| `--gpu 1` | pin to physical GPU 1 (recommended on .80)    |
+| `--gpu 0` | pin to GPU 0 — OOMs while Strata holds it     |
+| `--gpu cpu` | force CPU                                   |
+| `--gpu auto` | let torch pick                             |
+
+### Key options
+
+| flag | default | meaning |
+|------|---------|---------|
+| `--conf` | 0.3 | detector confidence (model's own recommendation; it is high-precision) |
+| `--min-width` / `--min-height` | 100 | minimum face box size in px |
+| `--min-face-fraction` | 0.005 | face box area as fraction of frame — kills distant background extras; `0` disables |
+| `--edge-margin` | 0 | require faces this far inside frame edges; `0` accepts edge-touching faces (original hard-rejected them) |
+| `--scene-threshold` | 30.0 | PySceneDetect ContentDetector threshold (same as original) |
+| `--imgsz` | 1280 | inference size; model was trained at 1280, `640` is faster at some recall cost |
+| `--preview-dir` | off | also write annotated copies (green boxes) of every saved frame for eyeball QC |
+
+## Notes
+
+- Runs every frame of every scene until a face qualifies, then waits
+  for the next scene start. That is the original behavior and it is why
+  dense scenes don't flood the output.
+- The detector ships as `yolov8x6` — big and a little slow (it was
+  trained to be precise, not cheap). On a 3090 at 1280px expect roughly
+  8-12 fps; an hour of 24fps anime is ~15 min of inference per GPU.
+- **License**: the model weights are AGPL-3.0 (Fuyucchi); ultralytics is
+  AGPL-3.0 too. Fine for personal/research dataset building; if this
+  ever ships in a product, revisit.
+
+## Verification
+
+Run `--preview-dir` and eyeball the `_ann.png` files: every saved frame
+should show a green box on a real face. Compare against the original
+facesaver output to quantify how many old saves had no face at all.

anime_facesaver.py

@@ -0,0 +1,336 @@
+#!/usr/bin/env python3
+"""anime-facesaver: save one still per scene from videos that contain an
+anime face.
+
+Fork of facesaver (https://github.com/.../facesaver) fixing its core bug:
+the original used stock COCO yolov11l with class 0 ('person') as a proxy
+for 'face', so it happily saved stills of headless bodies, hands and
+shoes. This version uses a dedicated anime-face detector
+(yolov8x6_animeface, AGPL-3.0, by Fuyucchi) and only saves frames where
+an actual face of sufficient size is found.
+
+Scene-change gate preserved from the original: at most one still is saved
+per detected scene, which keeps the output dataset from exploding with
+near-duplicates.
+"""
+
+import argparse
+import os
+import sys
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(
+        description=(
+            "Detect anime faces in videos and save one still per scene "
+            "where a qualifying face appears."
+        ),
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--input-dir", "-I",
+        required=True,
+        help="Directory containing input video files.",
+    )
+    parser.add_argument(
+        "--output-dir", "-O",
+        required=True,
+        help="Directory to save screenshot outputs.",
+    )
+    parser.add_argument(
+        "--gpu", "-g",
+        default="1",
+        help=(
+            "GPU index to pin to (e.g. '1'), or 'cpu' to force CPU. "
+            "Use 'auto' to let torch pick the first visible device. "
+            "On the shared .80 box, GPU0 is usually held by the inference "
+            "server — pinning to GPU1 avoids OOM."
+        ),
+    )
+    parser.add_argument(
+        "--model",
+        default=os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                             "models", "yolov8x6_animeface.pt"),
+        help="Path to the anime-face YOLO weights.",
+    )
+    parser.add_argument(
+        "--download-model",
+        action="store_true",
+        help="Download the yolov8x6_animeface weights if --model is missing.",
+    )
+    parser.add_argument(
+        "--min-width", "-w",
+        type=int,
+        default=100,
+        help="Minimum face bounding box width (px) to trigger a save.",
+    )
+    parser.add_argument(
+        "--min-height", "-m",
+        type=int,
+        default=100,
+        help="Minimum face bounding box height (px) to trigger a save.",
+    )
+    parser.add_argument(
+        "--min-face-fraction",
+        type=float,
+        default=0.005,
+        help="Minimum face box area as a fraction of frame area (kills "
+             "background-extra faces in wide shots; 0 disables).",
+    )
+    parser.add_argument(
+        "--conf",
+        type=float,
+        default=0.3,
+        help="Detector confidence threshold (model ships recommending 0.3; "
+             "it is high-precision, raise for a stricter dataset).",
+    )
+    parser.add_argument(
+        "--iou",
+        type=float,
+        default=0.5,
+        help="NMS IoU threshold.",
+    )
+    parser.add_argument(
+        "--edge-margin",
+        type=int,
+        default=0,
+        help="Require the face box to stay this many px inside all frame "
+             "edges. 0 (default) allows cropped faces; set >0 to reject "
+             "faces touching the crop lines.",
+    )
+    parser.add_argument(
+        "--scene-threshold",
+        type=float,
+        default=30.0,
+        help="PySceneDetect ContentDetector threshold (same as original).",
+    )
+    parser.add_argument(
+        "--imgsz",
+        type=int,
+        default=1280,
+        help="Inference image size. The model was trained at 1280; lower "
+             "it (640) for speed at some recall cost.",
+    )
+    parser.add_argument(
+        "--preview-dir",
+        default=None,
+        help="Optional: also copy each SAVED frame plus an annotated "
+             "version (face boxes drawn) here for eyeball QC.",
+    )
+    return parser.parse_args()
+
+
+def resolve_gpu(gpu_arg):
+    """Pin CUDA before torch initializes. Returns torch device string."""
+    if gpu_arg.lower() == "cpu":
+        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+        return "cpu"
+    if gpu_arg.lower() == "auto":
+        return None  # let ultralytics/torch decide
+    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_arg
+    return "cuda"
+
+
+def ensure_directory(directory):
+    if directory and not os.path.exists(directory):
+        os.makedirs(directory)
+
+
+def download_model(path):
+    """Fetch yolov8x6_animeface.pt from Hugging Face (Fuyucchi/yolov8_animeface)."""
+    url = ("https://huggingface.co/Fuyucchi/yolov8_animeface/"
+           "resolve/main/yolov8x6_animeface.pt")
+    ensure_directory(os.path.dirname(path))
+    tmp = path + ".part"
+    print(f"Downloading anime-face model to {path} ...")
+    import urllib.request
+
+    def _hook(n, bs, total):
+        if total > 0 and n % 20000 == 0:
+            pct = min(100.0, n * bs * 100.0 / total)
+            print(f"\r  {pct:5.1f}%", end="", flush=True)
+
+    urllib.request.urlretrieve(url, tmp, reporthook=_hook)
+    print()
+    os.replace(tmp, path)
+    print(f"Saved model ({os.path.getsize(path) / 1e6:.0f} MB).")
+
+
+def face_qualifies(box, frame_shape, min_width, min_height,
+                   min_face_fraction, edge_margin):
+    """Original is_full_face(), corrected: the box must actually be a face
+    (enforced by the detector class), size gates stay, hard
+    'must-not-touch-edge' rule becomes an optional margin."""
+    x1, y1, x2, y2 = box
+    frame_h, frame_w = frame_shape[:2]
+
+    w, h = x2 - x1, y2 - y1
+    if w < min_width or h < min_height:
+        return False
+
+    if min_face_fraction > 0:
+        if (w * h) < (frame_w * frame_h * min_face_fraction):
+            return False
+
+    if edge_margin > 0:
+        if (x1 < edge_margin or y1 < edge_margin
+                or x2 > frame_w - edge_margin
+                or y2 > frame_h - edge_margin):
+            return False
+
+    return True
+
+
+def process_video(video_path, output_dir, args, model, device, draw,
+                  output_start=0):
+    from scenedetect import open_video, SceneManager, ContentDetector
+
+    video_name = os.path.splitext(os.path.basename(video_path))[0]
+
+    try:
+        import cv2
+        video = open_video(video_path)
+        scene_manager = SceneManager()
+        scene_manager.add_detector(
+            ContentDetector(threshold=args.scene_threshold))
+    except Exception as e:
+        print(f"Error initializing video for scene detection in "
+              f"{video_path}: {e}")
+        return output_start
+
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        print(f"Error opening video file {video_path}")
+        return output_start
+
+    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    try:
+        scene_manager.detect_scenes(video=video)
+        scene_list = scene_manager.get_scene_list()
+        scene_starts = {scene[0].get_frames() for scene in scene_list}
+    except Exception as e:
+        print(f"Error detecting scenes in {video_path}: {e}")
+        cap.release()
+        return output_start
+
+    print(f"{video_name}: {frame_count} frames, "
+          f"{len(scene_starts)} scene starts")
+
+    output_count = output_start
+    face_in_scene = False
+    frame_idx = 0
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        if frame_idx in scene_starts:
+            face_in_scene = False
+        # membership test on a set; pop matched starts to keep it cheap
+        # (scenes strictly increase, so set stays small either way)
+
+        if not face_in_scene:
+            try:
+                results = model.predict(frame, conf=args.conf,
+                                        iou=args.iou, imgsz=args.imgsz,
+                                        device=device, verbose=False)
+                for result in results:
+                    if result.boxes is None:
+                        continue
+                    boxes = result.boxes.xyxy.cpu().numpy()
+                    for box in boxes:
+                        if face_qualifies(box, frame.shape, args.min_width,
+                                          args.min_height,
+                                          args.min_face_fraction,
+                                          args.edge_margin):
+                            out = os.path.join(
+                                output_dir,
+                                f"{video_name}_face_{output_count:04d}.png")
+                            cv2.imwrite(out, frame)
+                            print(f"  saved {os.path.basename(out)} "
+                                  f"(frame {frame_idx}, "
+                                  f"box {int(box[2]-box[0])}x"
+                                  f"{int(box[3]-box[1])})")
+                            output_count += 1
+                            face_in_scene = True
+                            if args.preview_dir:
+                                annotated = draw(frame.copy(),
+                                                 result.boxes.xyxy.cpu())
+                                cv2.imwrite(
+                                    os.path.join(
+                                        args.preview_dir,
+                                        os.path.basename(out).replace(
+                                            ".png", "_ann.png")),
+                                    annotated)
+                            break
+                    if face_in_scene:
+                        break
+            except Exception as e:
+                print(f"Error during face detection in {video_path} "
+                      f"at frame {frame_idx}: {e}")
+
+        frame_idx += 1
+
+    cap.release()
+    print(f"Processed {video_name}: "
+          f"{output_count - output_start} screenshots saved.")
+    return output_count
+
+
+def main():
+    args = parse_arguments()
+
+    if not os.path.isdir(args.input_dir):
+        print(f"Error: Input directory '{args.input_dir}' does not exist.")
+        sys.exit(1)
+
+    if args.download_model and not os.path.exists(args.model):
+        download_model(args.model)
+    if not os.path.exists(args.model):
+        print(f"Error: model weights not found at {args.model}\n"
+              f"Re-run with --download-model, or place the weights there "
+              f"(see README).")
+        sys.exit(1)
+
+    device = resolve_gpu(args.gpu)
+
+    import torch
+    from ultralytics import YOLO
+    import cv2
+
+    if device == "cuda" and not torch.cuda.is_available():
+        print("CUDA requested but not visible after pinning "
+              f"(--gpu {args.gpu}). Falling back to CPU.")
+        device = "cpu"
+    print(f"Device: {device}"
+          + (f" (physical GPU {args.gpu})" if device == "cuda" else ""))
+
+    model = YOLO(args.model)
+    names = model.names
+    print(f"Model loaded. Classes: {names}")
+
+    def draw(img, boxes):
+        for (a, b, c, d) in boxes.int().tolist():
+            cv2.rectangle(img, (a, b), (c, d), (0, 255, 0), 2)
+        return img
+
+    ensure_directory(args.output_dir)
+    if args.preview_dir:
+        ensure_directory(args.preview_dir)
+
+    video_extensions = ('.mp4', '.avi', '.mov', '.mkv')
+    grand_total = 0
+    for filename in sorted(os.listdir(args.input_dir)):
+        if filename.lower().endswith(video_extensions):
+            video_path = os.path.join(args.input_dir, filename)
+            print(f"Processing video: {video_path}")
+            grand_total += process_video(
+                video_path, args.output_dir, args, model, device, draw)
+
+    print(f"Done. {grand_total} stills saved to {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,5 @@
+numpy
+opencv-python
+ultralytics
+scenedetect
+torch