diff --git a/src/jabs/project/pose_attribute_cache.py b/src/jabs/project/pose_attribute_cache.py new file mode 100644 index 00000000..db24d8da --- /dev/null +++ b/src/jabs/project/pose_attribute_cache.py @@ -0,0 +1,88 @@ +"""Persistent cache of per-video pose attributes to skip the load-time pose scan. + +The project-load scan opens every pose HDF5 file to read a handful of small, +intrinsic attributes (frame count, identity count, static objects, lixit +keypoint count, cm-per-pixel flag). Those values change only when the pose file +itself changes, so they are cached in ``jabs/cache/pose_attribute_cache.json`` +keyed by video filename and gated by a cheap ``stat`` token of the pose file. A +subsequent load then rescans only the videos whose pose file is new or changed. + +The cache is an optimization only: a missing, unreadable, or schema-mismatched +cache simply triggers a full rescan, and a failed write is logged and ignored. +""" + +import json +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) + +# Bump when the cached entry schema changes so stale caches are ignored wholesale. +SCHEMA_VERSION = 1 + + +def pose_token(pose_path: Path) -> str: + """Return a cheap change-detection token for a pose file. + + Uses a single ``stat`` (size and modification time) rather than opening or + hashing the file, so computing the token stays far cheaper than the scan it + guards. + + Args: + pose_path: Path to the pose HDF5 file. + + Returns: + A ``":"`` token string. + """ + st = pose_path.stat() + return f"{st.st_size}:{st.st_mtime_ns}" + + +def load(cache_path: Path | None) -> dict[str, dict]: + """Load the per-video attribute map from the cache file. + + Args: + cache_path: Path to the cache JSON file, or ``None`` when caching is + disabled (``use_cache=False``). + + Returns: + Mapping of video filename to its cached entry. Returns an empty mapping + when caching is disabled, or the file is missing, unreadable, or written + by a different schema version. + """ + if cache_path is None or not cache_path.exists(): + return {} + try: + with cache_path.open("r") as f: + data = json.load(f) + except (OSError, ValueError): + logger.warning("Could not read pose attribute cache %s; rescanning", cache_path) + return {} + if not isinstance(data, dict) or data.get("schema_version") != SCHEMA_VERSION: + return {} + videos = data.get("videos") + return videos if isinstance(videos, dict) else {} + + +def save(cache_path: Path | None, videos: dict[str, dict]) -> None: + """Atomically write the per-video attribute map to the cache file. + + A write failure is logged and swallowed: the cache is an optimization, so + failing to persist it must never break project loading. + + Args: + cache_path: Path to the cache JSON file, or ``None`` when caching is + disabled (in which case this is a no-op). + videos: Mapping of video filename to its cached entry. + """ + if cache_path is None: + return + payload = {"schema_version": SCHEMA_VERSION, "videos": videos} + tmp = cache_path.with_suffix(".json.tmp") + try: + cache_path.parent.mkdir(parents=True, exist_ok=True) + with tmp.open("w") as f: + json.dump(payload, f, indent=2, sort_keys=True) + tmp.replace(cache_path) + except OSError: + logger.warning("Could not write pose attribute cache %s", cache_path, exc_info=True) diff --git a/src/jabs/project/project.py b/src/jabs/project/project.py index a9f7699f..5db5c6a6 100644 --- a/src/jabs/project/project.py +++ b/src/jabs/project/project.py @@ -33,6 +33,7 @@ open_pose_file, ) +from . import pose_attribute_cache from .feature_manager import FeatureManager from .parallel_workers import ( BinaryFeatureLoadJobSpec, @@ -61,6 +62,78 @@ from jabs.core.utils.process_pool_manager import ProcessPoolManager +def _is_int(value: object) -> bool: + """Return True for a genuine int (JSON bools are ints in Python; reject them).""" + return isinstance(value, int) and not isinstance(value, bool) + + +def _is_str_list(value: object) -> bool: + """Return True for a list whose every element is a string.""" + return isinstance(value, list) and all(isinstance(item, str) for item in value) + + +# Pose-derived fields cached per video, mapped to a predicate validating the +# cached value's type (matching the corresponding VideoScanResult field). An +# entry is trusted only when it carries a matching token and pose filename plus +# every field present with the expected type; a parseable-but-malformed entry +# (e.g. a hand-edited or corrupted cache) is treated as a miss and rescanned. +_POSE_CACHE_FIELD_VALIDATORS: dict[str, Callable[[object], bool]] = { + "hdf5_frame_count": _is_int, + "identity_count": _is_int, + "static_objects": _is_str_list, + "lixit_keypoints": _is_int, + "has_cm_per_pixel": lambda value: isinstance(value, bool), +} + + +def _pose_cache_entry_matches(entry: object, token: str, pose_file: str) -> bool: + """Return True when a cache entry is usable for a video's current pose file. + + Validates field types as well as presence so a parseable-but-malformed entry + is treated as a miss (triggering a rescan) rather than reconstructed into + wrong metadata or a reconstruction-time error. + """ + return ( + isinstance(entry, dict) + and entry.get("token") == token + and entry.get("pose_file") == pose_file + and all( + field in entry and validator(entry[field]) + for field, validator in _POSE_CACHE_FIELD_VALIDATORS.items() + ) + ) + + +def _scan_result_from_pose_cache(video: str, entry: dict) -> VideoScanResult: + """Reconstruct a VideoScanResult from a cached entry. + + ``video_frame_count`` is left ``None``: video frame counts are validated on + demand rather than at load, so they are never cached. + """ + return VideoScanResult( + video=video, + hdf5_frame_count=entry["hdf5_frame_count"], + video_frame_count=None, + identity_count=entry["identity_count"], + static_objects=list(entry["static_objects"]), + lixit_keypoints=entry["lixit_keypoints"], + has_cm_per_pixel=entry["has_cm_per_pixel"], + ) + + +def _pose_cache_entry(result: VideoScanResult, pose_file: str, token: str) -> dict: + """Build a cache entry from a fresh scan result.""" + return { + "token": token, + "pose_file": pose_file, + "hdf5_frame_count": result["hdf5_frame_count"], + "identity_count": result["identity_count"], + "static_objects": list(result["static_objects"]), + "lixit_keypoints": result["lixit_keypoints"], + "has_cm_per_pixel": result["has_cm_per_pixel"], + } + + class Project: """Represents a JABS project, managing data, settings, and operations for a project directory. @@ -215,18 +288,99 @@ def _run_video_scan( if not jobs: return {} + # The opt-in up-front check reads video frame counts, which are not + # cached, so it always performs a full scan. + if enable_video_check: + return self._scan_jobs(jobs, process_pool) + + return self._run_cached_video_scan(jobs, process_pool) + + @staticmethod + def _scan_jobs( + jobs: list[VideoScanJobSpec], + process_pool: "ProcessPoolManager | None", + ) -> dict[str, VideoScanResult]: + """Run scan jobs in parallel when a pool is available, else sequentially. + + Args: + jobs: Scan jobs to execute. + process_pool: Optional shared pool; when ``None`` the jobs run + sequentially in the calling process. + + Returns: + Mapping from video filename to its scan result. + """ if process_pool is not None: future_to_video = { process_pool.submit(scan_video_metadata, job): job["video"] for job in jobs } results: dict[str, VideoScanResult] = {} for future in as_completed(future_to_video): - result: VideoScanResult = future.result() + try: + result: VideoScanResult = future.result() + except Exception: + logger.error( + "Failed to scan pose metadata for %s", + future_to_video[future], + exc_info=True, + ) + raise results[result["video"]] = result return results return {job["video"]: scan_video_metadata(job) for job in jobs} + def _run_cached_video_scan( + self, + jobs: list[VideoScanJobSpec], + process_pool: "ProcessPoolManager | None", + ) -> dict[str, VideoScanResult]: + """Scan only new/changed pose files, reusing cached attributes otherwise. + + Pose attributes are intrinsic to the pose file, so an entry whose cached + ``stat`` token still matches is reused without opening the file. Only + videos with a missing or stale entry are scanned; the cache is then + rewritten when anything changed (a rescan happened, or the set of videos + differs from what was cached). + + Args: + jobs: Scan jobs for every video with a locatable pose file. + process_pool: Optional shared pool for parallelizing the scan of the + uncached videos. + + Returns: + Mapping from video filename to its scan result. + """ + cache_path = self._paths.pose_attribute_cache_file + cached = pose_attribute_cache.load(cache_path) + + results: dict[str, VideoScanResult] = {} + tokens: dict[str, str] = {} + to_scan: list[VideoScanJobSpec] = [] + for job in jobs: + video = job["video"] + token = pose_attribute_cache.pose_token(job["pose_path"]) + tokens[video] = token + entry = cached.get(video) + if _pose_cache_entry_matches(entry, token, job["pose_path"].name): + results[video] = _scan_result_from_pose_cache(video, entry) + else: + to_scan.append(job) + + if to_scan: + results.update(self._scan_jobs(to_scan, process_pool)) + + updated = { + job["video"]: _pose_cache_entry( + results[job["video"]], job["pose_path"].name, tokens[job["video"]] + ) + for job in jobs + } + if to_scan or set(updated) != set(cached): + pose_attribute_cache.save(cache_path, updated) + + return results + def _validate_pose_files(self): """Ensure all videos have corresponding pose files.""" err = False diff --git a/src/jabs/project/project_paths.py b/src/jabs/project/project_paths.py index 95a51d6b..5b29f8d6 100644 --- a/src/jabs/project/project_paths.py +++ b/src/jabs/project/project_paths.py @@ -85,6 +85,16 @@ def cache_dir(self) -> Path | None: """Get the path to the cache directory.""" return self._cache_dir + @property + def pose_attribute_cache_file(self) -> Path | None: + """Get the path to the per-video pose-attribute cache file. + + Returns ``None`` when the project has no cache directory + (``use_cache=False``), in which case pose attributes are not persisted + and every load performs a full pose scan. + """ + return self._cache_dir / "pose_attribute_cache.json" if self._cache_dir else None + @property def session_dir(self) -> Path: """Get the path to the session directory.""" diff --git a/tests/project/test_pose_attribute_cache.py b/tests/project/test_pose_attribute_cache.py new file mode 100644 index 00000000..23dff83a --- /dev/null +++ b/tests/project/test_pose_attribute_cache.py @@ -0,0 +1,270 @@ +"""Tests for the per-video pose attribute cache (KLAUS-506).""" + +import json +import os +import shutil +from pathlib import Path + +import pytest + +from jabs.project import Project, pose_attribute_cache +from jabs.project.project import _pose_cache_entry_matches + +DATA_DIR = Path(__file__).parent.parent / "data" +SAMPLE_POSE = DATA_DIR / "sample_pose_est_v3.h5" +CACHE_RELPATH = Path("jabs") / "cache" / "pose_attribute_cache.json" + + +def _make_project(project_dir: Path, video_names: list[str]) -> Path: + """Create a minimal project dir: stub ``.avi`` videos + copied v3 pose files.""" + project_dir.mkdir(parents=True, exist_ok=True) + for name in video_names: + (project_dir / name).touch() + base = Path(name).with_suffix("").name + shutil.copy(SAMPLE_POSE, project_dir / f"{base}_pose_est_v3.h5") + return project_dir + + +def _open(project_dir: Path) -> Project: + """Open a project on the deferred (cached) scan path.""" + return Project(project_dir, enable_video_check=False, enable_session_tracker=False) + + +def _spy_scan(monkeypatch) -> list[str]: + """Spy on ``scan_video_metadata`` as Project uses it; returns scanned videos.""" + import jabs.project.project as project_mod + + real = project_mod.scan_video_metadata + calls: list[str] = [] + + def spy(job): + calls.append(job["video"]) + return real(job) + + monkeypatch.setattr(project_mod, "scan_video_metadata", spy) + return calls + + +def _bump_mtime(path: Path) -> None: + """Advance a file's modification time by one second (changes its token).""" + st = path.stat() + os.utime(path, ns=(st.st_atime_ns, st.st_mtime_ns + 1_000_000_000)) + + +# --- pose_attribute_cache module unit tests -------------------------------- + + +def test_pose_token_changes_with_mtime(tmp_path): + """The stat token changes when the file's modification time changes.""" + f = tmp_path / "x.h5" + f.write_bytes(b"abc") + before = pose_attribute_cache.pose_token(f) + _bump_mtime(f) + assert pose_attribute_cache.pose_token(f) != before + + +def test_load_none_returns_empty(): + """Caching disabled (None path) loads an empty map.""" + assert pose_attribute_cache.load(None) == {} + + +def test_load_missing_returns_empty(tmp_path): + """A missing cache file loads an empty map.""" + assert pose_attribute_cache.load(tmp_path / "nope.json") == {} + + +def test_load_corrupt_returns_empty(tmp_path): + """An unparseable cache file is treated as empty (triggers rescan).""" + p = tmp_path / "c.json" + p.write_text("{ not valid json") + assert pose_attribute_cache.load(p) == {} + + +def test_load_invalid_encoding_returns_empty(tmp_path): + """A cache file with invalid text encoding is treated as empty (rescan). + + ``UnicodeDecodeError`` is a subclass of ``ValueError``, so the decode failure + raised while reading is caught like any other unreadable cache. + """ + p = tmp_path / "c.json" + p.write_bytes(b"\xff\xfe not valid utf-8") + assert pose_attribute_cache.load(p) == {} + + +def test_load_schema_mismatch_returns_empty(tmp_path): + """A cache written by a different schema version is ignored.""" + p = tmp_path / "c.json" + p.write_text( + json.dumps( + {"schema_version": pose_attribute_cache.SCHEMA_VERSION + 1, "videos": {"a.avi": {}}} + ) + ) + assert pose_attribute_cache.load(p) == {} + + +def test_save_load_roundtrip(tmp_path): + """A saved map is returned verbatim by load.""" + p = tmp_path / "c.json" + videos = {"v.avi": {"token": "1:2", "pose_file": "v_pose_est_v3.h5", "hdf5_frame_count": 10}} + pose_attribute_cache.save(p, videos) + assert pose_attribute_cache.load(p) == videos + + +def test_save_none_is_noop(): + """Saving with a None path does nothing and does not raise.""" + pose_attribute_cache.save(None, {"v.avi": {}}) + + +# --- _pose_cache_entry_matches unit tests ---------------------------------- + + +def _valid_entry() -> dict: + """Return a well-formed cache entry for token ``"1:2"`` / pose ``"v.h5"``.""" + return { + "token": "1:2", + "pose_file": "v.h5", + "hdf5_frame_count": 10, + "identity_count": 2, + "static_objects": ["lixit", "food_hopper"], + "lixit_keypoints": 3, + "has_cm_per_pixel": True, + } + + +def test_matches_well_formed_entry(): + """A complete, correctly-typed entry with matching token/pose is a hit.""" + assert _pose_cache_entry_matches(_valid_entry(), "1:2", "v.h5") + + +def test_matches_empty_static_objects(): + """An empty static_objects list is valid (no static objects in the pose file).""" + entry = _valid_entry() + entry["static_objects"] = [] + assert _pose_cache_entry_matches(entry, "1:2", "v.h5") + + +@pytest.mark.parametrize( + ("token", "pose_file"), + [("9:9", "v.h5"), ("1:2", "other.h5")], + ids=["stale-token", "different-pose-file"], +) +def test_no_match_on_token_or_pose_mismatch(token, pose_file): + """A mismatched token or pose filename is a miss.""" + assert not _pose_cache_entry_matches(_valid_entry(), token, pose_file) + + +def test_no_match_when_not_dict(): + """A non-dict entry (e.g. None for a missing video) is a miss.""" + assert not _pose_cache_entry_matches(None, "1:2", "v.h5") + + +def test_no_match_on_missing_field(): + """An entry missing a required field is a miss.""" + entry = _valid_entry() + del entry["identity_count"] + assert not _pose_cache_entry_matches(entry, "1:2", "v.h5") + + +@pytest.mark.parametrize( + ("field", "bad_value"), + [ + ("hdf5_frame_count", "10"), + ("identity_count", 2.0), + ("lixit_keypoints", True), + ("has_cm_per_pixel", 1), + ("static_objects", {"lixit": 1}), + ("static_objects", "lixit"), + ("static_objects", [1, 2]), + ], + ids=[ + "count-as-str", + "count-as-float", + "int-field-as-bool", + "bool-field-as-int", + "static-objects-as-dict", + "static-objects-as-str", + "static-objects-non-str-items", + ], +) +def test_no_match_on_wrong_type(field, bad_value): + """A parseable-but-wrong-typed field is a miss so the video is rescanned.""" + entry = _valid_entry() + entry[field] = bad_value + assert not _pose_cache_entry_matches(entry, "1:2", "v.h5") + + +# --- Project-level caching behavior ---------------------------------------- + + +def test_second_load_uses_cache_no_scan(tmp_path, monkeypatch): + """A second load of an unchanged project rescans no pose files.""" + project_dir = _make_project(tmp_path / "proj", ["v1.avi"]) + _open(project_dir) # first load populates the cache + assert (project_dir / CACHE_RELPATH).exists() + + calls = _spy_scan(monkeypatch) + _open(project_dir) + assert calls == [] + + +def test_changed_pose_file_is_rescanned(tmp_path, monkeypatch): + """A pose file whose token changed is rescanned on the next load.""" + project_dir = _make_project(tmp_path / "proj", ["v1.avi"]) + _open(project_dir) + _bump_mtime(project_dir / "v1_pose_est_v3.h5") + + calls = _spy_scan(monkeypatch) + _open(project_dir) + assert calls == ["v1.avi"] + + +def test_new_video_only_scans_new(tmp_path, monkeypatch): + """Adding a video rescans only the new video; cached ones are reused.""" + project_dir = _make_project(tmp_path / "proj", ["v1.avi"]) + _open(project_dir) + _make_project(project_dir, ["v2.avi"]) # add a second video + pose + + calls = _spy_scan(monkeypatch) + _open(project_dir) + assert calls == ["v2.avi"] + + +def test_malformed_cache_entry_is_rescanned(tmp_path, monkeypatch): + """A parseable cache whose entry has a wrong-typed field is rescanned.""" + project_dir = _make_project(tmp_path / "proj", ["v1.avi"]) + _open(project_dir) # populate the cache + + cache_file = project_dir / CACHE_RELPATH + data = json.loads(cache_file.read_text()) + # Corrupt a field's type: identity_count stored as a string. + data["videos"]["v1.avi"]["identity_count"] = "not-an-int" + cache_file.write_text(json.dumps(data)) + + calls = _spy_scan(monkeypatch) + _open(project_dir) + assert calls == ["v1.avi"] + + +def test_use_cache_false_writes_no_cache(tmp_path): + """With use_cache=False the project loads but persists no cache file.""" + project_dir = _make_project(tmp_path / "proj", ["v1.avi"]) + project = Project( + project_dir, + use_cache=False, + enable_video_check=False, + enable_session_tracker=False, + ) + assert project.video_manager.videos == ["v1.avi"] + assert not (project_dir / CACHE_RELPATH).exists() + + +def test_cached_load_matches_fresh_scan(tmp_path): + """Values reconstructed from the cache match a full scan of the same project.""" + project_dir = _make_project(tmp_path / "proj", ["v1.avi", "v2.avi"]) + fresh = _open(project_dir) # full scan (no cache yet) + cached = _open(project_dir) # cache hits + + assert cached.total_project_identities == fresh.total_project_identities + assert cached.feature_manager.min_pose_version == fresh.feature_manager.min_pose_version + assert cached.feature_manager.static_objects == fresh.feature_manager.static_objects + assert cached.feature_manager.is_cm_unit == fresh.feature_manager.is_cm_unit