diff --git a/src/nimbus_inference/nimbus.py b/src/nimbus_inference/nimbus.py index 77c57ed..c2bb64f 100644 --- a/src/nimbus_inference/nimbus.py +++ b/src/nimbus_inference/nimbus.py @@ -403,5 +403,8 @@ def _stitch_tiles(self, tiles, padding): for b_ in range(b): stitched[b_, :, i * h : (i + 1) * h, j * w : (j + 1) * w] = tiles[i, j, b_] # remove padding - stitched = stitched[:, :, padding[0] : -padding[1], padding[2] : -padding[3]] + h_end = None if padding[1] == 0 else -padding[1] + w_end = None if padding[3] == 0 else -padding[3] + + stitched = stitched[:, :, padding[0]:h_end, padding[2]:w_end] return stitched diff --git a/tests/test_nimbus.py b/tests/test_nimbus.py index 2bd626e..a87c060 100644 --- a/tests/test_nimbus.py +++ b/tests/test_nimbus.py @@ -63,6 +63,19 @@ def test_tile_input(): assert tiled_input.shape == (3,3,1,2,512,512) assert padding == [192, 192, 192, 192] +def test_stitch_tiles_handles_zero_end_padding(): + nimbus = Nimbus(dataset="", output_dir="") + + # Shape: h_t, w_t, batch, channels, h, w + tiles = np.ones((1, 1, 1, 1, 10, 10)) + + # bottom and right padding are zero. + # Old slicing uses -0, which becomes 0 and collapses dimensions. + padding = [2, 0, 3, 0] + + stitched = nimbus._stitch_tiles(tiles, padding) + + assert stitched.shape == (1, 1, 8, 7) def test_tile_and_stitch(): # tests _tile_and_stitch which chains _tile_input, model.forward and _stitch_tiles