"""Benchmark / compare NSFW models on YOUR OWN photos. python -m nsfwtag.bench "" [-r] [--models a,b,c] [--limit N] [--threshold 0.85] Scores every image with each selected model and writes, into the sample folder: * bench_scores.csv — path + one nsfw-score column per model * bench_report.html — thumbnails with per-model scores, sortable, disagreements highlighted (open it to eyeball which model over-flags) …plus a console summary (flag counts, agreement, biggest disagreements). Every score is normalized to nsfw = 1 - P(safe-class) so binary and multi-class models are directly comparable. Nothing leaves the machine. Models (need `pip install timm` for the timm ones; each downloads once): falconsai Falconsai/nsfw_image_detection (current, transformers, binary) adamcodd AdamCodd/vit-base-nsfw-detector (transformers, binary) marqo Marqo/nsfw-image-detection-384 (timm, binary, tiny) owen-xs OwenElliott/image-safety-classifier-xs (timm, 3-class SFW/NSFW/NSFL) """ import argparse import csv import gc import html import json import os import sys import webbrowser from pathlib import Path from urllib.parse import quote from .exif import read_marks from .scoring import discover_images BATCH = 16 # label (lowercased) treated as "safe"; nsfw score = 1 - sum(P over these) SAFE_LABELS = {"sfw", "safe", "normal", "neutral"} MODELS = { # name -> (kind, hf_id) "falconsai": ("hf", "Falconsai/nsfw_image_detection"), "adamcodd": ("hf", "AdamCodd/vit-base-nsfw-detector"), "marqo": ("timm", "Marqo/nsfw-image-detection-384"), "owen-xs": ("timm", "OwenElliott/image-safety-classifier-xs"), } def _device(): import torch return "mps" if torch.backends.mps.is_available() else "cpu" def _unsafe_indices(labels): """Given lowercased class labels, return (safe_idx, use_safe). If any safe label is present we score 1 - P(safe); otherwise P(labels that look unsafe).""" safe = [i for i, l in enumerate(labels) if l in SAFE_LABELS] if safe: return safe, True unsafe = [i for i, l in enumerate(labels) if any(k in l for k in ("nsfw", "nsfl", "porn", "sexy", "hentai", "explicit"))] return unsafe, False def load_hf(model_id): import torch from transformers import AutoImageProcessor, AutoModelForImageClassification dev = _device() proc = AutoImageProcessor.from_pretrained(model_id) model = AutoModelForImageClassification.from_pretrained(model_id).to(dev).eval() labels = [model.config.id2label[i].lower() for i in range(len(model.config.id2label))] idx, use_safe = _unsafe_indices(labels) def score(imgs): inputs = proc(images=imgs, return_tensors="pt").to(dev) with torch.no_grad(): p = model(**inputs).logits.softmax(-1) part = p[:, idx].sum(-1) if idx else p[:, :1].sum(-1) * 0 return (1 - part if use_safe else part).cpu().tolist() return score, model def load_timm(model_id): import timm import torch dev = _device() model = timm.create_model(f"hf_hub:{model_id}", pretrained=True).to(dev).eval() cfg = timm.data.resolve_model_data_config(model) tf = timm.data.create_transform(**cfg, is_training=False) labels = [l.lower() for l in model.pretrained_cfg["label_names"]] idx, use_safe = _unsafe_indices(labels) def score(imgs): batch = torch.stack([tf(im) for im in imgs]).to(dev) with torch.no_grad(): p = model(batch).softmax(-1) part = p[:, idx].sum(-1) if idx else p[:, :1].sum(-1) * 0 return (1 - part if use_safe else part).cpu().tolist() return score, model def score_all(names, imgs): """Return {model_name: {path: nsfw_score}}; loads one model at a time.""" from PIL import Image, ImageFile ImageFile.LOAD_TRUNCATED_IMAGES = True dev = _device() out = {} for name in names: kind, mid = MODELS[name] print(f"\n[{name}] loading {mid} on {dev} …", file=sys.stderr) try: score, model = (load_hf if kind == "hf" else load_timm)(mid) except Exception as e: print(f" SKIP {name}: {e}", file=sys.stderr) continue res = {} for i in range(0, len(imgs), BATCH): pil, paths = [], [] for p in imgs[i:i + BATCH]: try: pil.append(Image.open(p).convert("RGB")); paths.append(p) except Exception as e: print(f"\n error {p}: {e}", file=sys.stderr) if not pil: continue for p, s in zip(paths, score(pil)): res[str(p)] = float(s) print(f"\r {min(i + BATCH, len(imgs))}/{len(imgs)}", end="", file=sys.stderr, flush=True) print(file=sys.stderr) out[name] = res del score, model gc.collect() return out def _chip_color(s): hue = round((1 - s) * 130) # green (safe) -> red (unsafe) return f"hsl({hue} 70% 40%)" def write_report(rows, names, target, threshold, gt): """rows: list of (path, {name: score}). gt: {path: 1(nsfw)/0(sfw)}. Writes bench_report.html in target. Chips that disagree with your manual tag get a red ring; sortable by error count.""" target = Path(target) cards = [] for p, scores in rows: rel = quote(os.path.relpath(p, target)) vals = [scores.get(n) for n in names if scores.get(n) is not None] spread = (max(vals) - min(vals)) if len(vals) > 1 else 0.0 gtv = gt.get(p) # 1 nsfw / 0 sfw / None errors, chips = 0, [] for n in names: s = scores.get(n) if s is None: chips.append(f'{n} —'); continue hot = s >= threshold err = gtv is not None and (hot != (gtv == 1)) errors += err cls = "chip" + (" hot" if hot else "") + (" err" if err else "") chips.append(f'{n} {s:.2f}') gtbadge = (f'you: {"nsfw" if gtv==1 else "sfw"}' if gtv is not None else "") data = html.escape(json.dumps({n: scores.get(n) for n in names}), quote=True) cards.append( f'
' f'{gtbadge}' f'
{"".join(chips)}
' f'
{html.escape(Path(p).name)}
') opts = [] if gt: opts.append('') opts.append('') opts += [f'' for n in names] opts.append('') doc = _REPORT.replace("%%CARDS%%", "".join(cards)).replace("%%SORTOPTS%%", "".join(opts)) \ .replace("%%N%%", str(len(rows))).replace("%%THR%%", f"{threshold:.2f}") out = target / "bench_report.html" out.write_text(doc, encoding="utf-8") return out _REPORT = """ NSFW model bench (%%N%%)
NSFW model bench%%N%% images · white ring = ≥ %%THR%% · red ring = disagrees with your tag
%%CARDS%%
""" def summarize(rows, names, threshold): print(f"\n=== {len(rows)} images · flagged = score ≥ {threshold} ===") print(f"{'model':<12}{'flagged':>10}{'mean':>8}{'median':>8}") for n in names: vals = sorted(s[n] for _, s in rows if s.get(n) is not None) if not vals: print(f"{n:<12}{'(skipped)':>10}"); continue flagged = sum(v >= threshold for v in vals) mean = sum(vals) / len(vals) med = vals[len(vals) // 2] print(f"{n:<12}{f'{flagged}/{len(vals)}':>10}{mean:>8.3f}{med:>8.3f}") ok = [n for n in names if any(s.get(n) is not None for _, s in rows)] if len(ok) > 1: print("\npairwise agreement (same side of threshold):") for i in range(len(ok)): for j in range(i + 1, len(ok)): a, b = ok[i], ok[j] both = [(s[a], s[b]) for _, s in rows if s.get(a) is not None and s.get(b) is not None] agree = sum((x >= threshold) == (y >= threshold) for x, y in both) print(f" {a:<10} vs {b:<10} {agree}/{len(both)} ({100*agree/max(1,len(both)):.0f}%)") if len(ok) > 1: print("\nbiggest disagreements (open the HTML to judge):") spread = sorted(rows, key=lambda r: -(max(v for v in (r[1].get(n) for n in ok) if v is not None) - min(v for v in (r[1].get(n) for n in ok) if v is not None))) for p, s in spread[:8]: print(" " + Path(p).name + " " + " ".join(f"{n}={s[n]:.2f}" for n in ok if s.get(n) is not None)) def metrics(rows, names, gt, threshold): """Per-model precision/recall/F1 vs. your manual tags, at the given threshold, plus each model's F1-optimal threshold.""" labs = [gt[p] for p, _ in rows if p in gt] npos, nneg = sum(labs), len(labs) - sum(labs) single = npos == 0 or nneg == 0 # a known single-class folder (--assume) print(f"\n=== vs. ground truth · {npos} nsfw + {nneg} sfw · flag = score ≥ {threshold} ===") print(f"{'model':<11}{'acc':>7}{'prec':>7}{'recall':>8}{'F1':>6}{'FP':>5}{'FN':>5}{' best-thr':>11}{'F1@best':>8}") for n in names: pairs = [(gt[p], s[n]) for p, s in rows if p in gt and s.get(n) is not None] if not pairs: print(f"{n:<11}{'(skipped)':>7}"); continue def conf(t): tp = fp = fn = tn = 0 for y, sc in pairs: pred = sc >= t if y == 1 and pred: tp += 1 elif y == 0 and pred: fp += 1 elif y == 1: fn += 1 else: tn += 1 return tp, fp, fn, tn def f1(t): tp, fp, fn, _ = conf(t) pr = tp / (tp + fp) if tp + fp else 0 rc = tp / (tp + fn) if tp + fn else 0 return 2 * pr * rc / (pr + rc) if pr + rc else 0 tp, fp, fn, tn = conf(threshold) pr = tp / (tp + fp) if tp + fp else 0 rc = tp / (tp + fn) if tp + fn else 0 f = 2 * pr * rc / (pr + rc) if pr + rc else 0 acc = (tp + tn) / len(pairs) if single: bt_col = f"{'—':>11}{'—':>8}" # threshold tuning needs both classes else: cand = sorted({round(sc, 3) for _, sc in pairs} | {threshold}) bt = max(cand, key=f1) bt_col = f"{bt:>11.2f}{f1(bt):>8.2f}" print(f"{n:<11}{acc:>7.2f}{pr:>7.2f}{rc:>8.2f}{f:>6.2f}{fp:>5}{fn:>5}{bt_col}") if single and npos: print("single-class NSFW folder: recall = catch rate, FN = missed. (precision/FP need safe images too — " "run the bench on a mixed/SFW folder for those.)") elif single: print("single-class SFW folder: FP = false alarms, acc = correct-reject rate. (recall needs NSFW images.)") else: print("prec = of what it flagged, how much you'd tagged nsfw (fewer false alarms) · " "recall = of your nsfw, how much it caught · FP/FN at the current threshold") def combine(csv_paths, threshold): """Pool the labeled rows from several bench_scores.csv files and print a precision/recall/F1 threshold sweep per model — the real cross-dataset verdict.""" per_model, seen = {}, set() pos = neg = 0 for cp in csv_paths: with open(cp, newline="", encoding="utf-8") as fh: r = csv.DictReader(fh) models = [c for c in r.fieldnames if c not in ("path", "manual")] for row in r: man = row.get("manual", "") if man not in ("nsfw", "sfw") or row["path"] in seen: continue seen.add(row["path"]) y = 1 if man == "nsfw" else 0 pos += y; neg += (1 - y) for m in models: v = row.get(m, "") if v != "": per_model.setdefault(m, []).append((y, float(v))) if not pos or not neg: sys.exit("Need both nsfw and sfw labels pooled — include a mixed/SFW run and an NSFW run.") print(f"pooled {pos} nsfw + {neg} sfw from {len(csv_paths)} file(s)") def prf(pairs, t): tp = sum(y == 1 and s >= t for y, s in pairs) fp = sum(y == 0 and s >= t for y, s in pairs) fn = sum(y == 1 and s < t for y, s in pairs) pr = tp / (tp + fp) if tp + fp else 0 rc = tp / (tp + fn) if tp + fn else 0 return pr, rc, (2 * pr * rc / (pr + rc) if pr + rc else 0), fp, fn ranking = [] for m, pairs in per_model.items(): cand = sorted({round(s, 3) for _, s in pairs}) best = max(cand, key=lambda t: prf(pairs, t)[2]) ranking.append((prf(pairs, best)[2], best, m)) print(f"\n[{m}] best-F1 threshold = {best:.2f}") print(f" {'thr':>5}{'prec':>7}{'recall':>8}{'F1':>6}{'FP':>6}{'FN':>6}") for t in sorted(set([0.5, 0.7, 0.85, 0.9, 0.95, 0.99, round(best, 2)])): pr, rc, f, fp, fn = prf(pairs, t) print(f" {t:>5.2f}{pr:>7.2f}{rc:>8.2f}{f:>6.2f}{fp:>6}{fn:>6}{' <- best F1' if abs(t - best) < 5e-3 else ''}") ranking.sort(reverse=True) print("\nranking by best F1: " + " ".join(f"{m}={f:.2f}@{t:.2f}" for f, t, m in ranking)) def main(): ap = argparse.ArgumentParser(prog="nsfwtag.bench", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("target", nargs="?", help="folder of photos to score with each model") ap.add_argument("--combine", nargs="+", metavar="CSV", help="pool these bench_scores.csv files (using their manual labels) and print a precision/recall/threshold sweep per model — no scoring") ap.add_argument("-r", "--recursive", action="store_true", help="descend into subfolders") ap.add_argument("--models", default=",".join(MODELS), help=f"comma list (default: all) — {', '.join(MODELS)}") ap.add_argument("--limit", type=int, default=0, help="cap number of images scored (quick pass)") ap.add_argument("--threshold", type=float, default=0.85, help="flag cutoff for the report (default 0.85)") ap.add_argument("--all", action="store_true", help="score ALL images, not just your manually-tagged ones") ap.add_argument("--assume", choices=["nsfw", "sfw"], help="treat EVERY image in the folder as this label (a known folder) instead of reading EXIF tags") args = ap.parse_args() if args.combine: combine(args.combine, args.threshold) return if not args.target: ap.error("a target folder is required (or use --combine)") target = Path(args.target).expanduser() if not target.is_dir(): sys.exit(f"Not a folder: {target}\n Check the path/spelling (that folder does not exist).") names = [n.strip() for n in args.models.split(",") if n.strip()] bad = [n for n in names if n not in MODELS] if bad: sys.exit(f"Unknown model(s): {bad}. Available: {list(MODELS)}") imgs = discover_images(target, recursive=args.recursive) if not imgs: has_subdirs = any(p.is_dir() for p in target.iterdir()) hint = " Images look to be in subfolders — add -r to scan recursively." if (has_subdirs and not args.recursive) else "" sys.exit(f"No images found directly in {target}.\n{hint}".rstrip()) if args.assume: # whole folder is a known class if args.limit: imgs = imgs[:args.limit] gt = {str(p): (1 if args.assume == "nsfw" else 0) for p in imgs} print(f"assuming all {len(imgs)} images are {args.assume.upper()} (folder is ground truth).", file=sys.stderr) else: # ground truth from your manual nsfw/sfw EXIF tags (one batched exiftool read) print(f"found {len(imgs)} images — reading your manual tags …", file=sys.stderr) marks = read_marks([str(p) for p in imgs]) gt = {p: 1 for p in marks["nsfw"]} gt.update({p: 0 for p in marks["sfw"] if p not in marks["nsfw"]}) if gt and not args.all: imgs = [p for p in imgs if str(p) in gt] # evaluate only what you've tagged print(f"evaluating on {len(imgs)} tagged images " f"({sum(gt.values())} nsfw / {len(gt) - sum(gt.values())} sfw). Use --all to score everything.", file=sys.stderr) elif not gt: print("no manual nsfw/sfw tags found — falling back to a relative model comparison.", file=sys.stderr) if args.limit: imgs = imgs[:args.limit] print(f"scoring {len(imgs)} images with {len(names)} model(s): {names}", file=sys.stderr) scores = score_all(names, imgs) names = [n for n in names if n in scores] # keep only models that loaded if not names: sys.exit("No models loaded (need: pip install torch transformers timm).") rows = [(str(p), {n: scores[n].get(str(p)) for n in names}) for p in imgs] csv_path = target / "bench_scores.csv" with open(csv_path, "w", newline="", encoding="utf-8") as fh: w = csv.writer(fh); w.writerow(["path", "manual", *names]) for p, s in rows: man = {1: "nsfw", 0: "sfw"}.get(gt.get(p), "") w.writerow([p, man, *[("" if s[n] is None else f"{s[n]:.4f}") for n in names]]) if any(p in gt for p, _ in rows): metrics(rows, names, gt, args.threshold) else: summarize(rows, names, args.threshold) report = write_report(rows, names, target, args.threshold, gt) print(f"\nscores -> {csv_path}\nreport -> {report}", file=sys.stderr) webbrowser.open(report.resolve().as_uri()) if __name__ == "__main__": main()