Refactored extra padding system and updated examples

the-lay · the-lay · Dec 14, 2021 · Sep 27, 2021 · Sep 29, 2021 · Nov 8, 2021
commit c230be4bb205ba22041a498143d803d75033a6ac
diff --git a/docs/index.html b/docs/index.html
diff --git a/docs/tiler_teaser.png b/docs/tiler_teaser.png
diff --git a/examples/2d_overlap_tile.py b/examples/2d_overlap_tile.py
@@ -12,69 +12,68 @@
 
 # Loading image
 # Photo by Christian Holzinger on Unsplash: https://unsplash.com/photos/CUY_YHhCFl4
-image = np.array(Image.open('example_image.jpg'))  # 1280x1920x3
+image = np.array(Image.open("example_image.jpg"))  # 1280x1920x3
 
-# Padding image
-# Overlap tile strategy assumes we only use the small, non-overlapping, middle part of tile.
-# Assuming we want tiles of size 128x128 and we want to only use middle 64x64,
-# we should pad the image by 32 from each side, using reflect mode
-padded_image = np.pad(image, ((32, 32), (32, 32), (0, 0)), mode='reflect')
-
-# Specifying tiling parameters
+# Overlap tile strategy assumes we only use the small, non-overlapping, middle part of a tile.
+# Assuming we want to use centre 64x64 square, we should specify
 # The overlap should be 0.5, 64 or explicitly (64, 64, 0)
-tiler = Tiler(data_shape=padded_image.shape, tile_shape=(128, 128, 3),
-              overlap=(64, 64, 0), channel_dimension=2)
+tiler = Tiler(
+    data_shape=image.shape,
+    tile_shape=(128, 128, 3),
+    overlap=(64, 64, 0),
+    channel_dimension=2,
+)
+
+# Calculate and apply extra padding, as well as adjust tiling parameters
+new_shape, padding = tiler.calculate_padding()
+tiler.recalculate(data_shape=new_shape)
+padded_image = np.pad(image, padding, mode="reflect")
 
 # Specifying merging parameters
 # You can define overlap-tile window explicitly, i.e.
-# window = np.zeros((128, 128, 3))
-# window[32:-32, 32:-32, :] = 1
-# merger = Merger(tiler=tiler, window=window)
-# or you can use overlap-tile window which will do that automatically based on tiler.overlap
-merger = Merger(tiler=tiler, window='overlap-tile')
+# >>> window = np.zeros((128, 128, 3))
+# >>> window[32:-32, 32:-32, :] = 1
+# >>> merger = Merger(tiler=tiler, window=window)
+# or you can use window="overlap-tile"
+# it will automatically calculate such window based on tiler.overlap and applied padding
+merger = Merger(tiler=tiler, window="overlap-tile")
 
 # Let's define a function that will be applied to each tile
-def process(patch: np.ndarray, sanity_check: bool = True) -> np.ndarray:
+# For this example, let's black out the sides that should be "cropped" by window function
+# as a way to confirm that only the middle parts are being merged
+def process(patch: np.ndarray) -> np.ndarray:
+    patch[:32, :, :] = 0
+    patch[-32:, :, :] = 0
+    patch[:, :32, :] = 0
+    patch[:, -32:, :] = 0
+    return patch
 
-    # One example can be a sanity check
-    # Make the parts that should be removed black
-    # There should not appear any black spots in the final merged image
-    if sanity_check:
-        patch[:32, :, :] = 0
-        patch[-32:, :, :] = 0
-        patch[:, :32, :] = 0
-        patch[:, -32:, :] = 0
-        return patch
-
-    # Another example can be to just modify the whole patch
-    # Using PIL, we adjust the color balance
-    enhancer = ImageEnhance.Color(Image.fromarray(patch))
-    return np.array(enhancer.enhance(5.0))
 
 # Iterate through all the tile and apply the processing function
 # as well as add them back to the merger
