Multi-channel training and prediction

viscy/cli/cli.py

@@ -6,7 +6,7 @@
 from lightning.pytorch.loggers import TensorBoardLogger
 
 from viscy.light.data import HCSDataModule
-from viscy.light.engine import PhaseToNuc25D, VSTrainer
+from viscy.light.engine import VSTrainer, VSUNet
 
 
 class VSLightningCLI(LightningCLI):
@@ -45,7 +45,7 @@ def add_arguments_to_parser(self, parser):
 def main():
     torch.set_float32_matmul_precision("high")
     _ = VSLightningCLI(
-        model_class=PhaseToNuc25D,
+        model_class=VSUNet,
         datamodule_class=HCSDataModule,
         trainer_class=VSTrainer,
     )

viscy/light/engine.py

@@ -72,7 +72,7 @@ def export(
         logging.info(f"ONNX exported at {export_path}")
 
 
-class PhaseToNuc25D(LightningModule):
+class VSUNet(LightningModule):
     def __init__(
         self,
         model_config: dict = {},
@@ -144,7 +144,7 @@ def validation_step(self, batch, batch_idx):
         target = batch["target"]
         pred = self.forward(source)
         loss = self.loss_function(pred, target)
-        self.log("val_loss", loss, batch_size=self.batch_size, sync_dist=True)
+        self.log("loss/validate", loss, batch_size=self.batch_size, sync_dist=True)
         if batch_idx < self.log_num_samples:
             self.validation_step_outputs.append(
                 self._detach_sample((source, target, pred))
@@ -186,18 +186,20 @@ def configure_optimizers(self):
 
     @staticmethod
     def _detach_sample(imgs: Sequence[torch.Tensor]):
-        return [
-            np.squeeze(img[0].detach().cpu().numpy().max(axis=(0, 1))) for img in imgs
-        ]
+        return [np.squeeze(img[0].detach().cpu().numpy().max(axis=(1))) for img in imgs]
 
     def _log_samples(self, key: str, imgs: Sequence[Sequence[np.ndarray]]):
         images_grid = []
         for sample_images in imgs:
             images_row = []
-            for im, cm_name in zip(sample_images, ["gray"] + ["inferno"] * 2):
-                im = rescale_intensity(im, out_range=(0, 1))
-                rendered_im = get_cmap(cm_name)(im, bytes=True)[..., :3]
-                images_row.append(rendered_im)
+            for i, image in enumerate(sample_images):
+                cm_name = "gray" if i == 0 else "inferno"
+                if image.ndim == 2:
+                    image = image[np.newaxis]
+                for channel in image:
+                    channel = rescale_intensity(channel, out_range=(0, 1))
+                    render = get_cmap(cm_name)(channel, bytes=True)[..., :3]
+                    images_row.append(render)
             images_grid.append(np.concatenate(images_row, axis=1))
         grid = np.concatenate(images_grid, axis=0)
         self.logger.experiment.add_image(

viscy/light/predict_writer.py

@@ -3,54 +3,78 @@
 from typing import Literal, Sequence
 
 import torch
-from iohub.ngff import _pad_shape, open_ome_zarr
+from iohub.ngff import ImageArray, _pad_shape, open_ome_zarr
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import BasePredictionWriter
+from numpy.typing import DTypeLike
 
 from viscy.light.data import Sample
 
 
+def _resize_image(image: ImageArray, t_index: int, z_index: int):
+    """Resize image array if incoming T and Z index is not within bounds."""
+    if image.shape[0] <= t_index or image.shape[2] <= z_index:
+        logging.debug(
+            f"Resizing image '{image.name}' {image.shape} for T={t_index}, Z={z_index}."
+        )
+        image.resize(
+            max(t_index + 1, image.shape[0]),
+            image.channels,
+            max(z_index + 1, image.shape[1]),
+            *image.shape[-2:],
+        )
+
+
 class HCSPredictionWriter(BasePredictionWriter):
+    """Callback to store virtual staining predictions as HCS OME-Zarr.
+
+    :param str output_store: Path to the zarr store to store output
+    :param bool write_input: Write the source and target channels too
+        (must be writing to a new store),
+        defaults to False
+    :param Literal['batch', 'epoch', 'batch_and_epoch'] write_interval:
+        When to write, defaults to "batch"
+    """
+
     def __init__(
         self,
         output_store: str,
         write_input: bool = False,
         write_interval: Literal["batch", "epoch", "batch_and_epoch"] = "batch",
     ) -> None:
-        """Callback to store virtual staining predictions as HCS OME-Zarr.
-
-        :param str output_store: Path to the zarr store to store output
-        :param bool write_input: Write the source and target channels too
-            (must be writing to a new store),
-            defaults to False
-        :param Literal['batch', 'epoch', 'batch_and_epoch'] write_interval:
-            When to write, defaults to "batch"
-        """
         super().__init__(write_interval)
-        source_channel = "source"
-        target_channel = "target"
-        prediction_channel = "prediction"
-        if os.path.exists(output_store):
-            if write_input:
+        self.output_store = output_store
+        self.write_input = write_input
+
+    def on_predict_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        source_channel: list[str] = trainer.datamodule.source_channel
+        target_channel: list[str] = trainer.datamodule.target_channel
+        prediction_channel = [ch + "_prediction" for ch in target_channel]
+        if os.path.exists(self.output_store):
+            if self.write_input:
                 raise FileExistsError(
                     "Cannot write input to an existing store. Aborting."
                 )
             else:
-                self.plate = open_ome_zarr(output_store, mode="r+")
-                self.plate.append_channel(prediction_channel, resize_arrays=True)
+                self.plate = open_ome_zarr(self.output_store, mode="r+")
+                for _, pos in self.plate.positions():
+                    for ch in prediction_channel:
+                        pos.append_channel(ch, resize_arrays=True)
         else:
-            channel_names = [prediction_channel]
-            if write_input:
-                channel_names += [source_channel, target_channel]
+            channel_names = prediction_channel
+            if self.write_input:
+                channel_names = source_channel + target_channel + channel_names
             self.plate = open_ome_zarr(
-                output_store, layout="hcs", mode="a", channel_names=channel_names
+                self.output_store, layout="hcs", mode="a", channel_names=channel_names
             )
         logging.info(f"Writing prediction to: '{self.plate.zgroup.store.path}'.")
-        if write_input:
-            self.source_index = self.plate.get_channel_index(source_channel)
-            self.target_index = self.plate.get_channel_index(target_channel)
-        self.prediction_index = self.plate.get_channel_index(prediction_channel)
-        self.write_input = write_input
+        if self.write_input:
+            self.source_index = self._get_channel_indices(source_channel)
+            self.target_index = self._get_channel_indices(target_channel)
+        self.prediction_index = self._get_channel_indices(prediction_channel)
+
+    def _get_channel_indices(self, channel_names: list[str]) -> list[int]:
+        return [self.plate.get_channel_index(ch) for ch in channel_names]
 
     def write_on_batch_end(
         self,
@@ -62,6 +86,7 @@ def write_on_batch_end(
         batch_idx: int,
         dataloader_idx: int,
     ) -> None:
+        logging.debug(f"Writing batch {batch_idx}.")
         for sample_index, _ in enumerate(batch["index"][0]):
             self.write_sample(batch, prediction[sample_index], sample_index)
 
@@ -71,36 +96,37 @@ def on_predict_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
     def write_sample(
         self, batch: Sample, sample_prediction: torch.Tensor, sample_index: int
     ) -> None:
+        logging.debug(f"Writing sample {sample_index}.")
         sample_prediction = sample_prediction.cpu().numpy()
         img_name, t_index, z_index = [batch["index"][i][sample_index] for i in range(3)]
         t_index = int(t_index)
         z_index = int(z_index)
-        if img_name not in self.plate.zgroup:
-            _, row_name, col_name, pos_name, arr_name = img_name.split("/")
-            position = self.plate.create_position(row_name, col_name, pos_name)
-            shape = [1] + list(sample_prediction.shape)
-            shape[1] = len(position.channel_names)
-            image = position.create_zeros(
-                arr_name,
-                shape=shape,
-                dtype=sample_prediction.dtype,
-                chunks=_pad_shape(tuple(shape[-2:]), 5),
-            )
-        else:
-            image = self.plate[img_name]
-        if image.shape[0] <= t_index or image.shape[1] <= z_index:
-            image.resize(
-                max(t_index + 1, image.shape[0]),
-                image.channels,
-                max(z_index + 1, image.shape[1]),
-                *image.shape[-2:],
-            )
+        image = self._create_image(
+            img_name, sample_prediction.shape, sample_prediction.dtype
+        )
+        _resize_image(image, t_index, z_index)
         if self.write_input:
             # FIXME: should write center sclice of source
             image[t_index, self.source_index, z_index] = batch["source"][
                 sample_index
-            ].cpu()[0, 0]
+            ].cpu()[:, 0]
             image[t_index, self.target_index, z_index] = batch["target"][
                 sample_index
-            ].cpu()[0, 0]
-        image[t_index, self.prediction_index, z_index] = sample_prediction[0, 0]
+            ].cpu()[:, 0]
+        # write C1YX
+        image.oindex[t_index, self.prediction_index, z_index] = sample_prediction[:, 0]
+
+    def _create_image(self, img_name: str, shape: tuple[int], dtype: DTypeLike):
+        if img_name in self.plate.zgroup:
+            return self.plate[img_name]
+        logging.debug(f"Creating image '{img_name}'")
+        _, row_name, col_name, pos_name, arr_name = img_name.split("/")
+        position = self.plate.create_position(row_name, col_name, pos_name)
+        shape = [1] + list(shape)
+        shape[1] = len(position.channel_names)
+        return position.create_zeros(
+            arr_name,
+            shape=shape,
+            dtype=dtype,
+            chunks=_pad_shape(tuple(shape[-2:]), 5),
+        )

viscy/unet/networks/Unet25D.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from viscy.unet.networks.layers import ConvBlock3D
+from viscy.unet.networks.layers.ConvBlock3D import ConvBlock3D
 
 
 class Unet25d(nn.Module):