2D FCMAE

pyproject.toml

@@ -26,7 +26,7 @@ dynamic = ["version"]
 metrics = [
     "cellpose==2.1.0",
     "scikit-learn>=1.1.3",
-    "torchmetrics[detection]>=1.0.0",
+    "torchmetrics[detection]>=1.3.1",
     "ptflops>=0.7",
 ]
 visual = ["ipykernel", "graphviz", "torchview"]

tests/unet/test_fcmae.py

@@ -6,6 +6,7 @@
     MaskedConvNeXtV2Block,
     MaskedConvNeXtV2Stage,
     MaskedMultiscaleEncoder,
+    PixelToVoxelShuffleHead,
     generate_mask,
     masked_patchify,
     masked_unpatchify,
@@ -104,6 +105,13 @@ def test_masked_multiscale_encoder():
         assert afeat.shape[2] == afeat.shape[3] == xy_size // stride
 
 
+def test_pixel_to_voxel_shuffle_head():
+    head = PixelToVoxelShuffleHead(240, 3, out_stack_depth=5, xy_scaling=4)
+    x = torch.rand(2, 240, 16, 16)
+    y = head(x)
+    assert y.shape == (2, 3, 5, 64, 64)
+
+
 def test_fcmae():
     x = torch.rand(2, 3, 5, 128, 128)
     model = FullyConvolutionalMAE(3, 3)
@@ -113,3 +121,16 @@ def test_fcmae():
     y, m = model(x, mask_ratio=0.6)
     assert y.shape == x.shape
     assert m.shape == (2, 1, 128, 128)
+
+
+def test_fcmae_head_conv():
+    x = torch.rand(2, 3, 5, 128, 128)
+    model = FullyConvolutionalMAE(
+        3, 3, head_conv=True, head_conv_expansion_ratio=4, head_conv_pool=True
+    )
+    y, m = model(x)
+    assert y.shape == x.shape
+    assert m is None
+    y, m = model(x, mask_ratio=0.6)
+    assert y.shape == x.shape
+    assert m.shape == (2, 1, 128, 128)

viscy/data/combined.py

@@ -79,7 +79,6 @@ class ConcatDataModule(LightningDataModule):
     The concatenated data module will have the same
     batch size and number of workers as the first data module.
     Each element will be sampled uniformly regardless of their original data module.
-
     :param Sequence[LightningDataModule] data_modules: data modules to concatenate
     """
 
@@ -93,9 +92,11 @@ def __init__(self, data_modules: Sequence[LightningDataModule]):
                 raise ValueError("Inconsistent number of workers")
             if dm.batch_size != self.batch_size:
                 raise ValueError("Inconsistent batch size")
+        self.prepare_data_per_node = True
 
     def prepare_data(self):
         for dm in self.data_modules:
+            dm.trainer = self.trainer
             dm.prepare_data()
 
     def setup(self, stage: Literal["fit", "validate", "test", "predict"]):

viscy/data/ctmc_v1.py

@@ -10,9 +10,8 @@
 
 
 class CTMCv1ValidationDataset(SlidingWindowDataset):
-    subsample_rate: int = 30
 
-    def __len__(self) -> int:
+    def __len__(self, subsample_rate: int = 30) -> int:
         # sample every 30th frame in the videos
         return super().__len__() // self.subsample_rate
 

viscy/data/hcs.py

@@ -191,8 +191,6 @@ def __getitem__(self, index: int) -> Sample:
             sample_images["norm_meta"] = norm_meta
         if self.transform:
             sample_images = self.transform(sample_images)
-        # if isinstance(sample_images, list):
-        #     sample_images = sample_images[0]
         if "weight" in sample_images:
             del sample_images["weight"]
         sample = {

viscy/data/livecell.py

@@ -3,9 +3,11 @@
 
 import torch
 from lightning.pytorch import LightningDataModule
-from monai.transforms import Compose, Transform
+from monai.transforms import Compose, MapTransform
+from pycocotools.coco import COCO
 from tifffile import imread
 from torch.utils.data import DataLoader, Dataset
+from torchvision.ops import box_convert
 
 from viscy.data.typing import Sample
 
@@ -15,10 +17,10 @@ class LiveCellDataset(Dataset):
     LiveCell dataset.
 
     :param list[Path] images: List of paths to single-page, single-channel TIFF files.
-    :param Transform | Compose transform: Transform to apply to the dataset
+    :param MapTransform | Compose transform: Transform to apply to the dataset
     """
 
-    def __init__(self, images: list[Path], transform: Transform | Compose) -> None:
+    def __init__(self, images: list[Path], transform: MapTransform | Compose) -> None:
         self.images = images
         self.transform = transform
 
@@ -32,36 +34,100 @@ def __getitem__(self, idx: int) -> Sample:
         return {"source": image, "target": image}
 
 
