Merge pull request #413 from KumarLabJax/feature/klaus-506-pose-attribute-cache

Cache per-video pose attributes to skip load-time pose scan (KLAUS-506)

Author: gbeane

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 ``"<size>:<mtime_ns>"`` 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)

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

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."""