alex
940c59d6c9
Ultralytics select_device() rewrites CUDA_VISIBLE_DEVICES on first predict() when the CUDA context is uninitialized: --gpu 1 silently became a two-GPU cuda:0,1 context that OOMed on GPU0 (Strata's card). Initializing torch and touching cuda once after setting CVD locks the pin; both scripts now use device cuda:0 within the pin.
116 lines
4.0 KiB
Python
116 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Audit an existing still collection with the anime-face detector:
|
|
how many saved images actually contain a qualifying face?
|
|
|
|
Useful for measuring how much junk the old person-proxy pipeline left
|
|
behind, and as a QC pass over any facesaver output.
|
|
|
|
python audit_stills.py -I <stills_dir> --gpu 1 [--write-report report.csv]
|
|
"""
|
|
|
|
import argparse
|
|
import csv
|
|
import os
|
|
import sys
|
|
|
|
|
|
def parse_args():
|
|
p = argparse.ArgumentParser(
|
|
description=(__doc__ or "Audit stills for anime faces.").split("\n")[0])
|
|
p.add_argument("-I", "--input-dir", required=True,
|
|
help="Directory of still images to audit.")
|
|
p.add_argument("--model",
|
|
default=os.path.join(
|
|
os.path.dirname(os.path.abspath(__file__)),
|
|
"models", "yolov8x6_animeface.pt"))
|
|
p.add_argument("--gpu", default="1",
|
|
help="GPU index, 'cpu', or 'auto'.")
|
|
p.add_argument("--conf", type=float, default=0.3)
|
|
p.add_argument("--imgsz", type=int, default=1280)
|
|
p.add_argument("--batch", type=int, default=8)
|
|
p.add_argument("--write-report", default=None,
|
|
help="Optional CSV: filename, n_faces, max_w, max_h, "
|
|
"ok (1/0).")
|
|
return p.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
if args.gpu.lower() == "cpu":
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
|
device = "cpu"
|
|
elif args.gpu.lower() != "auto":
|
|
# Lock the pin with a warm CUDA alloc BEFORE ultralytics runs
|
|
# select_device(), which would otherwise rewrite
|
|
# CUDA_VISIBLE_DEVICES and un-pin us onto GPU0.
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
torch.zeros(1, device="cuda")
|
|
device = "cuda:0"
|
|
else:
|
|
device = "cpu"
|
|
else:
|
|
device = None
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from ultralytics import YOLO
|
|
|
|
exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp")
|
|
files = sorted(f for f in os.listdir(args.input_dir)
|
|
if f.lower().endswith(exts))
|
|
if not files:
|
|
print("No images found.")
|
|
sys.exit(1)
|
|
print(f"Auditing {len(files)} images with "
|
|
f"{os.path.basename(args.model)} on {device or 'auto'}")
|
|
|
|
model = YOLO(args.model)
|
|
|
|
rows = []
|
|
no_face = 0
|
|
# stream in chunks so 1000s of 1080p pngs don't eat RAM
|
|
for i in range(0, len(files), args.batch):
|
|
chunk = files[i:i + args.batch]
|
|
imgs, kept = [], []
|
|
for f in chunk:
|
|
img = cv2.imread(os.path.join(args.input_dir, f))
|
|
if img is None:
|
|
print(f" WARN unreadable: {f}")
|
|
continue
|
|
imgs.append(img)
|
|
kept.append(f)
|
|
results = model.predict(imgs, conf=args.conf, imgsz=args.imgsz,
|
|
device=device, verbose=False)
|
|
for f, res in zip(kept, results):
|
|
n, mw, mh = 0, 0, 0
|
|
if res.boxes is not None and len(res.boxes):
|
|
xyxy = res.boxes.xyxy.cpu().numpy()
|
|
n = len(xyxy)
|
|
for (x1, y1, x2, y2) in xyxy:
|
|
mw = max(mw, int(x2 - x1))
|
|
mh = max(mh, int(y2 - y1))
|
|
ok = 1 if n > 0 else 0
|
|
if not ok:
|
|
no_face += 1
|
|
print(f" NO FACE: {f}")
|
|
rows.append((f, n, mw, mh, ok))
|
|
done = min(i + args.batch, len(files))
|
|
print(f" {done}/{len(files)}", flush=True)
|
|
|
|
print()
|
|
print(f"Total: {len(rows)} with face: {len(rows) - no_face} "
|
|
f"NO face: {no_face} ({100.0 * no_face / max(1, len(rows)):.1f}%)")
|
|
|
|
if args.write_report:
|
|
with open(args.write_report, "w", newline="") as fh:
|
|
w = csv.writer(fh)
|
|
w.writerow(["filename", "n_faces", "max_w", "max_h", "ok"])
|
|
w.writerows(rows)
|
|
print(f"Report: {args.write_report}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|