openmoderate 1.0.1__py3-none-any.whl

openmoderate/__init__.py

@@ -0,0 +1,43 @@
+"""openmoderate — open-source image and video content moderation.
+
+Quick start:
+
+```python
+from openmoderate import Moderator
+
+result = Moderator().analyze("upload.jpg")
+if not result.is_safe:
+    print(result.violations)
+```
+
+Everything importable from this top-level package is public API and follows semantic
+versioning. Submodules that are not re-exported here (and names starting with ``_``) may
+change in minor releases.
+"""
+
+from openmoderate.categories import DEFAULT_CATEGORIES, SAFE_PROMPTS, Category
+from openmoderate.config import ModerationConfig
+from openmoderate.io import ImageLoadError
+from openmoderate.moderator import LabelSpec, Moderator
+from openmoderate.result import Label, ModerationResult
+from openmoderate.training.head import LinearHead
+from openmoderate.video import VideoLoadError, VideoModerationResult, VideoSegment
+
+__version__ = "1.0.1"
+
+__all__ = [
+    "DEFAULT_CATEGORIES",
+    "SAFE_PROMPTS",
+    "Category",
+    "ImageLoadError",
+    "Label",
+    "LabelSpec",
+    "LinearHead",
+    "ModerationConfig",
+    "ModerationResult",
+    "Moderator",
+    "VideoLoadError",
+    "VideoModerationResult",
+    "VideoSegment",
+    "__version__",
+]

openmoderate/categories.py

@@ -0,0 +1,203 @@
+"""Default moderation taxonomy.
+
+Each category is described by natural-language prompts. The zero-shot engine compares an
+image against every category's prompts *and* a pool of neutral "safe" prompts, so a
+category only scores high when it explains the image better than ordinary content does.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Category:
+    name: str
+    prompts: tuple[str, ...]
+    threshold: float = 0.5
+    parent: str | None = None
+    description: str = ""
+    # Category-specific contrast prompts: benign look-alikes that should *not* count
+    # for this category but may legitimately count for another (e.g. cigarettes are
+    # tobacco, not drugs).
+    negatives: tuple[str, ...] = ()
+
+
+SAFE_PROMPTS: tuple[str, ...] = (
+    "a photo of a person wearing normal clothes",
+    "a photo of people at work or in a meeting",
+    "a portrait photo of a smiling person",
+    "a photo of a landscape or nature scene",
+    "a photo of a city street",
+    "a photo of food on a plate",
+    "a photo of an animal or pet",
+    "a photo of a building or room interior",
+    "a photo of everyday objects on a table",
+    "a screenshot of a document or website",
+    "a product photo on a white background",
+    "a photo of people playing sports",
+    "a family photo",
+    "a cartoon or illustration for children",
+    "a person at the beach wearing a swimsuit",
+    # Look-alikes: benign content that resembles a violating category. Without these the
+    # engine confuses national flags with hate symbols, kitchen knives with weapons, etc.
+    "the national flag of a country",
+    "a company logo or brand icon",
+    "abstract shapes and a simple geometric graphic",
+    "clipart, an icon or an emoji",
+    "kitchen knives and cooking utensils",
+    "a medical photo of prescription medicine",
+    "a doctor or nurse in a hospital",
+    "a soldier or police officer in uniform at a parade",
+    "people dancing at a party or concert",
+    "a red painting or red colored object",
+    "a glass of milk",
+    "a glass of orange juice or soft drink",
+    "a cup of coffee or tea",
+)
+
+DEFAULT_CATEGORIES: tuple[Category, ...] = (
+    Category(
+        "nudity",
+        (
+            "a photo of a naked person",
+            "a nude photo showing genitals or bare breasts",
+            "explicit nudity",
+        ),
+        threshold=0.5,
+        parent="explicit",
+        description="Exposed genitals, buttocks or female nipples.",
+    ),
+    Category(
+        "sexual_activity",
+        (
+            "a pornographic image of people having sex",
+            "an explicit sexual act",
+            "hardcore pornography",
+        ),
+        threshold=0.5,
+        parent="explicit",
+        description="Depictions of sexual acts.",
+    ),
+    Category(
+        "suggestive",
+        (
+            "a sexually suggestive photo of a person in lingerie",
+            "an erotic provocative pose",
+        ),
+        threshold=0.6,
+        parent="suggestive",
+        description="Revealing clothing or sexually suggestive poses without explicit nudity.",
+    ),
+    Category(
+        "violence",
+        (
+            "a photo of people fighting violently",
+            "a person being beaten or attacked",
+            "a violent assault",
+        ),
+        threshold=0.5,
+        parent="violence",
+        description="Physical violence between people.",
+        negatives=("a person holding a gun at a shooting range", "a sports match"),
+    ),
+    Category(
+        "gore",
+        (
+            "a graphic photo of blood and gore",
+            "a mutilated dead body",
+            "a gruesome bloody wound",
+        ),
+        threshold=0.45,
+        parent="violence",
+        description="Graphic blood, injuries, corpses.",
+    ),
+    Category(
+        "weapons",
+        (
+            "a photo of a gun",
+            "a person holding a firearm or rifle",
+            "a photo of a knife used as a weapon",
+            "a revolver or pistol",
+        ),
+        threshold=0.5,
+        parent="violence",
+        description="Firearms, knives and other weapons.",
+        negatives=("a kitchen knife for cooking food",),
+    ),
+    Category(
+        "hate_symbols",
+        (
+            "a nazi swastika flag",
+            "a hate symbol or extremist flag",
+            "a ku klux klan hood and robe",
+            "a nazi swastika armband or banner",
+            "people wearing ku klux klan robes and pointed hoods",
+            "white supremacist propaganda",
+        ),
+        threshold=0.5,
+        parent="hate",
+        description="Symbols and imagery associated with hate groups.",
+    ),
+    Category(
+        "drugs",
+        (
+            "a photo of illegal drugs",
+            "cocaine lines on a table",
+            "a person injecting drugs with a needle",
+            "pills and drug paraphernalia",
+            "a photo of marijuana buds",
+            "a rolled cannabis joint",
+            "a glass bong for smoking marijuana",
+            "a bag of white powder drugs",
+            "a person smoking marijuana",
+        ),
+        threshold=0.5,
+        parent="drugs",
+        description="Illegal drugs and drug use.",
+        negatives=("a person smoking a cigarette", "a vaccination or medical injection"),
+    ),
+    Category(
+        "self_harm",
+        (
+            "a photo of self harm cuts on a wrist",
+            "a person attempting suicide",
+        ),
+        threshold=0.45,
+        parent="self_harm",
+        description="Self-injury or suicide imagery.",
+    ),
+    Category(
+        "alcohol",
+        (
+            "a photo of alcoholic drinks and beer bottles",
+            "a person drinking alcohol",
+            "a glass of beer",
+            "a glass of red wine",
+            "a cocktail drink in a glass",
+            "a bottle of whisky or liquor",
+        ),
+        threshold=0.6,
+        parent="substances",
+        description="Alcoholic beverages (often allowed; off by default via high threshold).",
+    ),
+    Category(
+        "tobacco",
+        (
+            "a person smoking a cigarette",
+            "a photo of cigarettes or vaping",
+            "a cigar",
+            "a pack of cigarettes",
+        ),
+        threshold=0.6,
+        parent="substances",
+        description="Tobacco and vaping products.",
+    ),
+)
+
+
+def get_category(name: str, categories: tuple[Category, ...] = DEFAULT_CATEGORIES) -> Category:
+    for category in categories:
+        if category.name == name:
+            return category
+    raise KeyError(f"Unknown category: {name!r}")

openmoderate/cli.py

@@ -0,0 +1,390 @@
+"""``openmoderate`` command-line interface.
+
+openmoderate scan ./uploads [--json] [--fail-on-violation]
+openmoderate train ./dataset -o my_head.omh
+openmoderate evaluate ./dataset [--head my_head.omh]
+openmoderate categories
+openmoderate download
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections.abc import Sequence
+from pathlib import Path
+
+from openmoderate import __version__
+
+
+def _parse_thresholds(items: list[str] | None) -> dict[str, float]:
+    out = {}
+    for item in items or []:
+        name, sep, value = item.partition("=")
+        if not sep:
+            raise SystemExit(f"--threshold expects NAME=VALUE, got {item!r}")
+        out[name.strip()] = float(value)
+    return out
+
+
+def _add_policy_args(p: argparse.ArgumentParser) -> None:
+    p.add_argument("--config", help="Policy YAML file")
+    p.add_argument("--head", help="Trained head (.omh) from `openmoderate train`")
+    p.add_argument(
+        "--head-weight", type=float, default=1.0, help="Blend for labels in both (default 1.0)"
+    )
+    p.add_argument(
+        "--threshold", "-t", action="append", metavar="NAME=VALUE", help="Override a threshold"
+    )
+    p.add_argument("--only", nargs="+", metavar="CATEGORY", help="Only check these categories")
+    p.add_argument("--exclude", nargs="+", metavar="CATEGORY", help="Skip these categories")
+    p.add_argument("--device", default="auto", help="auto, cpu, cuda or mps")
+
+
+def _build_moderator(args: argparse.Namespace):
+    from openmoderate import Moderator
+
+    return Moderator(
+        args.config,
+        thresholds=_parse_thresholds(args.threshold),
+        categories=args.only,
+        exclude=args.exclude,
+        head=args.head,
+        head_weight=args.head_weight,
+        device=args.device,
+    )
+
+
+def cmd_scan(args: argparse.Namespace) -> int:
+    from openmoderate.io import ImageLoadError, iter_image_files
+    from openmoderate.video import is_video_path
+
+    sources: list[str] = []
+    videos: list[str] = []
+    for target in args.paths:
+        if target.startswith(("http://", "https://")):
+            sources.append(target)
+            continue
+        path = Path(target)
+        if path.is_file():
+            (videos if is_video_path(path) else sources).append(str(path))
+            continue
+        found_images = [str(f) for f in iter_image_files(path)]
+        found_videos = sorted(str(f) for f in path.rglob("*") if is_video_path(f))
+        if not found_images and not found_videos:
+            print(f"warning: no images or videos found at {target}", file=sys.stderr)
+        sources.extend(found_images)
+        videos.extend(found_videos)
+
+    mod = _build_moderator(args)
+    flagged = errors = 0
+    for video in videos:
+        try:
+            res = mod.analyze_video(
+                video, fps=args.fps, max_duration=args.max_duration, min_frames=args.min_frames
+            )
+        except (ImportError, ValueError) as exc:
+            errors += 1
+            _emit_error(args, video, exc)
+            continue
+        flagged += not res.is_safe
+        _emit_video_result(args, res)
+    for i in range(0, len(sources), args.batch_size):
+        chunk = sources[i : i + args.batch_size]
+        try:
+            results = mod.analyze_batch(chunk)
+        except ImageLoadError:
+            # Fall back to one-by-one so a single bad file doesn't sink the batch.
+            results = []
+            for src in chunk:
+                try:
+                    results.append(mod.analyze(src))
+                except ImageLoadError as exc:
+                    errors += 1
+                    _emit_error(args, src, exc)
+        for res in results:
+            flagged += not res.is_safe
+            _emit_result(args, res)
+
+    if not args.json:
+        total = len(sources) + len(videos)
+        print(f"\n{total} scanned, {flagged} flagged, {errors} unreadable", file=sys.stderr)
+    if args.fail_on_violation and flagged:
+        return 1
+    return 2 if errors else 0
+
+
+def _emit_result(args: argparse.Namespace, res) -> None:
+    if args.json:
+        print(json.dumps(res.to_dict()))
+        return
+    verdict = "SAFE   " if res.is_safe else "FLAGGED"
+    details = ", ".join(f"{v.name} {v.confidence:.0%}" for v in res.violations)
+    print(f"{verdict}  {res.source}  {details}".rstrip())
+
+
+def _emit_video_result(args: argparse.Namespace, res) -> None:
+    if args.json:
+        print(json.dumps(res.to_dict()))
+        return
+    verdict = "SAFE   " if res.is_safe else "FLAGGED"
+    note = " (truncated)" if res.truncated else ""
+    print(
+        f"{verdict}  {res.source}  [video {res.duration:.1f}s, {res.frames_analyzed} frames]{note}"
+    )
+    for seg in res.segments:
+        print(
+            f"           {_clock(seg.start)}-{_clock(seg.end)}  {seg.name} {seg.max_confidence:.0%}"
+        )
+
+
+def _clock(seconds: float) -> str:
+    minutes, secs = divmod(seconds, 60)
+    return f"{int(minutes)}:{secs:04.1f}"
+
+
+def _emit_error(args: argparse.Namespace, src: str, exc: Exception) -> None:
+    if args.json:
+        print(json.dumps({"source": src, "error": str(exc)}))
+    else:
+        print(f"ERROR    {src}  {exc}", file=sys.stderr)
+
+
+def cmd_train(args: argparse.Namespace) -> int:
+    from openmoderate.config import ModerationConfig
+    from openmoderate.engine.clip_engine import ClipEngine
+    from openmoderate.training import FolderDataset, TrainConfig, train_head
+
+    cfg = ModerationConfig.from_yaml(args.config) if args.config else ModerationConfig()
+    dataset = FolderDataset.from_folder(args.dataset, safe_label=args.safe_label)
+    print(f"Dataset: {len(dataset)} images", file=sys.stderr)
+    summary = dataset.summary()
+    for label, count in summary.items():
+        print(f"  {label:20} {count}", file=sys.stderr)
+    n_safe = summary["(no label)"]
+    if n_safe < max(50, max(summary[label] for label in dataset.labels)):
+        print(
+            f"\nwarning: only {n_safe} images in '{args.safe_label}/'. The head learns what is "
+            "NOT a violation from this folder, so fill it with a broad sample of your normal "
+            "traffic (people, products, scenes...). A small or narrow safe set causes false "
+            "positives that validation metrics will not reveal.",
+            file=sys.stderr,
+        )
+
+    engine = ClipEngine(cfg.model_name, cfg.pretrained, device=args.device, cache_dir=cfg.cache_dir)
+    embeddings = dataset.embed(engine, verbose=True, trim=cfg.trim_borders)
+    parents = {c.name: c.parent for c in cfg.categories if c.name in dataset.labels}
+    report = train_head(
+        embeddings,
+        dataset.targets,
+        dataset.labels,
+        TrainConfig(
+            epochs=args.epochs,
+            learning_rate=args.lr,
+            val_split=args.val_split,
+            threshold=args.threshold,
+        ),
+        model_name=cfg.model_name,
+        pretrained=cfg.pretrained,
+        parents=parents,
+    )
+    if report.val_metrics:
+        print(f"\nValidation ({report.n_val} held-out images):", file=sys.stderr)
+        _print_metrics(report.val_metrics)
+    else:
+        print("\n(Not enough images for a validation split; no metrics.)", file=sys.stderr)
+    report.head.save(args.output)
+    print(f"\nSaved head to {args.output} (trained on all {len(dataset)} images)", file=sys.stderr)
+    print(f'Use it with: Moderator(head="{args.output}")', file=sys.stderr)
+    return 0
+
+
+def _print_metrics(metrics: dict[str, dict[str, float]]) -> None:
+    print(f"  {'label':20} {'precision':>9} {'recall':>7} {'f1':>6} {'n':>5}", file=sys.stderr)
+    for label, m in metrics.items():
+        print(
+            f"  {label:20} {m['precision']:9.2f} {m['recall']:7.2f} {m['f1']:6.2f} "
+            f"{m['support']:5d}",
+            file=sys.stderr,
+        )
+
+
+def cmd_evaluate(args: argparse.Namespace) -> int:
+    import numpy as np
+
+    from openmoderate.training import FolderDataset, compute_metrics
+
+    dataset = FolderDataset.from_folder(args.dataset, safe_label=args.safe_label)
+    mod = _build_moderator(args)
+    names = mod.category_names
+    missing = [label for label in dataset.labels if label not in names]
+    if missing:
+        print(
+            f"warning: dataset labels not checked by this moderator, skipped: {missing}",
+            file=sys.stderr,
+        )
+    labels = [label for label in dataset.labels if label in names]
+    if not labels:
+        raise SystemExit("No dataset label matches a moderator category; nothing to evaluate.")
+
+    embeddings = dataset.embed(mod.engine, verbose=not args.json, trim=mod.config.trim_borders)
+    scores = mod.score_all(embeddings)
+    cols = [names.index(label) for label in labels]
+    targets = dataset.targets[:, [dataset.labels.index(label) for label in labels]]
+    all_thresholds = mod.thresholds
+    thresholds = [all_thresholds[label] for label in labels]
+    metrics = compute_metrics(targets, scores[:, cols], labels, thresholds)
+
+    # Overall: is an image flagged at all vs. does it carry any violation label.
+    flagged_any = np.any(scores >= np.array([all_thresholds[n] for n in names]), axis=1)
+    overall = compute_metrics(
+        (dataset.targets.sum(axis=1, keepdims=True) > 0).astype(float),
+        flagged_any[:, None].astype(float),
+        ["any_violation"],
+        [0.5],
+    )
+    if args.json:
+        print(json.dumps({"labels": metrics, **overall}, indent=2))
+    else:
+        print(f"\nEvaluated {len(dataset)} images:", file=sys.stderr)
+        _print_metrics({**metrics, **overall})
+    return 0
+
+
+def cmd_categories(args: argparse.Namespace) -> int:
+    from openmoderate import Moderator
+
+    class _NoEngine:  # categories don't need the model
+        pass
+
+    mod = Moderator(args.config, head=args.head, engine=_NoEngine())  # type: ignore[arg-type]
+    zero_shot = {c.name for c in mod.config.categories}
+    trained = set(mod.head.labels) if mod.head else set()
+    for spec in mod.label_specs:
+        kind = "+".join(
+            k for k, s in (("zero-shot", zero_shot), ("trained", trained)) if spec.name in s
+        )
+        print(f"{spec.name:20} {spec.parent or '-':12} {spec.threshold:5.2f}  {kind}")
+    return 0
+
+
+def cmd_download(args: argparse.Namespace) -> int:
+    from openmoderate.config import ModerationConfig
+    from openmoderate.engine.clip_engine import ClipEngine
+
+    cfg = ModerationConfig.from_yaml(args.config) if args.config else ModerationConfig()
+    print(f"Downloading {cfg.model_name} ({cfg.pretrained})...", file=sys.stderr)
+    ClipEngine(cfg.model_name, cfg.pretrained, device="cpu", cache_dir=cfg.cache_dir).encode_texts(
+        ["warm up"]
+    )
+    print("Done. The model is cached and openmoderate can now run offline.", file=sys.stderr)
+    return 0
+
+
+def cmd_serve(args: argparse.Namespace) -> int:
+    import os
+
+    try:
+        import uvicorn
+    except ImportError:
+        raise SystemExit(
+            "`serve` needs extra dependencies: pip install 'openmoderate[server]'"
+        ) from None
+
+    # The app reads its settings from the environment so multi-worker mode works too.
+    for env, value in (
+        ("OPENMODERATE_CONFIG", args.config),
+        ("OPENMODERATE_HEAD", args.head),
+        ("OPENMODERATE_DEVICE", args.device),
+    ):
+        if value:
+            os.environ[env] = value
+    if args.allow_urls:
+        os.environ["OPENMODERATE_ALLOW_URLS"] = "true"
+    if args.host not in ("127.0.0.1", "localhost") and not os.getenv("OPENMODERATE_API_KEY"):
+        print(
+            f"warning: serving on {args.host} without OPENMODERATE_API_KEY; anyone who can "
+            "reach this port can use the API.",
+            file=sys.stderr,
+        )
+    uvicorn.run(
+        "openmoderate.server.app:app",
+        host=args.host,
+        port=args.port,
+        workers=args.workers,
+        log_level=args.log_level,
+    )
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="openmoderate", description="Open-source image content moderation."
+    )
+    parser.add_argument("--version", action="version", version=f"openmoderate {__version__}")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p = sub.add_parser("scan", help="Scan images, folders or URLs")
+    p.add_argument("paths", nargs="+", help="Image files, folders (recursive) or URLs")
+    p.add_argument("--json", action="store_true", help="Output one JSON object per line")
+    p.add_argument(
+        "--fail-on-violation", action="store_true", help="Exit with code 1 if anything is flagged"
+    )
+    p.add_argument("--batch-size", type=int, default=16)
+    p.add_argument("--fps", type=float, default=1.0, help="Video frames sampled per second")
+    p.add_argument("--max-duration", type=float, help="Only analyse the first N seconds of videos")
+    p.add_argument(
+        "--min-frames", type=int, default=1, help="Ignore video segments shorter than N frames"
+    )
+    _add_policy_args(p)
+    p.set_defaults(func=cmd_scan)
+
+    p = sub.add_parser("train", help="Train a head on a folder dataset (dataset/<label>/*.jpg)")
+    p.add_argument("dataset", help="Folder with one sub-folder per label, plus safe/")
+    p.add_argument("-o", "--output", default="head.omh", help="Output file (default head.omh)")
+    p.add_argument("--safe-label", default="safe", help="Folder name for safe images")
+    p.add_argument("--epochs", type=int, default=300)
+    p.add_argument("--lr", type=float, default=0.05)
+    p.add_argument("--val-split", type=float, default=0.2)
+    p.add_argument("--threshold", type=float, default=0.5, help="Decision threshold to store")
+    p.add_argument("--config", help="Policy YAML (for model choice and category parents)")
+    p.add_argument("--device", default="auto")
+    p.set_defaults(func=cmd_train)
+
+    p = sub.add_parser("evaluate", help="Measure precision/recall on a labelled folder dataset")
+    p.add_argument("dataset")
+    p.add_argument("--safe-label", default="safe")
+    p.add_argument("--json", action="store_true")
+    _add_policy_args(p)
+    p.set_defaults(func=cmd_evaluate)
+
+    p = sub.add_parser("categories", help="List categories and thresholds")
+    p.add_argument("--config")
+    p.add_argument("--head")
+    p.set_defaults(func=cmd_categories)
+
+    p = sub.add_parser("download", help="Pre-download model weights (for offline/Docker use)")
+    p.add_argument("--config")
+    p.set_defaults(func=cmd_download)
+
+    p = sub.add_parser("serve", help="Run the REST API server")
+    p.add_argument("--host", default="127.0.0.1", help="Bind address (0.0.0.0 for all)")
+    p.add_argument("--port", type=int, default=8000)
+    p.add_argument("--workers", type=int, default=1, help="Processes (each loads the model)")
+    p.add_argument("--config", help="Policy YAML file")
+    p.add_argument("--head", help="Trained head (.omh)")
+    p.add_argument("--device", help="auto, cpu, cuda or mps")
+    p.add_argument("--allow-urls", action="store_true", help="Allow {'url': ...} requests")
+    p.add_argument("--log-level", default="info")
+    p.set_defaults(func=cmd_serve)
+    return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())