Skip to content

Multi-class merging with arbitrary class axis placement (resolves #20) #30

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
the-lay wants to merge 1 commit into
base: modernize
Choose a base branch
from
Open

Multi-class merging with arbitrary class axis placement (resolves #20) #30

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
104 changes: 80 additions & 24 deletions src/tiler/merger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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`.
Expand All @@ -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"""
Expand All @@ -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
Expand Down Expand Up @@ -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.

Expand All @@ -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)

Expand Down Expand Up @@ -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)):
Expand All @@ -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
Expand Down Expand Up @@ -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)]

Expand All @@ -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.
Expand All @@ -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)
Expand Down
125 changes: 125 additions & 0 deletions tests/test_merger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Loading