One-Class Porn Detection

One-class detection of pornographic images using vision embeddings and outlier detection.

How It Works

1. Image Embedding — Images are converted to high-dimensional embeddings using a vision model (CLIP, SigLIP2, DINOv3, C-RADIOv4 (incl. distilled), EUPE, AIMv2, MetaCLIP2, SAM3, Sapiens2, or TIPSv2).

2. One-Class / Outlier Detection — Embeddings are fed to PyOD models trained on normal (pornographic) images. A reject-rate sweep evaluates each model at multiple thresholds to optimize accuracy (or another metric).

Installation

Requires uv and Python 3.13.

git clone <repository-url>
cd one-class-porn-detection
uv sync

CLI Entry Points

Train

uv run train \
  --train-dir <path> \
  --val-dir <path> \
  --embedding-model clip

Sweeps over all configured PyOD models × reject rates, writes a sweep CSV and summary JSON to a timestamped subdirectory under trained_models/<embedding_model>/<reducer>/.

Test

uv run test \
  --test-dir <path> \
  --results-dir <path> \
  --sweep-csv <path/to/sweep.csv> \
  --embedding-model clip

By default, selects the best config from the sweep CSV by --select-metric (default: accuracy, validation metric from training) and evaluates only that single config on the test set. Use --all to test every config in the sweep CSV.

For the evaluated config(s), writes to a timestamped subdirectory under --results-dir:

Metric convention. The stored score is the raw outlier score (higher = more anomalous = Non_porn/label 1), but Porn (the inlier, label 0) is the positive class for all metrics. Threshold metrics (F1(0)/Precision(0)/ Recall(0)) and average precision treat Porn as positive; average precision and the PR/ROC curves negate the scores at evaluation time to orient Porn as the detection target. ROC-AUC is invariant to the positive-class choice, so it is computed on the raw scores (numerically equal to the Porn-positive AUC). This aligns average precision with the threshold metrics, which already frame Porn (label 0) as the detection target.

Key CLI Arguments

Train (uv run train)

Test (uv run test)

Embedding Models

Detection Methods

10 PyOD models are configured by default:

Embedding Cache

Embeddings are cached as .npz files (keyed by MD5 hash of each image). Subsequent runs only compute embeddings for new images. Use --no-cache to bypass.

A standalone scripts/cleanup_cache.py script is provided for removing cached entries matching a path filter.

Using Trained Models in Python

The OneClassDetector class provides an sklearn-like interface for inference:

from one_class_porn_detection import OneClassDetector

# Load from a sweep CSV row
detector = OneClassDetector.from_sweep_row(
    "trained_models/clip/none/sweep.csv",
    model="iforest",
    reject_rate=0.01,
)

# Or load from individual artifact files
detector = OneClassDetector.load(
    model_path="trained_models/clip/none/iforest_clip_none_base.joblib",
    threshold=0.123456,
    embedding_model="clip",
)

# Predict on a directory of images (handles embedding + normalization)
labels = detector.predict(image_dir="path/to/images/")
# 0 = inlier, 1 = outlier (as defined in classes.toml)

# Or on pre-computed embeddings
labels = detector.predict(embeddings=X)

# Get raw outlier scores
scores = detector.score_samples(image_dir="path/to/images/")

Plotting Results

A standalone script, scripts/plot_results.py, generates diagnostic plots from the *_predictions.csv files produced by test:

uv run python scripts/plot_results.py --results-dir <path>
# or single CSV
uv run python scripts/plot_results.py --predictions-csv <path>

Per-config plots (6 files per config):

Combined overlays (when multiple CSVs are found, unless --no-multi):

Selecting Extreme Samples

A standalone script, scripts/select_samples.py, copies the highest- and lowest-scoring images per class from a predictions CSV into named subdirectories (e.g. outdir/porn/min, outdir/porn/max, outdir/non_porn/min, outdir/non_porn/max). Useful for inspecting which images the model is most/least confident about.

uv run python scripts/select_samples.py results.csv --data-dir datasets/test_easy -n 10
uv run python scripts/select_samples.py results.csv --data-dir datasets/test_easy \
    -n 10 --outdir selected_samples

Scores follow the repo convention: higher = more anomalous = outlier.

Baseline: ShieldGemma 2 (VLM)

A zero-shot baseline that is separate from the trained one-class pipeline. Instead of learning embeddings + outlier detection, it asks the google/shieldgemma-2-4b-it vision-language model to classify each image against a sexually-explicit safety policy and records the model's probability that the image violates the policy (i.e. is pornographic). The model is gated, so HF_TOKEN must be set.

Usage

uv run python scripts/eval_shieldgemma.py \
  --test-dir datasets/test_easy \
  --results-dir baseline_results

Score convention

