Cleanup

scripts/split_nsfw_test_image.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""Split an image into two lossless, random-looking PNG fixture shares."""
+
+from __future__ import annotations
+
+import argparse
+from io import BytesIO
+from pathlib import Path
+import secrets
+
+from PIL import Image
+
+
+def split_image(input_path: Path, output_dir: Path, index: int | None = None) -> tuple[Path, Path]:
+    if index is None:
+        index = 1
+        while any((output_dir / f"image_{index}{part}.png").exists() for part in ("a", "b")):
+            index += 1
+    if index < 1:
+        raise ValueError("Image index must be a positive integer")
+
+    paths = (
+        output_dir / f"image_{index}a.png",
+        output_dir / f"image_{index}b.png",
+    )
+    for path in paths:
+        if path.exists():
+            raise FileExistsError(f"Refusing to overwrite {path}; choose another --index")
+
+    # RGBA shares preserve every decoded pixel, including transparency. For animated
+    # inputs, use the first frame, matching the dashboard's static image upload behavior.
+    with Image.open(input_path) as source:
+        image = source.convert("RGBA")
+    pixels = image.tobytes()
+    first = secrets.token_bytes(len(pixels))
+    second = bytes(pixel ^ mask for pixel, mask in zip(pixels, first))
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    for path, share in zip(paths, (first, second)):
+        # Exclusive creation also prevents overwriting a file created during splitting.
+        encoded = BytesIO()
+        Image.frombytes("RGBA", image.size, share).save(encoded, format="PNG")
+        with path.open("xb") as output:
+            output.write(encoded.getvalue())
+    return paths
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("input_image", type=Path)
+    parser.add_argument("--output-dir", type=Path, default=Path("."))
+    parser.add_argument("--index", type=int, help="Fixture number (default: first unused positive index)")
+    args = parser.parse_args()
+    try:
+        paths = split_image(args.input_image, args.output_dir, args.index)
+    except (OSError, ValueError) as error:
+        parser.exit(1, f"Could not split image: {error}\n")
+    for path in paths:
+        print(path)
+
+
+if __name__ == "__main__":
+    main()

src/web/image_uploads.rs

@@ -12,6 +12,9 @@ const NSFW_THRESHOLD: f32 = 0.5;
 const MODEL_BYTES: &[u8] =
     include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/assets/nsfw/model.onnx"));
 
+#[cfg(test)]
+mod fixture_tests;
+
 // An empty configuration keeps the dashboard's existing URL-only behavior.
 pub struct ImageUploads(pub Option<Arc<ImageUploadService>>);
 

src/web/image_uploads/fixture_tests.rs

@@ -0,0 +1,131 @@
+use std::{
+    collections::BTreeSet,
+    fs,
+    io::Cursor,
+    path::{Path, PathBuf},
+};
+
+use image::{ImageFormat, RgbaImage};
+
+use super::{prepare_image, ImageError, MODEL_BYTES};
+
+const NSFW_IMAGES: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/nsfw");
+const NON_NSFW_IMAGES: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/non_nsfw");
+
+fn fixture_pairs(directory: &Path) -> Vec<(PathBuf, PathBuf)> {
+    let mut indices = BTreeSet::new();
+    for entry in fs::read_dir(directory).expect("Could not read image fixture directory") {
+        let entry = entry.expect("Could not read image fixture entry");
+        let path = entry.path();
+        if path.extension().and_then(|extension| extension.to_str()) != Some("png") {
+            continue;
+        }
+        let name = path.file_name().unwrap().to_str().expect("Fixture name is not valid UTF-8");
+        let index = name
+            .strip_prefix("image_")
+            .and_then(|name| name.strip_suffix("a.png").or_else(|| name.strip_suffix("b.png")))
+            .filter(|index| index.parse::<u32>().is_ok_and(|index| index > 0))
+            .unwrap_or_else(|| {
+                panic!("Invalid fixture name {name}; expected image_1a.png/image_1b.png")
+            });
+        indices.insert(index.to_string());
+    }
+
+    assert!(
+        !indices.is_empty(),
+        "No paired fixtures in {}. Add images with scripts/split_nsfw_test_image.py.",
+        directory.display()
+    );
+    indices
+        .into_iter()
+        .map(|index| {
+            let first = directory.join(format!("image_{index}a.png"));
+            let second = directory.join(format!("image_{index}b.png"));
+            assert!(
+                first.is_file() && second.is_file(),
+                "Both fixture parts must exist: {} and {}",
+                first.display(),
+                second.display()
+            );
+            (first, second)
+        })
+        .collect()
+}
+
+fn recombine_parts(first: &RgbaImage, second: &RgbaImage) -> Result<RgbaImage, &'static str> {
+    if first.dimensions() != second.dimensions() {
+        return Err("Fixture parts must have identical dimensions");
+    }
+    let pixels = first.as_raw().iter().zip(second.as_raw()).map(|(a, b)| a ^ b).collect();
+    RgbaImage::from_raw(first.width(), first.height(), pixels).ok_or("Invalid fixture pixel data")
+}
+
+fn check_directory(directory: &str, expected_nsfw: bool) {
+    let pairs = fixture_pairs(Path::new(directory));
+    let model = nsfw::create_model(MODEL_BYTES).expect("Could not load bundled NSFW model");
+    for (first, second) in pairs {
+        let a = image::open(&first).expect("Could not decode first fixture part").to_rgba8();
+        let b = image::open(&second).expect("Could not decode second fixture part").to_rgba8();
+        let reconstructed = recombine_parts(&a, &b).expect("Could not recombine fixture parts");
+        let mut bytes = Cursor::new(Vec::new());
+        reconstructed
+            .write_to(&mut bytes, ImageFormat::Png)
+            .expect("Could not encode reconstructed fixture");
+
+        // Use the production path, including flattening transparency and the same
+        // combined category threshold. Never write the reconstructed image to disk.
+        let result = prepare_image(Some(&model), bytes.get_ref());
+        if expected_nsfw {
+            assert_eq!(
+                result.as_ref().err(),
+                Some(&ImageError::Nsfw),
+                "Expected {} to be blocked as NSFW, got {:?}",
+                first.display(),
+                result.as_ref().err()
+            );
+        } else {
+            assert!(
+                result.is_ok(),
+                "Expected {} to be accepted, got {:?}",
+                first.display(),
+                result.as_ref().err()
+            );
+        }
+        // The toggle must allow either category, while retaining image validation.
+        assert!(
+            prepare_image(None, bytes.get_ref()).is_ok(),
+            "Expected {} to be accepted with moderation disabled",
+            first.display()
+        );
+    }
+}
+
+#[test]
+#[ignore = "Add local split fixtures to tests/fixtures/nsfw, then run with --ignored"]
+fn nsfw_directory_images_are_blocked() {
+    check_directory(NSFW_IMAGES, true);
+}
+
+#[test]
+#[ignore = "Add local split fixtures to tests/fixtures/non_nsfw, then run with --ignored"]
+fn non_nsfw_directory_images_are_allowed() {
+    check_directory(NON_NSFW_IMAGES, false);
+}
+
+#[test]
+fn recombining_parts_restores_every_pixel_and_alpha_channel() {
+    let original = [0, 17, 254, 255, 45, 99, 128, 0];
+    let random = vec![18, 201, 50, 34, 166, 76, 11, 220];
+    let masked = original.iter().zip(&random).map(|(pixel, mask)| pixel ^ mask).collect();
+    let first = RgbaImage::from_raw(2, 1, random).unwrap();
+    let second = RgbaImage::from_raw(2, 1, masked).unwrap();
+    assert_eq!(recombine_parts(&first, &second).unwrap().as_raw(), &original);
+}
+
+#[test]
+fn rejects_parts_with_different_dimensions() {
+    assert_eq!(
+        recombine_parts(&RgbaImage::new(1, 1), &RgbaImage::new(2, 1)),
+        Err("Fixture parts must have identical dimensions")
+    );
+}

tests/fixtures/README.md

@@ -4,6 +4,27 @@ They accept the same `[1, 224, 224, 3]` float input as the production nsfw model
 other always returns Porn = 1. They exercise the real nsfw inference path without requiring
 image fixtures.
 
+## Image fixtures for the bundled model
+
+Add image pairs to these project directories:
+
+* `tests/fixtures/nsfw/` — images expected to be rejected by moderation.
+* `tests/fixtures/non_nsfw/` — images expected to be accepted by moderation.
+
+Install the splitter's dependency with `python3 -m pip install Pillow`, then run from the
+project root:
+
+```sh
+python3 scripts/split_nsfw_test_image.py /path/to/input.jpg --output-dir tests/fixtures/nsfw
+python3 scripts/split_nsfw_test_image.py /path/to/safe.jpg --output-dir tests/fixtures/non_nsfw
+```
+
+To run the tests:
+
+```sh
+SQLX_OFFLINE=true cargo test --locked directory_images -- --ignored
+```
+
 ## Deterministic model fixtures
 
 To regenerate them with Python's `onnx` package: