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:
@@ -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()
|
||||
Reference in New Issue
Block a user