fix type errors in cascaded model

Author: henrykrumb

ncalab/models/basicNCA.py

@@ -31,7 +31,7 @@ def __init__(
         use_alive_mask: bool = False,
         immutable_image_channels: bool = True,
         num_learned_filters: int = 2,
-        filter_padding: str = "circular",
+        filter_padding: str = "reflect",
         use_laplace: bool = False,
         kernel_size: int = 3,
         pad_noise: bool = False,

ncalab/models/cascadeNCA.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Dict, Optional, Tuple, List
 
 import numpy as np
 import torch  # type: ignore[import-untyped]
@@ -99,12 +99,13 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> Prediction:
         assert len(self.scales) > 0
         assert len(self.models) > 0
         assert len(self.steps) > 0
+        prediction = None
         x_scaled = downscale(x, self.scales[0])
         for i, (model, scale, scale_steps) in enumerate(
             zip(self.models, self.scales, self.steps)
         ):
             steps = scale_steps + np.random.randint(
-                -int(scale_steps * 0.2), int(scale_steps * 0.2)
+                -int(scale_steps * 0.2), int(scale_steps * 0.2) + 1
             )
             if steps <= 0:
                 steps = 1
@@ -117,12 +118,13 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> Prediction:
                     self.scales[i + 1],
                 )
         # TODO prediction has incorrect number of steps
-        return prediction
+        return unwrap(prediction)
 
     def record_steps(self, x: torch.Tensor):
         # TODO let "Prediction" class record steps
         step_outputs = []
         x_scaled = downscale(x, self.scales[0])
+        prediction = None
         for i, (model, scale, scale_steps) in enumerate(
             zip(self.models, self.scales, self.steps)
         ):
@@ -132,15 +134,15 @@ def record_steps(self, x: torch.Tensor):
                 step_outputs.append(upscale(prediction.output_image, scale))
                 x_in = prediction.output_image
             if i < len(self.scales) - 1:
-                x_scaled = upscale(prediction.output_image, scale / self.scales[i + 1])
+                x_scaled = upscale(unwrap(prediction).output_image, scale / self.scales[i + 1])
                 # replace input with downscaled variant of original image
                 x_scaled[:, : model.num_image_channels, :, :] = downscale(
                     x[:, : model.num_image_channels, :, :],
                     self.scales[i + 1],
                 )
         return step_outputs
 
-    def validate(self, image: torch.Tensor, label: torch.Tensor, steps: int = 1):
+    def validate(self, image: torch.Tensor, label: torch.Tensor, steps: int = 1) -> Optional[Tuple[Dict[str, float], Prediction]]:
         """
         Validation method.
 
@@ -175,4 +177,4 @@ def validate(self, image: torch.Tensor, label: torch.Tensor, steps: int = 1):
                     image[:, : model.num_image_channels, :, :],
                     self.scales[i + 1],
                 )
-        return metrics, prediction
+        return unwrap(metrics), unwrap(prediction)

ncalab/models/segmentationNCA.py

@@ -27,6 +27,7 @@ def __init__(
         num_learned_filters: int = 2,
         pad_noise: bool = True,
         autostepper: Optional[AutoStepper] = None,
+        filter_padding: str = "reflect",
         **kwargs,
     ):
         """
@@ -38,6 +39,7 @@ def __init__(
         :param hidden_size [int]: Number of neurons in hidden layer. Defaults to 128.
         :param learned_filters [int]: Number of learned filters. If 0, use sobel. Defaults to 2.
         :param pad_noise [bool]: Whether to pad input images with noise. Defaults to True.
+        :param filter_padding [str]: Padding type to use. Might affect reliance on spatial cues. Defaults to "circular".
         """
         self.num_classes = num_classes
         super(SegmentationNCAModel, self).__init__(
@@ -54,6 +56,7 @@ def __init__(
             num_learned_filters=num_learned_filters,
             pad_noise=pad_noise,
             autostepper=autostepper,
+            filter_padding=filter_padding,
             **kwargs,
         )
 
@@ -92,15 +95,15 @@ def metrics(self, pred: torch.Tensor, label: torch.Tensor) -> Dict[str, float]:
         """
         outputs = pred[:, self.num_image_channels + self.num_hidden_channels :, :, :]
         tp, fp, fn, tn = smp.metrics.get_stats(
-            outputs.cpu(),
+            outputs.cpu().float(),
             label[:, None, :, :].cpu().long(),
             mode="binary",
             threshold=0.1,
         )
-        tp = tp.squeeze()
-        fp = fp.squeeze()
-        fn = fn.squeeze()
-        tn = tn.squeeze()
+        tp = tp.squeeze().long()
+        fp = fp.squeeze().long()
+        fn = fn.squeeze().long()
+        tn = tn.squeeze().long()
         iou_score = smp.metrics.iou_score(
             tp,
             fp,