+class LiveCellTestDataset(Dataset):
+    """
+    LiveCell dataset.
+
+    :param list[Path] images: List of paths to single-page, single-channel TIFF files.
+    :param MapTransform | Compose transform: Transform to apply to the dataset
+    """
+
+    def __init__(
+        self,
+        image_dir: Path,
+        transform: MapTransform | Compose,
+        annotations: Path,
+        load_target: bool = False,
+        load_labels: bool = False,
+    ) -> None:
+        self.image_dir = image_dir
+        self.transform = transform
+        self.coco = COCO(str(annotations))
+        self.image_ids = list(self.coco.imgs.keys())
+        self.load_target = load_target
+        self.load_labels = load_labels
+
+    def __len__(self) -> int:
+        return len(self.image_ids)
+
+    def __getitem__(self, idx: int) -> Sample:
+        image_id = self.image_ids[idx]
+        file_name = self.coco.imgs[image_id]["file_name"]
+        image_path = self.image_dir / file_name
+        image = imread(image_path)[None, None]
+        image = torch.from_numpy(image).to(torch.float32)
+        sample = Sample(source=image)
+        if self.load_target:
+            sample["target"] = image
+        if self.load_labels:
+            anns = self.coco.loadAnns(self.coco.getAnnIds(image_id)) or []
+            boxes = [torch.tensor(ann["bbox"]).to(torch.float32) for ann in anns]
+            masks = [
+                torch.from_numpy(self.coco.annToMask(ann)).to(torch.bool)
+                for ann in anns
+            ]
+            dets = {
+                "boxes": box_convert(torch.stack(boxes), in_fmt="xywh", out_fmt="xyxy"),
+                "labels": torch.zeros(len(anns)).to(torch.uint8),
+                "masks": torch.stack(masks),
+            }
+            sample["detections"] = dets
+            sample["file_name"] = file_name
+        self.transform(sample)
+        return sample
+
+
 class LiveCellDataModule(LightningDataModule):
     def __init__(
         self,
-        train_val_images: Path,
-        train_annotations: Path,
-        val_annotations: Path,
-        train_transforms: list[Transform],
-        val_transforms: list[Transform],
+        train_val_images: Path | None = None,
+        test_images: Path | None = None,
+        train_annotations: Path | None = None,
+        val_annotations: Path | None = None,
+        test_annotations: Path | None = None,
+        train_transforms: list[MapTransform] = [],
+        val_transforms: list[MapTransform] = [],
+        test_transforms: list[MapTransform] = [],
         batch_size: int = 16,
         num_workers: int = 8,
     ) -> None:
         super().__init__()
         self.train_val_images = Path(train_val_images)
         if not self.train_val_images.is_dir():
             raise NotADirectoryError(str(train_val_images))
+        self.test_images = Path(test_images)
+        if not self.test_images.is_dir():
+            raise NotADirectoryError(str(test_images))
         self.train_annotations = Path(train_annotations)
         if not self.train_annotations.is_file():
             raise FileNotFoundError(str(train_annotations))
         self.val_annotations = Path(val_annotations)
         if not self.val_annotations.is_file():
             raise FileNotFoundError(str(val_annotations))
+        self.test_annotations = Path(test_annotations)
+        if not self.test_annotations.is_file():
+            raise FileNotFoundError(str(test_annotations))
         self.train_transforms = Compose(train_transforms)
         self.val_transforms = Compose(val_transforms)
+        self.test_transforms = Compose(test_transforms)
         self.batch_size = batch_size
         self.num_workers = num_workers
 
     def setup(self, stage: str) -> None:
-        if stage != "fit":
-            raise NotImplementedError("Only fit stage is supported")
-        self._setup_fit()
+        if stage == "fit":
+            self._setup_fit()
+        elif stage == "test":
+            self._setup_test()
 
     def _parse_image_names(self, annotations: Path) -> list[Path]:
         with open(annotations) as f:
@@ -80,6 +146,14 @@ def _setup_fit(self) -> None:
             transform=self.val_transforms,
         )
 
+    def _setup_test(self) -> None:
+        self.test_dataset = LiveCellTestDataset(
+            self.test_images,
+            transform=self.test_transforms,
+            annotations=self.test_annotations,
+            load_labels=True,
+        )
+
     def train_dataloader(self) -> DataLoader:
         return DataLoader(
             self.train_dataset,
@@ -96,3 +170,8 @@ def val_dataloader(self) -> DataLoader:
             num_workers=self.num_workers,
             persistent_workers=bool(self.num_workers),
         )
+
+    def test_dataloader(self) -> DataLoader:
+        return DataLoader(
+            self.test_dataset, batch_size=self.batch_size, num_workers=self.num_workers
+        )

viscy/evaluation/evaluation_metrics.py

@@ -9,7 +9,7 @@
 from monai.metrics.regression import compute_ssim_and_cs
 from scipy.optimize import linear_sum_assignment
 from skimage.measure import label, regionprops
-from torchmetrics.detection import MeanAveragePrecision
+from torchmetrics.detection.mean_ap import MeanAveragePrecision
 from torchvision.ops import masks_to_boxes
 
 
@@ -172,7 +172,12 @@ def mean_average_precision(
         :py:class:`torchmetrics.detection.MeanAveragePrecision`
     :return dict[str, torch.Tensor]: COCO-style metrics
     """
-    map_metric = MeanAveragePrecision(box_format="xyxy", iou_type="segm", **kwargs)
+    defaults = dict(
+        iou_type="segm", box_format="xyxy", max_detection_thresholds=[1, 100, 10000]
+    )
+    if not kwargs:
+        kwargs = {}
+    map_metric = MeanAveragePrecision(**(defaults | kwargs))
     map_metric.update(
         [labels_to_detection(pred_labels)], [labels_to_detection(target_labels)]
     )

viscy/light/engine.py

@@ -146,6 +146,7 @@ def __init__(
         self.log_batches_per_epoch = log_batches_per_epoch
         self.log_samples_per_batch = log_samples_per_batch
         self.training_step_outputs = []
+        self.validation_losses = []
         self.validation_step_outputs = []
         # required to log the graph
         if architecture == "2D":
@@ -170,32 +171,49 @@ def __init__(
     def forward(self, x: Tensor) -> Tensor:
         return self.model(x)
 
-    def training_step(self, batch: Sample, batch_idx: int):
-        source = batch["source"]
-        target = batch["target"]
-        pred = self.forward(source)
-        loss = self.loss_function(pred, target)
+    def training_step(self, batch: Sample | Sequence[Sample], batch_idx: int):
+        losses = []
+        batch_size = 0
+        if not isinstance(batch, Sequence):
+            batch = [batch]
+        for b in batch:
+            source = b["source"]
+            target = b["target"]
+            pred = self.forward(source)
+            loss = self.loss_function(pred, target)
+            losses.append(loss)
+            batch_size += source.shape[0]
+            if batch_idx < self.log_batches_per_epoch:
+                self.training_step_outputs.extend(
+                    self._detach_sample((source, target, pred))
+                )
+        loss_step = torch.stack(losses).mean()
         self.log(
             "loss/train",
-            loss,
+            loss_step.to(self.device),
             on_step=True,
             on_epoch=True,
             prog_bar=True,
             logger=True,
             sync_dist=True,
+            batch_size=batch_size,
         )
-        if batch_idx < self.log_batches_per_epoch:
-            self.training_step_outputs.extend(
-                self._detach_sample((source, target, pred))
-            )
-        return loss
+        return loss_step
 
     def validation_step(self, batch: Sample, batch_idx: int, dataloader_idx: int = 0):
-        source = batch["source"]
-        target = batch["target"]
+        source: Tensor = batch["source"]
+        target: Tensor = batch["target"]
         pred = self.forward(source)
         loss = self.loss_function(pred, target)
-        self.log("loss/validate", loss, sync_dist=True, add_dataloader_idx=False)
+        if dataloader_idx + 1 > len(self.validation_losses):
+            self.validation_losses.append([])
+        self.validation_losses[dataloader_idx].append(loss.detach())
+        self.log(
+            f"loss/val/{dataloader_idx}",
+            loss.to(self.device),
+            sync_dist=True,
+            batch_size=source.shape[0],
+        )
         if batch_idx < self.log_batches_per_epoch:
             self.validation_step_outputs.extend(
                 self._detach_sample((source, target, pred))
@@ -305,8 +323,16 @@ def on_train_epoch_end(self):
         self.training_step_outputs = []
 
     def on_validation_epoch_end(self):
+        super().on_validation_epoch_end()
         self._log_samples("val_samples", self.validation_step_outputs)
         self.validation_step_outputs = []
+        # average within each dataloader
+        loss_means = [torch.tensor(losses).mean() for losses in self.validation_losses]
+        self.log(
+            "loss/validate",
+            torch.tensor(loss_means).mean().to(self.device),
+            sync_dist=True,
+        )
 
     def on_test_start(self):
         """Load CellPose model for segmentation."""
@@ -382,7 +408,6 @@ class FcmaeUNet(VSUNet):
     def __init__(self, fit_mask_ratio: float = 0.0, **kwargs):
         super().__init__(architecture="fcmae", **kwargs)
         self.fit_mask_ratio = fit_mask_ratio
-        self.validation_losses = []
 
     def forward(self, x: Tensor, mask_ratio: float = 0.0):
         return self.model(x, mask_ratio)
@@ -434,13 +459,3 @@ def validation_step(self, batch: Sample, batch_idx: int, dataloader_idx: int = 0
             self.validation_step_outputs.extend(
                 self._detach_sample((source, target * mask.unsqueeze(2), pred))
             )
-
-    def on_validation_epoch_end(self):
-        super().on_validation_epoch_end()
-        # average within each dataloader
-        loss_means = [torch.tensor(losses).mean() for losses in self.validation_losses]
-        self.log(
-            "loss/validate",
-            torch.tensor(loss_means).mean().to(self.device),
-            sync_dist=True,
-        )