-for tile_id, tile in tiler(padded_image):
+for tile_id, tile in tiler(padded_image, progress_bar=True):
     processed_tile = process(tile)
     merger.add(tile_id, processed_tile)
 
-# Merger.merge() returns unpadded from tiler image, but we still need to unpad line#21
-final_image = merger.merge().astype(np.uint8)
-final_unpadded_image = final_image[32:-32, 32:-32, :]
+# Merge processed tiles
+final_image = merger.merge(extra_padding=padding, dtype=image.dtype)
+
+print(f"Sanity check: {np.all(image == final_image)}")
 
 # Show the final merged image, weights and number of times each pixel was seen in tiles
-fig, ax = plt.subplots(3, 2, sharex=True, sharey=True)
-ax[0, 0].set_title('Original image')
+fig, ax = plt.subplots(3, 2)
+ax[0, 0].set_title("Original image")
 ax[0, 0].imshow(image)
-ax[0, 1].set_title('Final unpadded image')
-ax[0, 1].imshow(final_unpadded_image)
+ax[0, 1].set_title("Final merged image")
+ax[0, 1].imshow(final_image)
 
-ax[1, 0].set_title('Padded image')
+ax[1, 0].set_title("Padded image")
 ax[1, 0].imshow(padded_image)
-ax[1, 1].set_title('Merged image')
-ax[1, 1].imshow(final_image)
+ax[1, 1].set_title("Overlap-tile window")
+ax[1, 1].imshow(merger.window)
 
-ax[2, 0].set_title('Weights sum')
+ax[2, 0].set_title("Weights sum")
 ax[2, 0].imshow(merger.weights_sum[:, :, 0], vmin=0, vmax=merger.weights_sum.max())
-ax[2, 1].set_title('Pixel visits')
+ax[2, 1].set_title("Pixel visits")
 ax[2, 1].imshow(merger.data_visits[:, :, 0], vmin=0, vmax=merger.data_visits.max())
 plt.show()
diff --git a/examples/3d_overlap_tile.py b/examples/3d_overlap_tile.py
@@ -16,23 +16,27 @@
 volume[:75] *= np.linspace(3, 10, 75)[:, None, None]
 volume[75:] *= np.linspace(10, 3, 75)[:, None, None]
 
-# Let's assume we want to use tiles of size 48x48x48 and only the middle 20x20x20 for the final image
-# That means we need to pad the image by 14 from each side
-# To extrapolate missing context let's use reflect mode
-padded_volume = np.pad(volume, 14, mode='reflect')
+# Let's assume we want to use tiles of size 48x48x48 and only the middle 20x20x20 contribute to the final image
+# The overlap then should be 28 (48-20) voxels
+tiler = Tiler(
+    data_shape=volume.shape,
+    tile_shape=(48, 48, 48),
+    overlap=(28, 28, 28),
+)
 
-# Specifying tiling
-# The overlap should be 28 voxels
-tiler = Tiler(data_shape=padded_volume.shape,
-              tile_shape=(48, 48, 48),
-              overlap=(28, 28, 28))
+# Calculate and apply extra padding, as well as adjust tiling parameters
+new_shape, padding = tiler.calculate_padding()
+tiler.recalculate(data_shape=new_shape)
+padded_volume = np.pad(volume, padding, mode="reflect")
 
-# Window function for merging
-window = np.zeros((48, 48, 48))
-window[14:-14, 14:-14, 14:-14] = 1
-
-# Specifying merging
-merger = Merger(tiler=tiler, window=window)
+# Specifying merging parameters
+# You can define overlap-tile window explicitly, i.e.
+# >>> window = np.zeros((48, 48, 48))
+# >>> window[14:-14, 14:-14, 14:-14] = 1
+# >>> merger = Merger(tiler=tiler, window=window)
+# or you can use window="overlap-tile"
+# it will automatically calculate such window based on tiler.overlap and applied padding
+merger = Merger(tiler=tiler, window="overlap-tile")
 
 # Let's define a function that will be applied to each tile
 # For this example, let's black out the sides that should be "cropped" by window function
