occlude 1.0.0__py3-none-any.whl

occlude/__init__.py

@@ -0,0 +1,3 @@
+"""OCCLUDE — blur immodestly dressed people in videos."""
+
+__version__ = "1.0.0"

occlude/__main__.py

@@ -0,0 +1,5 @@
+"""Allow `python -m occlude` invocation."""
+from occlude.cli import main
+
+if __name__ == "__main__":
+    main()

occlude/cli.py

@@ -0,0 +1,195 @@
+"""OCCLUDE — CLI entry point.
+
+Usage:
+    occlude --input <video> [--output <path>] [--blur-strength N]
+
+Detects immodestly dressed people frame-by-frame and writes a clean
+video with the original audio preserved. See OCCLUDE_SPEC.md.
+"""
+from __future__ import annotations
+
+import os
+
+# Diagnostic instrumentation: traps os._exit + sys.exit + fatal signals and
+# dumps all-thread stacks to /tmp so a silent death pinpoints its origin.
+# Activated only when OCCLUDE_DEBUG=1 to keep `--help` side-effect-free.
+if os.getenv("OCCLUDE_DEBUG") == "1":
+    import sys as _diag_sys
+    import time as _diag_time
+    import traceback as _diag_tb
+    import threading as _diag_threading
+    import faulthandler as _diag_faulthandler
+    import signal as _diag_signal
+    import atexit as _diag_atexit
+
+    _DIAG_LOG = "/tmp/occlude_diag.log"
+    _FAULT_LOG = "/tmp/occlude_faulthandler.log"
+
+    def _diag_log(msg: str) -> None:
+        with open(_DIAG_LOG, "a") as f:
+            f.write(f"[{_diag_time.strftime('%H:%M:%S')}] {msg}\n")
+            f.flush()
+
+    _real_os_exit = os._exit
+
+    def _traced_os_exit(code):  # noqa: ANN001
+        with open(_DIAG_LOG, "a") as f:
+            f.write(
+                f"\n========== os._exit({code}) at "
+                f"{_diag_time.strftime('%H:%M:%S')} ==========\n"
+            )
+            f.write(f"main thread: {_diag_threading.main_thread().ident}\n")
+            f.write(
+                f"current thread: {_diag_threading.current_thread().ident} "
+                f"({_diag_threading.current_thread().name})\n"
+            )
+            f.write("--- python stack at exit ---\n")
+            _diag_tb.print_stack(file=f)
+            f.write("--- all threads ---\n")
+            for tid, frame in _diag_sys._current_frames().items():
+                f.write(f"\n>>> thread {tid}\n")
+                _diag_tb.print_stack(frame, file=f)
+            f.flush()
+        _real_os_exit(code)
+
+    os._exit = _traced_os_exit
+
+    _real_sys_exit = _diag_sys.exit
+
+    def _traced_sys_exit(code=0):  # noqa: ANN001
+        with open(_DIAG_LOG, "a") as f:
+            f.write(
+                f"[{_diag_time.strftime('%H:%M:%S')}] sys.exit({code}) "
+                "called from:\n"
+            )
+            _diag_tb.print_stack(file=f)
+            f.flush()
+        _real_sys_exit(code)
+
+    _diag_sys.exit = _traced_sys_exit
+
+    _diag_fault_fh = open(_FAULT_LOG, "w")
+    _diag_faulthandler.enable(file=_diag_fault_fh, all_threads=True)
+    for _diag_sig in (
+        _diag_signal.SIGTERM,
+        _diag_signal.SIGINT,
+        _diag_signal.SIGHUP,
+        _diag_signal.SIGUSR1,
+        _diag_signal.SIGUSR2,
+        _diag_signal.SIGPIPE,
+        _diag_signal.SIGQUIT,
+    ):
+        try:
+            _diag_faulthandler.register(
+                _diag_sig, file=_diag_fault_fh, all_threads=True, chain=True
+            )
+        except Exception as _e:  # noqa: BLE001
+            _diag_log(f"faulthandler.register({_diag_sig}) failed: {_e}")
+
+    def _diag_on_atexit() -> None:
+        _diag_log("atexit fired (orderly python shutdown)")
+
+    _diag_atexit.register(_diag_on_atexit)
+    _diag_log(
+        f"=== occlude started, pid={os.getpid()}, "
+        f"argv={_diag_sys.argv} ==="
+    )
+
+import argparse
+import shutil
+import sys
+from pathlib import Path
+
+from rich.console import Console
+
+
+def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        prog="occlude",
+        description="Blur immodestly dressed people in a video.",
+    )
+    parser.add_argument(
+        "--input", required=True, type=Path,
+        help="path to the input video file",
+    )
+    parser.add_argument(
+        "--output", type=Path, default=None,
+        help="output video path (default: <input_stem>_occluded.mp4 next to input)",
+    )
+    # Default sourced from video.DEFAULT_BLUR_KERNEL so a single
+    # constant controls the kernel size — historical 51 was too soft
+    # on 1280×720 footage (subjects still recognizable on multi-person
+    # video). Currently 199.
+    from occlude.pipeline.video import DEFAULT_BLUR_KERNEL
+    parser.add_argument(
+        "--blur-strength", type=int, default=DEFAULT_BLUR_KERNEL,
+        help=f"Gaussian blur kernel size (must be odd, default {DEFAULT_BLUR_KERNEL})",
+    )
+    parser.add_argument(
+        "--frame-stride", type=int, default=1,
+        help="run perception every Nth frame; skipped frames re-render the previous frame's blur (default 1 = every frame)",
+    )
+    return parser.parse_args(argv)
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _parse_args(argv)
+
+    from occlude.ui.ascii_art import get_header_panel
+    console = Console()
+    console.print(get_header_panel())
+
+    input_path: Path = args.input
+    if not input_path.exists() or not input_path.is_file():
+        print(f"error: input file not found: {input_path}", file=sys.stderr)
+        return 1
+
+    if args.blur_strength <= 0 or args.blur_strength % 2 == 0:
+        print(
+            f"error: --blur-strength must be a positive odd integer, got {args.blur_strength}",
+            file=sys.stderr,
+        )
+        return 1
+
+    if args.frame_stride < 1:
+        print(
+            f"error: --frame-stride must be >= 1, got {args.frame_stride}",
+            file=sys.stderr,
+        )
+        return 1
+
+    if shutil.which("ffmpeg") is None:
+        print(
+            "error: ffmpeg not found on PATH. Install via: brew install ffmpeg",
+            file=sys.stderr,
+        )
+        return 1
+
+    output_path: Path = args.output or (
+        input_path.parent / f"{input_path.stem}_occluded.mp4"
+    )
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Imported here so `--help` and arg validation don't pay the
+    # multi-second model-loading cost.
+    from occlude.pipeline.video import VideoProcessor
+
+    try:
+        processor = VideoProcessor(
+            blur_kernel=args.blur_strength,
+            frame_stride=args.frame_stride,
+        )
+        processor.process(input_path, output_path)
+    except KeyboardInterrupt:
+        print("\ncancelled by user.", file=sys.stderr)
+        return 130
+    except Exception as e:
+        print(f"error: {e}", file=sys.stderr)
+        return 1
+
+    print(f"done. output: {output_path}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

occlude/pipeline/__init__.py

@@ -0,0 +1,21 @@
+"""Pipeline package — perception, rules, blur, video.
+
+Stage 4 lives in :mod:`occlude.pipeline.perception` and produces the data
+structure the Stage 5 rule layer consumes.
+"""
+
+from occlude.pipeline.perception import (
+    SEG_LABELS,
+    TARGET_LABELS,
+    Perception,
+    Perceiver,
+    Person,
+)
+from occlude.pipeline.rules import Decision, RuleEngine
+from occlude.pipeline.video import VideoProcessor, Tracker, blur_region
+
+__all__ = [
+    "SEG_LABELS", "TARGET_LABELS", "Perception", "Perceiver", "Person",
+    "Decision", "RuleEngine",
+    "VideoProcessor", "Tracker", "blur_region",
+]

occlude/pipeline/perception.py

@@ -0,0 +1,301 @@
+"""Stage 4 — Perception.
+
+Wraps the YOLO person detector, the SegFormer body-part segmenter, and
+the InsightFace gender classifier as a single callable that maps an
+image to a list of per-person observations. This is the data structure
+the Stage 5 rule layer consumes.
+
+The class loads all three models once on construction; calling the
+instance on an image runs detection, then segments + classifies each
+person crop.
+"""
+import gc
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Generator, Protocol, runtime_checkable
+
+import cv2
+import numpy as np
+import onnxruntime
+import torch
+from insightface.app import FaceAnalysis
+from insightface.model_zoo import model_zoo as _imz
+from PIL import Image
+from transformers import AutoModelForSemanticSegmentation, SegformerImageProcessor
+from ultralytics import YOLO
+
+
+class _NoArenaPickableSession(_imz.PickableInferenceSession):
+    """InsightFace `PickableInferenceSession` with the ONNX Runtime CPU
+    memory arena disabled.
+
+    By default ORT keeps a per-session arena that retains freed buffers
+    for reuse. Empirically, on multi-person video, each
+    `face_app.get(bgr)` call leaks ~45 MB/person into compressed-VM
+    pages that macOS's Activity Monitor tallies into the Jetsam
+    footprint metric (vmmap MALLOC_LARGE virtual = 4.6 GB at frame 25
+    of laughing_people.mp4 with arena on, ~250 MB total with arena
+    off). Disabling the arena trades a small inference slowdown for
+    bounded steady-state memory.
+    """
+
+    def __init__(self, model_path, **kwargs):
+        if "sess_options" not in kwargs:
+            opts = onnxruntime.SessionOptions()
+            opts.enable_cpu_mem_arena = False
+            # mem_pattern caches per-input-shape allocation plans.
+            # InsightFace receives variable-sized BGR crops per
+            # detection call, so each unique shape grew the cache.
+            # Disabling forces fresh allocations and bounds the heap.
+            opts.enable_mem_pattern = False
+            kwargs["sess_options"] = opts
+        super().__init__(model_path, **kwargs)
+
+YOLO_MODEL_ID = "yolov8n.pt"
+SEG_MODEL_ID = "mattmdjaga/segformer_b2_clothes"
+INSIGHT_MODEL_ID = "buffalo_l"
+PERSON_CLASS_ID = 0
+DEFAULT_PERSON_CONF = 0.40
+INSIGHT_DET_SIZE = (640, 640)
+
+# 18 SegFormer classes from `mattmdjaga/segformer_b2_clothes`. Index
+# into this list maps to the integer label in the predicted mask.
+SEG_LABELS = [
+    "Background", "Hat", "Hair", "Sunglasses", "Upper-clothes",
+    "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe",
+    "Face", "Left-leg", "Right-leg", "Left-arm", "Right-arm",
+    "Bag", "Scarf",
+]
+
+# Subset the modesty rule layer cares about. Surfaced here (not buried
+# in test scripts) so Stage 5 can import the same set.
+TARGET_LABELS = frozenset({
+    "Hair", "Hat", "Scarf", "Upper-clothes", "Pants", "Skirt", "Dress",
+    "Face", "Left-leg", "Right-leg", "Left-arm", "Right-arm",
+})
+
+
+@dataclass
+class Person:
+    """Everything the rule layer needs to know about one detected person."""
+
+    # YOLO bbox in source-image pixel coordinates.
+    bbox: tuple[int, int, int, int]
+    # YOLO person-class confidence.
+    det_conf: float
+    # Cropped RGB image at the bbox.
+    crop: Image.Image
+    # SegFormer pixel-wise label map, same H×W as `crop`. Values index
+    # into SEG_LABELS.
+    seg_mask: np.ndarray
+    # InsightFace gender: 'M', 'F', or None when no face was found.
+    gender: str | None
+    # InsightFace *face detection* confidence — answers "is there a
+    # face here?", not "is the gender prediction reliable." Stage 3
+    # (`docs/04-gender-classifier.md`, Finding 2) is explicit about
+    # this distinction. 0.0 when no face was detected.
+    face_det_score: float
+    # Boolean mask per SEG_LABELS entry, same H×W as seg_mask.
+    # Populated by Perception.detect_and_segment so the rule layer
+    # never needs to know integer label indices — only string names.
+    label_masks: dict[str, np.ndarray]
+
+
+@runtime_checkable
+class Perceiver(Protocol):
+    """Interface the video pipeline depends on.
+
+    One production adapter exists: :class:`Perception`.  Fakes that
+    satisfy this protocol can be injected into :class:`VideoProcessor`
+    for testing the tracking / temporal-smoothing logic without loading
+    any model weights.
+    """
+
+    def detect_and_segment(self, image: Image.Image) -> list[Person]: ...
+    def classify(self, crop: Image.Image) -> tuple[str | None, float]: ...
+
+
+def _pick_device() -> torch.device:
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+class Perception:
+    def __init__(self, person_conf: float = DEFAULT_PERSON_CONF) -> None:
+        self.person_conf = person_conf
+        self.device = _pick_device()
+
+        self.detector = YOLO(YOLO_MODEL_ID)
+
+        self.seg_processor = SegformerImageProcessor.from_pretrained(SEG_MODEL_ID)
+        self.seg_model = (
+            AutoModelForSemanticSegmentation.from_pretrained(SEG_MODEL_ID)
+            .to(self.device)
+            .eval()
+        )
+        # fp16: logits are immediately argmax'd so precision loss is benign.
+        # CPU fp16 support in PyTorch is incomplete; restrict to MPS/CUDA.
+        if self.device.type in ("mps", "cuda"):
+            self.seg_model = self.seg_model.half()
+
+        # torch.compile fuses transformer attention kernels on top of fp16.
+        # dynamic=True treats the batch dim as symbolic so frames with
+        # different person counts (1 person, then 3, then 2 …) don't each
+        # trigger a recompile. First batch incurs one-time JIT compilation
+        # (~10-30 frames of wall-time warm-up at 1 fps baseline); all
+        # subsequent batches run fused. On MPS, unsupported ops fall back to
+        # eager silently — no correctness risk. Silently skip if the backend
+        # raises (e.g. an inductor limitation on this torch build).
+        if self.device.type in ("mps", "cuda"):
+            try:
+                self.seg_model = torch.compile(self.seg_model, dynamic=True)
+            except Exception:
+                pass
+
+        # Two scoped tweaks for memory bounding:
+        #
+        # 1. `allowed_modules` filters the buffalo_l bundle from 5
+        #    sub-models down to the 2 we actually consume. We only
+        #    read `face.gender` and `face.det_score`, so landmarks
+        #    (2d_106, 3d_68) and recognition embeddings are pure
+        #    overhead — fewer models per call = less per-call CPU
+        #    allocation.
+        # 2. Swap insightface's PickableInferenceSession for our
+        #    arena/mem_pattern-disabled subclass during construction,
+        #    then restore — keeps the patch scoped to this instance.
+        _orig = _imz.PickableInferenceSession
+        _imz.PickableInferenceSession = _NoArenaPickableSession
+        try:
+            self.face_app = FaceAnalysis(
+                name=INSIGHT_MODEL_ID,
+                providers=["CPUExecutionProvider"],
+                allowed_modules=["detection", "genderage"],
+            )
+            self.face_app.prepare(ctx_id=0, det_size=INSIGHT_DET_SIZE)
+        finally:
+            _imz.PickableInferenceSession = _orig
+
+    @staticmethod
+    def make_label_masks(seg_mask: np.ndarray) -> dict[str, np.ndarray]:
+        """Convert an integer seg_mask to named boolean arrays.
+
+        Isolates the model's integer label convention here so the rule
+        layer only deals with string keys and never imports SEG_LABELS.
+        """
+        return {name: (seg_mask == i) for i, name in enumerate(SEG_LABELS)}
+
+    def __call__(self, image: Image.Image) -> list[Person]:
+        """Single-image entry point: detect, segment, classify all in one.
+
+        Used by test scripts. The video pipeline uses
+        :meth:`detect_and_segment` + :meth:`classify` separately so it
+        can cache gender per IoU-tracked person and avoid the
+        ~200 MB/frame heap growth that running InsightFace ONNX every
+        frame produces (see docs/07-video-pipeline.md Finding 6).
+        """
+        people = self.detect_and_segment(image)
+        for person in people:
+            person.gender, person.face_det_score = self.classify(person.crop)
+        return people
+
+    def detect_and_segment(self, image: Image.Image) -> list[Person]:
+        """Detect persons, segment each, return Person list with
+        ``gender=None`` and ``face_det_score=0.0``.
+
+        The video pipeline calls this once per frame, then only invokes
+        :meth:`classify` for newly-detected tracks.
+        """
+        det = self.detector.predict(
+            source=image,
+            classes=[PERSON_CLASS_ID],
+            conf=self.person_conf,
+            verbose=False,
+        )[0]
+        boxes = det.boxes
+        if boxes is None or len(boxes) == 0:
+            return []
+
+        xyxy = boxes.xyxy.cpu().numpy()
+        confs = boxes.conf.cpu().numpy()
+
+        crops: list[Image.Image] = []
+        meta: list[tuple[int, int, int, int, float]] = []
+        for (x1f, y1f, x2f, y2f), conf in zip(xyxy, confs):
+            x1, y1, x2, y2 = int(x1f), int(y1f), int(x2f), int(y2f)
+            crops.append(image.crop((x1, y1, x2, y2)))
+            meta.append((x1, y1, x2, y2, float(conf)))
+
+        # Batched segmentation: one forward pass over all person crops
+        # in this frame instead of N sequential calls. The processor
+        # always resizes inputs to 512×512 so the batch tensor is
+        # (N, 3, 512, 512) regardless of crop sizes — same per-call
+        # path as the single-image version, just amortizes GPU pipeline
+        # overhead. On a 6-person 1280×720 frame this is ~2× the
+        # throughput of sequential calls on Apple MPS.
+        seg_masks = self._segment_batch(crops)
+
+        people: list[Person] = []
+        for (x1, y1, x2, y2, conf), crop, seg_mask in zip(meta, crops, seg_masks):
+            people.append(Person(
+                bbox=(x1, y1, x2, y2),
+                det_conf=conf,
+                crop=crop,
+                seg_mask=seg_mask,
+                gender=None,
+                face_det_score=0.0,
+                label_masks=Perception.make_label_masks(seg_mask),
+            ))
+        return people
+
+    def classify(self, crop: Image.Image) -> tuple[str | None, float]:
+        """Run face detection + gender classification on a person crop."""
+        return self._classify(crop)
+
+    @contextmanager
+    def _cleanup_device_memory(self) -> Generator[None, None, None]:
+        try:
+            yield
+        finally:
+            gc.collect()
+            if self.device.type == "mps":
+                torch.mps.empty_cache()
+            elif self.device.type == "cuda":
+                torch.cuda.empty_cache()
+
+    def _segment_batch(self, crops: list[Image.Image]) -> list[np.ndarray]:
+        """One forward pass for all person crops in a frame. Argmax on
+        device, INTER_NEAREST upsample on CPU, per-batch empty_cache.
+        """
+        if not crops:
+            return []
+        with self._cleanup_device_memory():
+            inputs = self.seg_processor(images=list(crops), return_tensors="pt").to(self.device)
+            inputs["pixel_values"] = inputs["pixel_values"].to(
+                dtype=next(self.seg_model.parameters()).dtype
+            )
+            with torch.no_grad():
+                outputs = self.seg_model(**inputs)
+            # (B, 128, 128) int64 on device — argmax before CPU move keeps
+            # the 18-channel logits tensor off CPU heap (Finding 7).
+            pred_small = outputs.logits.argmax(dim=1)
+            del outputs, inputs
+            pred_small_cpu = pred_small.detach().to("cpu").numpy().astype(np.int32)
+            del pred_small
+            results: list[np.ndarray] = []
+            for i, crop in enumerate(crops):
+                results.append(
+                    cv2.resize(pred_small_cpu[i], crop.size, interpolation=cv2.INTER_NEAREST)
+                )
+        return results
+
+    def _classify(self, crop: Image.Image) -> tuple[str | None, float]:
+        bgr = np.array(crop)[:, :, ::-1]
+        faces = self.face_app.get(bgr)
+        if not faces:
+            return None, 0.0
+        face = max(faces, key=lambda f: f.det_score)
+        gender = "M" if int(face.gender) == 1 else "F"
+        return gender, float(face.det_score)