create class-based API for tensorboard visuals

Author: henrykrumb

ncalab/models/basicNCA.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Callable, Optional, Dict, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import numpy as np
 
@@ -10,6 +10,7 @@
 from ..autostepper import AutoStepper
 from ..prediction import Prediction
 from ..utils import pad_input
+from ..visualization import Visual
 
 
 class BasicNCAModel(nn.Module):
@@ -23,7 +24,7 @@ def __init__(
         num_image_channels: int,
         num_hidden_channels: int,
         num_output_channels: int,
-        plot_function: Optional[Callable] = None,
+        plot_function: Optional[Visual] = None,
         validation_metric: Optional[str] = None,
         fire_rate: float = 0.5,
         hidden_size: int = 128,
@@ -46,7 +47,7 @@ def __init__(
         :param num_output_channels [int]: Number of output channels.
         :param fire_rate [float]: Fire rate for stochastic weight update. Defaults to 0.5.
         :param hidden_size [int]: Number of neurons in hidden layer. Defaults to 128.
-        :param use_alive_mask [bool]: Whether to use alive masking during training. Defaults to False.
+        :param use_alive_mask [bool]: Whether to use alive masking (channel 3) during training. Defaults to False.
         :param immutable_image_channels [bool]: If image channels should be fixed during inference, which is the case for most segmentation or classification problems. Defaults to True.
         :param num_learned_filters [int]: Number of learned filters. If zero, use two sobel filters instead. Defaults to 2.
         :param filter_padding [str]: Padding type to use. Might affect reliance on spatial cues. Defaults to "circular".
@@ -80,16 +81,17 @@ def __init__(
         self.plot_function = plot_function
         self.validation_metric = validation_metric
 
+        # define input filters
         self._define_filters(num_learned_filters)
 
         # define model structure
-        self._define_network()
+        self.network = self._define_network().to(self.device)
 
     def _define_network(self):
         input_vector_size = self.num_channels * (self.num_filters + 1)
         if self.use_temporal_encoding:
             input_vector_size += 1
-        self.network = nn.Sequential(
+        network = nn.Sequential(
             nn.Conv2d(
                 in_channels=input_vector_size,
                 out_channels=self.hidden_size,
@@ -107,11 +109,11 @@ def _define_network(self):
                 padding=0,
                 kernel_size=1,
             ),
-        ).to(self.device)
-
+        )
         # initialize final layer with 0
         with torch.no_grad():
-            self.network[-1].weight.data.fill_(0)
+            network[-1].weight.data.fill_(0)
+        return network
 
     def _define_filters(self, num_learned_filters: int):
         """
@@ -134,9 +136,9 @@ def _define_filters(self, num_learned_filters: int):
                         padding_mode=self.filter_padding,
                         groups=self.num_channels,
                         bias=False,
-                    ).to(self.device)
+                    )
                 )
-            self.filters = nn.ModuleList(filters)
+            self.filters = nn.ModuleList(filters).to(self.device)
         else:
             sobel_x = np.outer([1, 2, 1], [-1, 0, 1]) / 8.0
             sobel_y = sobel_x.T
@@ -185,9 +187,13 @@ def _perceive_with(x, weight):
         perception = [x]
         perception.extend([_perceive_with(x, w) for w in self.filters])
         if self.use_temporal_encoding:
+            normalization = 100
+            if self.autostepper is not None:
+                normalization = self.autostepper.max_steps
             perception.append(
                 torch.mul(
-                    torch.ones((x.shape[0], 1, x.shape[2], x.shape[3])), step / 100
+                    torch.ones((x.shape[0], 1, x.shape[2], x.shape[3])),
+                    step / normalization,
                 ).to(self.device)
             )
         dx = torch.cat(perception, 1)
@@ -243,7 +249,6 @@ def forward(
                     x = x.permute(1, 0, 2, 3)  # C B W H --> B C W H
             return Prediction(self, steps, x)
 
-
         for step in range(self.autostepper.max_steps):
             if self.autostepper.check(step):
                 return Prediction(self, step, x)
@@ -347,3 +352,6 @@ def validate(
         prediction = self.predict(image.to(self.device), steps=steps)
         metrics = self.metrics(prediction.output_image, label.to(self.device))
         return metrics, prediction
+
+    def to_dict(self) -> Dict[str, Any]:
+        return dict()

ncalab/models/cascadeNCA.py

@@ -92,9 +92,7 @@ def __init__(self, backbone: BasicNCAModel, scales: List[int], steps: List[int])
         models = [backbone for _ in scales]
         self.models = nn.ModuleList(models)
 
-    def forward(
-        self, x: torch.Tensor, *args, **kwargs
-    ) -> Prediction:
+    def forward(self, x: torch.Tensor, *args, **kwargs) -> Prediction:
         """
         :param x [torch.Tensor]: Input image tensor, BCWH.
         :param steps [int]: Unused, as steps are defined in constructor.

ncalab/models/classificationNCA.py

@@ -1,12 +1,14 @@
-from typing import Dict
+from typing import Dict, Optional
 
 import torch  # type: ignore[import-untyped]
 import torch.nn.functional as F  # type: ignore[import-untyped]
 
 import torchmetrics
 import torchmetrics.classification
 
+from ..autostepper import AutoStepper
 from .basicNCA import BasicNCAModel
+from ..visualization import VisualBinaryImageClassification, VisualMultiImageClassification
 
 
 class ClassificationNCAModel(BasicNCAModel):
@@ -17,39 +19,66 @@ def __init__(
         num_hidden_channels: int,
         num_classes: int,
         fire_rate: float = 0.8,
+        hidden_size: int = 128,
         use_alive_mask: bool = False,
         pixel_wise_loss: bool = False,
+        num_learned_filters: int = 2,
         filter_padding: str = "reflect",
+        use_laplace: bool = False,
+        kernel_size: int = 3,
         pad_noise: bool = False,
+        autostepper: Optional[AutoStepper] = None,
+        use_temporal_encoding: bool = False,
         **kwargs,
     ):
         """
-        :param device [torch.device]: Compute device.
+        Constructor.
+
+        :param device [device]: Pytorch device descriptor.
         :param num_image_channels [int]: _description_
         :param num_hidden_channels [int]: _description_
         :param num_classes [int]: _description_
-        :param fire_rate [float]: _description_. Defaults to 0.8.
-        :param use_alive_mask [bool]: _description_. Defaults to False.
+        :param fire_rate [float]: Fire rate for stochastic weight update. Defaults to 0.8.
+        :param hidden_size [int]: Number of neurons in hidden layer. Defaults to 128.
+        :param use_alive_mask [bool]: Whether to use alive masking (channel 3) during training. Defaults to False.
         :param pixel_wise_loss [bool]: Whether a prediction per pixel is desired, like in self-classifying MNIST. Defaults to False.
-        :param filter_padding [str]: _description_. Defaults to "reflect".
-        :param pad_noise [bool]: _description_. Defaults to False.
+        :param num_learned_filters [int]: Number of learned filters. If zero, use two sobel filters instead. Defaults to 2.
+        :param filter_padding [str]: Padding type to use. Might affect reliance on spatial cues. Defaults to "circular".
+        :param pad_noise [bool]: Whether to pad input image tensor with noise in hidden / output channels
         """
         super(ClassificationNCAModel, self).__init__(
-            device,
-            num_image_channels,
-            num_hidden_channels,
-            num_classes,
+            device=device,
+            num_image_channels=num_image_channels,
+            num_hidden_channels=num_hidden_channels,
+            num_output_channels=num_classes,
             fire_rate=fire_rate,
+            hidden_size=hidden_size,
             use_alive_mask=use_alive_mask,
             immutable_image_channels=True,
             plot_function=None,
             validation_metric="accuracy_micro",
             filter_padding=filter_padding,
+            use_laplace=use_laplace,
+            kernel_size=kernel_size,
             pad_noise=pad_noise,
-            **kwargs,
+            autostepper=autostepper,
+            use_temporal_encoding=use_temporal_encoding,
         )
-        self.num_classes = num_classes
+        self._num_classes = num_classes
         self.pixel_wise_loss = pixel_wise_loss
+        if num_classes < 2:
+            self.plot_function = VisualBinaryImageClassification()
+        else:
+            self.plot_function = VisualMultiImageClassification()
+
+    @property
+    def num_classes(self) -> int:
+        return self._num_classes
+
+    @num_classes.setter
+    def num_classes(self, x: int):
+        self._num_classes = x
+        self.num_output_channels = x
 
     def classify(
         self, image: torch.Tensor, steps: int = 100, reduce: bool = False
@@ -181,7 +210,9 @@ def metrics(self, pred: torch.Tensor, label: torch.Tensor) -> Dict[str, float]:
         ]
         y_prob = class_channels
         y_prob = torch.mean(y_prob, dim=(2, 3))
-        y_true = label.squeeze(1)
+        y_true = label
+        if len(y_true.shape) == 2:
+            y_true = label.squeeze(1)
 
         accuracy_macro_metric.update(y_prob, y_true)
         accuracy_micro_metric.update(y_prob, y_true)

ncalab/models/depthNCA.py

@@ -2,7 +2,7 @@
 
 from .basicNCA import BasicNCAModel, AutoStepper
 
-from ..visualization import show_batch_depth
+from ..visualization import VisualDepthEstimation
 
 import torch  # type: ignore[import-untyped]
 import torch.nn as nn  # type: ignore[import-untyped]
@@ -83,7 +83,7 @@ def __init__(
             device,
             num_image_channels,
             num_hidden_channels,
-            plot_function=show_batch_depth,
+            plot_function=VisualDepthEstimation(),
             validation_metric="ssim",
             num_output_channels=1,
             fire_rate=fire_rate,

ncalab/models/growingNCA.py

@@ -7,7 +7,7 @@
 
 from .basicNCA import AutoStepper, BasicNCAModel
 from ..prediction import Prediction
-from ..visualization import show_batch_growing
+from ..visualization import VisualGrowing
 
 
 class GrowingNCAModel(BasicNCAModel):
@@ -40,7 +40,7 @@ def __init__(
             device,
             num_image_channels,
             num_hidden_channels,
-            plot_function=show_batch_growing,
+            plot_function=VisualGrowing(),
             num_output_channels=0,
             fire_rate=fire_rate,
             hidden_size=hidden_size,
@@ -99,11 +99,7 @@ def grow(
                     prediction = self.forward(x, steps=1)  # type: ignore[assignment]
                     step_outs.append(
                         np.clip(
-                            prediction.image_channels
-                            .squeeze(0)
-                            .detach()
-                            .cpu()
-                            .numpy(),
+                            prediction.image_channels.squeeze(0).detach().cpu().numpy(),
                             0,
                             1,
                         )
@@ -112,8 +108,6 @@ def grow(
                 return step_outs
             else:
                 prediction = self.forward(x, steps=steps)  # type: ignore[assignment]
-            out_np = (
-                prediction.image_channels.detach().cpu().numpy().squeeze(0)
-            )
+            out_np = prediction.image_channels.detach().cpu().numpy().squeeze(0)
             out_np = np.clip(out_np, 0, 1)
             return out_np

ncalab/models/segmentationNCA.py

@@ -5,7 +5,7 @@
 
 from .basicNCA import AutoStepper, BasicNCAModel
 from ..losses import DiceBCELoss
-from ..visualization import show_batch_binary_segmentation
+from ..visualization import VisualBinaryImageSegmentation
 
 
 class SegmentationNCAModel(BasicNCAModel):
@@ -46,7 +46,7 @@ def __init__(
             num_image_channels,
             num_hidden_channels,
             num_output_channels=num_classes,
-            plot_function=show_batch_binary_segmentation,
+            plot_function=VisualBinaryImageSegmentation(),
             validation_metric="Dice",
             fire_rate=fire_rate,
             hidden_size=hidden_size,
@@ -58,7 +58,6 @@ def __init__(
             **kwargs,
         )
 
-
     def loss(self, image: torch.Tensor, label: torch.Tensor) -> Dict[str, torch.Tensor]:
         """
         Compute Dice + BCE loss.
@@ -85,16 +84,14 @@ def loss(self, image: torch.Tensor, label: torch.Tensor) -> Dict[str, torch.Tens
         loss = loss_segmentation
         return {"total": loss}
 
-    def metrics(self, pred: torch.Tensor, label: torch.Tensor):
+    def metrics(self, pred: torch.Tensor, label: torch.Tensor) -> Dict[str, float]:
         """
         Return dict of standard evaluation metrics.
 
         :param pred [torch.Tensor]: Predicted image.
         :param label [torch.Tensor]: Ground truth label.
         """
-        outputs = pred[
-            :, self.num_image_channels + self.num_hidden_channels :, :, :
-        ]
+        outputs = pred[:, self.num_image_channels + self.num_hidden_channels :, :, :]
         tp, fp, fn, tn = smp.metrics.get_stats(
             outputs.cpu(),
             label[:, None, :, :].cpu().long(),