Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 88 additions & 0 deletions src/jabs/project/pose_attribute_cache.py
Original file line number Diff line number Diff line change
@@ -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)
125 changes: 125 additions & 0 deletions src/jabs/project/project.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
open_pose_file,
)

from . import pose_attribute_cache
from .feature_manager import FeatureManager
from .parallel_workers import (
BinaryFeatureLoadJobSpec,
Expand Down Expand Up @@ -61,6 +62,57 @@
from jabs.core.utils.process_pool_manager import ProcessPoolManager


# Pose-derived fields cached per video; an entry is trusted only when it carries
# a matching token and pose filename plus all of these fields.
_POSE_CACHE_FIELDS = (
"hdf5_frame_count",
"identity_count",
"static_objects",
"lixit_keypoints",
"has_cm_per_pixel",
)


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."""
return (
isinstance(entry, dict)
and entry.get("token") == token
and entry.get("pose_file") == pose_file
and all(field in entry for field in _POSE_CACHE_FIELDS)
)
Comment on lines +76 to +83


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.

Expand Down Expand Up @@ -211,6 +263,28 @@ 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
Expand All @@ -223,6 +297,57 @@ def _run_video_scan(

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
Expand Down
10 changes: 10 additions & 0 deletions src/jabs/project/project_paths.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Loading
Loading