Fix zero-padding edge case in tile stitching

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

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