56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
"""Characterization parity: the extracted safety logic matches the donor on the
|
|
same fixture files (real exiftool-written keywords)."""
|
|
|
|
import shutil
|
|
import subprocess
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from photo_pipeline.integrations import exiftool
|
|
from photo_pipeline.services import safety
|
|
|
|
EXIFTOOL = shutil.which("exiftool")
|
|
pytestmark = pytest.mark.skipif(EXIFTOOL is None, reason="exiftool not installed")
|
|
|
|
|
|
def _jpeg(path, seed=1):
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
arr = np.random.default_rng(seed).integers(0, 256, (32, 32, 3), dtype=np.uint8)
|
|
Image.fromarray(arr).save(path, quality=90)
|
|
return str(path)
|
|
|
|
|
|
def _tag(path, keyword):
|
|
subprocess.run(
|
|
["exiftool", "-m", "-overwrite_original", f"-Keywords={keyword}", f"-Subject={keyword}", path],
|
|
check=True,
|
|
capture_output=True,
|
|
)
|
|
|
|
|
|
def test_extracted_marks_and_partition_match_donor(tmp_path):
|
|
donor_exif = pytest.importorskip("nsfwtag.exif")
|
|
|
|
nsfw = _jpeg(tmp_path / "nsfw.jpg", 1)
|
|
sfw = _jpeg(tmp_path / "sfw.jpg", 2)
|
|
plain = _jpeg(tmp_path / "plain.jpg", 3)
|
|
_tag(nsfw, "nsfw")
|
|
_tag(sfw, "sfw")
|
|
paths = [nsfw, sfw, plain]
|
|
|
|
# Donor behavior.
|
|
donor_marks = donor_exif.read_marks(paths)
|
|
donor_nsfw = donor_exif.read_tagged(paths, "nsfw")
|
|
|
|
# Extracted behavior on the same files.
|
|
keyword_sets = exiftool.read_keyword_sets(paths)
|
|
extracted_marks = safety.marks_from_keywords(keyword_sets)
|
|
_, skipped = safety.partition_nsfw(keyword_sets)
|
|
|
|
assert extracted_marks == donor_marks
|
|
assert set(skipped) == donor_nsfw
|
|
assert extracted_marks["nsfw"] == {nsfw}
|
|
assert extracted_marks["sfw"] == {sfw}
|