Add grid parameter to H5FileImageWriter constructor

ahcore/callbacks/h5_callback.py

@@ -130,6 +130,10 @@ def _batch_end(
             current_dataset: TiledWsiDataset
             current_dataset, _ = total_dataset.index_to_dataset(self._dataset_index)  # type: ignore
             slide_image = current_dataset.slide_image
+            if stage == "validate":
+                grid = current_dataset._grids[0][0]  # pylint: disable=protected-access
+            else:
+                grid = None  # During inference we don't have a grid around ROI
 
             data_description: DataDescription = pl_module.data_description  # type: ignore
             inference_grid: GridDescription = data_description.inference_grid
@@ -159,6 +163,7 @@ def _batch_end(
                 is_compressed_image=False,
                 progress=None,
                 precision=InferencePrecision(self._precision),
+                grid=grid,
             )
             new_process = Process(target=new_writer.consume, args=(self.generator(new_queue), child_conn))
             new_process.start()

ahcore/writers.py

@@ -62,10 +62,11 @@ def __init__(
         progress: Optional[Any] = None,
         extra_metadata: Optional[dict[str, Any]] = None,
         precision: InferencePrecision | None = None,
+        grid: Grid | None = None,
     ) -> None:
-        self._grid: Optional[Grid] = None
+        self._grid = grid
         self._grid_coordinates: Optional[npt.NDArray[np.int_]] = None
-        self._grid_offset: tuple[int, int] | None = None
+        self._grid_offset: npt.NDArray[np.int_] | None = None
         self._filename: Path = filename
         self._size: tuple[int, int] = size
         self._mpp: float = mpp
@@ -103,7 +104,7 @@ def init_writer(self, first_coordinates: GenericArray, first_batch: GenericArray
         self._current_index = 0
         # The grid can be smaller than the actual image when slide bounds are given.
         # As the grid should cover the image, the offset is given by the first tile.
-        self._grid_offset = list(first_coordinates[0])
+        self._grid_offset = np.array(first_coordinates[0])
 
         self._coordinates_dataset = h5file.create_dataset(
             "coordinates",
@@ -117,16 +118,19 @@ def init_writer(self, first_coordinates: GenericArray, first_batch: GenericArray
         # TODO: This would also support multiple grids
         # TODO: One would need to collect the coordinates and based on the first and the last
         # TODO: of a single grid, one can determine the grid, and the empty indices.
-        grid = Grid.from_tiling(
-            self._grid_offset,
-            size=self._size,
-            tile_size=self._tile_size,
-            tile_overlap=self._tile_overlap,
-            mode=TilingMode.overflow,
-            order=GridOrder.C,
-        )
+        if self._grid is None:  # During validation, the grid is passed as a parameter
+            grid = Grid.from_tiling(
+                self._grid_offset,
+                size=self._size,
+                tile_size=self._tile_size,
+                tile_overlap=self._tile_overlap,
+                mode=TilingMode.overflow,
+                order=GridOrder.C,
+            )
+        else:
+            grid = self._grid
+
         num_tiles = len(grid)
-        self._grid = grid
         self._tile_indices = h5file.create_dataset(
             "tile_indices",
             shape=(num_tiles,),