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.
This commit is contained in:
2026-10-18 09:02:55 -04:00
commit 01eb729c2f
4 changed files with 432 additions and 0 deletions
+8
View File
@@ -0,0 +1,8 @@
.venv/
__pycache__/
*.pyc
models/
output/
preview/
*.pt
*.part
+83
View File
@@ -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.
+336
View File
@@ -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()
+5
View File
@@ -0,0 +1,5 @@
numpy
opencv-python
ultralytics
scenedetect
torch