Fix MONAI-ITK orientation. Improve Writer docstrings.

Makefile

@@ -17,6 +17,7 @@ setup: poetry-download
 	source ~/.bashrc
 	source ~/.profile
 	poetry self add poetry-bumpversion@latest
+	poetry self add poetry-plugin-export@latest
 
 
 .PHONY: poetry-download
@@ -65,7 +66,7 @@ mypy:
 .PHONY: check-safety
 check-safety:
 	poetry check
-	poetry run safety check --full-report
+	poetry export | poetry run safety check --stdin
 	poetry run bandit -ll --recursive lighter tests
 
 .PHONY: lint

lighter/callbacks/writer/base.py

@@ -57,15 +57,14 @@ def writers(self) -> Dict[str, Callable]:
     def write(self, tensor: torch.Tensor, id: int) -> None:
         """
         Method to define how a tensor should be saved. The input tensor will be a single tensor without
-        the batch dimension. If the batch dimension is needed, apply `tensor.unsqueeze(0)` before saving,
-        either in this method or in the particular writer function.
+        the batch dimension.
 
         For each supported format, there should be a corresponding writer function registered in `self.writers`
         A specific writer function can be retrieved using `self.get_writer(self.format)`.
 
         Args:
-            tensor (torch.Tensor): Tensor to be saved. It will be a single tensor without the batch dimension.
-            id (int): Identifier for the tensor, can be used for naming or indexing.
+            tensor (torch.Tensor): tensor, without the batch dimension, to be saved.
+            id (int): identifier for the tensor, can be used for naming files or adding table records.
         """
 
     def setup(self, trainer: Trainer, pl_module: LighterSystem, stage: str) -> None:

lighter/callbacks/writer/file.py

@@ -3,7 +3,6 @@
 from functools import partial
 from pathlib import Path
 
-import monai
 import torch
 import torchvision
 from monai.data import metatensor_to_itk_image
@@ -16,14 +15,16 @@
 
 class LighterFileWriter(LighterBaseWriter):
     """
-    Writer for writing predictions to files. Supports multiple formats, and
-    additional custom formats can be added either through `additional_writers`
-    argument at initialization, or by calling `add_writer` method after initialization.
+    Writer for writing predictions to files. Supports writing tensors, images, videos, and ITK images.
+    To write to other formats, a custom writer function can be provided to the `writer` argument or,
+    for a more permanent solution, it can be added to the `self.writers` dictionary.
 
     Args:
         directory (Union[str, Path]): The directory where the files should be written.
-        writer (Union[str, Callable]): Name of the writer function registered in `self.writers`, or a custom writer function.
+        writer (Union[str, Callable]): Name of the writer function registered in `self.writers` or a custom writer function.
             Available writers: "tensor", "image", "video", "itk_nrrd", "itk_seg_nrrd", "itk_nifti".
+            A custom writer function must take two arguments: `path` and `tensor`, and write the tensor to the specified path.
+            `tensor` is a single tensor without the batch dimension.
     """
 
     def __init__(self, directory: Union[str, Path], writer: Union[str, Callable]) -> None:
@@ -45,9 +46,9 @@ def write(self, tensor: torch.Tensor, id: Union[int, str]) -> None:
         Write the tensor to the specified path in the given format.
 
         Args:
-            tensor (Tensor): The tensor to be written.
-            id (Union[int, str]): The identifier for naming.
-            format (str): Format in which tensor should be written.
+            tensor (Tensor): tensor, without the batch dimension, to be written.
+            id (Union[int, str]): identifier, used for file-naming.
+            format (str): format in which tensor should be written.
         """
         # Determine the path for the file based on prediction count. The suffix must be added by the writer function.
         path = self.directory / str(id)
@@ -82,10 +83,5 @@ def write_video(path, tensor):
 
 def write_itk_image(path: str, tensor: torch.Tensor, suffix) -> None:
     path = path.with_suffix(suffix)
-
-    # TODO: Remove this code when fixed https://github.com/Project-MONAI/MONAI/issues/6985
-    if tensor.meta["space"] == "RAS":
-        tensor.affine = monai.data.utils.orientation_ras_lps(tensor.affine)
-
     itk_image = metatensor_to_itk_image(tensor, channel_dim=0, dtype=tensor.dtype)
     OPTIONAL_IMPORTS["itk"].imwrite(itk_image, str(path), True)

lighter/callbacks/writer/table.py

@@ -17,8 +17,9 @@ class LighterTableWriter(LighterBaseWriter):
 
     Args:
         directory (Path): The directory where the CSV will be saved.
-        writer (Union[str, Callable]): Name of the writer function registered in `self.writers`, or a custom writer function.
-            Available writers: "tensor".
+        writer (Union[str, Callable]): Name of the writer function registered in `self.writers` or a custom writer function.
+            Available writers: "tensor". A custom writer function must take a single argument: `tensor`, and return the record
+            to be saved in the CSV file. The tensor will be a single tensor without the batch dimension.
     """
 
     def __init__(self, directory: Union[str, Path], writer: Union[str, Callable]) -> None:
@@ -36,12 +37,11 @@ def write(self, tensor: Any, id: Union[int, str]) -> None:
         Write the tensor as a table record using the specified writer.
 
         Args:
-            tensor (Any): The tensor to be written.
-            id (Union[int, str]): The identifier used as the key for the record.
+            tensor (Any): tensor, without the batch dimension, to be recorded.
+            id (Union[int, str]): identifier, used as the key for the record.
         """
         column = "pred"
         record = self.writer(tensor)
-
         self.csv_records.setdefault(id, {})[column] = record
 
     def on_predict_epoch_end(self, trainer: Trainer, pl_module: LighterSystem) -> None: