diff --git a/deeptrack/optical/augmentations.py b/deeptrack/optical/augmentations.py index faa2fd370..4e4927535 100644 --- a/deeptrack/optical/augmentations.py +++ b/deeptrack/optical/augmentations.py @@ -95,6 +95,7 @@ import numpy as np from deeptrack import utils, TORCH_AVAILABLE +from deeptrack.backend.core import DeepTrackDataDict from deeptrack.features import Feature from deeptrack.types import PropertyLike from deeptrack.optical.scatterers import ScatteredVolume, ScatteredField @@ -462,6 +463,37 @@ def __init__( self.feature = self.add_feature(feature) self.counter = 0 self.cache = [] + self._cache_dependency_data = [] + + @staticmethod + def _copy_data_dict(data: DeepTrackDataDict) -> DeepTrackDataDict: + """Create a shallow copy of a node data dictionary.""" + + copied = DeepTrackDataDict() + for key, data_object in data.dict.items(): + copied.create_index(key) + copied[key].store(data_object.current_value()) + if not data_object.is_valid(): + copied[key].invalidate() + + return copied + + def _snapshot_feature_data(self: Reuse) -> dict[Any, DeepTrackDataDict]: + """Snapshot cached values of the wrapped feature graph.""" + + return { + dependency: self._copy_data_dict(dependency.data) + for dependency in self.feature.recurse_dependencies() + } + + def _restore_feature_data( + self: Reuse, + snapshot: dict[Any, DeepTrackDataDict], + ) -> None: + """Restore cached values of the wrapped feature graph.""" + + for dependency, data in snapshot.items(): + dependency.data = self._copy_data_dict(data) def get( self: Reuse, @@ -503,11 +535,14 @@ def get( if recompute: output = self.feature(data) self.cache.append(output) + self._cache_dependency_data.append(self._snapshot_feature_data()) if len(self.cache) > storage: self.cache.pop(0) + self._cache_dependency_data.pop(0) else: index = self.counter % storage output = self.cache[index] + self._restore_feature_data(self._cache_dependency_data[index]) self.counter += 1 diff --git a/tests/test_augmentations.py b/tests/test_augmentations.py index a0a785734..560098d64 100644 --- a/tests/test_augmentations.py +++ b/tests/test_augmentations.py @@ -167,6 +167,37 @@ def get(self, data=None, **kwargs): self.assertEqual(counter["i"], 3) self.assertEqual(out6, 3) + def test_Reuse_preserves_take_properties_on_cached_output(self): + + config.set_backend("numpy") + + class ImageWithPosition(features.Feature): + __distributed__ = False + + def __init__(self): + super().__init__(position=lambda: np.array([1, 2])) + self.counter = 0 + + def get(self, data=None, position=None, **kwargs): + self.counter += 1 + return np.ones((2, 2)) * self.counter + + base_feature = ImageWithPosition() + reuse = augmentations.Reuse(base_feature, uses=2, storage=1) + label = features.TakeProperties(reuse, "position") + combined = reuse & label + + first = combined.update()() + second = combined.update()() + + self.assertIsInstance(first, list) + self.assertIsInstance(second, list) + self.assertEqual(len(first), 2) + self.assertEqual(len(second), 2) + np.testing.assert_array_equal(first[1], np.array([1, 2])) + np.testing.assert_array_equal(second[1], np.array([1, 2])) + np.testing.assert_array_equal(label(), np.array([1, 2])) + def test_FlipLR(self): backends = ["numpy"]