@@ -46,20 +50,25 @@ def process(patch: np.ndarray) -> np.ndarray:
     patch[:, :, -14:] = 0
     return patch
 
-# Iterate through all the tiles and apply the processing function and merge everything back
+
+# Apply the processing function to each tile and add it to the merger
 for tile_id, tile in tiler(padded_volume, progress_bar=True):
     processed_tile = process(tile)
     merger.add(tile_id, processed_tile)
 
-final_volume = merger.merge()
-final_unpadded_volume = final_volume[14:-14, 14:-14, 14:-14]
+# Merge processed tiles
+final_volume = merger.merge(extra_padding=padding)
+
+print(f"Sanity check: {np.all(volume == final_volume)}")
 
-# Show all the
-with napari.gui_qt():
-    v = napari.Viewer()
-    v.add_image(volume, name='Original volume')
-    v.add_image(padded_volume, name='Padded volume')
-    v.add_image(final_volume, name='Final volume')
-    v.add_image(final_unpadded_volume, name='Final unpadded volume')
-    v.add_image(merger.weights_sum, name='Merger weights sum')
-    v.add_image(merger.data_visits, name='Merger data visits')
+# Show all the produced volumes
+v = napari.Viewer()
+# v.add_image(merger.window, name="window")
+# v.add_image(tile, name="tile")
+# v.add_image(processed_tile, name="processed_tile")
+v.add_image(volume, name="Original volume")
+v.add_image(padded_volume, name="Padded volume")
+v.add_image(final_volume, name="Final merged volume")
+v.add_image(merger.weights_sum, name="Merger weights sum")
+v.add_image(merger.data_visits, name="Merger data visits")
+v.show(block=True)
diff --git a/examples/batch_processing.py b/examples/batch_processing.py
@@ -4,7 +4,7 @@
 
 # Loading image
 # Photo by Christian Holzinger on Unsplash: https://unsplash.com/photos/CUY_YHhCFl4
-image = np.array(Image.open('example_image.jpg'))  # 1280x1920x3
+image = np.array(Image.open("example_image.jpg"))  # 1280x1920x3
 
 # Setup Tiler and Merger
 tiler = Tiler(data_shape=image.shape, tile_shape=(200, 200, 3), channel_dimension=2)
@@ -13,25 +13,25 @@
 # Example 1: process all tiles one by one, i.e. batch_size=0
 for tile_i, tile in tiler(image, batch_size=0):
     merger.add(tile_i, tile)
-result_bs0 = merger.merge().astype(np.uint8)
+result_bs0 = merger.merge(dtype=image.dtype)
 
 # Example 2: process all tiles in batches of 1, i.e. batch_size=1
 merger.reset()
 for batch_i, batch in tiler(image, batch_size=1):
     merger.add_batch(batch_i, 1, batch)
-result_bs1 = merger.merge().astype(np.uint8)
+result_bs1 = merger.merge(dtype=image.dtype)
 
 # Example 3: process all tiles in batches of 10, i.e. batch_size=10
 merger.reset()
 for batch_i, batch in tiler(image, batch_size=10):
     merger.add_batch(batch_i, 10, batch)
-result_bs10 = merger.merge().astype(np.uint8)
+result_bs10 = merger.merge(dtype=image.dtype)
 
 # Example 4: process all tiles in batches of 10, but drop the batch that has <batch_size tiles, drop_last=True
 merger.reset()
 for batch_i, batch in tiler(image, batch_size=10, drop_last=True):
     merger.add_batch(batch_i, 10, batch)
-result_bs10 = merger.merge().astype(np.uint8)
+result_bs10 = merger.merge(dtype=image.dtype)
 
 assert np.all(result_bs0 == result_bs1)
 assert np.all(result_bs0 == result_bs10)
