Vendor pad shape function (#189)

* vendor the private function from iohub

* add test for the added function

* test full shape case

tests/translation/test_predict_writer.py

@@ -0,0 +1,8 @@
+from viscy.translation.predict_writer import _pad_shape
+
+
+def test_pad_shape():
+    assert _pad_shape((2, 3), 3) == (1, 2, 3)
+    assert _pad_shape((4, 5), 4) == (1, 1, 4, 5)
+    full_shape = tuple(range(1, 6))
+    assert _pad_shape(full_shape, 5) == full_shape

viscy/translation/predict_writer.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 import torch
-from iohub.ngff import ImageArray, Plate, Position, _pad_shape, open_ome_zarr
+from iohub.ngff import ImageArray, Plate, Position, open_ome_zarr
 from iohub.ngff_meta import TransformationMeta
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import BasePredictionWriter
@@ -17,6 +17,15 @@
 _logger = logging.getLogger("lightning.pytorch")
 
 
+def _pad_shape(shape: tuple[int, ...], target: int = 5) -> tuple[int, ...]:
+    """
+    Pad shape tuple to a target length.
+    Vendored from ``iohub.ngff.nodes._pad_shape()``.
+    """
+    pad = target - len(shape)
+    return (1,) * pad + shape
+
+
 def _resize_image(image: ImageArray, t_index: int, z_slice: slice) -> None:
     """Resize image array if incoming (1, C, Z, Y, X) stack is not within bounds."""
     if image.shape[0] <= t_index or image.shape[2] < z_slice.stop: