519 lines
18 KiB
Python
519 lines
18 KiB
Python
#!/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(
|
|
"--max-face-fraction",
|
|
type=float,
|
|
default=0.30,
|
|
help="Maximum face box area as a fraction of frame area (kills "
|
|
"detector meltdowns on full-frame effects like ED feather "
|
|
"storms; 0 disables).",
|
|
)
|
|
parser.add_argument(
|
|
"--min-aspect",
|
|
type=float,
|
|
default=0.45,
|
|
help="Minimum face box width/height ratio.",
|
|
)
|
|
parser.add_argument(
|
|
"--max-aspect",
|
|
type=float,
|
|
default=1.6,
|
|
help="Maximum face box width/height ratio (a face wider than 1.6x "
|
|
"itself is almost always an effect blob).",
|
|
)
|
|
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(
|
|
"--probe-stride",
|
|
type=int,
|
|
default=1,
|
|
help="Within a scene still needing a face, run detection every Nth "
|
|
"frame instead of every frame (1 = original behavior, free on "
|
|
"CUDA). For slow devices (iGPU ~5s/img) use e.g. 25: decode "
|
|
"stays frame-perfect, only the expensive inference is thinned.",
|
|
)
|
|
parser.add_argument(
|
|
"--probe-max-per-scene",
|
|
type=int,
|
|
default=4,
|
|
help="With --probe-stride>1, probes before giving up on a scene and "
|
|
"waiting for the next scene start (default 4).",
|
|
)
|
|
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(
|
|
"--all-files",
|
|
action="store_true",
|
|
help="Also scan OP/ED files (name matches 'opening'/'ending'). "
|
|
"Skipped by default: music-video files are dominated by "
|
|
"stylized graphics (feather storms, transformation "
|
|
"emblems) that false-positive face detectors; character "
|
|
"close-ups for training datasets come from episodes.",
|
|
)
|
|
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.
|
|
|
|
NOTE: ultralytics' select_device() rewrites CUDA_VISIBLE_DEVICES on
|
|
first predict() if the CUDA context is not yet initialized, silently
|
|
un-pinning us onto GPU0. We force-initialize a tiny CUDA allocation
|
|
here so our CVD pin is locked in before ultralytics can touch it.
|
|
"""
|
|
if gpu_arg.lower() == "cpu":
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
|
return "cpu"
|
|
if gpu_arg.lower() == "igpu":
|
|
# Intel iGPU via OpenVINO (Level Zero) — no CUDA, no torch pinning.
|
|
# OVFaceDetector takes over model loading in main().
|
|
return "igpu"
|
|
if gpu_arg.lower() == "auto":
|
|
return None # let ultralytics/torch decide
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_arg
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
torch.zeros(1, device="cuda") # lock the context to our pin
|
|
else:
|
|
return "cpu"
|
|
return "cuda:0" # within our CVD, physical GPU N is always cuda:0
|
|
|
|
|
|
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, max_face_fraction,
|
|
min_aspect, max_aspect, 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.
|
|
Upper-bound gates (max area, aspect) kill detector meltdowns on
|
|
full-frame stylized effects (ED feather storms etc.)."""
|
|
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
|
|
|
|
area_frac = (w * h) / float(frame_w * frame_h)
|
|
if min_face_fraction > 0 and area_frac < min_face_fraction:
|
|
return False
|
|
if max_face_fraction > 0 and area_frac > max_face_fraction:
|
|
return False
|
|
|
|
if h > 0:
|
|
aspect = w / float(h)
|
|
if aspect < min_aspect or aspect > max_aspect:
|
|
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
|
|
|
|
|
|
class AvFrameReader:
|
|
"""Forward-only frame reader with a cv2.VideoCapture-compatible
|
|
surface (read/release). PyAV ships its own ffmpeg build with libdav1d:
|
|
the ffmpeg bundled inside opencv-python has no software AV1 decoder,
|
|
and Alder Lake-N has no AV1 hw decode, so BD AV1 rips (e.g. SavI0r)
|
|
decode 0 frames through cv2. scenedetect likewise prefers its pyav
|
|
backend once av is installed, fixing scene detection on the same files.
|
|
"""
|
|
|
|
def __init__(self, path):
|
|
import av
|
|
self._c = av.open(path)
|
|
self._vs = self._c.streams.video[0]
|
|
self._vs.thread_type = "AUTO"
|
|
self._iter = self._c.decode(self._vs)
|
|
|
|
@property
|
|
def frame_count(self):
|
|
vs = self._vs
|
|
if vs.frames:
|
|
return int(vs.frames)
|
|
# MKV headers often leave stream.frames unset; estimate from the
|
|
# container duration (it is this file's total decode work anyway).
|
|
dur = self._c.duration # microseconds, or None
|
|
if dur and vs.average_rate:
|
|
return int(float(dur) / 1e6 * float(vs.average_rate))
|
|
return 0
|
|
|
|
def read(self):
|
|
try:
|
|
frame = next(self._iter)
|
|
except StopIteration:
|
|
return False, None
|
|
return True, frame.to_ndarray(format="bgr24")
|
|
|
|
def release(self):
|
|
try:
|
|
self._c.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def open_reader(video_path):
|
|
"""PyAV when available (AV1-capable), else plain cv2."""
|
|
import cv2
|
|
try:
|
|
return AvFrameReader(video_path)
|
|
except ImportError:
|
|
pass
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
return None
|
|
return cap
|
|
|
|
|
|
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
|
|
try:
|
|
import av # noqa: F401
|
|
# scenedetect defaults to backend="opencv", which cannot decode
|
|
# AV1 in this env; pyav ships its own ffmpeg with libdav1d.
|
|
video = open_video(video_path, backend="pyav")
|
|
except ImportError:
|
|
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 = open_reader(video_path)
|
|
if cap is None:
|
|
print(f"Error opening video file {video_path}")
|
|
return output_start
|
|
|
|
frame_count = getattr(cap, "frame_count", 0)
|
|
if not frame_count and hasattr(cap, "get"):
|
|
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
|
|
try:
|
|
print(f"{video_name}: detecting scenes "
|
|
"(silent for several minutes on CPU decode)...", flush=True)
|
|
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
|
|
|
|
# Probe schedule. probe_stride=1 reproduces the original every-frame
|
|
# loop (right on CUDA, where predict costs ~0.2 s). On slow devices
|
|
# (iGPU ~5 s/img) predicting every frame is hopeless: decode is ~300 fps,
|
|
# inference 0.2/s. With stride>1 we decode through frames for free and
|
|
# predict only at scene starts, retrying every stride frames up to
|
|
# probe_max attempts per scene, then waiting for the next scene start.
|
|
probe_stride = getattr(args, "probe_stride", 1)
|
|
probe_max = getattr(args, "probe_max_per_scene", 4)
|
|
next_probe = 0
|
|
probes_used = 0
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
if frame_idx in scene_starts:
|
|
face_in_scene = False
|
|
next_probe = frame_idx
|
|
probes_used = 0
|
|
|
|
if face_in_scene:
|
|
frame_idx += 1
|
|
continue
|
|
if probe_stride > 1 and (frame_idx < next_probe
|
|
or probes_used >= probe_max):
|
|
frame_idx += 1
|
|
continue
|
|
|
|
if not face_in_scene:
|
|
probes_used += 1
|
|
next_probe = frame_idx + probe_stride
|
|
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.max_face_fraction,
|
|
args.min_aspect, args.max_aspect,
|
|
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 == "igpu":
|
|
# Intel iGPU path: OpenVINO replaces ultralytics' torch predict loop.
|
|
# Ultralytics import stays (shared pre/post code paths), but weights
|
|
# load as ONNX through OVFaceDetector instead of torch.
|
|
from ov_detector import OVFaceDetector
|
|
onnx_path = os.path.splitext(args.model)[0] + ".onnx"
|
|
if not os.path.exists(onnx_path):
|
|
print(f"Error: --gpu igpu needs an exported ONNX at {onnx_path}\n"
|
|
" (one-time: YOLO('<pt>').export(format='onnx', "
|
|
"dynamic=True, simplify=True, opset=17))")
|
|
sys.exit(1)
|
|
model = OVFaceDetector(onnx_path, device="GPU", imgsz=args.imgsz)
|
|
names = model.names
|
|
print("Device: igpu (OpenVINO, Intel iGPU via Level Zero)")
|
|
print(f"Model loaded (OpenVINO GPU). Classes: {names}")
|
|
else:
|
|
if device == "cuda:0" 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 not filename.lower().endswith(video_extensions):
|
|
continue
|
|
if (not args.all_files
|
|
and any(k in filename.lower()
|
|
for k in ("opening", "ending", " op", " ed"))):
|
|
print(f"Skipping OP/ED file (use --all-files to include): "
|
|
f"{filename}")
|
|
continue
|
|
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()
|