ShieldGemma returns probabilities[:,0] = P(violates policy) = P(porn) and [:,1] = P(safe). To stay compatible with the one-class pipeline (where a higher score = label 1 = Non_porn/outlier), the score column is score = 1 - P(violates) = P(safe). Predictions use y_pred = (score > threshold) with --threshold (default 0.5), so an image the model deems pornographic (P(violates) > 0.5) is predicted as inlier (0).

Outputs

Writes a timestamped subdirectory under --results-dir with the same schema as the test CLI, consumable by scripts/plot_results.py:

Flags

VLM inference is slow, so a lightweight per-image results cache is kept under --cache-dir as <model>_<policy-hash>_scores.csv (keyed by image MD5 + policy text). Re-runs skip already-scored images; use --no-cache to bypass. The default policy is ShieldGemma's "No Sexually Explicit content" description; use --policy to supply a custom one (this changes the cache file so different policies are scored independently).

Baseline: LlavaGuard (VLM, SGLang)

A zero-shot baseline that is separate from the trained one-class pipeline. Instead of learning embeddings + outlier detection, it asks the AIML-TUDA/LlavaGuard-v1.2-0.5B-OV vision-language model to assess each image against the O3 (Sexual Content) and O4 (Nudity Content) safety-policy categories and records the model's Safe/Unsafe rating.

This script is a thin HTTP client that talks to a running SGLang server (OpenAI-compatible /v1/chat/completions vision endpoint). It does not load the model itself — no sglang/torch/transformers dependency — only requests + pillow + eval libs (declared via PEP 723 inline metadata). Launch the server first (from the model card):

CUDA_VISIBLE_DEVICES=0 python -m sglang.launch_server \
    --model-path AIML-TUDA/LlavaGuard-v1.2-0.5B-OV --port 10000

then run the evaluation script:

uv run python scripts/eval_llavaguard.py \
  --test-dir datasets/test_easy \
  --results-dir baseline_results

If you customised --served-model-name on the server, pass the same value to --served-model-name here. The non-HF variant (LlavaQwenForCausalLM, model_type: "llava") is the one SGLang implements natively; the -hf variant is not SGLang-compatible.

LlavaGuard is a generative VLM: it emits a JSON {"rating": "Safe"|"Unsafe", ...} and has no probability head. The rating is therefore mapped to a discrete binary score (1.0 for Safe, 0.0 for Unsafe). Because of this, ROC/PR curves and threshold sweeps are degenerate by design (there is no continuous ranking between images). Generation uses the official SGLang model-card hyperparameters (temperature=0.2, top_p=0.95, top_k=50, max_new_tokens=500; no beam search); sampling_seed=--seed is passed server-side for reproducibility.

Score convention

LlavaGuard returns a Safe/Unsafe rating, where Unsafe means the image violates the sexual/nudity policy (i.e. is pornographic). To stay compatible with the one-class pipeline (where a higher score = label 1 = Non_porn/outlier), the score column is score = 1 - P(Unsafe) = P(Safe) (so Safe -> 1.0, Unsafe -> 0.0). Predictions use a fixed boundary y_pred = (score > 0.5) (not configurable), so an image rated Unsafe (pornographic, score 0) is predicted as inlier (0) and Safe (non-pornographic, score 1) as outlier (1).

Partial results are checkpointed to {tag}_checkpoint.csv every --checkpoint-interval images (default 1000) so an interrupted inference run can be resumed with --resume (pass the exact previous run directory as --results-dir). Without --resume every run scores all (subsampled) images fresh, overwriting any existing checkpoint in the output directory. Images whose response does not contain a parseable Safe/Unsafe rating are logged and excluded from evaluation.

Outputs

Writes a timestamped subdirectory under --results-dir with the same schema as the test CLI, consumable by scripts/plot_results.py:

Flags

Baseline: LlavaGuard (VLM, HuggingFace)

A zero-shot baseline that is separate from the trained one-class pipeline. Instead of learning embeddings + outlier detection, it asks the AIML-TUDA/LlavaGuard-v1.2-0.5B-OV-hf vision-language model to assess each image against the O3 (Sexual Content) and O4 (Nudity Content) safety-policy categories and records the model's Safe/Unsafe rating.

Unlike the SGLang variant above, this script loads the model locally via transformers (LlavaOnevisionForConditionalGeneration / AutoModelForImageTextToText) — no external server is needed, but it requires a CUDA-capable environment with the full HuggingFace stack installed.

LlavaGuard is a generative VLM: it emits a JSON {"rating": "Safe"|"Unsafe", ...} and has no probability head. The rating is therefore mapped to a discrete binary score (1.0 for Safe, 0.0 for Unsafe). Because of this, ROC/PR curves and threshold sweeps are degenerate by design (there is no continuous ranking between images). Generation uses stochastic sampling (do_sample=True, temperature=0.2, top_p=0.95, top_k=50, num_beams=2, max_new_tokens=200); --seed governs both subsampling and the PyTorch/TorchVision RNG for reproducibility.

