diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1a60936..b7188a6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -15,7 +15,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12", "3.13"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 @@ -27,7 +27,7 @@ jobs: python -m pip install --upgrade pip pip install -e ".[dev]" - name: Lint - run: ruff check src tests main.py + run: ruff check src tests scripts main.py - name: Test run: pytest -q diff --git a/README.md b/README.md index 5b96a90..9eaf30f 100644 --- a/README.md +++ b/README.md @@ -89,7 +89,7 @@ pyframe clip.gif --prescreen --backend aws # cascade: local gate then AW pyframe a.gif b.gif c.png --json # batch, machine-readable ``` -Exit code: `0` clean, `1` NSFW (per `--fail-on`), `2` bad input, `3` backend not installed, so it drops straight into a shell gate: `pyframe upload.gif || reject`. Equivalent module form: `python -m pyframe clip.gif`. +Exit code: `0` clean, `1` NSFW (per `--fail-on`), `2` bad input, `3` backend not installed, `4` could not classify, so it drops straight into a shell gate: `pyframe upload.gif || reject`. A broken backend exits `4` rather than `0` — if nothing scored, nothing was cleared. Equivalent module form: `python -m pyframe clip.gif`. ### Options @@ -98,12 +98,12 @@ Exit code: `0` clean, `1` NSFW (per `--fail-on`), `2` bad input, `3` backend not | `--backend` | `auto` | `local`, `aws`, or `local:` | | `--model` | model default | HuggingFace model id (local backend) | | `--region` | `us-east-1` | AWS region (aws backend) | -| `--max-frames` | `10` | frames to extract from a GIF/video | +| `--max-frames` | `10` | frames to extract from a GIF/video (must be ≥ 1) | | `--min-confidence` | backend default | NSFW threshold (0-1); `0.5` local, `0.8` aws | | `--sampler` | `motion` | `motion` (bucketing) or `dense` (uniform) | | `--prescreen` | off | enable the two-stage cascade | | `--escalate-threshold` | `0.15` | cascade gate (low = recall-safe) | -| `--max-escalations` | `2` | hard cap on precise (AWS) calls per file | +| `--max-escalations` | `2` | hard cap on precise (AWS) calls per file (must be ≥ 1) | | `--screen-fps` | `2.0` | soft-screen sample rate | | `--use-merged` / `--frames-per-batch` | off / `2` | merge frames into a grid before classifying | | `--json` / `--fail-on` | off / `nsfw` | output format / exit-code policy | @@ -117,6 +117,8 @@ Exit code: `0` clean, `1` NSFW (per `--fail-on`), `2` bad input, `3` backend not **Single-pass** (default): extract `max_frames` via motion bucketing, then classify each with one backend. +Decoding is split in two so memory tracks the sample rather than the clip: one pass reads the whole timeline as per-frame metadata (index, timestamp, motion), sampling picks from that, and a second pass fetches pixels for the selected frames only. A 900-frame 640x480 video scans in ~12 MB of frame memory instead of ~444 MB, so long videos are bounded rather than fatal. See [performance](docs/performance.md#memory) for the measurements and what the second pass costs. + **Cascade** (`--prescreen`): a free local model densely soft-screens the whole clip; if any frame scores above `--escalate-threshold` (a deliberately *low* recall gate), the most-suspicious frames are merged into grids and sent to the precise backend, capped at `--max-escalations` calls per file (default 2) so a heavily-flagged clip can never cost more than a single-pass scan. Clean media short-circuits to ~$0 and never hits the expensive backend. Because the soft-screen looks at *content* (not motion), it won't discard a unique suspicious frame the way motion bucketing can, and it fails *open*: a decode/inference error escalates rather than silently clearing. ## Cost diff --git a/docs/output.md b/docs/output.md index 7d527a0..bfce3bf 100644 --- a/docs/output.md +++ b/docs/output.md @@ -68,7 +68,7 @@ A clean image scanned with the local backend: | `source` | string | The input path exactly as passed in. | | `media_kind` | string | `"image"` or `"animation"` (GIF/video). | | `verdict` | string | Overall category: `clean`, `uncertain`, `nsfw`, or `error`. See [Verdict values](#verdict-values). | -| `is_nsfw` | bool | The authoritative pass/fail: `true` if any classified frame met the NSFW threshold. Branch on this. | +| `is_nsfw` | bool | The authoritative pass/fail: `true` if and only if `verdict` is `nsfw`. Branch on this. | | `max_score` | float | Highest NSFW score (0..1) across the classified frames, rounded to 4 dp. | | `worst_frame` | object \| null | The single highest-scoring frame (a [frame object](#frame-object)), or `null` if nothing was classified. | | `frames` | array | The [frame objects](#frame-object) that were classified. In a short-circuited clean cascade these are the soft-screen frames. | @@ -114,7 +114,10 @@ boolean version; `verdict` adds an "uncertain" band: | `nsfw` | `max_score >= min_confidence` (threshold default: 0.5 local, 0.8 aws) | | `uncertain` | `uncertain_threshold <= max_score < min_confidence` (default `uncertain_threshold` 0.3) | | `clean` | `max_score < uncertain_threshold` | -| `error` | every classified frame failed to score | +| `error` | every classified frame failed to score (CLI exit `4`) | + +Media that decodes to zero frames is not `clean` — it raises `MediaDecodeError` (CLI +exit `2`), since nothing ever looked at it. ## Single-pass vs cascade @@ -139,8 +142,15 @@ exit code encodes the outcome so it slots into shell gates: |------|---------| | `0` | clean | | `1` | NSFW (subject to `--fail-on`) | -| `2` | bad input (unsupported type, decode error, missing file) | +| `2` | bad input (unsupported type, decode error, missing file, unknown `--backend`, out-of-range option) | | `3` | backend not installed (missing optional extra) | +| `4` | could not classify: every frame failed to score (`verdict` is `error`) | + +Code `4` exists so a broken backend can't be mistaken for a clean file. If credentials +expire or the model fails to load, nothing was ever cleared, so `pyframe upload.gif || reject` +must reject rather than accept. The reason is printed to stderr. `--fail-on never` still +exits `0` in every case, including this one — it's the explicit "don't gate me, I only +want the JSON" escape hatch. ```bash pyframe upload.gif --backend local || echo "rejected" diff --git a/docs/performance.md b/docs/performance.md index b9b3186..fcb7826 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -39,8 +39,33 @@ Per-core throughput is the unit that transfers between machines, not a box total ## Memory -~0.5 GB resident per worker (model weights + buffers). Memory is not the bottleneck for -this path. +~0.5 GB resident per worker (model weights + buffers). Decoding adds only the sample, not +the clip: frames are selected against per-frame metadata first, and pixels are fetched in +a second pass for the selected frames alone. Peak decode memory is therefore bounded by +`max_frames` (or by `frames_per_batch * max_escalations` on an escalation), not by the +length of the media. + +Measured on a 900 frame 640x480 clip, `max_frames=10`, backend stubbed: + +| | peak RSS | growth over baseline | +|---|---|---| +| holding every decoded frame | 504 MB | 444 MB | +| two pass decode | 82 MB | 12 MB | + +The clip holds 829 MB of raw frame data, so the old shape scaled with the file: a 60 +second 1080p30 video is roughly 11 GB and does not complete. + +### What the second pass costs + +Two sequential decodes instead of one, unconditionally. Unwanted frames are skipped with +`grab()` rather than `read()`, and a third pass happens only when the cascade escalates, +so clean media stops at two. + +On the same clip that is 3.4s to 6.0s of wall clock, because a stubbed backend makes the +run purely decode bound. That ratio is the worst case, not the typical one: the per-stage +table above measures inference at ~91% of a real GIF scan, so a second decode moves a +much smaller share when a model is actually running. It is the right trade either way, +since the alternative on a long video is not a faster scan but an exhausted machine. ## Decode: file path vs in-memory (`scan_bytes`) diff --git a/pyproject.toml b/pyproject.toml index ef45000..b7aaf4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "pyframe-gif-video-image-moderation" -version = "0.4.0" +version = "0.5.0" description = "Two-stage NSFW moderation for GIFs, videos, and images via local HuggingFace models and/or AWS Rekognition." readme = "README.md" requires-python = ">=3.10" @@ -28,6 +28,7 @@ classifiers = [ "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", "Operating System :: OS Independent", "Topic :: Multimedia :: Graphics", "Topic :: Scientific/Engineering :: Image Recognition", @@ -65,3 +66,13 @@ only-include = ["src/pyframe", "tests", "README.md", "LICENSE"] [tool.pytest.ini_options] testpaths = ["tests"] + +[tool.ruff] +line-length = 110 +target-version = "py310" + +[tool.ruff.lint] +# Ruff's current defaults, pinned explicitly so a future release changing them can't +# silently widen or narrow what CI enforces. Widening the set (I, UP, B) is worth doing +# but touches files unrelated to any bug, so it belongs in its own change. +select = ["E4", "E7", "E9", "F"] diff --git a/src/pyframe/__init__.py b/src/pyframe/__init__.py index aa7961d..266767c 100644 --- a/src/pyframe/__init__.py +++ b/src/pyframe/__init__.py @@ -9,7 +9,19 @@ UnsupportedMediaError, ) from .image_utils import merge_images_to_grid, merge_to_grid -from .media import Frame, MediaKind, iter_frames, iter_frames_from_bytes, media_kind +from .media import ( + Frame, + FrameLike, + FrameMeta, + MediaKind, + iter_frame_meta, + iter_frame_meta_from_bytes, + iter_frames, + iter_frames_at, + iter_frames_from_bytes, + iter_frames_from_bytes_at, + media_kind, +) from .pipe import Pipe, scan, scan_bytes from .results import Label, ScanResult, Severity, Verdict from .scanner import Scanner @@ -20,9 +32,9 @@ try: __version__ = version("pyframe-gif-video-image-moderation") except PackageNotFoundError: - __version__ = "0.4.0" + __version__ = "0.5.0" except Exception: - __version__ = "0.4.0" + __version__ = "0.5.0" __all__ = [ "Pipe", @@ -39,9 +51,15 @@ "load_backend", "clear_backend_cache", "Frame", + "FrameMeta", + "FrameLike", "MediaKind", "iter_frames", "iter_frames_from_bytes", + "iter_frame_meta", + "iter_frame_meta_from_bytes", + "iter_frames_at", + "iter_frames_from_bytes_at", "media_kind", "merge_to_grid", "merge_images_to_grid", diff --git a/src/pyframe/backends/base.py b/src/pyframe/backends/base.py index e80be39..caadaec 100644 --- a/src/pyframe/backends/base.py +++ b/src/pyframe/backends/base.py @@ -1,7 +1,7 @@ from __future__ import annotations import abc -from typing import Sequence +from collections.abc import Iterable from ..media import Frame from ..results import Label, Verdict @@ -53,5 +53,11 @@ def classify(self, frame: Frame, *, min_confidence: float = 0.8) -> Verdict: timestamp=frame.timestamp, ) - def classify_batch(self, frames: Sequence[Frame], *, min_confidence: float = 0.8) -> list[Verdict]: + def classify_batch(self, frames: Iterable[Frame], *, min_confidence: float = 0.8) -> list[Verdict]: + """Score frames in order. + + `frames` may be a lazy iterable that decodes as it is consumed, so an override + should iterate it once and not index, len() or retain it. Materialising it puts + the whole sample in memory, which on a long video is the thing this avoids. + """ return [self.classify(f, min_confidence=min_confidence) for f in frames] diff --git a/src/pyframe/cli.py b/src/pyframe/cli.py index 4cd652b..4b38152 100644 --- a/src/pyframe/cli.py +++ b/src/pyframe/cli.py @@ -26,6 +26,14 @@ def _print_human(result: ScanResult) -> None: head += ", escalated)" if result.escalated else ", short-circuit clean)" print(head) + if result.verdict is Severity.ERROR: + reason = next((f.error for f in result.frames if f.error), None) + if reason: + # Flush first: stdout is block-buffered when piped but stderr isn't, so + # without this the reason lands above the line it explains. + sys.stdout.flush() + print(f" could not classify: {reason}", file=sys.stderr) + for frame in result.flagged_frames: names = ", ".join(label.name for label in frame.labels) or "flagged" print(f" t={frame.timestamp:.2f}s {frame.score:.2f} {names}") @@ -53,7 +61,7 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--frames-per-batch", type=int, default=2) parser.add_argument("--prescreen", action="store_true", help="enable the two-stage cascade") parser.add_argument("--escalate-threshold", type=float, default=0.15, help="cascade gate (low = recall-safe)") - parser.add_argument("--max-escalations", type=int, default=2, help="max precise (AWS) calls per file") + parser.add_argument("--max-escalations", type=int, default=2, help="max precise (AWS) calls per file (>= 1)") parser.add_argument("--screen-fps", type=float, default=2.0, help="soft-screen sample rate") parser.add_argument("--json", action="store_true", help="machine-readable output") parser.add_argument("--fail-on", choices=("nsfw", "uncertain", "never"), default="nsfw") @@ -90,6 +98,10 @@ def main() -> int: except BackendUnavailableError as exc: print(exc, file=sys.stderr) return 3 + except ValueError as exc: + # Unknown --backend, or an out-of-range knob: bad input, not a missing extra. + print(f"error: {exc}", file=sys.stderr) + return 2 rc = 0 for path in args.paths: @@ -105,7 +117,12 @@ def main() -> int: else: _print_human(result) - if args.fail_on == "nsfw" and result.is_nsfw: + if result.verdict is Severity.ERROR and args.fail_on != "never": + # Nothing scored, so nothing was cleared. Exiting 0 here would make + # `pyframe upload.gif || reject` accept every upload while the backend is + # down. --fail-on never stays the explicit "don't gate me" escape hatch. + rc = max(rc, 4) + elif args.fail_on == "nsfw" and result.is_nsfw: rc = max(rc, 1) elif args.fail_on == "uncertain" and result.verdict in (Severity.NSFW, Severity.UNCERTAIN): rc = max(rc, 1) diff --git a/src/pyframe/media.py b/src/pyframe/media.py index 16d848e..3480441 100644 --- a/src/pyframe/media.py +++ b/src/pyframe/media.py @@ -1,9 +1,22 @@ +"""Decoding media into frames, in two separable halves. + +Sampling decisions need only per-frame metadata (index, timestamp, motion); backends +need only the pixels of the frames finally selected. Nothing needs both at once, so +decoding is split: `iter_frame_meta` walks the whole file holding one frame at a time, +and `iter_frames_at` re-walks it to materialise just the selected frames. That keeps +peak memory proportional to the sample rather than to the length of the media. + +`iter_frames` (everything, with pixels) remains for callers that want it, and is built +on the same private generators so the two halves cannot drift apart. +""" + from __future__ import annotations import os +from collections.abc import Iterator, Sequence from dataclasses import dataclass from enum import Enum -from typing import Iterator +from typing import Protocol import cv2 import numpy as np @@ -22,6 +35,27 @@ class MediaKind(str, Enum): ANIMATION = "animation" # gif or video +class FrameLike(Protocol): + """What the samplers read. Frame and FrameMeta both satisfy it structurally.""" + + index: int + timestamp: float + motion_score: float + + +@dataclass(frozen=True, slots=True) +class FrameMeta: + """A frame's position and motion, without its pixels. + + Tens of bytes against several megabytes for a decoded Frame, which is what lets the + sampling pass hold the whole timeline at once. + """ + + index: int + timestamp: float # seconds from start + motion_score: float = 0.0 + + @dataclass class Frame: index: int @@ -63,15 +97,25 @@ def _read_image(source: str | os.PathLike) -> "np.ndarray": raise MediaDecodeError(f"Could not read image {source}: {exc}") from exc -def iter_frames(source: str | os.PathLike) -> Iterator[Frame]: +def _motion_gray(bgr: "np.ndarray") -> "np.ndarray": + return cv2.cvtColor(cv2.resize(bgr, (64, 64)), cv2.COLOR_BGR2GRAY) + + +def _motion_score(gray: "np.ndarray", prev_gray) -> float: + if prev_gray is None: + return 0.0 + return float(np.sum(cv2.absdiff(gray, prev_gray))) + + +def _require_file(source: str | os.PathLike) -> None: if not os.path.exists(source): raise FileNotFoundError(f"File not found: {source}") - kind = media_kind(source) - if kind is MediaKind.IMAGE: - yield Frame(index=0, timestamp=0.0, image=_read_image(source), motion_score=0.0) - return +def _iter_video(source: str | os.PathLike) -> Iterator[tuple[FrameMeta, "np.ndarray"]]: + """Every frame of a cv2-decodable file as (meta, BGR). One frame alive at a time; + the caller decides what to keep. Single source of truth for index, timestamp and + motion, so the metadata pass and the pixel pass cannot disagree.""" cap = cv2.VideoCapture(str(source)) if not cap.isOpened(): raise MediaDecodeError(f"Could not open {source}") @@ -87,38 +131,40 @@ def iter_frames(source: str | os.PathLike) -> Iterator[Frame]: ok, frame = cap.read() if not ok: break - small = cv2.resize(frame, (64, 64)) - gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY) - if prev_gray is None: - motion = 0.0 - else: - motion = float(np.sum(cv2.absdiff(gray, prev_gray))) + gray = _motion_gray(frame) + motion = _motion_score(gray, prev_gray) prev_gray = gray - yield Frame(index=index, timestamp=index / fps, image=frame, motion_score=motion) + yield FrameMeta(index=index, timestamp=index / fps, motion_score=motion), frame index += 1 finally: cap.release() -def iter_frames_from_bytes(data: bytes) -> Iterator[Frame]: - """Decode a GIF / static image from memory (no disk). Pillow only; for video - bytes use the path-based API. Motion + timestamps match iter_frames.""" +def _open_bytes(data: bytes): + """Open in-memory media and probe its frame count. Shared by all three bytes + entry points so the 'no video bytes' message stays identical across them.""" import io from PIL import Image try: img = Image.open(io.BytesIO(data)) - n_frames = getattr(img, "n_frames", 1) + return img, getattr(img, "n_frames", 1) except Exception as exc: raise MediaDecodeError( f"could not decode bytes in memory: {exc} " "(video bytes are not supported by scan_bytes; use the path-based API)" ) from exc + +def _iter_bytes(data: bytes) -> Iterator[tuple[FrameMeta, "np.ndarray"]]: + """The in-memory counterpart of _iter_video, via Pillow. GIF timestamps accumulate + per-frame durations rather than dividing by a container fps.""" + img, n_frames = _open_bytes(data) + if n_frames <= 1: rgb = np.asarray(img.convert("RGB")) - yield Frame(index=0, timestamp=0.0, image=cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)) + yield FrameMeta(index=0, timestamp=0.0), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) return prev_gray = None @@ -127,8 +173,137 @@ def iter_frames_from_bytes(data: bytes) -> Iterator[Frame]: img.seek(index) rgb = np.asarray(img.convert("RGB")) frame = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) - gray = cv2.cvtColor(cv2.resize(frame, (64, 64)), cv2.COLOR_BGR2GRAY) - motion = 0.0 if prev_gray is None else float(np.sum(cv2.absdiff(gray, prev_gray))) + gray = _motion_gray(frame) + motion = _motion_score(gray, prev_gray) prev_gray = gray - yield Frame(index=index, timestamp=timestamp, image=frame, motion_score=motion) + yield FrameMeta(index=index, timestamp=timestamp, motion_score=motion), frame timestamp += (img.info.get("duration") or 100) / 1000.0 # per-frame GIF duration (ms) + + +def _as_frame(meta: FrameMeta, image: "np.ndarray") -> Frame: + return Frame( + index=meta.index, + timestamp=meta.timestamp, + image=image, + motion_score=meta.motion_score, + ) + + +def _ordered_unique(metas: Sequence[FrameLike]) -> list[FrameLike]: + by_index = {m.index: m for m in metas} + return [by_index[i] for i in sorted(by_index)] + + +def iter_frames(source: str | os.PathLike) -> Iterator[Frame]: + """Every frame with its pixels. Holds one frame at a time itself, but `list()` of it + is proportional to the length of the media; prefer iter_frame_meta + iter_frames_at + when you only need a sample.""" + _require_file(source) + + if media_kind(source) is MediaKind.IMAGE: + yield Frame(index=0, timestamp=0.0, image=_read_image(source), motion_score=0.0) + return + + for meta, image in _iter_video(source): + yield _as_frame(meta, image) + + +def iter_frame_meta(source: str | os.PathLike) -> Iterator[FrameMeta]: + """Pass one: the whole timeline, no pixels retained. Raises what iter_frames raises, + at the same point, except that a still image is not decoded until it is fetched.""" + _require_file(source) + + if media_kind(source) is MediaKind.IMAGE: + yield FrameMeta(index=0, timestamp=0.0) + return + + for meta, _image in _iter_video(source): + yield meta + + +def iter_frames_at(source: str | os.PathLike, metas: Sequence[FrameLike]) -> Iterator[Frame]: + """Pass two: re-decode and yield only the frames named by `metas`, in index order. + + Skips unwanted frames with cap.grab() rather than seeking. CAP_PROP_POS_FRAMES is an + approximate keyframe seek on long-GOP codecs, VFR containers and GIF, so it can land + on a neighbouring frame; this pass must return exactly the frame pass one measured. + Do not replace the walk with a seek. + + timestamp and motion_score are carried from `metas`, never recomputed: motion is a + diff against the previous *decoded* frame, so recomputing it here would measure + across the skipped gap instead. + """ + wanted = _ordered_unique(metas) + if not wanted: + return + + _require_file(source) + if media_kind(source) is MediaKind.IMAGE: + image = _read_image(source) + for meta in wanted: + yield _as_frame(meta, image) + return + + cap = cv2.VideoCapture(str(source)) + if not cap.isOpened(): + raise MediaDecodeError(f"Could not open {source}") + + pending = iter(wanted) + target = next(pending) + index = 0 + try: + while True: + if index != target.index: + if not cap.grab(): + break + index += 1 + continue + ok, image = cap.read() + if not ok: + break + yield _as_frame(target, image) + index += 1 + try: + target = next(pending) + except StopIteration: + return + finally: + cap.release() + + raise MediaDecodeError( + f"{source}: second decode pass ended at frame {index}, before the selected " + f"frame {target.index}. The file may decode non-deterministically." + ) + + +def iter_frames_from_bytes(data: bytes) -> Iterator[Frame]: + """Decode a GIF / static image from memory (no disk). Pillow only; for video + bytes use the path-based API. Motion + timestamps match iter_frames.""" + for meta, image in _iter_bytes(data): + yield _as_frame(meta, image) + + +def iter_frame_meta_from_bytes(data: bytes) -> Iterator[FrameMeta]: + """Pass one for the in-memory path.""" + for meta, _image in _iter_bytes(data): + yield meta + + +def iter_frames_from_bytes_at(data: bytes, metas: Sequence[FrameLike]) -> Iterator[Frame]: + """Pass two for the in-memory path. Pillow replays the GIF from frame 0, so this + walks forward and keeps the matches rather than seeking.""" + wanted = {m.index: m for m in metas} + if not wanted: + return + + for meta, image in _iter_bytes(data): + target = wanted.pop(meta.index, None) + if target is None: + continue + yield _as_frame(target, image) + if not wanted: + return + + raise MediaDecodeError( + f"second decode pass ended before the selected frame(s) {sorted(wanted)}" + ) diff --git a/src/pyframe/sampling.py b/src/pyframe/sampling.py index 8a59e77..cfffe80 100644 --- a/src/pyframe/sampling.py +++ b/src/pyframe/sampling.py @@ -13,15 +13,20 @@ from __future__ import annotations -from typing import Iterable, Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence +from typing import TypeVar -from .media import Frame +from .media import FrameLike + +# Samplers only read index, timestamp and motion_score, so they work on bare metadata +# as well as on decoded frames. Binding the element type preserves which one came in. +F = TypeVar("F", bound=FrameLike) class MotionBucketSampler: # Highest-motion frame per equal-width bucket. Lossy/content-blind: cost lever # only, never the cascade gate (motion is uncorrelated with NSFW content). - def select(self, frames: Sequence[Frame], budget: int) -> list[Frame]: + def select(self, frames: Sequence[F], budget: int) -> list[F]: n = len(frames) if n == 0: return [] @@ -29,7 +34,7 @@ def select(self, frames: Sequence[Frame], budget: int) -> list[Frame]: return list(frames) chunk = n / budget - chosen: list[Frame] = [] + chosen: list[F] = [] for i in range(budget): start = int(i * chunk) end = n if i == budget - 1 else int((i + 1) * chunk) @@ -46,7 +51,7 @@ class DenseUniformSampler: def __init__(self, target_fps: float = 2.0): self.target_fps = max(target_fps, 0.01) - def select(self, frames: Sequence[Frame]) -> list[Frame]: + def select(self, frames: Sequence[F]) -> list[F]: n = len(frames) if n <= 1: return list(frames) @@ -57,17 +62,23 @@ def select(self, frames: Sequence[Frame]) -> list[Frame]: source_fps = (n - 1) / duration stride = max(1, round(source_fps / self.target_fps)) - return list(frames[::stride]) + selected = list(frames[::stride]) + # A strided slice lands on the last frame only when (n - 1) % stride == 0, so up + # to stride-1 frames at the end of every clip would sit outside the floor. An + # event that runs to the end of the clip has to stay catchable. + if selected[-1].index != frames[-1].index: + selected.append(frames[-1]) + return selected class SuspicionSampler: # Keep the most-suspicious frames in a window (screen score, then motion). def select( self, - frames: Sequence[Frame], + frames: Sequence[F], budget: int, scores: Mapping[int, float] | None = None, - ) -> list[Frame]: + ) -> list[F]: n = len(frames) if n == 0: return [] diff --git a/src/pyframe/scanner.py b/src/pyframe/scanner.py index 6ae9c3b..f0cb841 100644 --- a/src/pyframe/scanner.py +++ b/src/pyframe/scanner.py @@ -4,8 +4,16 @@ from .backends import Backend, load_backend from .config import Config +from .errors import MediaDecodeError from .image_utils import merge_to_grid -from .media import MediaKind, iter_frames, iter_frames_from_bytes, media_kind +from .media import ( + MediaKind, + iter_frame_meta, + iter_frame_meta_from_bytes, + iter_frames_at, + iter_frames_from_bytes_at, + media_kind, +) from .results import ScanResult, Severity, Verdict from .sampling import ( DenseUniformSampler, @@ -28,6 +36,22 @@ def __init__(self, precise: Backend, *, screen: Backend | None = None, config: C @classmethod def from_config(cls, config: Config) -> "Scanner": + # Validate before load_backend: constructing a backend pulls ~0.5 GB of weights, + # which is a lot of work to do before rejecting the config. + # + # A budget below 1 doesn't disable escalation, it removes the cap: the suspicion + # sampler treats a non-positive budget as "keep everything", so every flagged + # frame would be escalated. Reject it rather than guess which of "never escalate" + # or "escalate once" was meant. + if config.prescreen.enabled and config.prescreen.max_escalations < 1: + raise ValueError( + f"max_escalations must be >= 1, got {config.prescreen.max_escalations}" + ) + # Same trap on the single-pass side: the samplers read a non-positive budget as + # "keep everything", which would materialise the whole clip. + if config.max_frames < 1: + raise ValueError(f"max_frames must be >= 1, got {config.max_frames}") + precise = load_backend(config.backend, model=config.model, region=config.region) screen = None if config.prescreen.enabled: @@ -37,44 +61,59 @@ def from_config(cls, config: Config) -> "Scanner": def scan(self, source) -> ScanResult: start = time.perf_counter() kind = media_kind(source) - frames = list(iter_frames(source)) - return self._scan_frames(str(source), kind, frames, start) + # Pass one holds the whole timeline as metadata; pixels are fetched later, and + # only for the frames sampling actually selects. + metas = list(iter_frame_meta(source)) + + def fetch(selected): + return iter_frames_at(source, selected) + + return self._scan_frames(str(source), kind, metas, fetch, start) def scan_bytes(self, data, *, label: str = "") -> ScanResult: """Scan a GIF/image decoded from memory, no disk touched.""" start = time.perf_counter() - frames = list(iter_frames_from_bytes(data)) - kind = MediaKind.ANIMATION if len(frames) > 1 else MediaKind.IMAGE - return self._scan_frames(label, kind, frames, start) + metas = list(iter_frame_meta_from_bytes(data)) + kind = MediaKind.ANIMATION if len(metas) > 1 else MediaKind.IMAGE - def _scan_frames(self, source, kind, frames, start) -> ScanResult: - if kind is MediaKind.IMAGE: - verdicts = self.precise.classify_batch(frames, min_confidence=self.min_confidence) - return self._aggregate(source, kind, verdicts, [], len(frames), start) + def fetch(selected): + return iter_frames_from_bytes_at(data, selected) + + return self._scan_frames(label, kind, metas, fetch, start) + + def _scan_frames(self, source, kind, metas, fetch, start) -> ScanResult: + # Nothing to look at is not the same as nothing to find: aggregating zero frames + # would report a confident "clean" for media no backend ever saw. + if not metas: + raise MediaDecodeError(f"decoded 0 frames from {source}") - if not frames: - return self._aggregate(source, kind, [], [], 0, start) + if kind is MediaKind.IMAGE: + verdicts = self.precise.classify_batch(fetch(metas), min_confidence=self.min_confidence) + return self._aggregate(source, kind, verdicts, [], len(metas), start) if self.config.prescreen.enabled and self.screen is not None: - return self._cascade(source, kind, frames, start) - return self._single_pass(source, kind, frames, start) + return self._cascade(source, kind, metas, fetch, start) + return self._single_pass(source, kind, metas, fetch, start) - def _single_pass(self, source, kind, frames, start) -> ScanResult: + def _single_pass(self, source, kind, metas, fetch, start) -> ScanResult: cfg = self.config if cfg.sampler == "dense": - selected = DenseUniformSampler(cfg.prescreen.screen_fps).select(frames) + selected = DenseUniformSampler(cfg.prescreen.screen_fps).select(metas) if len(selected) > cfg.max_frames: selected = MotionBucketSampler().select(selected, cfg.max_frames) else: - selected = self._motion_select_with_floor(frames) + selected = self._motion_select_with_floor(metas) + # The only materialisation on this path, and every branch above caps `selected` + # at max_frames. + frames = list(fetch(selected)) if cfg.use_merged: - verdicts = self._classify_merged(selected) + verdicts = self._classify_merged(frames) else: - verdicts = self.precise.classify_batch(selected, min_confidence=self.min_confidence) - return self._aggregate(source, kind, verdicts, [], len(frames), start) + verdicts = self.precise.classify_batch(frames, min_confidence=self.min_confidence) + return self._aggregate(source, kind, verdicts, [], len(metas), start) - def _motion_select_with_floor(self, frames): + def _motion_select_with_floor(self, metas): # Recall floor for the default (motion) sampler. The uniform-by-time sample at # screen_fps bounds the sampling stride, so no NSFW event longer than that stride # can fall entirely between selected frames. Motion is content-blind (it can keep @@ -83,7 +122,7 @@ def _motion_select_with_floor(self, frames): # cf. Ding, Sener, and Yao, arXiv:2210.10352 (temporal coverage as a prior, and # the decoupling of motion from static semantic content). cfg = self.config - floor = DenseUniformSampler(cfg.prescreen.screen_fps).select(frames) + floor = DenseUniformSampler(cfg.prescreen.screen_fps).select(metas) if len(floor) >= cfg.max_frames: # The floor already fills the budget; motion only decides what to drop, # exactly as the `dense` path trims its own uniform sample. @@ -92,19 +131,24 @@ def _motion_select_with_floor(self, frames): # the highest-motion frames the floor did not already include. have = {f.index for f in floor} extra = sorted( - (f for f in frames if f.index not in have), + (m for m in metas if m.index not in have), key=lambda f: f.motion_score, reverse=True, ) selected = floor + extra[: cfg.max_frames - len(floor)] return sorted(selected, key=lambda f: f.index) - def _cascade(self, source, kind, frames, start) -> ScanResult: + def _cascade(self, source, kind, metas, fetch, start) -> ScanResult: cfg = self.config pc = cfg.prescreen - screen_frames = DenseUniformSampler(pc.screen_fps).select(frames) - screen_verdicts = self.screen.classify_batch(screen_frames, min_confidence=pc.escalate_threshold) + screen_metas = DenseUniformSampler(pc.screen_fps).select(metas) + # The screen set is screen_fps x duration, not max_frames, so on a long clip it + # is most of the timeline. classify_batch only iterates, so handing it the lazy + # fetch keeps one decoded frame alive at a time instead of all of them. + screen_verdicts = self.screen.classify_batch( + fetch(screen_metas), min_confidence=pc.escalate_threshold + ) scores = {v.frame_index: v.score for v in screen_verdicts} flagged = [ @@ -114,7 +158,7 @@ def _cascade(self, source, kind, frames, start) -> ScanResult: ] if not flagged: return self._aggregate( - source, kind, [], screen_verdicts, len(frames), start, escalated=False, windows=0 + source, kind, [], screen_verdicts, len(metas), start, escalated=False, windows=0 ) # Keep the most-suspicious flagged frames, capped so we make at most @@ -122,27 +166,32 @@ def _cascade(self, source, kind, frames, start) -> ScanResult: per_batch = max(1, cfg.frames_per_batch) frame_budget = pc.max_escalations * per_batch flagged_set = set(flagged) - flagged_frames = [f for f in frames if f.index in flagged_set] - selected = SuspicionSampler().select(flagged_frames, frame_budget, scores) + flagged_metas = [m for m in metas if m.index in flagged_set] + selected = SuspicionSampler().select(flagged_metas, frame_budget, scores) # Always fill at least one full grid (send both even if only one frame flagged). - selected = self._ensure_min_frames(selected, frames, scores, per_batch) + selected = self._ensure_min_frames(selected, metas, scores, per_batch) - # Send the top suspicious frames to the precise backend as merged grids. - precise = self._classify_merged(selected) + # Send the top suspicious frames to the precise backend as merged grids. A third + # decode, and only when something flagged: clean media stops after two. + precise = self._classify_merged(list(fetch(selected))) - windows = group_flagged_into_windows(flagged, len(frames), pc.group_gap, pc.window_pad) + windows = group_flagged_into_windows(flagged, len(metas), pc.group_gap, pc.window_pad) return self._aggregate( - source, kind, precise, screen_verdicts, len(frames), start, + source, kind, precise, screen_verdicts, len(metas), start, escalated=True, windows=len(windows), ) - def _ensure_min_frames(self, selected, frames, scores, minimum): + def _ensure_min_frames(self, selected, metas, scores, minimum): if len(selected) >= minimum: return selected have = {f.index for f in selected} + # Same key as SuspicionSampler: screen score first, motion only as a tiebreak. + # These are different units -- scores are 0..1, motion is a pixel-diff sum up to + # ~1e6 -- so one flat key would rank any moving frame above a screened frame that + # scored 0.99. extra = sorted( - (f for f in frames if f.index not in have), - key=lambda f: scores.get(f.index, f.motion_score), + (m for m in metas if m.index not in have), + key=lambda f: (scores.get(f.index, -1.0), f.motion_score), reverse=True, ) if not extra: @@ -176,9 +225,13 @@ def _aggregate( worst = max(primary, key=lambda v: v.score) if primary else None max_score = worst.score if worst else 0.0 - is_nsfw = any(v.is_nsfw for v in classified) errored = bool(primary) and all(v.error for v in primary) severity = Severity.from_score(max_score, self.min_confidence, cfg.uncertain_threshold, errored=errored) + # Derive is_nsfw from the severity rather than computing it separately, so the + # two can't disagree. They used to: a short-circuited cascade scores max_score + # off the screen verdicts, which no classified frame ever backed, and the result + # could read verdict=nsfw with is_nsfw=False. + is_nsfw = severity is Severity.NSFW cost = len(classified) * self.precise.cost_per_image if self.screen is not None: diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..dbe5474 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,130 @@ +import json +import sys + +import numpy as np +import pytest +from PIL import Image + +from pyframe import cli +from pyframe import scanner as scanner_mod +from pyframe.backends.base import Backend + + +class ScriptedBackend(Backend): + # Scores every frame the same, or fails on every frame, so a CLI test can pin one + # exact verdict without downloading a model. + cost_per_image = 0.0 + default_min_confidence = 0.5 + + def __init__(self, score=0.0, fail_with=None, name="fake"): + self.name = name + self.score = score + self.fail_with = fail_with + + def _score(self, image): + if self.fail_with is not None: + raise RuntimeError(self.fail_with) + return self.score, [], None + + +@pytest.fixture +def gif(tmp_path): + path = tmp_path / "clip.gif" + imgs = [Image.fromarray(np.full((32, 32, 3), v, np.uint8)) for v in (10, 60, 120)] + imgs[0].save(path, save_all=True, append_images=imgs[1:], duration=80, loop=0) + return str(path) + + +def _run(monkeypatch, argv, backend=None): + """Drive cli.main() end to end. Omit `backend` to let the real load_backend run.""" + if backend is not None: + monkeypatch.setattr(scanner_mod, "load_backend", lambda *a, **k: backend) + monkeypatch.setattr(sys, "argv", ["pyframe", *argv]) + return cli.main() + + +def test_clean_media_exits_zero(monkeypatch, gif): + assert _run(monkeypatch, [gif], ScriptedBackend(score=0.01)) == 0 + + +def test_nsfw_media_exits_one(monkeypatch, gif): + assert _run(monkeypatch, [gif], ScriptedBackend(score=0.99)) == 1 + + +def test_errored_scan_exits_four(monkeypatch, capsys, gif): + # Every frame failing means nothing was cleared. Exiting 0 here would make + # `pyframe upload.gif || reject` accept every upload while the backend is down. + rc = _run(monkeypatch, [gif], ScriptedBackend(fail_with="credentials expired")) + + assert rc == 4 + assert "credentials expired" in capsys.readouterr().err + + +def test_fail_on_never_exits_zero_even_on_error(monkeypatch, gif): + # The explicit "don't gate me, I just want the JSON" escape hatch. + assert _run(monkeypatch, [gif, "--fail-on", "never"], ScriptedBackend(fail_with="boom")) == 0 + + +def test_fail_on_never_exits_zero_on_nsfw(monkeypatch, gif): + assert _run(monkeypatch, [gif, "--fail-on", "never"], ScriptedBackend(score=0.99)) == 0 + + +def test_fail_on_uncertain_gates_the_middle_band(monkeypatch, gif): + # 0.4 sits under the 0.5 threshold but over uncertain_threshold 0.3. + assert _run(monkeypatch, [gif, "--fail-on", "uncertain"], ScriptedBackend(score=0.4)) == 1 + assert _run(monkeypatch, [gif], ScriptedBackend(score=0.4)) == 0 + + +def test_missing_file_exits_two(monkeypatch, tmp_path): + assert _run(monkeypatch, [str(tmp_path / "nope.gif")], ScriptedBackend()) == 2 + + +def test_unsupported_type_exits_two(monkeypatch, tmp_path): + path = tmp_path / "notes.txt" + path.write_text("hello") + + assert _run(monkeypatch, [str(path)], ScriptedBackend()) == 2 + + +def test_unknown_backend_exits_two_without_a_traceback(monkeypatch, capsys, gif): + rc = _run(monkeypatch, [gif, "--backend", "bogus"]) # deliberately unpatched + + assert rc == 2 + assert "bogus" in capsys.readouterr().err + + +def test_max_escalations_of_zero_is_rejected(monkeypatch, capsys, gif): + # 0 used to mean "no cap": the suspicion sampler keeps every frame on a non-positive + # budget, so the flag whose job is bounding spend removed the bound instead. + rc = _run(monkeypatch, [gif, "--prescreen", "--max-escalations", "0"], ScriptedBackend()) + + assert rc == 2 + assert "max_escalations" in capsys.readouterr().err + + +def test_json_output_carries_the_documented_keys(monkeypatch, capsys, gif): + rc = _run(monkeypatch, [gif, "--json"], ScriptedBackend(score=0.01)) + payload = json.loads(capsys.readouterr().out) + + assert rc == 0 + for key in ( + "source", "media_kind", "verdict", "is_nsfw", "max_score", "worst_frame", + "frames", "backends_used", "frames_total", "frames_screened", + "frames_classified", "cost_usd", "prescreen_used", "escalated", "windows", + "elapsed_s", + ): + assert key in payload, key + + +def test_json_verdict_and_is_nsfw_agree(monkeypatch, capsys, gif): + _run(monkeypatch, [gif, "--json"], ScriptedBackend(score=0.99)) + payload = json.loads(capsys.readouterr().out) + + assert payload["verdict"] == "nsfw" + assert payload["is_nsfw"] is True + + +def test_batch_exit_code_is_the_worst_across_files(monkeypatch, gif, tmp_path): + rc = _run(monkeypatch, [gif, str(tmp_path / "nope.gif")], ScriptedBackend(score=0.99)) + + assert rc == 2 # bad input (2) outranks nsfw (1) diff --git a/tests/test_media.py b/tests/test_media.py new file mode 100644 index 0000000..41c3aee --- /dev/null +++ b/tests/test_media.py @@ -0,0 +1,198 @@ +import io + +import numpy as np +import pytest +from PIL import Image + +from pyframe.errors import MediaDecodeError, UnsupportedMediaError +from pyframe.media import ( + ANIMATION_EXTS, + IMAGE_EXTS, + FrameMeta, + MediaKind, + iter_frame_meta, + iter_frame_meta_from_bytes, + iter_frames, + iter_frames_at, + iter_frames_from_bytes, + iter_frames_from_bytes_at, + media_kind, +) + + +def _write_gif(path, fills=(10, 60, 250, 120, 30), size=32, duration=80): + imgs = [Image.fromarray(np.full((size, size, 3), v, np.uint8)) for v in fills] + imgs[0].save(path, save_all=True, append_images=imgs[1:], duration=duration, loop=0) + return str(path) + + +def _gif_bytes(fills=(10, 60, 250, 120, 30), size=16, duration=80): + imgs = [Image.fromarray(np.full((size, size, 3), v, np.uint8)) for v in fills] + buf = io.BytesIO() + imgs[0].save(buf, format="GIF", save_all=True, append_images=imgs[1:], duration=duration, loop=0) + return buf.getvalue() + + +# The cv2.VideoCapture path had no coverage at all, so nothing checked that the OpenCV +# major version CI actually resolves can still decode the format the package is named +# after. `opencv-python-headless>=4.8` currently resolves to OpenCV 5. + + +def test_iter_frames_decodes_a_gif_from_disk(tmp_path): + frames = list(iter_frames(_write_gif(tmp_path / "clip.gif"))) + + assert [f.index for f in frames] == [0, 1, 2, 3, 4] + assert frames[0].image.shape == (32, 32, 3) + assert frames[0].image.dtype == np.uint8 + + +def test_iter_frames_timestamps_start_at_zero_and_increase(tmp_path): + stamps = [f.timestamp for f in iter_frames(_write_gif(tmp_path / "clip.gif"))] + + assert stamps[0] == 0.0 + assert all(later > earlier for earlier, later in zip(stamps, stamps[1:])) + + +def test_iter_frames_scores_motion_against_the_previous_frame(tmp_path): + frames = list(iter_frames(_write_gif(tmp_path / "clip.gif"))) + + assert frames[0].motion_score == 0.0 # nothing to diff against + assert any(f.motion_score > 0 for f in frames[1:]) + + +def test_iter_frames_reads_a_still_image(tmp_path): + path = tmp_path / "still.png" + Image.fromarray(np.full((16, 24, 3), 200, np.uint8)).save(path) + + frames = list(iter_frames(str(path))) + + assert len(frames) == 1 + assert (frames[0].index, frames[0].timestamp, frames[0].motion_score) == (0, 0.0, 0.0) + assert frames[0].image.shape == (16, 24, 3) + + +def test_iter_frames_on_a_missing_file(tmp_path): + with pytest.raises(FileNotFoundError): + list(iter_frames(str(tmp_path / "nope.gif"))) + + +def test_iter_frames_on_a_gif_that_is_not_a_gif(tmp_path): + path = tmp_path / "fake.gif" + path.write_text("this is not a gif") + + with pytest.raises(MediaDecodeError): + list(iter_frames(str(path))) + + +def test_iter_frames_from_bytes_on_garbage(): + with pytest.raises(MediaDecodeError): + list(iter_frames_from_bytes(b"not an image")) + + +# The two-pass split is only safe if pass two returns exactly what a single full decode +# would have. These pin that, since a drift between the passes would silently moderate a +# different frame than the one sampling chose. + +MANY = tuple(range(0, 200, 8)) + + +def test_meta_pass_matches_the_full_decode(tmp_path): + path = _write_gif(tmp_path / "clip.gif", fills=MANY) + + full = [(f.index, f.timestamp, f.motion_score) for f in iter_frames(path)] + meta = [(m.index, m.timestamp, m.motion_score) for m in iter_frame_meta(path)] + + assert meta == full + + +def test_fetched_frames_match_the_full_decode(tmp_path): + path = _write_gif(tmp_path / "clip.gif", fills=MANY) + full = {f.index: f for f in iter_frames(path)} + wanted = [m for m in iter_frame_meta(path) if m.index in (0, 7, 13)] + + got = list(iter_frames_at(path, wanted)) + + assert [f.index for f in got] == [0, 7, 13] + for frame in got: + original = full[frame.index] + assert frame.timestamp == original.timestamp + # Carried from the meta, never recomputed: a diff against the previous *kept* + # frame would measure across the skipped gap instead. + assert frame.motion_score == original.motion_score + assert np.array_equal(frame.image, original.image) + + +def test_iter_frames_at_sorts_and_deduplicates(tmp_path): + path = _write_gif(tmp_path / "clip.gif") + metas = list(iter_frame_meta(path)) + scrambled = [metas[3], metas[1], metas[3], metas[0]] + + assert [f.index for f in iter_frames_at(path, scrambled)] == [0, 1, 3] + + +def test_iter_frames_at_with_nothing_wanted(tmp_path): + path = _write_gif(tmp_path / "clip.gif") + + assert list(iter_frames_at(path, [])) == [] + + +def test_iter_frames_at_raises_when_the_pass_ends_early(tmp_path): + path = _write_gif(tmp_path / "clip.gif") # five frames + + with pytest.raises(MediaDecodeError): + list(iter_frames_at(path, [FrameMeta(index=99, timestamp=9.9)])) + + +def test_meta_and_fetch_round_trip_for_a_still_image(tmp_path): + path = tmp_path / "still.png" + Image.fromarray(np.full((16, 24, 3), 200, np.uint8)).save(path) + + metas = list(iter_frame_meta(str(path))) + frames = list(iter_frames_at(str(path), metas)) + + assert [(m.index, m.timestamp) for m in metas] == [(0, 0.0)] + assert frames[0].image.shape == (16, 24, 3) + + +def test_bytes_meta_pass_matches_the_full_decode(): + # Guards the cumulative per-frame GIF durations, which are not reconstructible from + # an index and an fps the way the file path timestamps are. + data = _gif_bytes() + + full = [(f.index, f.timestamp, f.motion_score) for f in iter_frames_from_bytes(data)] + meta = [(m.index, m.timestamp, m.motion_score) for m in iter_frame_meta_from_bytes(data)] + + assert meta == full + + +def test_bytes_fetch_matches_the_full_decode(): + data = _gif_bytes() + full = {f.index: f for f in iter_frames_from_bytes(data)} + wanted = [m for m in iter_frame_meta_from_bytes(data) if m.index in (1, 4)] + + got = list(iter_frames_from_bytes_at(data, wanted)) + + assert [f.index for f in got] == [1, 4] + for frame in got: + assert frame.timestamp == full[frame.index].timestamp + assert np.array_equal(frame.image, full[frame.index].image) + + +@pytest.mark.parametrize("ext", sorted(IMAGE_EXTS)) +def test_media_kind_maps_every_image_extension(ext): + assert media_kind(f"x{ext}") is MediaKind.IMAGE + + +@pytest.mark.parametrize("ext", sorted(ANIMATION_EXTS)) +def test_media_kind_maps_every_animation_extension(ext): + assert media_kind(f"x{ext}") is MediaKind.ANIMATION + + +def test_media_kind_ignores_case(): + assert media_kind("X.GIF") is MediaKind.ANIMATION + + +@pytest.mark.parametrize("name", ["x.txt", "x.pdf", "noextension"]) +def test_media_kind_rejects_unknown_types(name): + with pytest.raises(UnsupportedMediaError): + media_kind(name) diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 0000000..780d83d --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,145 @@ +"""The decode path must hold a sample, not a timeline. + +These assert peak *live Frame count*, not bytes. Counting objects is exact and platform +independent, where tracemalloc may not observe cv2's buffers at all and ru_maxrss is a +high water mark reported in bytes on macOS and kilobytes on Linux. Each bound is paired +with a positive control that fails if the tracker ever stops working, because a memory +test that silently measures nothing is worse than no test. +""" + +import weakref + +import numpy as np +import pytest +from PIL import Image + +from pyframe import media as media_mod +from pyframe.backends.base import Backend +from pyframe.config import Config, PrescreenConfig +from pyframe.scanner import Scanner + +N_FRAMES = 240 + + +@pytest.fixture +def long_gif(tmp_path): + path = tmp_path / "long.gif" + imgs = [ + Image.fromarray(np.full((32, 32, 3), (i * 7) % 256, np.uint8)) for i in range(N_FRAMES) + ] + imgs[0].save(path, save_all=True, append_images=imgs[1:], duration=20, loop=0) + return str(path) + + +class FrameTracker: + """Counts live Frames by wrapping the constructor media.py resolves at call time.""" + + def __init__(self, monkeypatch): + self.live = self.peak = self.total = 0 + real = media_mod.Frame + + def factory(*args, **kwargs): + frame = real(*args, **kwargs) + self.live += 1 + self.total += 1 + self.peak = max(self.peak, self.live) + weakref.finalize(frame, self._released) + return frame + + monkeypatch.setattr(media_mod, "Frame", factory) + + def _released(self): + self.live -= 1 + + +class FlatBackend(Backend): + name = "fake" + cost_per_image = 0.0 + default_min_confidence = 0.5 + + def __init__(self, score=0.0, name="fake"): + self.name = name + self.score = score + + def _score(self, image): + return self.score, [], None + + +def test_iter_frames_holds_the_whole_clip(monkeypatch, long_gif): + # Positive control. If this does not peak at the full length then the tracker is + # broken and every bound below is vacuous. + tracker = FrameTracker(monkeypatch) + + frames = list(media_mod.iter_frames(long_gif)) + + assert len(frames) == N_FRAMES + assert tracker.peak == N_FRAMES + + +def test_meta_pass_holds_one_frame_at_a_time(monkeypatch, long_gif): + tracker = FrameTracker(monkeypatch) + + metas = list(media_mod.iter_frame_meta(long_gif)) + + assert len(metas) == N_FRAMES # the whole timeline is still measured + assert tracker.total == 0 # and not one Frame was built to do it + + +def test_single_pass_decodes_only_the_sample(monkeypatch, long_gif): + tracker = FrameTracker(monkeypatch) + config = Config(backend=FlatBackend(), max_frames=10) + + result = Scanner.from_config(config).scan(long_gif) + + assert result.frames_total == N_FRAMES # every frame was considered + assert tracker.total == 10 # only the selected ones were decoded + assert tracker.peak <= 12 + + +def test_cascade_streams_the_screen_pass(monkeypatch, long_gif): + # The screen set is screen_fps x duration, not max_frames, so this is the pass that + # would still blow up if only the single-pass path had been bounded. + tracker = FrameTracker(monkeypatch) + config = Config( + backend=FlatBackend(name="aws"), + screen_backend=FlatBackend(name="local"), + max_frames=10, + prescreen=PrescreenConfig(enabled=True, screen_fps=2.0), + ) + + result = Scanner.from_config(config).scan(long_gif) + + assert result.frames_screened > 1 + assert result.escalated is False + assert tracker.peak <= 4 + + +def test_cascade_escalation_stays_within_its_budget(monkeypatch, long_gif): + tracker = FrameTracker(monkeypatch) + config = Config( + backend=FlatBackend(score=0.99, name="aws"), + screen_backend=FlatBackend(score=0.99, name="local"), + max_frames=10, + frames_per_batch=2, + prescreen=PrescreenConfig(enabled=True, screen_fps=2.0, max_escalations=2), + ) + + result = Scanner.from_config(config).scan(long_gif) + + assert result.escalated is True + # Every screened frame flagged, yet only max_escalations x frames_per_batch are + # ever held at once. + assert tracker.peak <= 6 + + +def test_scan_bytes_is_bounded_too(monkeypatch, long_gif): + with open(long_gif, "rb") as fh: + data = fh.read() + tracker = FrameTracker(monkeypatch) + config = Config(backend=FlatBackend(), max_frames=10) + + result = Scanner.from_config(config).scan_bytes(data) + + assert result.frames_total == N_FRAMES + assert tracker.total == 10 + assert tracker.peak <= 12 diff --git a/tests/test_sampling.py b/tests/test_sampling.py index 2ecce88..f5d0324 100644 --- a/tests/test_sampling.py +++ b/tests/test_sampling.py @@ -1,6 +1,6 @@ import numpy as np -from pyframe.media import Frame +from pyframe.media import Frame, FrameMeta from pyframe.sampling import ( DenseUniformSampler, MotionBucketSampler, @@ -30,8 +30,21 @@ def test_motion_bucket_returns_all_when_under_budget(): def test_dense_sampler_respects_target_fps(): frames = _frames([0] * 20) # 20 frames over ~1.9s -> ~10 fps source selected = DenseUniformSampler(target_fps=2.0).select(frames) - assert len(selected) == 4 # stride 5 - assert selected[0].index == 0 + assert [f.index for f in selected] == [0, 5, 10, 15, 19] # stride 5, plus the tail + + +def test_dense_sampler_always_covers_the_tail(): + # frames[::stride] stops at 15 here, leaving 16-19 outside the recall floor. An NSFW + # event that runs to the end of the clip would have no sampled frame to catch it. + frames = _frames([0] * 20) + selected = DenseUniformSampler(target_fps=2.0).select(frames) + assert selected[-1].index == 19 + + +def test_dense_sampler_does_not_duplicate_an_aligned_tail(): + frames = _frames([0] * 21) # 21 frames, stride 5 -> 0,5,10,15,20 already ends on 20 + selected = DenseUniformSampler(target_fps=2.0).select(frames) + assert [f.index for f in selected] == [0, 5, 10, 15, 20] def test_dense_sampler_keeps_all_when_no_duration(): @@ -58,3 +71,30 @@ def test_group_windows_merges_overlap_after_padding(): def test_group_windows_empty(): assert group_flagged_into_windows([], n_frames=10, gap=2, pad=1) == [] + + +def _metas(motions, fps=10.0): + return [ + FrameMeta(index=i, timestamp=i / fps, motion_score=m) for i, m in enumerate(motions) + ] + + +def test_samplers_pick_the_same_frames_from_metadata_as_from_frames(): + # Sampling runs on metadata now and pixels are fetched afterwards, so the two must + # agree exactly or the scan would moderate a different frame than the one chosen. + motions = [(i * 37) % 100 for i in range(40)] + frames, metas = _frames(motions), _metas(motions) + scores = {i: (i % 7) / 10.0 for i in range(40)} + + assert ( + [f.index for f in DenseUniformSampler(2.0).select(frames)] + == [m.index for m in DenseUniformSampler(2.0).select(metas)] + ) + assert ( + [f.index for f in MotionBucketSampler().select(frames, 6)] + == [m.index for m in MotionBucketSampler().select(metas, 6)] + ) + assert ( + [f.index for f in SuspicionSampler().select(frames, 5, scores)] + == [m.index for m in SuspicionSampler().select(metas, 5, scores)] + ) diff --git a/tests/test_scanner.py b/tests/test_scanner.py index ee60945..0d1e7cd 100644 --- a/tests/test_scanner.py +++ b/tests/test_scanner.py @@ -1,9 +1,11 @@ import time import numpy as np +import pytest from pyframe.backends.base import Backend from pyframe.config import Config, PrescreenConfig +from pyframe.errors import MediaDecodeError from pyframe.media import Frame, MediaKind from pyframe.results import Severity from pyframe.scanner import Scanner @@ -33,10 +35,18 @@ def _scanner(precise, screen=None, **prescreen): return Scanner(precise, screen=screen, config=cfg) +def _fetch(frames): + """Stand in for the real two-pass decode: the scanner selects against metadata and + then asks for pixels. Frame satisfies FrameLike, so these tests pass the same list + as both. Lazy on purpose, to catch anything that indexes or len()s the fetch.""" + by_index = {f.index: f for f in frames} + return lambda selected: (by_index[m.index] for m in selected) + + def test_single_pass_flags_bright_frame(): frames = _frames([10] * 9 + [250]) scanner = _scanner(FakeBackend(cost=0.001)) - result = scanner._single_pass("clip.gif", MediaKind.ANIMATION, frames, time.perf_counter()) + result = scanner._single_pass("clip.gif", MediaKind.ANIMATION, frames, _fetch(frames), time.perf_counter()) assert result.is_nsfw assert result.verdict is Severity.NSFW assert result.cost_usd > 0 @@ -45,7 +55,7 @@ def test_single_pass_flags_bright_frame(): def test_cascade_short_circuits_clean_media(): frames = _frames([10] * 20) scanner = _scanner(FakeBackend("aws", cost=0.001), screen=FakeBackend("local"), enabled=True) - result = scanner._cascade("clip.gif", MediaKind.ANIMATION, frames, time.perf_counter()) + result = scanner._cascade("clip.gif", MediaKind.ANIMATION, frames, _fetch(frames), time.perf_counter()) assert not result.is_nsfw assert result.escalated is False assert result.frames_classified == 0 @@ -56,7 +66,7 @@ def test_cascade_escalates_top_suspicious_as_merged(): frames = _frames([10] * 20) frames[12].image[:] = 250 # one suspicious frame scanner = _scanner(FakeBackend("aws", cost=0.001), screen=FakeBackend("local"), enabled=True) - result = scanner._cascade("clip.gif", MediaKind.ANIMATION, frames, time.perf_counter()) + result = scanner._cascade("clip.gif", MediaKind.ANIMATION, frames, _fetch(frames), time.perf_counter()) assert result.is_nsfw assert result.escalated is True assert 0 < result.frames_classified <= 2 # merged grids, capped @@ -66,7 +76,7 @@ def test_cascade_escalates_top_suspicious_as_merged(): def test_cascade_caps_aws_calls_at_max_escalations(): frames = _frames([250] * 40) # every frame flags scanner = _scanner(FakeBackend("aws", 0.001), screen=FakeBackend("local"), enabled=True, max_escalations=2) - result = scanner._cascade("c.gif", MediaKind.ANIMATION, frames, time.perf_counter()) + result = scanner._cascade("c.gif", MediaKind.ANIMATION, frames, _fetch(frames), time.perf_counter()) assert result.escalated is True assert result.frames_classified <= 2 # hard cap, regardless of how many frames flag @@ -75,7 +85,7 @@ def test_cascade_pads_to_full_grid_when_one_frame_flagged(): frames = _frames([10] * 20) frames[5].image[:] = 250 # only one suspicious frame scanner = _scanner(FakeBackend("aws", 0.001), screen=FakeBackend("local"), enabled=True, max_escalations=1) - result = scanner._cascade("c.gif", MediaKind.ANIMATION, frames, time.perf_counter()) + result = scanner._cascade("c.gif", MediaKind.ANIMATION, frames, _fetch(frames), time.perf_counter()) assert result.escalated is True assert result.frames_classified == 1 # one merged grid (the flagged frame plus a neighbor) assert result.is_nsfw @@ -114,10 +124,92 @@ def _score(self, image): frames = _frames([10] * 12) scanner = _scanner(FakeBackend("aws", cost=0.001), screen=BrokenScreen(), enabled=True, fail_open=True) - result = scanner._cascade("clip.gif", MediaKind.ANIMATION, frames, time.perf_counter()) + result = scanner._cascade("clip.gif", MediaKind.ANIMATION, frames, _fetch(frames), time.perf_counter()) assert result.frames_classified > 0 # errors were escalated, not silently cleared +def test_all_frames_failing_reports_error_not_clean(): + class BrokenPrecise(Backend): + name = "aws" + cost_per_image = 0.001 + + def _score(self, image): + raise RuntimeError("credentials expired") + + frames = _frames([10] * 6) + scanner = _scanner(BrokenPrecise()) + result = scanner._single_pass("clip.gif", MediaKind.ANIMATION, frames, _fetch(frames), time.perf_counter()) + + assert result.verdict is Severity.ERROR + assert result.is_nsfw is False # an error is not a positive finding... + assert result.max_score == 0.0 # ...but the CLI must not read it as a clean bill either + + +def test_short_circuited_cascade_keeps_is_nsfw_and_verdict_in_lockstep(): + # escalate_threshold above min_confidence: the screen scores high enough to be NSFW + # but not high enough to escalate, so nothing is ever classified. is_nsfw used to be + # computed only from the classified frames, so it read False while verdict read nsfw + # -- and the CLI gates on is_nsfw. + frames = _frames([220] * 10) # 220/255 = 0.86, over min_confidence 0.8 + scanner = _scanner( + FakeBackend("aws", 0.001), screen=FakeBackend("local"), + enabled=True, escalate_threshold=0.95, + ) + result = scanner._cascade("clip.gif", MediaKind.ANIMATION, frames, _fetch(frames), time.perf_counter()) + + assert result.escalated is False + assert result.frames_classified == 0 + assert result.verdict is Severity.NSFW + assert result.is_nsfw is True + + +def test_media_that_decodes_to_nothing_raises_rather_than_reporting_clean(): + scanner = _scanner(FakeBackend()) + + with pytest.raises(MediaDecodeError): + scanner._scan_frames("clip.gif", MediaKind.ANIMATION, [], _fetch([]), time.perf_counter()) + + +def test_ensure_min_frames_fills_by_suspicion_not_motion(): + # Screen scores are 0..1 while motion_score is a pixel-diff sum reaching ~1e6, so a + # single flat sort key ranked any moving frame above a screened frame that had + # nearly flagged. + frames = _frames([10] * 3) + frames[1].motion_score = 0.0 # screened, scored just under the gate + frames[2].motion_score = 1_000_000.0 # never screened, merely busy + scores = {0: 0.9, 1: 0.4} + + selected = _scanner(FakeBackend())._ensure_min_frames([frames[0]], frames, scores, 2) + + assert [f.index for f in selected] == [0, 1] + + +def test_max_escalations_below_one_is_rejected(): + # A non-positive budget does not disable escalation, it uncaps it: SuspicionSampler + # returns every frame when budget <= 0. + cfg = Config(backend=FakeBackend(), prescreen=PrescreenConfig(enabled=True, max_escalations=0)) + + with pytest.raises(ValueError, match="max_escalations"): + Scanner.from_config(cfg) + + +def test_max_frames_below_one_is_rejected(): + # Same trap as max_escalations on the single pass side: the samplers read a + # non-positive budget as "keep everything", which would materialise the whole clip. + with pytest.raises(ValueError, match="max_frames"): + Scanner.from_config(Config(backend=FakeBackend(), max_frames=0)) + + +def test_max_escalations_of_one_still_caps_at_one_call(): + frames = _frames([250] * 40) # every frame flags + scanner = _scanner( + FakeBackend("aws", 0.001), screen=FakeBackend("local"), enabled=True, max_escalations=1 + ) + result = scanner._cascade("c.gif", MediaKind.ANIMATION, frames, _fetch(frames), time.perf_counter()) + + assert result.frames_classified == 1 + + def test_motion_sampler_always_includes_time_coverage_floor(): from pyframe.sampling import DenseUniformSampler diff --git a/tests/test_smoke.py b/tests/test_smoke.py index e7d5558..aa1f2e1 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -19,6 +19,34 @@ def test_public_api_surface(): assert hasattr(pyframe, name), name +def test_frame_still_constructs_positionally(): + # Frame is exported and its field order is public, so the metadata split must not + # have reordered or defaulted its way into a breaking constructor. + import numpy as np + + import pyframe + + frame = pyframe.Frame(3, 1.5, np.zeros((2, 2, 3), np.uint8)) + + assert (frame.index, frame.timestamp, frame.motion_score) == (3, 1.5, 0.0) + assert frame.image.shape == (2, 2, 3) + + +def test_two_pass_decode_names_are_exported(): + import pyframe + + for name in ( + "FrameMeta", + "FrameLike", + "iter_frame_meta", + "iter_frames_at", + "iter_frame_meta_from_bytes", + "iter_frames_from_bytes_at", + ): + assert hasattr(pyframe, name), name + assert name in pyframe.__all__, name + + def test_unsupported_media_raises(): import pytest