predtiler 0.0.1__py3-none-any.whl

predtiler/dataset.py

@@ -0,0 +1,53 @@
+
+from predtiler.tile_manager import TileIndexManager, TilingMode
+
+# class TilingDataset:
+#     def __init_subclass__(cls, parent_class=None, tile_manager=None, **kwargs):
+#         super().__init_subclass__(**kwargs)
+#         assert tile_manager is not None, 'tile_manager must be provided'
+#         cls.tile_manager = tile_manager
+#         if parent_class is not None:
+#             has_callable_method = callable(getattr(parent_class, 'patch_location', None))
+#             assert has_callable_method, f'{parent_class.__name__} must have a callable method with following signature: def patch_location(self, index)'
+#             cls.__bases__ = (parent_class,) + cls.__bases__
+
+#     def __len__(self):
+#         return self.tile_manager.total_grid_count()
+
+#     def patch_location(self, index):
+#         print('Calling patch_location')
+#         patch_loc_list = self.tile_manager.get_patch_location_from_dataset_idx(index)
+#         return patch_loc_list
+    
+
+# def get_tiling_dataset(dataset_class, tile_manager) -> type:
+#     class CorrespondingTilingDataset(TilingDataset, parent_class=dataset_class, tile_manager=tile_manager):
+#         pass
+    
+#     return CorrespondingTilingDataset
+
+def get_tiling_dataset(dataset_class, tile_manager) -> type:
+    has_callable_method = callable(getattr(dataset_class, 'patch_location', None))
+    assert has_callable_method, f'{dataset_class.__name__} must have a callable method with following signature: def patch_location(self, index)'
+
+    class TilingDataset(dataset_class):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.tile_manager = tile_manager
+
+        def __len__(self):
+            return self.tile_manager.total_grid_count()
+
+        def patch_location(self, index):
+            patch_loc_list = self.tile_manager.get_patch_location_from_dataset_idx(index)
+            return patch_loc_list
+    
+    return TilingDataset
+    
+
+
+
+def get_tile_manager(data_shape, tile_shape, patch_shape, tiling_mode=TilingMode.ShiftBoundary):
+    return TileIndexManager(data_shape, tile_shape, patch_shape, tiling_mode)
+
+

predtiler/tile_manager.py