Usage

uv run python scripts/eval_llavaguard_hf.py \
  --test-dir datasets/test_easy \
  --results-dir baseline_results

Score convention

LlavaGuard returns a Safe/Unsafe rating, where Unsafe means the image violates the sexual/nudity policy (i.e. is pornographic). To stay compatible with the one-class pipeline (where a higher score = label 1 = Non_porn/outlier), the score column is score = 1 - P(Unsafe) = P(Safe) (so Safe -> 1.0, Unsafe -> 0.0). Predictions use y_pred = (score > threshold) with --threshold (default 0.5), so an image rated Unsafe (pornographic) is predicted as inlier (0).

There is no results cache: every run scores all (subsampled) images fresh. Images whose response does not contain a parseable Safe/Unsafe rating are logged and excluded from evaluation.

Outputs

Writes a timestamped subdirectory under --results-dir with the same schema as the test CLI, consumable by scripts/plot_results.py:

Flags

Baseline: Falconsai NSFW Image Detection (ViT)

A zero-shot baseline that is separate from the trained one-class pipeline. Instead of learning embeddings + outlier detection, it asks a Falconsai fine-tuned Vision Transformer to classify each image as normal or nsfw and records the model's probability that the image is NSFW (i.e. pornographic). Two models are selectable via --model:

Both are ViTForImageClassification (ViT-base, 224×224, patch16, 85.8M params).

Usage

# Legacy public model (default)
uv run python scripts/eval_falconsai.py \
  --test-dir datasets/test_easy \
  --results-dir baseline_results

# 2026 gated model (requires HF_TOKEN)
uv run python scripts/eval_falconsai.py \
  --model nsfw_image_detection_26 \
  --test-dir datasets/test_easy \
  --results-dir baseline_results

Score convention

Falconsai returns P(nsfw) = P(porn) (label index resolved from model.config.id2label). To stay compatible with the one-class pipeline (where a higher score = label 1 = Non_porn/outlier), the score column is score = 1 - P(nsfw) = P(normal). Predictions use y_pred = (score > threshold) with --threshold (default 0.5), so an image the model deems pornographic (P(nsfw) > 0.5) is predicted as inlier (0).

Outputs

Writes a timestamped subdirectory under --results-dir with the same schema as the test CLI, consumable by scripts/plot_results.py:

Flags

A lightweight per-image results cache is kept under --cache-dir as <model>_scores.csv (keyed by image MD5). Re-runs skip already-scored images; use --no-cache to bypass. The cache file is per --model-name, so the two models are scored and cached independently.

Data Convention

Class-to-label mappings are defined in src/one_class_porn_detection/classes.toml. The default config:

[[classes]]
name = "Porn"
prefix = "Porn"
label = 0
role = "inlier"

[[classes]]
name = "Non_porn"
prefix = "Non_porn"
label = 1
role = "outlier"

Image filenames must be prefixed with the corresponding prefix field:

• Porn_* → inlier (label 0)
• Non_porn_* → outlier (label 1)

The video test pipeline (test-video) applies the same prefix convention to video filenames (e.g. Porn_clip.mp4, Non_porn_scene.mp4).

Override with --classes-config <path> to use a custom mapping.

Dataset Preparation

Raw images typically live in source folders (possibly with subfolders) organized by class. A standalone script, scripts/prepare_data.py, discovers, validates, splits, and prefix-tags images so they are consumable by the train/test CLIs described above.

Recipe

Splits and sources are declared in a fixed TOML recipe (scripts/data_recipe.toml). Each [[source]] maps a raw folder to a role (inlier → Porn_ prefix, outlier → Non_porn_ prefix), and each [[split]] lists one or more [[split.sample]] entries specifying a source and count:

[defaults]
seed = 42
link = "copy"

[[source]]
name = "porn"
role = "inlier"
root = "data/LSPD/porn"

[[source]]
name = "normal"
role = "outlier"
root = "data/LSPD/normal"

[[split]]
name = "train"
[[split.sample]]
source = "porn"
count = 100000

[[split]]
name = "val"
[[split.sample]]
source = "porn"
count = 50000
[[split.sample]]
source = "normal"
count = 50000

count = "all" uses every valid image from that source (e.g., for APD-2M). A group key reuses the same image set across multiple splits — used by test_easy and test_hard to share the same porn images while pairing them with different outlier classes:

[[split]]
name = "test_easy"
[[split.sample]]
source = "porn"
count = 50000
group = "test_inlier"
[[split.sample]]
source = "normal"
count = 50000

[[split]]
name = "test_hard"
[[split.sample]]
source = "porn"
count = 50000
group = "test_inlier"
[[split.sample]]
source = "sexy"
count = 50000

Prefixes (Porn_, Non_porn_) are read from classes.toml via LabelConfig, so the script never hardcodes them.

Usage

# Dry-run: validate counts against your data without copying
uv run python scripts/prepare_data.py --dry-run

# Full run
uv run python scripts/prepare_data.py \
  --recipe scripts/data_recipe.toml \
  --out-root datasets

What the script does

1. Discover — recursively scans each source root for image files.
2. Validate — each image is checked with PIL.Image.verify(); --strict-verify adds a full decode to catch truncated files. Corrupted images are skipped and logged.
3. Allocate — per-source valid images are shuffled deterministically, then exact counts are drawn split-by-split with a global disjointness guarantee (no image appears in two splits unless explicitly shared via group). If a source can't satisfy a requested count, the script raises a clear error.
4. Write — images are copied (or linked) into flat --out-root/<split>/ directories with the appropriate filename prefix. Collision-safe renaming prevents overwrites.
5. Manifests — per-split CSV files (source_path, output_name, prefix, source, role, group), a summary.json, and a corrupted.log are written next to the output.

Output layout

datasets/
  train/        Porn_*         (inlier only)
  val/          Porn_* + Non_porn_*
  test_easy/    Porn_* + Non_porn_*
  test_hard/    Porn_* + Non_porn_*
  test_ood/     Porn_* + Non_porn_*
  apd2m_test/   Porn_* + Non_porn_*
manifests/
  summary.json
  <split>.csv
  corrupted.log

These split directories plug directly into the train/test CLIs:

uv run train --train-dir datasets/train --val-dir datasets/val --embedding-model clip
uv run test  --test-dir datasets/test_easy --results-dir results --sweep-csv <path> --embedding-model clip

Data-Scaling Ablation

A shell driver, scripts/run_ablation.sh, trains on nested subsets of the inlier training set (permutation with --seed, so 100 ⊂ 500 ⊂ … ⊂ full within each seed), then tests each resulting model. Writes a manifest CSV linking seed + subset size → sweep CSV → test-results CSV, consumable by scripts/plot_ablation.py.

The first run extracts all train+val embeddings (slow); subsequent subset runs read from the NPZ cache and only fit the detector (fast). Val and test sets are fixed across all runs.

# train + test (default)
bash scripts/run_ablation.sh

# train only
STAGE=train bash scripts/run_ablation.sh

# test only (needs prior train)
STAGE=test bash scripts/run_ablation.sh

Testing runs against every split in TEST_SPLITS (default: test_easy, test_hard, test_ood), each expected as a sibling of TEST_DIR. Each split gets its own manifest and results directory.

Plotting Ablation Results

A standalone script, scripts/plot_ablation.py, plots metrics vs training-set size (log x-axis) from one or more manifest CSVs. ROC-AUC and average precision are recomputed post-hoc from each config's predictions CSV so test runs before and after the Porn-positive reframe plot on a consistent basis. Both PNG and SVG are written.

# single manifest (val + test)
uv run python scripts/plot_ablation.py --manifest ablation_cradiov4_h_test_easy_manifest.csv

# combined overlay across test splits
uv run python scripts/plot_ablation.py \
    --manifest ablation_cradiov4_h_test_easy_manifest.csv \
               ablation_cradiov4_h_test_hard_manifest.csv \
               ablation_cradiov4_h_test_ood_manifest.csv \
    --split-labels test_easy,test_hard,test_ood \
    --metrics accuracy,macro_f1,bac,roc_auc --csv

Five-Way Score Distribution

A standalone script, scripts/plot_score_dist_5way.py, produces a single figure (PNG + SVG) showing outlier-score distributions for five categories drawn from three test splits:

1. porn — inlier of test_easy (shared with test_hard via the test_inlier group, so taken from test_easy)
2. normal — outlier of test_easy
3. sexy — outlier of test_hard
4. drawing — outlier of test_ood
5. hentai — inlier of test_ood

The predictions CSV carries no sub-category column, but each split has a fixed inlier/outlier category pair, so split + y_true deterministically recovers the category. The trained threshold is identical across splits (same sweep CSV) and is marked on the plot.

# default: resolve from the three default ablation manifests
uv run python scripts/plot_score_dist_5way.py

# explicit predictions CSVs (fixed order: test_easy, test_hard, test_ood)
uv run python scripts/plot_score_dist_5way.py \
    --predictions-csv easy/hbos_rr0p01_predictions.csv \
                     hard/hbos_rr0p01_predictions.csv \
                     ood/hbos_rr0p01_predictions.csv