diff --git a/tests/test_merger.py b/tests/test_merger.py
@@ -1,8 +1,6 @@
 import unittest
 import numpy as np
 from tiler import Tiler, Merger
-from contextlib import redirect_stderr
-import os
 
 
 class TestMergingCommon(unittest.TestCase):
@@ -284,11 +282,16 @@ def test_merge(self):
             ),
         )
 
-        # Test explicit unpadding
+        # Test extra padding
         tiler = Tiler(data_shape=self.data.shape, tile_shape=(12,))
+        new_shape, padding = tiler.calculate_padding()
+        np.testing.assert_equal(new_shape, [112])
+        np.testing.assert_equal(padding, [(6, 6)])
+
+        tiler.recalculate(data_shape=new_shape)
+        padded_data = np.pad(self.data, pad_width=padding, mode="reflect")
         merger = Merger(tiler)
-        padded_data = tiler.apply_padding(self.data)
-        np.testing.assert_equal(tiler._padding, [(4, 4)])
         for t_id, t in tiler(padded_data):
             merger.add(t_id, t)
-        np.testing.assert_equal(merger.merge(unpad=True), self.data)
+        np.testing.assert_equal(merger.merge(), padded_data)
+        np.testing.assert_equal(merger.merge(extra_padding=padding), self.data)
diff --git a/tests/test_tiler.py b/tests/test_tiler.py
@@ -560,24 +560,28 @@ def test_tile_bbox_position(self):
             ([0, 90], [3, 100]), tiler2.get_tile_bbox_position(tile_id, True)
         )
 
-    def test_apply_padding(self):
-
-        t1 = Tiler(data_shape=self.data.shape, tile_shape=(3,), mode="reflect")
-
-        # without apply padding the padding is done based on data in the tile
-        # last tile has only one 99, but when reflect padded becomes 99, 99, 99
-        np.testing.assert_equal(t1.get_tile(self.data, 0), [0, 1, 2])
-        np.testing.assert_equal(t1.get_tile(self.data, len(t1) - 1), [99, 99, 99])
-
-        # with apply padding, padding should now correctly take into account data outside the tile
-        data = t1.apply_padding(self.data, mode="reflect")
-        np.testing.assert_equal(t1.get_tile(data, 0), [1, 0, 1])
-        np.testing.assert_equal(t1.get_tile(data, len(t1) - 1), [98, 99, 98])
-
-        # giving data with wrong (in this case old) data shape should raise an exception
-        with self.assertRaises(ValueError):
-            t1.apply_padding(self.data)
-
-        # re applying padding should show a warning
-        with self.assertWarns(Warning):
-            t1.apply_padding(data, mode="reflect")
+    def test_calculate_padding(self):
+
+        # no overlap, even
+        tiler = Tiler(data_shape=self.data.shape, tile_shape=(10,))
+        new_shape, padding = tiler.calculate_padding()
+        np.testing.assert_equal(new_shape, [110])
+        np.testing.assert_equal(padding, [(5, 5)])
+
+        # no overlap, odd
+        tiler = Tiler(data_shape=self.data.shape, tile_shape=(13,))
+        new_shape, padding = tiler.calculate_padding()
+        np.testing.assert_equal(new_shape, [113])
+        np.testing.assert_equal(padding, [(6, 7)])
+
+        # overlap, even
+        tiler = Tiler(data_shape=self.data.shape, tile_shape=(10,), overlap=0.2)
+        new_shape, padding = tiler.calculate_padding()
+        np.testing.assert_equal(new_shape, [108])
+        np.testing.assert_equal(padding, [(4, 4)])
+
+        # overlap, odd
+        tiler = Tiler(data_shape=self.data.shape, tile_shape=(10,), overlap=0.3)
+        new_shape, padding = tiler.calculate_padding()
+        np.testing.assert_equal(new_shape, [107])
+        np.testing.assert_equal(padding, [(3, 4)])