@@ -0,0 +1,210 @@
+from dataclasses import dataclass
+
+import numpy as np
+
+
+class TilingMode:
+    """
+    Enum for the tiling mode.
+    """
+    TrimBoundary = 0
+    PadBoundary = 1
+    ShiftBoundary = 2
+
+@dataclass
+class TileIndexManager:
+    data_shape: tuple
+    grid_shape: tuple
+    patch_shape: tuple
+    tiling_mode: TilingMode
+
+    def __post_init__(self):
+        assert len(self.data_shape) == len(self.grid_shape), f"Data shape:{self.data_shape} and grid size:{self.grid_shape} must have the same dimension"
+        assert len(self.data_shape) == len(self.patch_shape), f"Data shape:{self.data_shape} and patch shape:{self.patch_shape} must have the same dimension"
+        innerpad = np.array(self.patch_shape) - np.array(self.grid_shape)
+        for dim, pad in enumerate(innerpad):
+            if pad < 0:
+                raise ValueError(f"Patch shape:{self.patch_shape} must be greater than or equal to grid shape:{self.grid_shape} in dimension {dim}")
+            if pad % 2 != 0:
+                raise ValueError(f"Patch shape:{self.patch_shape} must have even padding in dimension {dim}")
+
+    def patch_offset(self):
+        return (np.array(self.patch_shape) - np.array(self.grid_shape))//2
+    
+    def get_individual_dim_grid_count(self, dim:int):
+        """
+        Returns the number of the grid in the specified dimension, ignoring all other dimensions.
+        """
+        assert dim < len(self.data_shape), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+
+        if self.grid_shape[dim]==1 and self.patch_shape[dim]==1:
+            return self.data_shape[dim]
+        elif self.tiling_mode == TilingMode.PadBoundary:
+            return int(np.ceil(self.data_shape[dim] / self.grid_shape[dim]))
+        elif self.tiling_mode == TilingMode.ShiftBoundary:
+            excess_size = self.patch_shape[dim] - self.grid_shape[dim]
+            return int(np.ceil((self.data_shape[dim] - excess_size) / self.grid_shape[dim]))
+        else:
+            excess_size = self.patch_shape[dim] - self.grid_shape[dim]
+            return int(np.floor((self.data_shape[dim] - excess_size) / self.grid_shape[dim]))
+    
+    def total_grid_count(self):
+        """
+        Returns the total number of grids in the dataset.
+        """
+        return self.grid_count(0) * self.get_individual_dim_grid_count(0)
+    
+    def grid_count(self, dim:int):
+        """
+        Returns the total number of grids for one value in the specified dimension.
+        """
+        assert dim < len(self.data_shape), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        if dim == len(self.data_shape)-1:
+            return 1
+        
+        return self.get_individual_dim_grid_count(dim+1) * self.grid_count(dim+1)
+    
+    def get_grid_index(self, dim:int, coordinate:int):
+        """
+        Returns the index of the grid in the specified dimension.
+        """
+        assert dim < len(self.data_shape), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        assert coordinate < self.data_shape[dim], f"Coordinate {coordinate} is out of bounds for data shape {self.data_shape}"
+
+        if self.grid_shape[dim]==1 and self.patch_shape[dim]==1:
+            return coordinate
+        elif self.tiling_mode == TilingMode.PadBoundary: #self.trim_boundary is False:
+            return np.floor(coordinate / self.grid_shape[dim])
+        elif self.tiling_mode == TilingMode.TrimBoundary:
+            excess_size = (self.patch_shape[dim] - self.grid_shape[dim])//2
+            # can be <0 if coordinate is in [0,grid_shape[dim]]
+            return max(0, np.floor((coordinate - excess_size) / self.grid_shape[dim]))
+        elif self.tiling_mode == TilingMode.ShiftBoundary:
+            excess_size = (self.patch_shape[dim] - self.grid_shape[dim])//2
+            if coordinate + self.grid_shape[dim] + excess_size == self.data_shape[dim]:
+                return self.get_individual_dim_grid_count(dim) - 1
+            else:
+                # can be <0 if coordinate is in [0,grid_shape[dim]]
+                return max(0, np.floor((coordinate - excess_size) / self.grid_shape[dim]))
+
+        else:
+            raise ValueError(f"Unsupported tiling mode {self.tiling_mode}")
+        
+    def dataset_idx_from_grid_idx(self, grid_idx:tuple):
+        """
+        Returns the index of the grid in the dataset.
+        """
+        assert len(grid_idx) == len(self.data_shape), f"Dimension indices {grid_idx} must have the same dimension as data shape {self.data_shape}"
+        index = 0
+        for dim in range(len(grid_idx)):
+            index += grid_idx[dim] * self.grid_count(dim)
+        return index
+    
+    def get_patch_location_from_dataset_idx(self, dataset_idx:int):
+        """
+        Returns the patch location of the grid in the dataset.
+        """
+        grid_location = self.get_location_from_dataset_idx(dataset_idx)
+        offset = self.patch_offset()
+        return tuple(np.array(grid_location) - np.array(offset))
+
+    
+    def get_dataset_idx_from_grid_location(self, location:tuple):
+        assert len(location) == len(self.data_shape), f"Location {location} must have the same dimension as data shape {self.data_shape}"
+        grid_idx = [self.get_grid_index(dim, location[dim]) for dim in range(len(location))]
+        return self.dataset_idx_from_grid_idx(tuple(grid_idx))
+    
+    def get_gridstart_location_from_dim_index(self, dim:int, dim_index:int):
+        """
+        Returns the grid-start coordinate of the grid in the specified dimension.
+        """
+        assert dim < len(self.data_shape), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        assert dim_index < self.get_individual_dim_grid_count(dim), f"Dimension index {dim_index} is out of bounds for data shape {self.data_shape}"
+
+        if self.grid_shape[dim]==1 and self.patch_shape[dim]==1:
+            return dim_index
+        elif self.tiling_mode == TilingMode.PadBoundary:
+            return dim_index * self.grid_shape[dim]
+        elif self.tiling_mode == TilingMode.TrimBoundary:
+            excess_size = (self.patch_shape[dim] - self.grid_shape[dim])//2
+            return dim_index * self.grid_shape[dim] + excess_size
+        elif self.tiling_mode == TilingMode.ShiftBoundary:
+            excess_size = (self.patch_shape[dim] - self.grid_shape[dim])//2
+            if dim_index < self.get_individual_dim_grid_count(dim) - 1:
+                return dim_index * self.grid_shape[dim] + excess_size
+            else:
+                # on boundary. grid should be placed such that the patch covers the entire data.
+                return self.data_shape[dim] - self.grid_shape[dim] - excess_size
+        else:
+            raise ValueError(f"Unsupported tiling mode {self.tiling_mode}")
+ 
+    def get_location_from_dataset_idx(self, dataset_idx:int):
+        """
+        Returns the start location of the grid in the dataset.
+        """
+        grid_idx = []
+        for dim in range(len(self.data_shape)):
+            grid_idx.append(dataset_idx // self.grid_count(dim))
+            dataset_idx = dataset_idx % self.grid_count(dim)
+        location = [self.get_gridstart_location_from_dim_index(dim, grid_idx[dim]) for dim in range(len(self.data_shape))]
+        return tuple(location)
+    
+    def on_boundary(self, dataset_idx:int, dim:int, only_end:bool=False):
+        """
+        Returns True if the grid is on the boundary in the specified dimension.
+        """
+        assert dim < len(self.data_shape), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        
+        if dim > 0:
+            dataset_idx = dataset_idx % self.grid_count(dim-1)
+
+        dim_index = dataset_idx // self.grid_count(dim)
+        if only_end:
+            return dim_index == self.get_individual_dim_grid_count(dim) - 1
+        
+        return dim_index == 0 or dim_index == self.get_individual_dim_grid_count(dim) - 1
+    
+    def next_grid_along_dim(self, dataset_idx:int, dim:int):
+        """
+        Returns the index of the grid in the specified dimension in the specified direction.
+        """
+        assert dim < len(self.data_shape), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        new_idx = dataset_idx + self.grid_count(dim)
+        if new_idx >= self.total_grid_count():
+            return None
+        return new_idx
+    
+    def prev_grid_along_dim(self, dataset_idx:int, dim:int):
+        """
+        Returns the index of the grid in the specified dimension in the specified direction.
+        """
+        assert dim < len(self.data_shape), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        new_idx = dataset_idx - self.grid_count(dim)
+        if new_idx < 0:
+            return None
+
+if __name__ == '__main__':
+    # data_shape =   (1, 5, 103, 103,2)
+    # grid_shape =   (1, 1, 16,16, 2)
+    # patch_shape =  (1, 3, 32, 32, 2)
+    data_shape = (5, 5, 64, 64, 2)
+    grid_shape = (1, 1, 8, 8, 2)
+    patch_shape = (1, 3, 16, 16, 2)
+    tiling_mode = TilingMode.ShiftBoundary
+    manager = TileIndexManager(data_shape, grid_shape, patch_shape, tiling_mode)
+    gc = manager.total_grid_count()
+    for i in range(gc):
+        loc = manager.get_location_from_dataset_idx(i)
+        print(i, loc)
+        inferred_i = manager.get_dataset_idx_from_grid_location(loc)
+        assert i == inferred_i, f"Index mismatch: {i} != {inferred_i}"
+    
+    for i in range(5):
+        print(manager.on_boundary(40, i))

predtiler/tile_stitcher.py

@@ -0,0 +1,65 @@
+from typing import List
+
+import numpy as np
+
+from predtiler.tile_manager import TilingMode
+
+
+def stitch_predictions(predictions:np.ndarray, manager):
+    """
+    Args:
+        predictions: N*C*H*W or N*C*D*H*W numpy array where N is the number of datasets, C is the number of channels, H is the height, W is the width, D is the depth.
+        manager:
+    """
+    
+    mng = manager
+    shape = list(mng.data_shape)
+    shape.append(predictions.shape[1])
+    print(shape)
+
+    output = np.zeros(shape, dtype=predictions.dtype)
+    for dset_idx in range(predictions.shape[0]):
+        # grid start, grid end
+        gs = np.array(mng.get_location_from_dataset_idx(dset_idx), dtype=int)
+        ge = gs + mng.grid_shape
+
+        # patch start, patch end
+        ps = gs - mng.patch_offset()
+        pe = ps + mng.patch_shape
+
+        # valid grid start, valid grid end
+        vgs = np.array([max(0,x) for x in gs], dtype=int)
+        vge = np.array([min(x,y) for x,y in zip(ge, mng.data_shape)], dtype=int)
+        assert np.all(vgs ==gs)
+        assert np.all(vge ==ge)
+        
+        if mng.tiling_mode == TilingMode.ShiftBoundary:
+            for dim in range(len(vgs)):
+                if ps[dim] == 0:
+                    vgs[dim] = 0
+                if pe[dim] == mng.data_shape[dim]:
+                    vge[dim]= mng.data_shape[dim]
+
+        # relative start, relative end. This will be used on pred_tiled
+        rs = vgs - ps
+        re = rs + ( vge - vgs)
+        
+        for ch_idx in range(predictions.shape[1]):
+            if len(output.shape) == 4:
+                # channel dimension is the last one.
+                output[vgs[0]:vge[0],
+                    vgs[1]:vge[1],
+                    vgs[2]:vge[2],
+                    ch_idx] = predictions[dset_idx][ch_idx,rs[1]:re[1], rs[2]:re[2]]
+            elif len(output.shape) == 5:
+                # channel dimension is the last one.
+                assert vge[0] - vgs[0] == 1, 'Only one frame is supported'
+                output[vgs[0],
+                    vgs[1]:vge[1],
+                    vgs[2]:vge[2],
+                    vgs[3]:vge[3],
+                    ch_idx] = predictions[dset_idx][ch_idx, rs[1]:re[1], rs[2]:re[2], rs[3]:re[3]]
+            else:
+                raise ValueError(f'Unsupported shape {output.shape}')
+            
+    return output

predtiler-0.0.1.dist-info/METADATA

@@ -0,0 +1,122 @@
+Metadata-Version: 2.3
+Name: predtiler
+Version: 0.0.1
+Summary: Converting your dataset class into a class that can be used for tiled prediction and eventually obtain stiched prediction.
+Project-URL: homepage, https://github.com/ashesh-0/PredTiler
+Project-URL: repository, https://github.com/ashesh-0/PredTiler
+Author: Ashesh
+License: MIT
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Typing :: Typed
+Requires-Python: >=3.9
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: sybil; extra == 'dev'
+Provides-Extra: examples
+Requires-Dist: jupyter; extra == 'examples'
+Requires-Dist: matplotlib; extra == 'examples'
+Description-Content-Type: text/markdown
+
+A lean wrapper around your dataset class to enable tiled prediction. 
+
+[![License](https://img.shields.io/badge/license-MIT-blue)](https://github.com/ashesh-0/PredTiler/blob/main/LICENSE)
+[![CI](https://github.com/ashesh-0/PredTiler/actions/workflows/ci.yml/badge.svg)](https://github.com/ashesh-0/PredTiler/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/ashesh-0/PredTiler/graph/badge.svg?token=M655MOS7EL)](https://codecov.io/gh/ashesh-0/PredTiler)
+
+## Objective
+This package subclasses the dataset class you use to train your network.
+With PredTiler, you can use your dataset class as is, and PredTiler will take care of the tiling logic for you. 
+It will automatically generate patches in such a way that they can be tiled with the overlap of `(patch_size - tile_size)//2`. 
+We also provide a function to stitch the tiles back together to get the final prediction. 
+
+In case you are facing issues, feel free to raise an issue and I will be happy to help you out ! 
+In future, I plan to add detailed instructions for:
+1. multi-channel data
+2. 3D data
+3. Data being a list of numpy arrays, each poissibly having different shapes.
+
+## Installation
+
+```bash
+pip install predtiler
+```
+
+## Usage
+To work with PredTiler, the only requirement is that your dataset class must have a **patch_location(self, index)** method that returns the location of the patch at the given index. 
+Your dataset class should only use the location information returned by this method to return the patch. 
+PredTiler will override this method to return the location of the patches needed for tiled prediction. 
+
+Note that your dataset class could be arbitrarily complex (augmentations, returning multiple patches, working with 3D data, etc.). The only requirement is that it should use the crop present at the location returned by **patch_location** method. Below is an example of a simple dataset class that can be used with PredTiler.
+
+```python
+class YourDataset:
+    def __init__(self, data_path, patch_size=64) -> None:
+        self.patch_size = patch_size
+        self.data = load_data(data_path) # shape: (N, H, W, C)
+
+    def patch_location(self, index:int)-> Tuple[int, int, int]:
+        # it just ignores the index and returns a random location
+        n_idx = np.random.randint(0,len(self.data))
+        h = np.random.randint(0, self.data.shape[1]-self.patch_size)
+        w = np.random.randint(0, self.data.shape[2]-self.patch_size)
+        return (n_idx, h, w)
+    
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        n_idx, h, w = self.patch_location(index)
+        # return the patch at the location (patch_size, patch_size)
+        return self.data[n_idx, h:h+self.patch_size, w:w+self.patch_size]
+```
+
+## Getting overlapping patches needed for tiled prediction
+To use PredTiler, we need to get a new class that wraps around your dataset class.
+For this we also need a tile manager that will manage the tiles.
+
+```python
+
+from predtiler.dataset import get_tiling_dataset, get_tile_manager
+patch_size = 256
+tile_size = 128
+data_shape = (10, 2048, 2048) # size of the data you are working with
+manager = get_tile_manager(data_shape=data_shape, tile_shape=(1,tile_size,tile_size), 
+                               patch_shape=(1,patch_size,patch_size))
+    
+dset_class = get_tiling_dataset(YourDataset, manager)
+```
+
+At this point, you can use the `dset_class` as you would use `YourDataset` class. 
+
+```python
+data_path = ... # path to your data
+dset = dset_class(data_path, patch_size=patch_size)
+```
+
+## Stitching the predictions
+The benefit of using PredTiler is that it will automatically generate the patches in such a way that they can be tiled with the overlap of `(patch_size - tile_size)//2`. This allows you to use your dataset class as is, without worrying about the tiling logic.
+
+```python
+model = ... # your model
+predictions = []
+for i in range(len(dset)):
+    inp = dset[i]
+    inp = torch.Tensor(inp)[None,None]
+    pred = model(inp)
+    predictions.append(pred[0].numpy())
+
+predictions = np.stack(predictions) # shape: (number_of_patches, C, patch_size, patch_size)
+stitched_pred = stitch_predictions(predictions, dset.tile_manager)
+```
+
+

predtiler-0.0.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+predtiler/dataset.py,sha256=tHgVZksqJ2h2sw212x2g8SOW4Tf2TDZ_f6ibTpQPUMM,2136
+predtiler/tile_manager.py,sha256=V-iunzWcDw4_GytvkzhlrG8hLepf2_iSTBA3Da5nWgE,9958
+predtiler/tile_stitcher.py,sha256=rAgEV0y_FyRzU9RC_W11AE9OaGOrm1CgWPA-hbTh2vE,2336
+predtiler-0.0.1.dist-info/METADATA,sha256=cXyF36cQHVKxttq9zTUdbksOs9H4Pd5YYhqbF8NRPfA,5317
+predtiler-0.0.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+predtiler-0.0.1.dist-info/licenses/LICENSE,sha256=QyujMlqgcoASx9pVONwI8UO7xMSntGUEgQdlMqwVv4w,1063
+predtiler-0.0.1.dist-info/RECORD,,

predtiler-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.25.0
+Root-Is-Purelib: true
+Tag: py3-none-any

predtiler-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 ashesh
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.