alex
89b0f708d7
Tested against real failure samples from the v2 run: full-frame feather storm (area .64) and wide effect blob (area .70) now rejected; normal and close-up faces pass. Note: stylized OP/ED emblems are geometrically face-shaped and NOT catchable this way — documented, CCIP gate is the answer for those.
381 lines
13 KiB
Python
381 lines
13 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(
|
|
"--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.
|
|
|
|
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() == "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
|
|
|
|
|
|
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.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 == "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 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()
|