- frame caching

This commit is contained in:
w.pomp
2026-08-04 14:35:24 +02:00
parent c975b1943f
commit 9f5f9debc5
5 changed files with 296 additions and 133 deletions
+1 -2
View File
@@ -3,7 +3,7 @@ name = "ndbioimage"
version = "0.2.0"
edition = "2024"
rust-version = "1.94.0"
authors = ["Wim Pomp <w.pomp@nki.nl>", "opencode"]
authors = ["Wim Pomp <w.pomp@nki.nl>"]
license = "MIT OR Apache-2.0"
description = "Read bio image formats using the bio-formats java package."
homepage = "https://git.pomppervova.nl/wim/ndbioimage/src/branch/rs"
@@ -24,7 +24,6 @@ color-eyre = { version = "0.6", optional = true }
console = { version = "0.16", optional = true }
downloader = { version = "0.2", optional = true, default-features = false, features = ["rustls-tls"] }
ffmpeg-sidecar = { version = "2", optional = true }
#ffmpreg = { version = "0.1", optional = true }
itertools = "0.15"
indexmap = { version = "2", features = ["serde"] }
indicatif = { version = "0.18", features = ["rayon"], optional = true }
-94
View File
@@ -1,94 +0,0 @@
use std::hash::{Hash, Hasher};
use std::path::{Path, PathBuf};
use std::sync::{LazyLock, RwLock};
use indexmap::{Equivalent, IndexMap};
use ndarray::IxDyn;
use crate::error::Error;
use crate::readers::ArrayT;
static CACHE: LazyLock<RwLock<IndexMap<ViewHash, ArrayT<IxDyn>>>> =
LazyLock::new(|| RwLock::new(IndexMap::new()));
#[derive(Debug, Clone, Eq, PartialEq)]
struct ViewHash {
reader: String,
path: PathBuf,
series: usize,
position: usize,
c: isize,
z: isize,
t: isize,
}
impl Hash for ViewHash {
fn hash<H: Hasher>(&self, state: &mut H) {
self.reader.as_str().hash(state);
self.path.as_path().hash(state);
self.series.hash(state);
self.position.hash(state);
self.c.hash(state);
self.z.hash(state);
self.t.hash(state);
}
}
impl ViewHash {
fn new(reader: String, path: PathBuf, series: usize, position: usize, c: isize, z: isize, t: isize) -> Self {
Self { reader, path, series, position, c, z, t }
}
}
struct ViewEquiv<'a> {
reader: &'a str,
path: &'a Path,
series: usize,
position: usize,
c: isize,
z: isize,
t: isize,
}
impl Hash for ViewEquiv<'_> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.reader.hash(state);
self.path.hash(state);
self.series.hash(state);
self.position.hash(state);
self.c.hash(state);
self.z.hash(state);
self.t.hash(state);
}
}
impl<'a> ViewEquiv<'a> {
fn new(reader: &'a str, path: &'a Path, series: usize, position: usize, c: isize, z: isize, t: isize) -> ViewEquiv<'a> {
Self { reader, path, series, position, c, z, t }
}
}
impl Equivalent<ViewEquiv<'_>> for ViewHash {
fn equivalent(&self, key: &ViewEquiv) -> bool {
(self.reader == key.reader) && (self.path == key.path) && (self.series == key.series) && (self.position == key.position) && (self.c == key.c) && (self.z == key.z) && (self.t == key.t)
}
}
impl Equivalent<ViewHash> for ViewEquiv<'_> {
fn equivalent(&self, key: &ViewHash) -> bool {
key.equivalent(self)
}
}
pub fn cache_get_or_insert(f: &dyn Fn() -> ArrayT<IxDyn>, reader: &str, path: &Path, series: usize, position: usize, c: isize, z: isize, t: isize) -> Result<ArrayT<IxDyn>, Error> {
if let Some(frame) = CACHE.read().unwrap().get(&ViewEquiv::new(reader, path, series, position, c, z, t)) {
Ok(frame.clone())
} else {
let a = f();
let mut cache = CACHE.write().unwrap();
cache.insert(ViewHash::new(reader.to_string(), path.to_path_buf(), series, position, c, z, t), a.clone());
// TODO: find an IndexMapDeque to pop efficiently at the other end
while cache.len() > 1024 {
cache.pop();
}
Ok(a)
}
}
+1 -1
View File
@@ -246,7 +246,7 @@ impl PyView {
#[getter]
fn reader_name(&self) -> String {
self.view.reader_name()
self.view.reader_name().to_string()
}
#[allow(unused_variables)]
+44 -27
View File
@@ -143,8 +143,8 @@ pub trait Reader: Clone + Sized + Debug + Send + Hash + Into<DynReader> {
where
P: AsRef<Path>;
fn reader_name(&self) -> String {
type_name::<Self>().to_string()
fn reader_name(&self) -> &'static str {
type_name::<Self>()
}
// TODO: read from file if present
@@ -257,83 +257,77 @@ impl_frame_cast! {
isize: INT32
}
impl<D, T> TryInto<Array<T, D>> for ArrayT<D>
where
D: Dimension,
T: FromPrimitive + Zero + 'static,
{
type Error = Error;
fn try_into(self) -> Result<Array<T, D>, Self::Error> {
macro_rules! try_into_array_body {
($self:expr, $map:ident) => {{
let mut err = Ok(());
let arr = match self {
ArrayT::I8(v) => v.mapv_into_any(|x| {
let arr = match $self {
ArrayT::I8(v) => v.$map(|x| {
T::from_i8(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::U8(v) => v.mapv_into_any(|x| {
ArrayT::U8(v) => v.$map(|x| {
T::from_u8(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::I16(v) => v.mapv_into_any(|x| {
ArrayT::I16(v) => v.$map(|x| {
T::from_i16(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::U16(v) => v.mapv_into_any(|x| {
ArrayT::U16(v) => v.$map(|x| {
T::from_u16(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::I32(v) => v.mapv_into_any(|x| {
ArrayT::I32(v) => v.$map(|x| {
T::from_i32(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::U32(v) => v.mapv_into_any(|x| {
ArrayT::U32(v) => v.$map(|x| {
T::from_u32(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::F32(v) => v.mapv_into_any(|x| {
ArrayT::F32(v) => v.$map(|x| {
T::from_f32(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::F64(v) | ArrayT::F128(v) => v.mapv_into_any(|x| {
ArrayT::F64(v) | ArrayT::F128(v) => v.$map(|x| {
T::from_f64(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::I64(v) => v.mapv_into_any(|x| {
ArrayT::I64(v) => v.$map(|x| {
T::from_i64(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::U64(v) => v.mapv_into_any(|x| {
ArrayT::U64(v) => v.$map(|x| {
T::from_u64(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::I128(v) => v.mapv_into_any(|x| {
ArrayT::I128(v) => v.$map(|x| {
T::from_i128(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
})
}),
ArrayT::U128(v) => v.mapv_into_any(|x| {
ArrayT::U128(v) => v.$map(|x| {
T::from_u128(x).unwrap_or_else(|| {
err = Err(Error::Cast(x.to_string(), type_name::<T>().to_string()));
T::zero()
@@ -344,6 +338,30 @@ where
Err(err) => Err(err),
Ok(()) => Ok(arr),
}
}};
}
impl<D, T> TryInto<Array<T, D>> for ArrayT<D>
where
D: Dimension,
T: FromPrimitive + Zero + 'static,
{
type Error = Error;
fn try_into(self) -> Result<Array<T, D>, Self::Error> {
try_into_array_body!(self, mapv_into_any)
}
}
impl<D, T> TryInto<Array<T, D>> for &ArrayT<D>
where
D: Dimension,
T: Clone + FromPrimitive + Zero + 'static,
{
type Error = Error;
fn try_into(self) -> Result<Array<T, D>, Self::Error> {
try_into_array_body!(self, mapv)
}
}
@@ -398,8 +416,8 @@ impl Reader for DynReader {
))
}
fn reader_name(&self) -> String {
let name = match self {
fn reader_name(&self) -> &'static str {
match self {
#[cfg(feature = "tiff")]
DynReader::Tiff(r) => r.reader_name(),
#[cfg(feature = "tiffseq")]
@@ -412,8 +430,7 @@ impl Reader for DynReader {
DynReader::BioFormatsJava(r) => r.reader_name(),
#[allow(unreachable_patterns)]
_ => unreachable!(),
};
format!("DynReader<{}>", name)
}
}
fn metadata(&self) -> Result<Ome, Error> {
+250 -9
View File
@@ -1,9 +1,9 @@
use crate::axes::{Ax, Axis, Operation, Shape, Slice, SliceInfoElemDef, slice_info};
use crate::error::Error;
use crate::metadata::Metadata;
use crate::readers::{Dimensions, DynReader, Reader};
use crate::readers::{Dimensions, DynReader, Frame, Reader};
use crate::stats::MinMax;
use indexmap::IndexMap;
use indexmap::{Equivalent, IndexMap};
use itertools::{Itertools, iproduct};
use ndarray::{
Array, Array0, Array1, Array2, ArrayD, Dimension, IntoDimension, Ix0, Ix1, Ix2, Ix5, IxDyn,
@@ -15,12 +15,13 @@ use serde::{Deserialize, Serialize};
use serde_with::serde_as;
use std::any::type_name;
use std::collections::{HashMap, HashSet};
use std::fmt::{Display, Formatter};
use std::fmt::{Debug, Display, Formatter};
use std::hash::{Hash, Hasher};
use std::iter::Sum;
use std::marker::PhantomData;
use std::ops::{AddAssign, Deref, Div};
use std::path::Path;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, OnceLock};
fn idx_bnd(idx: isize, bnd: isize) -> Result<isize, Error> {
if idx < -bnd {
@@ -62,9 +63,157 @@ impl<T> Number for T where
{
}
/// maximum number of frames held in the cache
const DEFAULT_FRAME_CACHE_SIZE: usize = 128;
/// identity of the reader a frame was read from
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct ReaderKey {
name: String,
path: PathBuf,
series: usize,
position: usize,
}
type FrameKey = (ReaderKey, usize, usize, usize);
/// borrowed view of [`FrameKey`] for cache lookups without allocation.
/// hashes byte-identically to [`FrameKey`] (str/String and Path/PathBuf hash
/// the same), so it only matches the same reader identity.
#[derive(Debug, PartialEq, Eq)]
struct FrameKeyRef<'a> {
name: &'a str,
path: &'a Path,
series: usize,
position: usize,
c: usize,
z: usize,
t: usize,
}
impl Hash for FrameKeyRef<'_> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.name.hash(state);
self.path.hash(state);
self.series.hash(state);
self.position.hash(state);
self.c.hash(state);
self.z.hash(state);
self.t.hash(state);
}
}
impl Equivalent<FrameKey> for FrameKeyRef<'_> {
fn equivalent(&self, key: &FrameKey) -> bool {
let (rk, c, z, t) = key;
self.name == rk.name
&& self.path == rk.path.as_path()
&& self.series == rk.series
&& self.position == rk.position
&& self.c == *c
&& self.z == *z
&& self.t == *t
}
}
impl FrameKeyRef<'_> {
fn to_owned(&self) -> FrameKey {
(
ReaderKey {
name: self.name.to_string(),
path: self.path.to_path_buf(),
series: self.series,
position: self.position,
},
self.c,
self.z,
self.t,
)
}
}
/// process-wide LRU cache of frames, shared between all views and threads
static GLOBAL_FRAME_CACHE: OnceLock<FrameCache> = OnceLock::new();
/// thread-safe LRU cache of frames read from the underlying reader
struct FrameCache {
inner: Mutex<FrameCacheInner>,
}
struct FrameCacheInner {
map: IndexMap<FrameKey, Arc<Frame>>,
capacity: usize,
}
impl Default for FrameCache {
fn default() -> Self {
Self::new(DEFAULT_FRAME_CACHE_SIZE)
}
}
impl FrameCache {
fn new(capacity: usize) -> Self {
Self {
inner: Mutex::new(FrameCacheInner {
map: IndexMap::with_capacity(capacity),
capacity,
}),
}
}
fn global() -> &'static FrameCache {
GLOBAL_FRAME_CACHE.get_or_init(FrameCache::default)
}
fn capacity(&self) -> usize {
self.inner.lock().unwrap().capacity
}
fn set_capacity(&self, capacity: usize) {
let mut inner = self.inner.lock().unwrap();
inner.capacity = capacity;
while inner.map.len() > inner.capacity {
inner.map.shift_remove_index(0);
}
}
fn get<Q>(&self, key: &Q) -> Option<Arc<Frame>>
where
Q: ?Sized + Hash + Equivalent<FrameKey>,
{
let mut inner = self.inner.lock().unwrap();
if let Some(idx) = inner.map.get_index_of(key) {
let (key, frame) = inner.map.shift_remove_index(idx).unwrap();
inner.map.insert(key, frame.clone());
Some(frame)
} else {
None
}
}
fn insert(&self, key: FrameKey, frame: Arc<Frame>) {
let mut inner = self.inner.lock().unwrap();
inner.map.insert(key, frame);
while inner.map.len() > inner.capacity {
inner.map.shift_remove_index(0);
}
}
#[cfg(test)]
fn len(&self) -> usize {
self.inner.lock().unwrap().map.len()
}
}
impl Debug for FrameCache {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FrameCache").finish_non_exhaustive()
}
}
/// sliceable view on an image file
#[serde_as]
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct View<D: Dimension, R: Reader = DynReader> {
reader: R,
/// same order as axes
@@ -176,6 +325,18 @@ impl<D: Dimension, R: Reader> View<D, R> {
self.reader.series()
}
/// set the maximum number of frames the process-wide cache may hold,
/// evicting the least recently used frames
pub fn with_cache_capacity(self, capacity: usize) -> Self {
FrameCache::global().set_capacity(capacity);
self
}
/// the maximum number of frames the process-wide cache may hold
pub fn cache_capacity(&self) -> usize {
FrameCache::global().capacity()
}
fn with_operations(mut self, operations: IndexMap<Axis, Operation>) -> Self {
self.operations = operations;
self
@@ -744,13 +905,13 @@ impl<D: Dimension, R: Reader> View<D, R> {
if let Some(i) = axes_out_idx[2] {
slice[i] = SliceInfoElem::Index(t)
};
let frame = self.reader.get_frame(
let frame = self.get_cached_frame(
(c % shape.c as isize) as usize,
(z % shape.z as isize) as usize,
(t % shape.t as isize) as usize,
)?;
let arr_frame: Array2<T> = frame.try_into()?;
let arr_frame: Array2<T> = frame.as_ref().try_into()?;
let arr_frame = match xy_dim {
0 => {
if op_xy.contains_key(&Axis::X) && op_xy.contains_key(&Axis::Y) {
@@ -888,6 +1049,28 @@ impl<D: Dimension, R: Reader> View<D, R> {
.collect())
}
fn frame_key(&self, c: usize, z: usize, t: usize) -> FrameKeyRef<'_> {
FrameKeyRef {
name: self.reader.reader_name(),
path: self.reader.path(),
series: self.reader.series(),
position: self.reader.position(),
c,
z,
t,
}
}
fn get_cached_frame(&self, c: usize, z: usize, t: usize) -> Result<Arc<Frame>, Error> {
let key = self.frame_key(c, z, t);
if let Some(frame) = FrameCache::global().get(&key) {
return Ok(frame);
}
let frame = Arc::new(self.reader.get_frame(c, z, t)?);
FrameCache::global().insert(key.to_owned(), frame.clone());
Ok(frame)
}
/// retrieve a single frame at czt, sliced accordingly
pub fn get_frame<T, N>(&self, c: N, z: N, t: N) -> Result<Array2<T>, Error>
where
@@ -1116,11 +1299,13 @@ to_bytes_vec_impl!(
mod tests {
use crate::axes::Axis;
use crate::error::Error;
use crate::readers::{DynReader, Reader};
use crate::readers::{DynReader, Frame, Reader};
use crate::stats::MinMax;
use crate::view::Item;
use crate::view::{FrameCache, Item, ReaderKey};
use ndarray::{Array, Array4, Array5, NewAxis};
use ndarray::{Array2, s};
use std::path::PathBuf;
use std::sync::Arc;
fn open(file: &str) -> Result<DynReader, Error> {
let path = std::env::current_dir()?
@@ -1388,4 +1573,60 @@ mod tests {
assert_eq!(a.shape(), [1280, 1280]);
Ok(())
}
#[test]
fn frame_cache() -> Result<(), Error> {
let cache = FrameCache::new(2);
let key = |c: usize, z: usize, t: usize| {
(
ReaderKey {
name: "test".to_string(),
path: PathBuf::from("test.tif"),
series: 0,
position: 0,
},
c,
z,
t,
)
};
cache.insert(
key(0, 0, 0),
Arc::new(Frame::from(Array2::from_elem((2, 2), 1u16))),
);
cache.insert(
key(0, 0, 1),
Arc::new(Frame::from(Array2::from_elem((2, 2), 2u16))),
);
assert!(cache.get(&key(0, 0, 0)).is_some());
cache.insert(
key(0, 0, 2),
Arc::new(Frame::from(Array2::from_elem((2, 2), 3u16))),
);
assert!(cache.get(&key(0, 0, 0)).is_some());
assert!(cache.get(&key(0, 0, 1)).is_none());
assert!(cache.get(&key(0, 0, 2)).is_some());
assert_eq!(cache.len(), 2);
Ok(())
}
#[test]
fn as_array_caches_frames() -> Result<(), Error> {
let file = "tiffseq/YTL1841B2-2-1_1hr_DMSO_galinduction_1";
let reader = open(file)?;
let view = reader.view();
let shape = view.shape();
let (c, z, t) = (shape[0] - 1, shape[1] - 1, shape[2] - 1);
let cache = FrameCache::global();
let a: Array5<usize> = view.clone().try_into()?;
assert!(cache.get(&view.frame_key(c, z, t)).is_some());
let b: Array5<usize> = view.clone().try_into()?;
assert_eq!(a, b);
// a fresh reader on the same file shares the same cache key
let view2 = open(file)?.view();
assert_eq!(view.frame_key(c, z, t), view2.frame_key(c, z, t));
let c2: Array5<usize> = view2.clone().try_into()?;
assert_eq!(a, c2);
Ok(())
}
}