diff --git a/README.md b/README.md index c1085de..2954f2d 100644 --- a/README.md +++ b/README.md @@ -77,7 +77,20 @@ for batch_id, batch in tiler(image, batch_size=10): # Final merging: applies tapering and optional unpadding final_image = merger.merge(unpad=True) # (3, 1920, 1080) ``` - + +Multi-class segmentation example — a model that consumes all channels and outputs class logits, +with the class axis placed wherever your framework puts it: +```python +# Input tiles are (32, 32, 4) channel-last, model output is (32, 32, 6) — 6 classes instead of channels +tiler = Tiler(data_shape=(1080, 1920, 4), tile_shape=(32, 32, 4), channel_dimension=-1) +merger = Merger(tiler, logits=6, logits_dim=-1, ignore_channels=True) + +for tile_id, tile in tiler(image): + merger.add(tile_id, model(tile)) + +segmentation = merger.merge(argmax=True) # (1080, 1920), argmax over the class axis +``` + Installation ------------- The latest release is available through pip: diff --git a/src/tiler/merger.py b/src/tiler/merger.py index e885c28..8069019 100644 --- a/src/tiler/merger.py +++ b/src/tiler/merger.py @@ -59,6 +59,8 @@ def __init__( tiler: Tiler, window: Union[None, str, np.ndarray] = None, logits: int = 0, + logits_dim: int = 0, + ignore_channels: bool = False, save_visits: bool = True, data_dtype: npt.DTypeLike = np.float32, weights_dtype: npt.DTypeLike = np.float32, @@ -76,7 +78,15 @@ def __init__( Must be one of `Merger.SUPPORTED_WINDOWS` or a numpy array with the same size as the tile. Default is None which creates a boxcar window (constant 1s). - logits (int): Specify whether to add logits dimensions in front of the data array. Default is `0`. + logits (int): Specify the number of classes/logits dimensions in the added tiles. Default is `0`. + + logits_dim (int): Specifies in which axis of the added tiles the logits are expected. + Negative indexing is allowed. Only used when `logits` is set. + Default is `0`, i.e. logits in front of the data, matching the behavior of earlier versions. + + ignore_channels (bool): If True, makes Merger expect tiles without the Tiler's channel dimension, + e.g. for models that consume all channels and output per-class or single-channel predictions. + Requires the Tiler to have `channel_dimension` set. Default is `False`. save_visits (bool): Specify whether to save which elements has been modified and how many times in `self.data_visits`. Can be disabled to save some memory. Default is `True`. @@ -94,12 +104,33 @@ def __init__( self.tiler = tiler """@private""" + # Channel dimension can be ignored, e.g. for models that consume channels and output classes + if ignore_channels and tiler.channel_dimension is None: + raise ValueError("ignore_channels requires the Tiler to have channel_dimension set.") + self.ignore_channels = ignore_channels + """@private""" + # Logits support if not isinstance(logits, int) or logits < 0: raise ValueError(f"Logits must be an integer 0 or a positive number ({logits}).") self.logits = int(logits) """@private""" + if self.logits: + if not isinstance(logits_dim, int): + raise ValueError(f"Logits dimension must be an integer ({logits_dim}).") + + # added tiles have one extra axis for logits, support negative indexing against that shape + out_n_dim = len(self._base_shape(self.tiler.tile_shape)) + 1 + if logits_dim >= out_n_dim or logits_dim < -out_n_dim: + raise ValueError(f"Logits dimension must be from {-out_n_dim} to {out_n_dim - 1} ({logits_dim}).") + if logits_dim < 0: + logits_dim = out_n_dim + logits_dim + elif logits_dim != 0: + raise ValueError("logits_dim is only used when logits is set.") + self.logits_dim = logits_dim + """@private""" + # Generate data and normalization arrays self.data = self.data_visits = self.weights_sum = None """@private""" @@ -114,6 +145,13 @@ def __init__( """@private""" self.set_window(window) + def _base_shape(self, shape) -> np.ndarray: + """Drops the channel axis from the given shape if `ignore_channels` is set.""" + shape = np.asarray(shape) + if self.ignore_channels: + shape = np.delete(shape, self.tiler.channel_dimension) + return shape + def _generate_window(self, window: str, shape: Union[Tuple, List]) -> np.ndarray: """Generate n-dimensional window according to the given shape. Adapted from: https://stackoverflow.com/a/53588640/1668421 @@ -183,6 +221,13 @@ def set_window(self, window: Union[None, str, np.ndarray] = None) -> None: else: raise ValueError(f"Unsupported type for window function ({type(window)}), expected str or np.ndarray.") + # The window that is actually applied in add(): the channel axis (always weighted 1s) is dropped + # when the added tiles do not carry it + if self.ignore_channels: + self._merge_window = np.take(self.window, 0, axis=self.tiler.channel_dimension) + else: + self._merge_window = self.window + def reset(self, save_visits: bool = True) -> None: """Reset data, weights and optional data_visits buffers. @@ -196,11 +241,11 @@ def reset(self, save_visits: bool = True) -> None: None """ - padded_data_shape = self.tiler._new_shape + padded_data_shape = self._base_shape(self.tiler._new_shape) # Image holds sum of all processed tiles multiplied by the window if self.logits: - self.data = np.zeros((self.logits, *padded_data_shape), dtype=self.data_dtype) + self.data = np.zeros(np.insert(padded_data_shape, self.logits_dim, self.logits), dtype=self.data_dtype) else: self.data = np.zeros(padded_data_shape, dtype=self.data_dtype) @@ -228,9 +273,11 @@ def add(self, tile_id: int, data: np.ndarray) -> None: ) data_shape = np.array(data.shape) - expected_tile_shape = ( - ((self.logits,) + tuple(self.tiler.tile_shape)) if self.logits > 0 else tuple(self.tiler.tile_shape) - ) + base_tile_shape = self._base_shape(self.tiler.tile_shape) + if self.logits > 0: + expected_tile_shape = tuple(np.insert(base_tile_shape, self.logits_dim, self.logits)) + else: + expected_tile_shape = tuple(base_tile_shape) if self.tiler.mode != "irregular": if not np.all(np.equal(data_shape, expected_tile_shape)): @@ -245,17 +292,24 @@ def add(self, tile_id: int, data: np.ndarray) -> None: # Select coordinates for data shape_diff = expected_tile_shape - data_shape - a, b = self.tiler.get_tile_bbox(tile_id, with_channel_dim=True) + if self.logits > 0: + # slicing and window are aligned with the tile axes, without the logits axis + spatial_diff = np.delete(shape_diff, self.logits_dim) + else: + spatial_diff = shape_diff + a, b = self.tiler.get_tile_bbox(tile_id, with_channel_dim=not self.ignore_channels) - sl = [slice(x, y - shape_diff[i]) for i, (x, y) in enumerate(zip(a, b))] - win_sl = [slice(None, -diff) if (diff > 0) else slice(None, None) for diff in shape_diff] + sl = [slice(x, y - spatial_diff[i]) for i, (x, y) in enumerate(zip(a, b))] + win_sl = [slice(None, -diff) if (diff > 0) else slice(None, None) for diff in spatial_diff] + win = self._merge_window[tuple(win_sl)] if self.logits > 0: - self.data[tuple([slice(None, None, None)] + sl)] += data * self.window[tuple(win_sl[1:])] - self.weights_sum[tuple(sl)] += self.window[tuple(win_sl[1:])] + data_sl = list(sl) + data_sl.insert(self.logits_dim, slice(None, None, None)) + self.data[tuple(data_sl)] += data * np.expand_dims(win, self.logits_dim) else: - self.data[tuple(sl)] += data * self.window[tuple(win_sl)] - self.weights_sum[tuple(sl)] += self.window[tuple(win_sl)] + self.data[tuple(sl)] += data * win + self.weights_sum[tuple(sl)] += win if self.data_visits is not None: self.data_visits[tuple(sl)] += 1 @@ -298,17 +352,15 @@ def _unpad(self, data: np.ndarray, extra_padding: Optional[List[Tuple[int, int]] ((before_1, after_1), … (before_N, after_N)) unique pad widths for each axis. Default is None. """ + base_data_shape = self._base_shape(self.tiler.data_shape) if extra_padding: - sl = [ - slice(pad_from, shape - pad_to) - for shape, (pad_from, pad_to) in zip(self.tiler.data_shape, extra_padding) - ] + sl = [slice(pad_from, shape - pad_to) for shape, (pad_from, pad_to) in zip(base_data_shape, extra_padding)] else: - sl = [slice(None, self.tiler.data_shape[i]) for i in range(len(self.tiler.data_shape))] + sl = [slice(None, shape) for shape in base_data_shape] - # if merger has logits dimension, add another slicing in front + # if merger has logits dimension, add another slicing for it if self.logits: - sl = [slice(None, None, None)] + sl + sl.insert(self.logits_dim, slice(None, None, None)) return data[tuple(sl)] @@ -330,8 +382,8 @@ def merge( ((before_1, after_1), … (before_N, after_N)) unique pad widths for each axis. Default is None. - argmax (bool): If argmax is True, the first dimension will be argmaxed. - Useful when merger is initialized with `logits=True`. + argmax (bool): If argmax is True, the logits dimension (`logits_dim`, first by default) + will be argmaxed. Useful when merger is initialized with `logits`. Default is False. normalize_by_weights (bool): If normalize is True, the accumulated data will be divided by weights. @@ -355,13 +407,17 @@ def merge( # ignoring should be more precise without atol # but can hide other errors with np.errstate(divide="ignore", invalid="ignore"): - data = np.nan_to_num(data / self.weights_sum) + weights = self.weights_sum + if self.logits: + # weights do not have the logits axis, expand for broadcasting + weights = np.expand_dims(weights, self.logits_dim) + data = np.nan_to_num(data / weights) if unpad: data = self._unpad(data, extra_padding) if argmax: - data = np.argmax(data, 0) + data = np.argmax(data, self.logits_dim) if dtype is not None: return data.astype(dtype) diff --git a/tests/test_merger.py b/tests/test_merger.py index 1c8dc2e..ba9e4fc 100644 --- a/tests/test_merger.py +++ b/tests/test_merger.py @@ -293,3 +293,128 @@ def test_merge(self): merger.add(t_id, t) np.testing.assert_equal(merger.merge(), padded_data) np.testing.assert_equal(merger.merge(extra_padding=padding), self.data) + + +class TestMulticlassMerging(unittest.TestCase): + """Multi-class (logits) merging with arbitrary class axis placement, see issue #20.""" + + def test_init_validation(self): + tiler = Tiler(data_shape=(100,), tile_shape=(10,)) + tiler_channel = Tiler(data_shape=(4, 96, 96), tile_shape=(4, 32, 32), channel_dimension=0) + + with self.assertRaises(ValueError): + Merger(tiler, logits=2, logits_dim="unsupported_type") + # 1D data + logits -> expected tile is 2D, so valid dims are -2..1 + with self.assertRaises(ValueError): + Merger(tiler, logits=2, logits_dim=2) + with self.assertRaises(ValueError): + Merger(tiler, logits=2, logits_dim=-3) + # logits_dim without logits is a no-op and most likely a mistake + with self.assertRaises(ValueError): + Merger(tiler, logits_dim=1) + # ignore_channels requires a channel dimension to ignore + with self.assertRaises(ValueError): + Merger(tiler, ignore_channels=True) + + # negative indexing is allowed + merger = Merger(tiler, logits=3, logits_dim=-1) + np.testing.assert_equal(merger.data.shape, (100, 3)) + merger = Merger(tiler_channel, logits=3, logits_dim=-1, ignore_channels=True) + np.testing.assert_equal(merger.data.shape, (96, 96, 3)) + + def test_logits_dim_last(self): + # 1D data, class axis last: model output (10, 3) per tile + tiler = Tiler(data_shape=(100,), tile_shape=(10,)) + merger = Merger(tiler, logits=3, logits_dim=-1) + + for tile_id in range(len(tiler)): + prediction = np.zeros((10, 3)) + prediction[:, tile_id % 3] = 1 + merger.add(tile_id, prediction) + + merged = merger.merge() + np.testing.assert_equal(merged.shape, (100, 3)) + argmaxed = merger.merge(argmax=True) + np.testing.assert_equal(argmaxed.shape, (100,)) + np.testing.assert_equal(argmaxed[:10], np.zeros(10)) + np.testing.assert_equal(argmaxed[10:20], np.ones(10)) + + def test_logits_dim_zero_matches_legacy_logits(self): + tiler = Tiler(data_shape=(100,), tile_shape=(10,)) + legacy = Merger(tiler, logits=3) + explicit = Merger(tiler, logits=3, logits_dim=0) + + for tile_id in range(len(tiler)): + prediction = np.zeros((3, 10)) + prediction[tile_id % 3] = 1 + legacy.add(tile_id, prediction) + explicit.add(tile_id, prediction) + + np.testing.assert_equal(legacy.merge(), explicit.merge()) + np.testing.assert_equal(legacy.merge(argmax=True), explicit.merge(argmax=True)) + + def test_ignore_channels_with_class_axis_last(self): + # The exact scenario from issue #20: channel-last input tiles, + # model replaces the channel axis with a class axis + tiler = Tiler(data_shape=(96, 96, 4), tile_shape=(32, 32, 4), channel_dimension=-1) + merger = Merger(tiler, logits=6, logits_dim=-1, ignore_channels=True) + + for tile_id in range(len(tiler)): + prediction = np.zeros((32, 32, 6)) + prediction[..., tile_id % 6] = 1 + merger.add(tile_id, prediction) + + merged = merger.merge() + np.testing.assert_equal(merged.shape, (96, 96, 6)) + + argmaxed = merger.merge(argmax=True) + np.testing.assert_equal(argmaxed.shape, (96, 96)) + # tiles are iterated row-major over the 3x3 mosaic + self.assertEqual(argmaxed[0, 0], 0) + self.assertEqual(argmaxed[0, 40], 1) + self.assertEqual(argmaxed[40, 0], 3) + + def test_ignore_channels_without_logits(self): + # model collapses channels into a single map: (4, 32, 32) in, (32, 32) out + tiler = Tiler(data_shape=(4, 96, 96), tile_shape=(4, 32, 32), channel_dimension=0) + merger = Merger(tiler, ignore_channels=True) + + for tile_id in range(len(tiler)): + merger.add(tile_id, np.full((32, 32), tile_id, dtype=np.float32)) + + merged = merger.merge() + np.testing.assert_equal(merged.shape, (96, 96)) + self.assertEqual(merged[0, 0], 0) + self.assertEqual(merged[40, 40], 4) + + def test_multiclass_with_window_and_overlap(self): + # normalization must stay correct when the window is non-trivial + # and the class axis is not the leading one + tiler = Tiler(data_shape=(64, 64), tile_shape=(16, 16), overlap=0.5) + merger = Merger(tiler, window="hann", logits=2, logits_dim=-1) + + for tile_id in range(len(tiler)): + merger.add(tile_id, np.ones((16, 16, 2))) + + merged = merger.merge() + np.testing.assert_equal(merged.shape, (64, 64, 2)) + visited = merger.merge(normalize_by_weights=True) + nonzero = visited[visited != 0] + np.testing.assert_allclose(nonzero, 1.0, rtol=1e-5) + + def test_multiclass_batch_add(self): + tiler = Tiler(data_shape=(100,), tile_shape=(10,)) + merger = Merger(tiler, logits=3, logits_dim=-1) + + batch = np.ones((5, 10, 3)) + merger.add_batch(0, 5, batch) + merger.add_batch(1, 5, batch) + np.testing.assert_equal(merger.merge().shape, (100, 3)) + + def test_reset_preserves_multiclass_shapes(self): + tiler = Tiler(data_shape=(100,), tile_shape=(10,)) + merger = Merger(tiler, logits=3, logits_dim=-1) + merger.add(0, np.ones((10, 3))) + merger.reset() + np.testing.assert_equal(merger.data.shape, (100, 3)) + self.assertEqual(np.count_nonzero(merger.data), 0)