Feature(Write predictions to disk): Write strategies cache samples, to mirror source files organisation

src/careamics/lightning/callbacks/prediction_writer_callback/README.md

@@ -0,0 +1,119 @@
+# `PredictionWriterCallback`
+
+## Class diagram for `PredictionWriteCallback` related classes
+```mermaid
+classDiagram
+    PredictionWriterCallback*--WriteStrategy : composition
+    WriteStrategy<--WriteTiles : implements
+    WriteStrategy<--WriteImage : implements
+    WriteStrategy<--WriteZarrTiles : implements
+    WriteTiles*--TileCache : composition
+    WriteTiles*--SampleCache : composition
+    WriteImage*--SampleCache : composition
+
+    class PredictionWriterCallback
+    PredictionWriterCallback : +bool writing_predictions
+    PredictionWriterCallback : +WriteStrategy write_strategy
+    PredictionWriterCallback : +write_on_batch_end(...) 
+    PredictionWriterCallback : +on_predict_epoch_end(...)
+
+    class WriteStrategy
+    <<interface>> WriteStrategy
+    WriteStrategy : +write_batch(...)*
+    WriteStrategy : +set_file_data(lists~str~ write_filenames, list~int~ n_samples_per_file)*
+    WriteStrategy : +reset()*
+
+    class WriteTiles
+    WriteTiles : +WriteFunc write_func
+
+    WriteTiles : +TileCache tile_cache
+    WriteTiles : +SampleCache sample_cache 
+    WriteTiles : +write_batch(...)
+    WriteTiles : +set_file_data(lists~str~ write_filenames, list~int~ n_samples_per_file)
+    WriteTiles : +reset()
+
+    class WriteImage
+    WriteImage : +WriteFunc write_func
+    WriteImage : +SampleCache sample_cache
+    WriteImage : +write_batch(...)
+    WriteImage : +set_file_data(lists~str~ write_filenames, list~int~ n_samples_per_file)
+    WriteImage : +reset()
+
+    class WriteZarrTiles
+    WriteZarrTiles : +write_batch(...) NotImplemented
+    WriteZarrTiles : +set_file_data(lists~str~ write_filenames, list~int~ n_samples_per_file) NotImplemented
+    WriteZarrTiles : +reset() NotImplemented
+
+    class TileCache
+    TileCache : +list~NDArray~ array_cache
+    TileCache : +list~TileInformation~ tile_info_cache
+    TileCache : +add(NDArray, list~TileInformation~ item)
+    TileCache : +has_last_tile() bool
+    TileCache : +pop_image_tiles() NDArray, list~TileInformation~
+    TileCache : +reset()
+
+    class SampleCache
+    SampleCache : +list~int~ n_samples_per_file
+    SampleCache : +Iterator n_samples_iter
+    SampleCache : +int n_samples
+    SampleCache : +sample_cache list~NDArray~
+    SampleCache : +add(NDArray item)
+    SampleCache : +has_all_file_samples() bool
+    SampleCache : +pop_file_samples() list~NDArray~
+    SampleCache : +reset()
+```
+
+## Sequence diagram for writing tiles
+
+```mermaid
+sequenceDiagram
+    participant Trainer
+    participant PredictionWriterCallback
+    participant WriteTiles
+    participant TileCache
+
+    Trainer->>PredictionWriterCallback: write_on_batch_end(batch, ...)
+    activate PredictionWriterCallback
+    activate PredictionWriterCallback
+    PredictionWriterCallback->>WriteTiles: write_batch(batch, ...)
+    activate WriteTiles
+    activate WriteTiles
+    activate WriteTiles
+    WriteTiles->>TileCache: add(batch)
+    activate TileCache
+    WriteTiles ->> TileCache: has_last_tile()
+    TileCache -->> WriteTiles: True/False
+    deactivate TileCache
+    alt If does not have last tile
+        WriteTiles -->> PredictionWriterCallback: return
+        deactivate WriteTiles
+        PredictionWriterCallback -->> Trainer: return
+        deactivate PredictionWriterCallback
+    else  If has last tile
+        WriteTiles ->> TileCache: pop_image_tiles()
+        activate TileCache
+        TileCache -->> WriteTiles: tiles, tile_infos
+        deactivate TileCache
+        Note right of WriteTiles: Tiles are stitched to create prediction_image.
+        WriteTiles ->> SampleCache: add(prediction_image)
+        activate SampleCache
+        WriteTiles ->> SampleCache: has_all_file_samples()
+        SampleCache -->> WriteTiles: True/False
+        deactivate SampleCache
+        alt If does not have all samples
+            WriteTiles -->> PredictionWriterCallback: return
+            deactivate WriteTiles
+        else If has all samples
+            WriteTiles ->> SampleCache: pop_file_samples()
+            activate SampleCache
+            SampleCache -->> WriteTiles: samples
+            deactivate SampleCache
+            Note right of WriteTiles: Concatenated samples are written to disk.
+            WriteTiles -->> PredictionWriterCallback: return
+            deactivate WriteTiles
+        end
+        PredictionWriterCallback -->> Trainer: return
+        deactivate PredictionWriterCallback
+    end
+
+```

src/careamics/lightning/callbacks/prediction_writer_callback/__init__.py

@@ -5,14 +5,14 @@
     "create_write_strategy",
     "WriteStrategy",
     "WriteImage",
-    "CacheTiles",
+    "WriteTiles",
     "WriteTilesZarr",
     "select_write_extension",
     "select_write_func",
 ]
 
 from .prediction_writer_callback import PredictionWriterCallback
-from .write_strategy import CacheTiles, WriteImage, WriteStrategy, WriteTilesZarr
+from .write_strategy import WriteImage, WriteStrategy, WriteTiles, WriteTilesZarr
 from .write_strategy_factory import (
     create_write_strategy,
     select_write_extension,

src/careamics/lightning/callbacks/prediction_writer_callback/prediction_writer_callback.py

@@ -7,7 +7,6 @@
 
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import BasePredictionWriter
-from torch.utils.data import DataLoader
 
 from careamics.dataset import (
     IterablePredDataset,
@@ -127,7 +126,7 @@ def from_write_func_params(
         )
         return cls(write_strategy=write_strategy, dirpath=dirpath)
 
-    def _init_dirpath(self, dirpath):
+    def _init_dirpath(self, dirpath: Union[Path, str]):
         """
         Initialize directory path. Should only be called from `__init__`.
 
@@ -161,6 +160,7 @@ def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> Non
             Stage of training e.g. 'predict', 'fit', 'validate'.
         """
         super().setup(trainer, pl_module, stage)
+        # TODO: move directory creation to hook on_predict_epoch_start
         if stage == "predict":
             # make prediction output directory
             logger.info("Making prediction output directory.")
@@ -202,25 +202,6 @@ def write_on_batch_end(
         if not self.writing_predictions:
             return
 
-        dataloaders: Union[DataLoader, list[DataLoader]] = trainer.predict_dataloaders
-        dataloader: DataLoader = (
-            dataloaders[dataloader_idx]
-            if isinstance(dataloaders, list)
-            else dataloaders
-        )
-        dataset: ValidPredDatasets = dataloader.dataset
-        if not (
-            isinstance(dataset, IterablePredDataset)
-            or isinstance(dataset, IterableTiledPredDataset)
-        ):
-            # Note: Error will be raised before here from the source type
-            # This is for extra redundancy of errors.
-            raise TypeError(
-                "Prediction dataset has to be `IterableTiledPredDataset` or "
-                "`IterablePredDataset`. Cannot be `InMemoryPredDataset` because "
-                "filenames are taken from the original file."
-            )
-
         self.write_strategy.write_batch(
             trainer=trainer,
             pl_module=pl_module,
@@ -231,3 +212,21 @@ def write_on_batch_end(
             dataloader_idx=dataloader_idx,
             dirpath=self.dirpath,
         )
+
+    def on_predict_epoch_end(
+        self, trainer: Trainer, pl_module: LightningModule
+    ) -> None:
+        """
+        Lightning hook called at the end of prediction.
+
+        Resets write_strategy.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            PyTorch Lightning trainer.
+        pl_module : LightningModule
+            PyTorch Lightning module.
+        """
+        # reset write_strategy to prevent bugs if trainer.predict is called twice.
+        self.write_strategy.reset()