Repository: instanseg/instanseg (language: python)

Issue:
Move the code from `instanseg/instanseg.py` to `instanseg/__init__.py`
When I work with the [editable install](https://setuptools.pypa.io/en/latest/userguide/development_mode.html) of `instanseg` I get this error when running `from instanseg.utils import utils`:

```
ModuleNotFoundError: No module named 'instanseg.utils'; 'instanseg' is not a package 
```

The reason is that Python tries to import from the `instanseg.py` file and not from the `instanseg` package (folder). This can be easily fixed by moving all the code that is in [ `instanseg/instanseg.py`](https://github.com/instanseg/instanseg/blob/main/instanseg/instanseg.py) to [`instanseg/__init__.py`](https://github.com/instanseg/instanseg/blob/main/instanseg/__init__.py_) and removing the` instanseg/instanseg.py` file, which is also the best practice. 

Everything else would work the same. In general, it is better to avoid creating a Python file inside a package with the same name as the folder because the folder gives the name to the package. When you want to import code with `from instanseg import InstanSeg` the code for the class `InstanSeg` should reside in the `__init__.py` file. 

I would already have a PR to open, but I cannot tell if my suggestion would break something else. If you tell me this is possible I can open the PR immediately. Thanks


Discussion / hints:

Fix patch:
diff --git a/instanseg/inference_class.py b/instanseg/inference_class.py
index 2ff3730..1e1dbb5 100644
--- a/instanseg/inference_class.py
+++ b/instanseg/inference_class.py
@@ -1,13 +1,13 @@
-from typing import Union, List, Optional
+from typing import Union, List, Optional, Tuple
 import numpy as np
 import torch
 from torch import nn
 from torch.nn.functional import interpolate
 from pathlib import Path, PosixPath
+from tiffslide import TiffSlide
+import zarr
 import os
 
-import sys
-import pdb
 
 def to_ndim(x: torch.Tensor, n: int) -> torch.Tensor:
     """
@@ -65,7 +65,7 @@ def _rescale_to_pixel_size(image: torch.Tensor,
     return to_ndim(image, original_dim)
     
 
-def display_colourized(mIF):
+def _display_colourized(mIF):
     from instanseg.utils.utils import _move_channel_axis, generate_colors
 
     mIF = _to_tensor_float32(mIF)
@@ -89,25 +89,15 @@ def __init__(self,
                  device: Optional[str] = None, 
                  image_reader: str = "tiffslide",
                  github_token: str = os.environ.get("GITHUB_TOKEN"),
-                 verbosity = 1, #0,1,2
-                ):
+                 verbosity: int = 1 #0,1,2
+                 ):
         
         """
-        Parameters:
-            model_type (Union[str,nn.Module])
-                The type of model to use. If a string is provided, the model will be downloaded. 
-                If the model is not public, it will look for a model in your bioimageio folder. 
-                If an nn.Module is provided, this model will be used.
-
-            device (Optional[str]): 
-                The device to run the model on. If None, the device will be chosen automatically.
-
-            image_reader (str): 
-                The image reader to use. Options are "tiffslide", "skimage.io", "bioio", "AICSImageIO".
-
-            verbosity (int): 
-                The verbosity level. 0 is silent, 1 is normal, 2 is verbose.
-        
+        :param model_type: The type of model to use. If a string is provided, the model will be downloaded. If the model is not public, it will look for a model in your bioimageio folder. If an nn.Module is provided, this model will be used.
+        :param device: The device to run the model on. If None, the device will be chosen automatically.
+        :param image_reader: The image reader to use. Options are "tiffslide", "skimage.io", "bioio", "AICSImageIO".
+        :param github_token: The GitHub API token to use to authenticate model downloads. May be necessary to avoid rate limits in some circumstances.
+        :param verbosity: The verbosity level. 0 is silent, 1 is normal, 2 is verbose.
         """
         from instanseg.utils.utils import download_model, _choose_device
 
@@ -127,7 +117,12 @@ def __init__(self,
         self.medium_image_threshold = 10000 * 10000 #max number of image pixels that could be loaded in RAM.
         self.prediction_tag = "_instanseg_prediction"
 
-    def read_image(self, image_str: str):
+    def read_image(self, image_str: str) -> Union[Tuple[str, float], Tuple[np.ndarray, float]]:
+        """
+        Read an image file from disk.
+        :param image_str: The path to the image.
+        :return: The image array if it can be safely read (or the path to the image if it cannot) and the pixel size in microns.
+        """
         if self.prefered_image_reader == "tiffslide":
             from tiffslide import TiffSlide
             slide = TiffSlide(image_str)
@@ -177,7 +172,12 @@ def read_image(self, image_str: str):
 
         return image_array, img_pixel_size
     
-    def read_pixel_size(self,image_str: str):
+    def read_pixel_size(self,image_str: str) -> float:
+        """
+        Read the pixel size from an image on disk.
+        :param image_str: The path to the image.
+        :return: The pixel size in microns.
+        """
         try:
             from tiffslide import TiffSlide
             slide = TiffSlide(image_str)
@@ -223,9 +223,12 @@ def read_pixel_size(self,image_str: str):
         return None
 
 
-    def read_slide(self, image_str: str):      
+    def read_slide(self, image_str: str):
+        """
+        Read a whole slide image from disk.
+        :param image_str: The path to the image.
+        """
         if self.prefered_image_reader == "tiffslide":
-            from tiffslide import TiffSlide
             slide = TiffSlide(image_str)
         # elif self.prefered_image_reader == "AICSImageIO":
         #     from aicsimageio import AICSImage
@@ -255,14 +258,22 @@ def _normalise(self, image: torch.Tensor) -> torch.Tensor:
 
         return image
 
-    def eval(self, image: Union[str, List[str]], 
-            pixel_size: Optional[float] = None,
-            save_output: bool = False,
-            save_overlay: bool = False,
-            save_geojson: bool = False,
-            **kwargs):
+    def eval(self,
+             image: Union[str, List[str]], 
+             pixel_size: Optional[float] = None,
+             save_output: bool = False,
+             save_overlay: bool = False,
+             save_geojson: bool = False,
+             **kwargs) -> Union[torch.Tensor, List[torch.Tensor], None]:
         """
         Evaluate the input image or list of images using the InstanSeg model.
+        :param image: The path to the image, or a list of such paths.
+        :param pixel_size: The pixel size in microns.
+        :param save_output: Controls whether the output is saved to disk (see :func:`save_output <instanseg.Instanseg.save_output>`).
+        :param save_overlay: Controls whether the output is saved to disk as an overlay (see :func:`save_output <instanseg.Instanseg.save_output>`).
+        :param save_geojson: Controls whether the geojson output labels are saved to disk (see :func:`save_output <instanseg.Instanseg.save_output>`).
+        :param kwargs: Passed to other eval methods, eg :func:`save_output <instanseg.Instanseg.eval_small_image>`, :func:`save_output <instanseg.Instanseg.eval_medium_image>`, :func:`save_output <instanseg.Instanseg.eval_whole_slide_image>` 
+        :return: A torch.Tensor of outputs if the input is a path to a single image, or a list of such outputs if the input is a list of paths, or None if the input is a whole slide image.
         """
 
         if isinstance(image, PosixPath):
@@ -314,12 +325,20 @@ def eval(self, image: Union[str, List[str]],
         
         return output
     
-    def save_output(self,image_path: str, 
+    def save_output(self,
+                    image_path: str, 
                     labels: torch.Tensor,
                     image_array: Optional[np.ndarray] = None,
                     save_overlay = False,
-                    save_geojson = False,):
-        
+                    save_geojson = False) -> None:
+        """
+        Save the output of InstanSeg to disk.
+        :param image_path: The path to the image, and where outputs will be saved.
+        :param labels: The output labels.
+        :param image_array: The image in array format. Required to save overlay.
+        :param save_overlay: Save the labels overlaid on the image.
+        :param save_geojson: Save the labels as a GeoJSON feature collection.
+        """
         import os
 
         if isinstance(image_path, str):
@@ -367,18 +386,26 @@ def save_output(self,image_path: str,
 
 
 
-    def eval_small_image(self,image: torch.Tensor, pixel_size: Optional[float] = None,
-                          normalise: bool = True,
-                          return_image_tensor: bool = True,
-                          target = "all_outputs", #or "nuclei" or "cells"
-                          rescale_output: bool = True,
-                          **kwargs):
+    def eval_small_image(self,
+                         image: torch.Tensor,
+                         pixel_size: Optional[float] = None,
+                         normalise: bool = True,
+                         return_image_tensor: bool = True,
+                         target: str = "all_outputs", #or "nuclei" or "cells"
+                         rescale_output: bool = True,
+                         **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         """
-        Evaluate the input image using the InstanSeg model.
+        Evaluate a small input image using the InstanSeg model.
+        
+        :param image:: The input image(s) to be evaluated.
+        :param pixel_size: The pixel size of the image, in microns. If not provided, it will be read from the image metadata.
+        :param normalise: Controls whether the image is normalised.
+        :param return_image_tensor: Controls whether the input image is returned as part of the output.
+        :param target: Controls what type of output is given, usually "all_outputs", "nuclei", or "cells".
+        :param rescale_output: Controls whether the outputs should be rescaled to the same coordinate space as the input (useful if the pixel size is different to that of the InstanSeg model being used).
+        :param kwargs: Passed to pytorch.
         
-        Args:
-            image (torch.Tensor): The input image(s) to be evaluated.
-            pixel_size (Optional[float]): The pixel size of the image. If not provided, it will be read from the image metadata.
+        :return: A tensor corresponding to the output targets specified, as well as the input image if requested.
         """
         from instanseg.utils.utils import percentile_normalize
 
@@ -428,23 +455,40 @@ def eval_small_image(self,image: torch.Tensor, pixel_size: Optional[float] = Non
         else:
             return instances.cpu()
 
-    def eval_medium_image(self,image: torch.Tensor, 
+    def eval_medium_image(self,
+                          image: torch.Tensor, 
                           pixel_size: Optional[float] = None, 
                           normalise: bool = True,
                           tile_size: int = 512,
                           batch_size: int = 1,
-                        return_image_tensor: bool = True,
-                        normalisation_subsampling_factor: int = 1,
-                        target = "all_outputs", #or "nuclei" or "cells"
-                        rescale_output: bool = True,
-                          **kwargs):
+                          return_image_tensor: bool = True,
+                          normalisation_subsampling_factor: int = 1,
+                          target: str = "all_outputs", #or "nuclei" or "cells"
+                          rescale_output: bool = True,
+                          **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Evaluate a medium input image using the InstanSeg model. The image will be split into tiles, and then inference and object merging will be handled internally.
+        
+        :param image:: The input image(s) to be evaluated.
+        :param pixel_size: The pixel size of the image, in microns. If not provided, it will be read from the image metadata.
+        :param normalise: Controls whether the image is normalised.
+        :param tile_size: The width/height of the tiles that the image will be split into.
+        :param batch_size: The number of tiles to be run simultaneously.
+        :param return_image_tensor: Controls whether the input image is returned as part of the output.
+        :param normalisation_subsampling_factor: The subsampling or downsample factor at which to calculate normalisation parameters.
+        :param target: Controls what type of output is given, usually "all_outputs", "nuclei", or "cells".
+        :param rescale_output: Controls whether the outputs should be rescaled to the same coordinate space as the input (useful if the pixel size is different to that of the InstanSeg model being used).
+        :param kwargs: Passed to pytorch.
         
+        :return: A tensor corresponding to the output targets specified, as well as the input image if requested.
+        """
+
         from instanseg.utils.utils import percentile_normalize, _move_channel_axis
 
     
         image = _to_tensor_float32(image)
         
-        from instanseg.utils.tiling import sliding_window_inference
+        from instanseg.utils.tiling import _sliding_window_inference
         original_shape = image.shape
         original_ndim = image.dim()
 
@@ -478,16 +522,15 @@ def eval_medium_image(self,image: torch.Tensor,
         else:
             target_segmentation = torch.tensor([1,1])
 
-        instances = sliding_window_inference(image,
-                    self.instanseg, 
-                    window_size = (tile_size,tile_size),
-                    sw_device = self.inference_device,
-                    device = 'cpu', 
-                    batch_size= batch_size,
-                    output_channels = output_dimension,
-                    show_progress= self.verbose,
-                    target_segmentation = target_segmentation,
-                    **kwargs).float()
+        instances = _sliding_window_inference(image,
+                                              self.instanseg,
+                                              window_size = (tile_size,tile_size),sw_device = self.inference_device,
+                                              device = 'cpu', 
+                                              batch_size= batch_size,
+                                              output_channels = output_dimension,
+                                              show_progress= self.verbose,
+                                              target_segmentation = target_segmentation,
+                                              **kwargs).float()
 
         instances = to_ndim(instances, 4)
         image = to_ndim(image, 4)
@@ -506,36 +549,219 @@ def eval_medium_image(self,image: torch.Tensor,
         else:
             return instances.cpu()
         
-    def eval_whole_slide_image(self, image: str, pixel_size: Optional[float] = None, 
-                               normalise: bool = True, 
-                               batch_size: int = 1,
+    def eval_whole_slide_image(self,
+                               image: str,
+                               pixel_size: Optional[float] = None, 
+                               normalise: bool = True,
+                               normalisation_subsampling_factor: int = 10,
                                tile_size: int = 512,
+                               overlap: int = 100,
+                               detection_size: int = 20, 
+                               batch_size: int = 1,
                                output_geojson: bool = False,
                                **kwargs):
-            from instanseg.utils.tiling import segment_image_larger_than_memory
-                                 
-            segment_image_larger_than_memory(instanseg_class= self,
-                                            image_path= image, 
-                                            memory_block_size = (int(self.medium_image_threshold**0.5),int(self.medium_image_threshold **0.5)),  #this is the size of the image that will be read in memory
-                                            inference_tile_size = (tile_size,tile_size), #this is the size of the image that will be passed to the model
-                                            batch_size= batch_size,
-                                            to_geojson= output_geojson, 
-                                            pixel_size =pixel_size,
-                                            prediction_tag = self.prediction_tag, 
-                                            sw_device = self.inference_device,
-                                            normalisation_subsling_factor=10,
-                                            **kwargs)
+            """
+            Evaluate a whole slide input image using the InstanSeg model. This function uses slideio to read an image and then segments it using the instanseg model. The segmentation is done in a tiled manner to avoid memory issues. 
+            
+            :param image:: The input image to be evaluated.
+            :param pixel_size: The pixel size of the image, in microns. If not provided, it will be read from the image metadata.
+            :param normalise: Controls whether the image is normalised.
+            :param tile_size: The width/height of the tiles that the image will be split into.
+            :param overlap: The overlap (in pixels) betwene tiles.
+            :param detection_size: The expected maximum size of detection objects.
+            :param batch_size: The number of tiles to be run simultaneously.
+            :param normalisation_subsampling_factor: The subsampling or downsample factor at which to calculate normalisation parameters.
+            :param target: Controls what type of output is given, usually "all_outputs", "nuclei", or "cells".
+            :param rescale_output: Controls whether the outputs should be rescaled to the same coordinate space as the input (useful if the pixel size is different to that of the InstanSeg model being used).
+            :param kwargs: Passed to pytorch.
+            :return: Returns a zarr file with the segmentation. The zarr file is saved in the same directory as the image with the same name but with the extension .zarr.
+            """
+
+            memory_block_size =  (int(self.medium_image_threshold**0.5), int(self.medium_image_threshold **0.5))
+            inference_tile_size = (tile_size,tile_size)
+
+            from itertools import product
+            from instanseg.utils.pytorch_utils import torch_fastremap, match_labels
+            from pathlib import Path
+            from tqdm import tqdm
+            from instanseg.utils.tiling import _chops, _remove_edge_labels, _zarr_to_json_export
+            
+            instanseg = self.instanseg
+
+            image, img_pixel_size = self.read_image(image)
+            slide = self.read_slide(image)
+
+            n_dim = 2 if instanseg.cells_and_nuclei else 1
+            model_pixel_size = instanseg.pixel_size
+
+            new_stem = Path(image).stem + self.prediction_tag
+            file_with_zarr_extension = Path(image).parent / (new_stem + ".zarr")
+
+
+            if img_pixel_size > 1 or img_pixel_size < 0.1:
+                import warnings
+                warnings.warn("The image pixel size {} is not in microns.".format(img_pixel_size))
+                if pixel_size is not None:
+                    img_pixel_size = pixel_size
+                else:
+                    raise ValueError("The image pixel size {} is not in microns.".format(img_pixel_size))
+            
+            scale_factor = model_pixel_size/img_pixel_size
+
+            dims = slide.dimensions
+            dims = (int(dims[1]/ scale_factor), int(dims[0]/scale_factor)) #The dimensions are opposite to numpy/torch/zarr dimensions.
+
+            pad2 = overlap + detection_size
+            pad = overlap
+
+            shape = memory_block_size
+            
+            chop_list = _chops(dims, shape, overlap=2*pad2)
+
+            chunk_shape = (n_dim,shape[0],shape[1])
+            store = zarr.DirectoryStore(file_with_zarr_extension) 
+            canvas = zarr.zeros((n_dim,dims[0],dims[1]), chunks=chunk_shape, dtype=np.int32, store=store, overwrite = True)
+
+            running_max = 0
+
+            total = len(chop_list[0]) * len(chop_list[1])
+            for _, ((i, window_i), (j, window_j)) in tqdm(enumerate(product(enumerate(chop_list[0]), enumerate(chop_list[1]))), total=total, colour = "green", desc = "Slide progress: "):
+
+            #  input_data = scene.read_block((int(window_j*scale_factor), int(window_i*scale_factor), int(shape[0]*scale_factor), int(shape[1]*scale_factor)), size = shape)
+
+                best_level = slide.get_best_level_for_downsample(scale_factor)
+                downsample_factor = slide.level_downsamples[best_level]
+
+                initial_pixel_size = img_pixel_size
+                itermediate_pixel_size = initial_pixel_size * downsample_factor
+                final_pixel_size = model_pixel_size
+
+                intermediate_to_final = final_pixel_size/itermediate_pixel_size
+                
+                # Calculate the size of the region needed at the base level to get the desired output size
+                intermediate_shape = (int(shape[0] * intermediate_to_final), int(shape[1] * intermediate_to_final))
+
+                input_data = slide.read_region((int(window_j*scale_factor), int(window_i*scale_factor)), best_level, (int(intermediate_shape[0]) , int(intermediate_shape[1])), as_array=True)
+            
+                input_tensor = self._to_tensor(input_data)
+
+                new_tile = self.eval_medium_image(input_tensor,
+                                                  pixel_size = itermediate_pixel_size,
+                                                  tile_size = inference_tile_size[0],
+                                                  batch_size = batch_size,
+                                                  return_image_tensor = False,
+                                                  normalise = normalise,
+                                                  normalisation_subsampling_factor = normalisation_subsampling_factor,
+                                                )
+                
             
+
+                from torch.nn.functional import interpolate
+                new_tile = interpolate(new_tile, size=shape[-2:], mode="nearest").int()[0]
+
+            # pdb.set_trace()
+
+                num_iter = new_tile.shape[0]
+
+                for n in range(num_iter):
+                    
+                    ignore_list = []
+                    if i == 0:
+                        ignore_list.append("top")
+                    if j == 0:
+                        ignore_list.append("left")
+                    if i == len(chop_list[0])-1:
+                        ignore_list.append("bottom")
+                    if j == len(chop_list[1])-1:
+                        ignore_list.append("right")
+
+                    if i == len(chop_list[0])-1 and j == len(chop_list[1])-1:
+                        tile1 = canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]]
+                        tile2 = _remove_edge_labels(new_tile[n,pad:shape[0],pad: shape[1]], ignore = ignore_list)
+
+                    elif i == len(chop_list[0])-1:
+                        tile1 = canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]]
+                        tile2 = _remove_edge_labels(new_tile[n,pad:shape[0],pad : shape[1]], ignore =ignore_list)
+
+                    elif j == len(chop_list[1])-1:
+                        tile1 = canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]]
+                        tile2 = _remove_edge_labels(new_tile[n,pad:shape[0],pad: shape[1]], ignore = ignore_list)
+
+                    elif i == 0 and j == 0:
+                        tile1 = canvas[n, window_i  :window_i + shape[0], window_j :window_j + shape[1]]
+                        tile2 = _remove_edge_labels(new_tile[n, :shape[0], : shape[1]], ignore = ignore_list)
+                    elif i == 0:
+                        tile1 = canvas[n, window_i  :window_i + shape[0], window_j + pad :window_j + shape[1]]
+                        tile2 = _remove_edge_labels(new_tile[n, :shape[0],pad : shape[1]], ignore = ignore_list)
+
+                    elif j == 0:
+                        tile1 = canvas[n, window_i  + pad:window_i + shape[0], window_j:window_j + shape[1]]
+                        tile2 = _remove_edge_labels(new_tile[n,pad :shape[0],: shape[1]], ignore = ignore_list)
+
+                    if j == 0 or i == 0 or j == len(chop_list[1])-1 or i == len(chop_list[0])-1:
+
+                        tile2 = torch_fastremap(tile2)
+                        tile2[tile2>0] = tile2[tile2>0] + running_max
+
+                        tile1_torch = torch.tensor(np.array(tile1), dtype = torch.int32)
+
+                        remapped = match_labels(tile1_torch, tile2, threshold = 0.1)[1]
+                        tile1_torch[remapped>0] = remapped[remapped>0].int()
+
+                        running_max = max(running_max, tile1_torch.max())
+
+                        if i == len(chop_list[0])-1 and j == len(chop_list[1])-1:
+                            canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
+                        elif i == len(chop_list[0])-1:
+                            canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
+                        elif j == len(chop_list[1])-1:
+                            canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
+                        elif i == 0 and j == 0:
+                            canvas[n, window_i  :window_i + shape[0], window_j :window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
+                        elif i == 0:
+                            canvas[n, window_i  :window_i + shape[0], window_j + pad :window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
+                        elif j == 0:
+                            canvas[n, window_i  + pad:window_i + shape[0], window_j:window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
+
+                    else:
+                        
+                        tile1 = canvas[n, window_i + pad:window_i + shape[0] - pad, window_j + pad:window_j + shape[1] - pad]
+                        tile2 = _remove_edge_labels(new_tile[n,pad:shape[0] -pad,pad: shape[1]-pad])
+                        
+                        tile2 = torch_fastremap(tile2)
+
+                        tile2[tile2>0] = tile2[tile2>0] + running_max
+
+                        tile1_torch = torch.tensor(np.array(tile1), dtype = torch.int32)
+                        remapped = match_labels(tile1_torch, tile2, threshold = 0.1)[1]
+
+                        tile1_torch[remapped>0] = remapped[remapped>0].int()
+                        running_max = max(running_max, tile1_torch.max())
+                    
+                        canvas[n, window_i + pad:window_i + shape[0] - pad, window_j + pad:window_j + shape[1] - pad] = tile1_torch.numpy().astype(np.int32)
+
+            if output_geojson:
+                print("Exporting to geojson")
+                _zarr_to_json_export(file_with_zarr_extension, detection_size = detection_size, size = shape[0], scale = scale_factor, n_dim = n_dim)
+                    
+
+
     
-    def display(self, image: torch.tensor,
-               instances: torch.Tensor):
-        
+    def display(self,
+                image: torch.tensor,
+                instances: torch.Tensor) -> np.ndarray:
+        """
+        Save the output of an InstanSeg model overlaid on the input.
+        See :func:`save_image_with_label_overlay <instanseg.utils.save_image_with_label_overlay>` for more details and return types.
+        :param image: The input image.
+        :param instances: The output labels.
+        """
         from instanseg.utils.utils import save_image_with_label_overlay
 
         if isinstance(image, torch.Tensor):
             image = image.cpu().detach().numpy()
 
-        im_for_display = display_colourized(image.squeeze())
+        im_for_display = _display_colourized(image.squeeze())
 
         output_dimension = instances.shape[1]
 
@@ -551,7 +777,7 @@ def display(self, image: torch.tensor,
         return image_overlay
     
 
-    def cluster_instances_by_mean_channel_intensity(self, image_tensor: torch.Tensor, labeled_output: torch.Tensor):
+    def _cluster_instances_by_mean_channel_intensity(self, image_tensor: torch.Tensor, labeled_output: torch.Tensor):
 
         #This is experimental code that is not yet implemented. You'll need to install rapids_singlecell, cuml and scanpy to run this code.
 
diff --git a/instanseg/utils/AI_utils.py b/instanseg/utils/AI_utils.py
index d764cc7..09d85d9 100755
--- a/instanseg/utils/AI_utils.py
+++ b/instanseg/utils/AI_utils.py
@@ -1,19 +1,12 @@
 import torch
 import torch
-import torchvision.transforms.functional as TF
-import sys
-import torch.nn as nn
-import pdb
 import numpy as np
 import matplotlib.pyplot as plt
 from tqdm.auto import tqdm
 from instanseg.utils.metrics import _robust_average_precision, _robust_f1_mean_calculator
 
-
 from instanseg.utils.augmentations import Augmentations
-from instanseg.utils.utils import timer
 import time
-import torchvision
 
 from instanseg.utils.utils import show_images
 import warnings
@@ -285,7 +278,7 @@ def objective(params={}):
                 return 1 - mean_f1
         
         elif val_images is not None and val_labels is not None:
-            from instanseg.utils.tiling import instanseg_padding, recover_padding
+            from instanseg.utils.tiling import _instanseg_padding, _recover_padding
             def objective(params={}):
                 pred_masks = []
                 gt_masks = []
@@ -300,9 +293,9 @@ def objective(params={}):
                     gt_mask = val_labels[i]
                     with torch.no_grad():
                         imgs = imgs.to(device)
-                        imgs, pad = instanseg_padding(imgs, min_dim = 32)
+                        imgs, pad = _instanseg_padding(imgs, min_dim = 32)
                         output = _model(imgs[None,])
-                        output = recover_padding(output, pad).squeeze(0)
+                        output = _recover_padding(output, pad).squeeze(0)
                         lab = postprocessing_fn(output.to(device), **params).cpu()
                         pred_masks.append(lab)
                         gt_masks.append(gt_mask)
diff --git a/instanseg/utils/augmentations.py b/instanseg/utils/augmentations.py
index 0d86e67..9637c7c 100755
--- a/instanseg/utils/augmentations.py
+++ b/instanseg/utils/augmentations.py
@@ -1,18 +1,15 @@
 from instanseg.utils.utils import show_images, _move_channel_axis
 import torchvision
-from skimage.draw import bezier_curve
 import torchvision.transforms.functional as TF
 import torch
 import random
 import numpy as np
-import skimage
 from torchvision.transforms import RandomCrop, Resize, RandomPerspective
-from monai.transforms import RandGaussianNoise, AdjustContrast, HistogramNormalize
+from monai.transforms import RandGaussianNoise, AdjustContrast
 from instanseg.utils.utils import percentile_normalize, generate_colors
 import warnings
 
 import time
-import pdb
 import fastremap
 
 time_dict = {}
@@ -190,8 +187,8 @@ def normalize_HE_stains(self, image: torch.Tensor, labels=None, amount=0, metada
         if metadata is not None and metadata["image_modality"] != "Brightfield":
             return image, labels
         
-        assert image.shape[0] == 3, pdb.set_trace()
-
+        assert image.shape[0] == 3
+        
         import torchstain
         if self.debug:
             orig = torch.clone(image)
@@ -209,7 +206,6 @@ def normalize_HE_stains(self, image: torch.Tensor, labels=None, amount=0, metada
         except:
             return image,labels
 
-        #    pdb.set_trace()
 
         out = _move_channel_axis(norm) / 255.
 
@@ -242,7 +238,6 @@ def randomJPEGcompression(self, image, labels=None, amount=0, metadata=None):
         image = image.clamp(0, 1) * 255
         image = image.byte()
 
-        #        pdb.set_trace()
 
         image = F.to_pil_image(image)
 
@@ -293,10 +288,6 @@ def HistogramNormalize(self, image, labels=None, amount=0, metadata=None):
 
         out = torch.stack([normalizer(c) for c in image])
 
-
-
-     #   pdb.set_trace()
-
         if self.debug:
             print("RandStdShiftIntensity")
             show_images([orig, out], titles=["Original", "Transformed"])
@@ -349,7 +340,7 @@ def rotate(self, image, labels, amount=0, metadata=None):
 
     def adjust_hue(self, image, labels, amount=0, metadata=None):
 
-        assert image.shape[0] == 3, pdb.set_trace()
+        assert image.shape[0] == 3
 
         if metadata is not None and metadata["image_modality"] != "Brightfield":
             return image, labels
@@ -391,7 +382,6 @@ def perspective(self, image, labels, amount=0, metadata=None):
         return out, labels
 
     def invert(self, image, labels, amount=0, metadata=None):
-        # pdb.set_trace()
 
         out = 1 - image
 
@@ -424,8 +414,7 @@ def pseudo_brightfield(self, image, labels=None, amount=0, c_nuclei=0, metadata=
             'eosin': torch.Tensor([0.2159, 0.8012, 0.5581])
         }
 
-        assert c_nuclei is not None, pdb.set_trace()
-
+        assert c_nuclei is not None
         im_nuclei = image[c_nuclei]
 
         my_stains = [torch.clone(stains['hematoxylin'])]
@@ -486,22 +475,19 @@ def colourize(self, image, labels=None, amount=0, c_nuclei=-1, random_seed=None,
         np.random.shuffle(colours)
         if np.min(image.shape) == 1:
             c_nuclei = 0  # If we have a greyscale image, we want to colourize the one and only nuclei channel
-        #   pdb.set_trace()
-
+        
         coloured_image = image[c_nuclei][None,] * torch.Tensor([0, 0, 1])[:, None, None]
         for i, image_channel in enumerate(image[np.arange(len(image)) != c_nuclei]):
             colour = colours[i]
             # greyscale_image_channel = image_channel[None,].expand(3, -1, -1)
             coloured_image_channel = image_channel * torch.Tensor(colour)[:, None, None]
             coloured_image += coloured_image_channel
-        assert not coloured_image.isnan().any(), pdb.set_trace()
-
+        assert not coloured_image.isnan().any()
         image = percentile_normalize(coloured_image, 0.1, subsampling_factor=5)
         if self.debug:
             print("Colourize")
             show_images([orig, image], titles=["Original", "Transformed"])
-        assert not coloured_image.isnan().any(), pdb.set_trace()
-
+        assert not coloured_image.isnan().any()
         return image, labels
 
     def pseudo_imc(self, image, labels, amount=0, metadata=None):
@@ -628,8 +614,6 @@ def channel_subsample(self, image, labels=None, max_channels=None, c_nuclei=None
                     np.random.shuffle(slice)
                     slice = np.append(slice, c_nuclei)
 
-                    #    pdb.set_trace()
-                    #   np.random.shuffle(slice)
                     c_nuclei = np.where(slice == c_nuclei)[0][0]
 
             image = image[slice]
@@ -680,8 +664,7 @@ def torch_rescale(self, image, labels=None, amount=0, current_pixel_size=None, r
             torch.manual_seed(random_seed)
 
         if labels is not None:
-            assert image.shape[-2:] == labels.shape[-2:], pdb.set_trace()
-
+            assert image.shape[-2:] == labels.shape[-2:]
 
 
         shape = (torch.tensor(image.shape[-2:]) * (1 + ((torch.rand(1) - 0.5) * amount))).int().tolist()
@@ -738,7 +721,7 @@ def torch_rescale(self, image, labels=None, amount=0, current_pixel_size=None, r
         if labels is not None:
             out_labels = resized_labels[:, i:i + h, j:j + w]
 
-            assert out_image.shape[-2:] == self.shape and out_labels.shape[-2:] == self.shape, pdb.set_trace()
+            assert out_image.shape[-2:] == self.shape and out_labels.shape[-2:] == self.shape
 
         if self.debug:
             print("Rescaling")
@@ -779,7 +762,6 @@ def __call__(self, image, labels, meta=None):
                 else:
                     observed_modality = meta["image_modality"]
             else:
-                pdb.set_trace()
                 observed_modality = _estimate_image_modality(image, labels)
             modality = observed_modality
         else:
@@ -870,8 +852,7 @@ def __call__(self, image, labels, meta=None):
 
                     image, labels = getattr(self, augmentation)(image, labels, amount=amount, metadata=metadata)
 
-                assert not image.isnan().any(), pdb.set_trace()
-
+                assert not image.isnan().any()
         
 
         image, labels = self.duplicate_grayscale_channels(image,
@@ -896,10 +877,6 @@ def __call__(self, image, labels, meta=None):
     label = tifffile.imread(r"../examples/LuCa1_label.tif")[:512,:512]
     meta = {"image_modality": "Fluorescence", "nuclei_channels": [7]}
 
-   # pdb.set_trace()
-    # img = tifffile.imread(r"../examples/HE_example.tif")
-    # label = tifffile.imread(r"../examples/HE_example_masks.tif")
-    # meta = {"image_modality": "Brightfield", "nuclei_channels": [0, 1, 2]}
 
     Augmenter = Augmentations(augmentation_dict=augmentation_dict, debug=False,dim_in = None)
 
diff --git a/instanseg/utils/biological_utils.py b/instanseg/utils/biological_utils.py
index 50ea7ae..41c8b9f 100644
--- a/instanseg/utils/biological_utils.py
+++ b/instanseg/utils/biological_utils.py
@@ -5,11 +5,8 @@
 import numpy as np
 import seaborn as sns
 import matplotlib.pyplot as plt
-import sys
 
-import pdb
 from instanseg.utils.pytorch_utils import torch_sparse_onehot, torch_fastremap, remap_values, torch_onehot
-from instanseg.utils.utils import show_images
 
 def get_intersection_over_union(label: torch.Tensor, return_lab: bool = True) -> torch.Tensor:
 
@@ -351,11 +348,8 @@ def show_umap_and_cluster(X_features):
 
 if __name__ == "__main__":
     from skimage import io
-    from instanseg.utils.tiling import sliding_window_inference
-    from instanseg.utils.pytorch_utils import torch_sparse_onehot, fast_sparse_dual_iou
+    from instanseg.utils.pytorch_utils import torch_sparse_onehot
     from instanseg.utils.utils import _choose_device
-    from instanseg.utils.utils import drag_and_drop_file
-    from pathlib import Path
 
     instanseg = torch.jit.load("../torchscripts/1793450.pt")
     device = _choose_device()
@@ -373,7 +367,7 @@ def show_umap_and_cluster(X_features):
 
     lab = instanseg(input_tensor[:,128:256,:128].unsqueeze(0).to(device))
 
-    # lab = sliding_window_inference(input_tensor, instanseg, window_size=(512, 512),
+    # lab = _sliding_window_inference(input_tensor, instanseg, window_size=(512, 512),
     #                                sw_device=device, device='cpu', output_channels=2)
     
     lab = resolve_cell_and_nucleus_boundaries(lab)
@@ -403,21 +397,8 @@ def show_umap_and_cluster(X_features):
                                         mixed_precision = False).to(device)
     out = super_model(input_tensor[None,][:,:,128:256,:128])
 
-
-
-    pdb.set_trace()
-
-
-
-
-
     export_to_torchscript(model_name, show_example=True, mixed_predicision=False)
 
-
-
-
-
-
     print("Calculating cellular features ...")
     X_cell, X_nuclei, X_cytoplasm = get_features_by_location(input_tensor, lab)
 
diff --git a/instanseg/utils/create_bioimageio_model.py b/instanseg/utils/create_bioimageio_model.py
index fda1c09..f3b0f55 100644
--- a/instanseg/utils/create_bioimageio_model.py
+++ b/instanseg/utils/create_bioimageio_model.py
@@ -1,12 +1,7 @@
 import os
-#from PIL import Image
-import tifffile
-
-# the imports for bioimage.io model export
 import bioimageio.core
 import numpy as np
 import torch
-import monai
 from aicsimageio import AICSImage
 
 from instanseg.utils.augmentations import Augmentations
@@ -176,7 +171,6 @@ def export_bioimageio(torchsript: torch.jit._script.RecursiveScriptModule,
 
     output_path = os.environ['INSTANSEG_BIOIMAGEIO_PATH']
 
-
     if output_name is None:
         output_name = model_name
     # create a directory to store bioimage.io model files
@@ -187,7 +181,6 @@ def export_bioimageio(torchsript: torch.jit._script.RecursiveScriptModule,
     model_pixel_size = torchsript.pixel_size
     print("Model pixel size: ", model_pixel_size)
 
-
     try:
         model,model_dict = load_model(model_name, path = os.environ['INSTANSEG_MODEL_PATH'])
     except:
@@ -249,9 +242,9 @@ def export_bioimageio(torchsript: torch.jit._script.RecursiveScriptModule,
         dim_out = output.shape[1]
     np.save(os.path.join(output_name, "test-output.npy"), output.cpu().numpy())
 
-    from instanseg.utils.utils import display_overlay
+    from instanseg.utils.utils import _display_overlay
 
-    cover = display_overlay(input_crop[0], output)
+    cover = _display_overlay(input_crop[0], output)
     show_images(cover, colorbar=False)
     show_images(cover, colorbar=False, save_str= os.path.join(output_name, "cover"))
 
@@ -260,7 +253,6 @@ def export_bioimageio(torchsript: torch.jit._script.RecursiveScriptModule,
     else:
         train_data = "Not specified"
 
-
     # create readme
     readme(output_name, model_dict)
 
@@ -348,7 +340,6 @@ def export_bioimageio(torchsript: torch.jit._script.RecursiveScriptModule,
 
     #unzip the folder
 
-    
     import zipfile
     import shutil
 
@@ -363,9 +354,4 @@ def export_bioimageio(torchsript: torch.jit._script.RecursiveScriptModule,
     
     make_archive(destination, input)
 
-
-    
-
-    
-
     shutil.rmtree(output_name)
diff --git a/instanseg/utils/data_download.py b/instanseg/utils/data_download.py
index 2fe7ed7..6ce29ba 100755
--- a/instanseg/utils/data_download.py
+++ b/instanseg/utils/data_download.py
@@ -1,7 +1,6 @@
 
-import sys
 from pathlib import Path
-from instanseg.utils.utils import show_images, _move_channel_axis
+from instanseg.utils.utils import _move_channel_axis
 import fastremap
 from tqdm import tqdm
 from skimage import io
@@ -10,7 +9,6 @@
 import skimage
 import numpy as np
 import os
-import subprocess
 
 import requests
 import zipfile
diff --git a/instanseg/utils/data_loader.py b/instanseg/utils/data_loader.py
index 5602e07..5adf342 100755
--- a/instanseg/utils/data_loader.py
+++ b/instanseg/utils/data_loader.py
@@ -1,7 +1,7 @@
 
-import torch
 import numpy as np
 import warnings
+
 def _keep_images(item, args):
     if args.source_dataset != "all" and item[
         'parent_dataset'] not in args.source_dataset:  # remove items that are not of the desired dataset
@@ -280,4 +280,4 @@ def get_loaders(train_images_local, train_labels_local, val_images_local, val_la
     test_loader = DataLoader(test_data, collate_fn=collate_fn, batch_size=args.batch_size, num_workers=args.num_workers,
                              sampler=test_sampler, persistent_workers=True)
 
-    return train_loader, test_loader
\ No newline at end of file
+    return train_loader, test_loader
diff --git a/instanseg/utils/loss/instanseg_loss.py b/instanseg/utils/loss/instanseg_loss.py
index c113d00..615ce0a 100644
--- a/instanseg/utils/loss/instanseg_loss.py
+++ b/instanseg/utils/loss/instanseg_loss.py
@@ -6,7 +6,7 @@
 
 from instanseg.utils.loss.lovasz_losses import binary_xloss
 from instanseg.utils.pytorch_utils import torch_fastremap, torch_onehot, remap_values, fast_iou, fast_sparse_iou, eccentricity_batch, connected_components
-from instanseg.utils.tiling import instanseg_padding, recover_padding
+from instanseg.utils.tiling import _instanseg_padding, _recover_padding
 
 import torch.nn.functional as F
 import torch.nn as nn
@@ -1109,9 +1109,9 @@ def TTA_postprocessing(self, img, model, transforms,
             for t in transforms:
                 with torch.cuda.amp.autocast():
                     augmented_image = t.augment_image(img)
-                    augmented_image, pad = instanseg_padding(augmented_image, extra_pad= 0, min_dim = 32)
+                    augmented_image, pad = _instanseg_padding(augmented_image, extra_pad= 0, min_dim = 32)
                     prediction = model(augmented_image)[:,i * dim_out:(i+1) * dim_out]
-                    prediction = recover_padding(prediction, pad)
+                    prediction = _recover_padding(prediction, pad)
                     mask_map = prediction[:,-1][None] 
                     mask_map = t.deaugment_mask(mask_map)
                 #  show_images(mask_map)
@@ -1287,7 +1287,7 @@ def forward(self, x: torch.Tensor,
         torch.clamp_max_(x, 3) #Safety check, please normalize inputs properly!
         torch.clamp_min_(x, -2)
 
-        x, pad = instanseg_padding(x, extra_pad=0)
+        x, pad = _instanseg_padding(x, extra_pad=0)
 
 
         with torch.no_grad():
@@ -1313,7 +1313,7 @@ def forward(self, x: torch.Tensor,
                     else:
                         x = x_full[image_index,dim_out:, :, :]
 
-                    x = recover_padding(x, pad)
+                    x = _recover_padding(x, pad)
 
                     height, width = x.size(1), x.size(2)
 
diff --git a/instanseg/utils/metrics.py b/instanseg/utils/metrics.py
index 79c415b..335a0d4 100755
--- a/instanseg/utils/metrics.py
+++ b/instanseg/utils/metrics.py
@@ -1,7 +1,4 @@
 import numpy as np
-import torch
-import pdb
-
 from stardist import matching
 import fastremap
 
diff --git a/instanseg/utils/model_loader.py b/instanseg/utils/model_loader.py
index ecff22c..fedbbe8 100755
--- a/instanseg/utils/model_loader.py
+++ b/instanseg/utils/model_loader.py
@@ -1,19 +1,5 @@
-import pdb
-import torch.nn as nn
-import torch
-import monai
-
-import torch.distributions as td
-import torch.nn as nn
-
-import torch.nn as nn
-import torch
-import torch.distributions as td
 import numpy as np
 
-from monai.networks.blocks.convolutions import Convolution, ResidualUnit
-
-
 
 def build_monai_model(model_str: str, build_model_dictionary: dict):
 
diff --git a/instanseg/utils/pytorch_utils.py b/instanseg/utils/pytorch_utils.py
index 494cc35..b66fe29 100644
--- a/instanseg/utils/pytorch_utils.py
+++ b/instanseg/utils/pytorch_utils.py
@@ -1,10 +1,7 @@
-import os
 import torch
 import torch.nn.functional as F
 
-from typing import Tuple, List, Union
-
-
+from typing import Tuple
 
 
 def remap_values(remapping: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
diff --git a/instanseg/utils/tiling.py b/instanseg/utils/tiling.py
index 9899ce9..4121023 100755
--- a/instanseg/utils/tiling.py
+++ b/instanseg/utils/tiling.py
@@ -2,7 +2,7 @@
 import numpy as np
 from tqdm import tqdm
 
-def edge_mask(labels, ignore=[None]):
+def _edge_mask(labels, ignore=[None]):
     labels = labels.squeeze()
     first_row = labels[0, :]
     last_row = labels[-1, :]
@@ -26,32 +26,12 @@ def edge_mask(labels, ignore=[None]):
     return torch.isin(labels, edges[edges > 0])
 
 
-def remove_edge_labels(labels, ignore=[None]):
-    return labels * ~edge_mask(labels, ignore=ignore)
+def _remove_edge_labels(labels, ignore=[None]):
+    return labels * ~_edge_mask(labels, ignore=ignore)
 
 
-def _to_shape(a, shape):
-    """Pad a tensor to a given shape."""
-    if len(a.shape) == 2:
-        a = a.unsqueeze(0)
-    y_, x_ = shape
-    y, x = a[0].shape[-2:]
-    y_pad = max(0, y_ - y)
-    x_pad = max(0, x_ - x)
-    return torch.nn.functional.pad(a, (x_pad // 2, x_pad // 2 + x_pad % 2, y_pad // 2, y_pad // 2 + y_pad % 2))
 
-
-def _to_shape_bottom_left(a, shape):
-    """Pad a tensor to a given shape."""
-    if len(a.shape) == 2:
-        a = a.unsqueeze(0)
-    y_, x_ = shape
-    y, x = a[0].shape[-2:]
-    y_pad = max(0, y_ - y)
-    x_pad = max(0, x_ - x)
-    return torch.nn.functional.pad(a, (0, x_pad, 0, y_pad))
-
-def chops(img_shape: tuple, shape: tuple, overlap: int = 0) -> tuple:
+def _chops(img_shape: tuple, shape: tuple, overlap: int = 0) -> tuple:
     """This function splits an image into desired windows and returns the indices of the windows"""
 
     if (torch.tensor(img_shape[-2:]) < torch.tensor(shape)).any():
@@ -79,8 +59,8 @@ def chops(img_shape: tuple, shape: tuple, overlap: int = 0) -> tuple:
     return h_index, v_index
 
 
-def tiles_from_chops(image: torch.Tensor, shape: tuple, tuple_index: tuple) -> list:
-    """This function takes an image, a shape, and a tuple of window indices (e.g., output of the function chops)
+def _tiles_from_chops(image: torch.Tensor, shape: tuple, tuple_index: tuple) -> list:
+    """This function takes an image, a shape, and a tuple of window indices (e.g., output of the function _chops)
     and returns a list of windows"""
     h_index, v_index = tuple_index
 
@@ -97,8 +77,8 @@ def tiles_from_chops(image: torch.Tensor, shape: tuple, tuple_index: tuple) -> l
     return tile_list
 
 
-def stitch(tiles: list, shape: tuple, chop_list: list, final_shape: tuple, offset : int):
-    """This function takes a list of tiles, a shape, and a tuple of window indices (e.g., outputed by the function chops)
+def _stitch(tiles: list, shape: tuple, chop_list: list, final_shape: tuple, offset : int):
+    """This function takes a list of tiles, a shape, and a tuple of window indices (e.g., outputed by the function _chops)
     and returns a stitched image"""
     from instanseg.utils.pytorch_utils import torch_fastremap, match_labels
 
@@ -124,31 +104,31 @@ def stitch(tiles: list, shape: tuple, chop_list: list, final_shape: tuple, offse
             if i == len(chop_list[0])-1 and j == len(chop_list[1])-1:
                 edge_window = True
                 tile1 = canvas[..., window_i + offset:window_i + shape[0], window_j + offset:window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[:,offset:shape[0],offset: shape[1]], ignore = ignore_list)
+                tile2 = _remove_edge_labels(new_tile[:,offset:shape[0],offset: shape[1]], ignore = ignore_list)
 
             elif i == len(chop_list[0])-1:
                 edge_window = True
                 tile1 = canvas[..., window_i + offset:window_i + shape[0], window_j + offset:window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[:,offset:shape[0],offset : shape[1]], ignore =ignore_list)
+                tile2 = _remove_edge_labels(new_tile[:,offset:shape[0],offset : shape[1]], ignore =ignore_list)
 
             elif j == len(chop_list[1])-1:
                 edge_window = True
                 tile1 = canvas[..., window_i + offset:window_i + shape[0], window_j + offset:window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[:,offset:shape[0],offset: shape[1]], ignore = ignore_list)
+                tile2 = _remove_edge_labels(new_tile[:,offset:shape[0],offset: shape[1]], ignore = ignore_list)
 
             if i == 0 and j == 0:
                 edge_window = True
                 tile1 = canvas[..., window_i  :window_i + shape[0], window_j :window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[:, :shape[0], : shape[1]], ignore = ignore_list)
+                tile2 = _remove_edge_labels(new_tile[:, :shape[0], : shape[1]], ignore = ignore_list)
             elif i == 0:
                 edge_window = True
                 tile1 = canvas[..., window_i  :window_i + shape[0], window_j + offset :window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[:, :shape[0],offset : shape[1]], ignore = ignore_list)
+                tile2 = _remove_edge_labels(new_tile[:, :shape[0],offset : shape[1]], ignore = ignore_list)
 
             elif j == 0:
                 edge_window = True
                 tile1 = canvas[..., window_i  + offset:window_i + shape[0], window_j:window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[:,offset :shape[0],: shape[1]], ignore = ignore_list)
+                tile2 = _remove_edge_labels(new_tile[:,offset :shape[0],: shape[1]], ignore = ignore_list)
 
             if edge_window:
 
@@ -163,7 +143,7 @@ def stitch(tiles: list, shape: tuple, chop_list: list, final_shape: tuple, offse
             else:
                 
                 tile1 = canvas[..., window_i + offset:window_i + shape[0] - offset, window_j + offset:window_j + shape[1] - offset]
-                tile2 = remove_edge_labels(new_tile[:,offset:shape[0] -offset,offset: shape[1]-offset])
+                tile2 = _remove_edge_labels(new_tile[:,offset:shape[0] -offset,offset: shape[1]-offset])
 
 
                 tile2 = torch_fastremap(tile2)
@@ -178,7 +158,7 @@ def stitch(tiles: list, shape: tuple, chop_list: list, final_shape: tuple, offse
     return canvas
 
 
-def zarr_to_json_export(path_to_zarr, cell_size = 30, size = 1024, scale = 1, n_dim = 1):
+def _zarr_to_json_export(path_to_zarr, detection_size = 30, size = 1024, scale = 1, n_dim = 1):
 
     import zarr
     import numpy as np
@@ -189,12 +169,12 @@ def zarr_to_json_export(path_to_zarr, cell_size = 30, size = 1024, scale = 1, n_
     output_path = str(path_to_zarr).replace(".zarr",".geojson")   
 
 
-    from instanseg.utils.tiling import chops
+    from instanseg.utils.tiling import _chops
 
-    chop_list = chops(z.shape, shape=(size,size), overlap=cell_size)
+    chop_list = _chops(z.shape, shape=(size,size), overlap=detection_size)
 
     from instanseg.utils.utils import labels_to_features
-    from instanseg.utils.tiling import remove_edge_labels
+    from instanseg.utils.tiling import _remove_edge_labels
     import json
 
     count = 0
@@ -218,7 +198,7 @@ def zarr_to_json_export(path_to_zarr, cell_size = 30, size = 1024, scale = 1, n_
                 
                 image = z[n,window_i:window_i+size, window_j:window_j+size]
               
-                image = remove_edge_labels(torch.tensor(image)).numpy()
+                image = _remove_edge_labels(torch.tensor(image)).numpy()
 
                 features = labels_to_features(image.astype(np.int32), object_type='detection', include_labels=True,
                                         classification=classes[n],offset=[window_j*scale,window_i*scale], downsample = scale)
@@ -236,191 +216,7 @@ def zarr_to_json_export(path_to_zarr, cell_size = 30, size = 1024, scale = 1, n_
         outfile.write(']')
 
 
-def segment_image_larger_than_memory(instanseg_class, # instanseg class method
-                                    image_path: str,  
-                                     memory_block_size: tuple = (3000,3000), #this is the size of the image that will be read in memory
-                                     inference_tile_size: tuple = (512,512), #this is the size of the image that will be passed to the model
-                                     overlap = 100,
-                                    cell_size = 20, 
-                                    to_geojson = False, 
-                                    batch_size = 3,
-                                    prediction_tag: str = "_instanseg_prediction",
-                                    pixel_size: float = None,
-                                    normalisation_subsampling_factor = 10,
-                                    **kwargs):
-    
-    """This function uses slideio to read an image and then segments it using the instanseg model. 
-    The segmentation is done in a tiled manner to avoid memory issues. 
-    The function returns a zarr file with the segmentation. The zarr file is saved in the same directory as the image with the same name but with the extension .zarr. 
-    The function also returns the zarr file object."""
-
-    import zarr
-    from itertools import product
-    from instanseg.utils.pytorch_utils import torch_fastremap, match_labels
-    from pathlib import Path
-    
-    instanseg = instanseg_class.instanseg
-
-    image_path, img_pixel_size = instanseg_class.read_image(image_path)
-    slide = instanseg_class.read_slide(image_path)
-
-    n_dim = 2 if instanseg.cells_and_nuclei else 1
-    model_pixel_size = instanseg.pixel_size
-
-    new_stem = Path(image_path).stem + prediction_tag
-    file_with_zarr_extension = Path(image_path).parent / (new_stem + ".zarr")
-
-
-    if img_pixel_size > 1 or img_pixel_size < 0.1:
-        import warnings
-        warnings.warn("The image pixel size {} is not in microns.".format(img_pixel_size))
-        if pixel_size is not None:
-            img_pixel_size = pixel_size
-        else:
-            raise ValueError("The image pixel size {} is not in microns.".format(img_pixel_size))
-    
-    scale_factor = model_pixel_size/img_pixel_size
-
-    dims = slide.dimensions
-    dims = (int(dims[1]/ scale_factor), int(dims[0]/scale_factor)) #The dimensions are opposite to numpy/torch/zarr dimensions.
-
-    pad2 = overlap + cell_size
-    pad = overlap
-
-    shape = memory_block_size
-    
-    chop_list = chops(dims, shape, overlap=2*pad2)
-
-    chunk_shape = (n_dim,shape[0],shape[1])
-    store = zarr.DirectoryStore(file_with_zarr_extension) 
-    canvas = zarr.zeros((n_dim,dims[0],dims[1]), chunks=chunk_shape, dtype=np.int32, store=store, overwrite = True)
-
-    running_max = 0
-
-    total = len(chop_list[0]) * len(chop_list[1])
-    for _, ((i, window_i), (j, window_j)) in tqdm(enumerate(product(enumerate(chop_list[0]), enumerate(chop_list[1]))), total=total, colour = "green", desc = "Slide progress: "):
-
-      #  input_data = scene.read_block((int(window_j*scale_factor), int(window_i*scale_factor), int(shape[0]*scale_factor), int(shape[1]*scale_factor)), size = shape)
-
-        best_level = slide.get_best_level_for_downsample(scale_factor)
-        downsample_factor = slide.level_downsamples[best_level]
-
-        initial_pixel_size = img_pixel_size
-        itermediate_pixel_size = initial_pixel_size * downsample_factor
-        final_pixel_size = model_pixel_size
-
-        intermediate_to_final = final_pixel_size/itermediate_pixel_size
-        
-        # Calculate the size of the region needed at the base level to get the desired output size
-        intermediate_shape = (int(shape[0] * intermediate_to_final), int(shape[1] * intermediate_to_final))
-
-        input_data = slide.read_region((int(window_j*scale_factor), int(window_i*scale_factor)), best_level, (int(intermediate_shape[0]) , int(intermediate_shape[1])), as_array=True)
-    
-        input_tensor = instanseg_class._to_tensor(input_data)
-
-        new_tile = instanseg_class.eval_medium_image(input_tensor,
-                                          pixel_size = itermediate_pixel_size,
-                                          tile_size = inference_tile_size[0],
-                                          batch_size = batch_size,
-                                          return_image_tensor = False,
-                                          normalisation_subsampling_factor = normalisation_subsampling_factor,
-                                          )
-        
-    
-
-        from torch.nn.functional import interpolate
-        new_tile = interpolate(new_tile, size=shape[-2:], mode="nearest").int()[0]
-
-       # pdb.set_trace()
-
-        num_iter = new_tile.shape[0]
-
-        for n in range(num_iter):
-            
-            ignore_list = []
-            if i == 0:
-                ignore_list.append("top")
-            if j == 0:
-                ignore_list.append("left")
-            if i == len(chop_list[0])-1:
-                ignore_list.append("bottom")
-            if j == len(chop_list[1])-1:
-                ignore_list.append("right")
-
-            if i == len(chop_list[0])-1 and j == len(chop_list[1])-1:
-                tile1 = canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[n,pad:shape[0],pad: shape[1]], ignore = ignore_list)
-
-            elif i == len(chop_list[0])-1:
-                tile1 = canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[n,pad:shape[0],pad : shape[1]], ignore =ignore_list)
-
-            elif j == len(chop_list[1])-1:
-                tile1 = canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[n,pad:shape[0],pad: shape[1]], ignore = ignore_list)
-
-            elif i == 0 and j == 0:
-                tile1 = canvas[n, window_i  :window_i + shape[0], window_j :window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[n, :shape[0], : shape[1]], ignore = ignore_list)
-            elif i == 0:
-                tile1 = canvas[n, window_i  :window_i + shape[0], window_j + pad :window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[n, :shape[0],pad : shape[1]], ignore = ignore_list)
-
-            elif j == 0:
-                tile1 = canvas[n, window_i  + pad:window_i + shape[0], window_j:window_j + shape[1]]
-                tile2 = remove_edge_labels(new_tile[n,pad :shape[0],: shape[1]], ignore = ignore_list)
-
-            if j == 0 or i == 0 or j == len(chop_list[1])-1 or i == len(chop_list[0])-1:
-
-                tile2 = torch_fastremap(tile2)
-                tile2[tile2>0] = tile2[tile2>0] + running_max
-
-                tile1_torch = torch.tensor(np.array(tile1), dtype = torch.int32)
-
-                remapped = match_labels(tile1_torch, tile2, threshold = 0.1)[1]
-                tile1_torch[remapped>0] = remapped[remapped>0].int()
-
-                running_max = max(running_max, tile1_torch.max())
-
-                if i == len(chop_list[0])-1 and j == len(chop_list[1])-1:
-                    canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
-                elif i == len(chop_list[0])-1:
-                    canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
-                elif j == len(chop_list[1])-1:
-                    canvas[n, window_i + pad:window_i + shape[0], window_j + pad:window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
-                elif i == 0 and j == 0:
-                    canvas[n, window_i  :window_i + shape[0], window_j :window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
-                elif i == 0:
-                    canvas[n, window_i  :window_i + shape[0], window_j + pad :window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
-                elif j == 0:
-                    canvas[n, window_i  + pad:window_i + shape[0], window_j:window_j + shape[1]] = tile1_torch.numpy().astype(np.int32)
-
-            else:
-                
-                tile1 = canvas[n, window_i + pad:window_i + shape[0] - pad, window_j + pad:window_j + shape[1] - pad]
-                tile2 = remove_edge_labels(new_tile[n,pad:shape[0] -pad,pad: shape[1]-pad])
-                
-                tile2 = torch_fastremap(tile2)
-
-                tile2[tile2>0] = tile2[tile2>0] + running_max
-
-                tile1_torch = torch.tensor(np.array(tile1), dtype = torch.int32)
-                remapped = match_labels(tile1_torch, tile2, threshold = 0.1)[1]
-
-                tile1_torch[remapped>0] = remapped[remapped>0].int()
-                running_max = max(running_max, tile1_torch.max())
-            
-                canvas[n, window_i + pad:window_i + shape[0] - pad, window_j + pad:window_j + shape[1] - pad] = tile1_torch.numpy().astype(np.int32)
-
-    if to_geojson:
-        print("Exporting to geojson")
-        zarr_to_json_export(file_with_zarr_extension, cell_size = cell_size, size = shape[0], scale = scale_factor, n_dim = n_dim)
-            
-
-    return canvas
-
-
-def instanseg_padding(img: torch.Tensor, extra_pad: int = 0, min_dim: int = 16, ensure_square: bool = False):
+def _instanseg_padding(img: torch.Tensor, extra_pad: int = 0, min_dim: int = 16, ensure_square: bool = False):
 
     is_square = img.shape[-2] == img.shape[-1]
     original_shape = img.shape[-2:]
@@ -447,7 +243,7 @@ def instanseg_padding(img: torch.Tensor, extra_pad: int = 0, min_dim: int = 16,
     return img, torch.stack((padx, pady))
 
 
-def recover_padding(x: torch.Tensor, pad: torch.Tensor):
+def _recover_padding(x: torch.Tensor, pad: torch.Tensor):
     # x must be 1,C,H,W or C,H,W
     squeeze = False
     if x.ndim == 3:
@@ -466,7 +262,7 @@ def recover_padding(x: torch.Tensor, pad: torch.Tensor):
 
 
 
-def sliding_window_inference(input_tensor, 
+def _sliding_window_inference(input_tensor, 
                              predictor, 
                              window_size=(512, 512), 
                              overlap = 80, 
@@ -484,8 +280,8 @@ def sliding_window_inference(input_tensor,
     input_tensor = input_tensor.to(device)
     predictor = predictor.to(sw_device)
  
-    tuple_index = chops(input_tensor.shape, shape=window_size, overlap=2 * (overlap + max_cell_size))
-    tile_list = tiles_from_chops(input_tensor, shape=window_size, tuple_index=tuple_index)
+    tuple_index = _chops(input_tensor.shape, shape=window_size, overlap=2 * (overlap + max_cell_size))
+    tile_list = _tiles_from_chops(input_tensor, shape=window_size, tuple_index=tuple_index)
  
  
     assert len(tile_list) > 0, "No tiles generated"
@@ -500,7 +296,7 @@ def sliding_window_inference(input_tensor,
             label_list = torch.cat([predictor(tile.to(sw_device),**kwargs).to(device) for tile in tqdm(batch_list, disable= not show_progress,leave = False, colour = "blue")])
  
    
-    lab = torch.cat([stitch([lab[i] for lab in label_list],
+    lab = torch.cat([_stitch([lab[i] for lab in label_list],
                             shape=window_size,
                             chop_list=tuple_index,
                             offset = overlap,
diff --git a/instanseg/utils/utils.py b/instanseg/utils/utils.py
index a134520..1de9d2d 100755
--- a/instanseg/utils/utils.py
+++ b/instanseg/utils/utils.py
@@ -2,7 +2,6 @@
 
 import json
 import os
-import pdb
 import fastremap
 import numpy as np
 import tifffile
@@ -14,7 +13,6 @@
 import rasterio.features
 from rasterio.transform import Affine
 from matplotlib.colors import LinearSegmentedColormap
-from matplotlib.cm import get_cmap
 import colorcet as cc
 import matplotlib as mpl
 import os
@@ -22,7 +20,8 @@
 import matplotlib.pyplot as plt
 import torch
 import tifffile
-from typing import Optional
+from typing import Optional, List
+import geojson
 
 def moving_average(x, w):
     """Moving average of an array x with window size w"""
@@ -41,25 +40,20 @@ def plot_average(train, test, clip=99, window_size=10):
     return fig
 
 
-def _to_shape(a, shape):
-    """Pad an array to a given shape."""
-    a = _move_channel_axis(a)
-    if len(np.shape(a)) == 2:
-        a = a[None,]
-    y_, x_ = shape
-    y, x = a[0].shape
-    y_pad = np.max([0, (y_ - y)])
-    x_pad = np.max([0, (x_ - x)])
-    return np.pad(a, ((0, 0), (y_pad // 2, y_pad // 2 + y_pad % 2),
-                      (x_pad // 2, x_pad // 2 + x_pad % 2)),
-                  mode='constant')
-
-
-def labels_to_features(lab: np.ndarray, object_type='annotation', connectivity: int = 4,
-                       transform: Affine = None, downsample: float = 1.0, include_labels=False,
-                       classification=None, offset=None):
+def labels_to_features(lab: np.ndarray,
+                       object_type='annotation',
+                       connectivity: int = 4,
+                       transform: Affine = None,
+                       downsample: float = 1.0,
+                       include_labels: bool = False,
+                       classification: str = None,
+                       offset: int = None) -> geojson.FeatureCollection:
     """
     Create a GeoJSON FeatureCollection from a labeled image.
+
+    :param classification: Optional classification added to output features.
+    :param offset: Offset added to x coordinates of output features.
+    :return: A GeoJSON FeatureCollection.
     """
     features = []
 
@@ -86,7 +80,7 @@ def labels_to_features(lab: np.ndarray, object_type='annotation', connectivity:
 
         # Just to show how a classification can be added
         if classification is not None:
-            props['classification'] = str(classification)
+            props['classification'] = classification
 
         if offset is not None:
             coordinates = obj[0]['coordinates']
@@ -94,32 +88,36 @@ def labels_to_features(lab: np.ndarray, object_type='annotation', connectivity:
                 [(int(x[0] + offset[0]), int(x[1] + offset[1])) for x in coordinates[0]]]
             obj[0]['coordinates'] = coordinates
 
-        # Wrap in a dict to effectively create a GeoJSON Feature
-        po = dict(type="Feature", geometry=obj[0], properties=props)
+        po = geojson.Feature(geometry = obj[0], properties=props)
 
         features.append(po)
+    return geojson.FeatureCollection(features)
 
-    return features
 
+def interp(image: np.ndarray, shape: float = None, scale:float = None):
 
-def interp(image: np.ndarray, shape=None, scale=None):
-    """Interpolate an image to a new shape or scale."""
     from scipy import interpolate
     x = np.array(range(image.shape[1]))
     y = np.array(range(image.shape[0]))
     interpolate_fn = interpolate.interp2d(x, y, image)
+
     if shape:
         x_new = np.linspace(0, image.shape[1] - 1, shape[1])
         y_new = np.linspace(0, image.shape[0] - 1, shape[0])
     elif scale:
         x_new = np.linspace(0, image.shape[1] - 1, int(np.floor(image.shape[1] * scale) + 1))
         y_new = np.linspace(0, image.shape[0] - 1, int(np.floor(image.shape[0] * scale) + 1))
+    
     znew = interpolate_fn(x_new, y_new)
     return znew
 
 
 
-def apply_cmap(x, fg_mask = None, cmap = "coolwarm_r", bg_intensity = 255, normalize = True):
+def apply_cmap(x,
+               fg_mask: np.ndarray = None,
+               cmap: str = "coolwarm_r",
+               bg_intensity: int = 255,
+               normalize: bool = True):
     """
     Apply a colormap to an image, with a background mask.
     x  and fg_mask should have the same shape, and be numpy arrays.
@@ -226,7 +224,14 @@ def show_images(*img_list, clip_pct=None, titles=None, save_str=False, n_cols=3,
 
 
 
-def display_as_grid(display_list, ncols,padding = 2, left_titles = None, top_titles = None, right_side = None, title_height = 20, fontsize = 12):
+def display_as_grid(display_list,
+                    ncols: int,
+                    padding: int = 2,
+                    left_titles: Optional[List[str]] = None,
+                    top_titles = None,
+                    right_side = None,
+                    title_height: int = 20,
+                    fontsize: float = 12):
 
     from instanseg.utils.augmentations import Augmentations
     Augmenter = Augmentations()
@@ -274,7 +279,7 @@ def _scale_area(size: float, pixel_size: float, do_round=True) -> float:
     return np.round(size_pixels) if do_round else size_pixels
 
 
-def _move_channel_axis(img: Union[np.ndarray, torch.Tensor], to_back: bool = False):
+def _move_channel_axis(img: Union[np.ndarray, torch.Tensor], to_back: bool = False) -> Union[np.ndarray, torch.Tensor]:
     if isinstance(img, np.ndarray):
         img = img.squeeze()
         if img.ndim != 3:
@@ -299,8 +304,11 @@ def _move_channel_axis(img: Union[np.ndarray, torch.Tensor], to_back: bool = Fal
         return img.movedim(ch, 0)
 
 
-def percentile_normalize(img: Union[np.ndarray, torch.Tensor], percentile=0.1, subsampling_factor: int = 1,
+def percentile_normalize(img: Union[np.ndarray, torch.Tensor],
+                         percentile: float = 0.1,
+                         subsampling_factor: int = 1,
                          epsilon: float = 1e-3):
+
     if isinstance(img, np.ndarray):
         assert img.ndim == 2 or img.ndim == 3, "Image must be 2D or 3D, got image of shape" + str(img.shape)
         img = np.atleast_3d(img)
@@ -341,16 +349,6 @@ def generate_colors(num_colors: int):
     return rgb_colors
 
 
-def tensor_or_np_copy(x):
-    if isinstance(x, torch.Tensor):
-        return x.clone()
-    else:
-        return x.copy()
-
-
-
-
-
 def export_annotations_and_images(output_dir, original_image, lab, base_name=None):
     from pathlib import Path
     import os
@@ -383,15 +381,13 @@ def export_annotations_and_images(output_dir, original_image, lab, base_name=Non
 
 
 import matplotlib.colors as mcolors
-def color_name_to_rgb(color_name):
+def color_name_to_rgb(color_name: str):
     """
     Convert a color name to its corresponding RGB values.
     
-    Args:
-    color_name (str): The name of the color.
+    :param color_name: The name of the color.
     
-    Returns:
-    tuple: A tuple containing the RGB values.
+    :return: A tuple containing the RGB values.
     """
     return mcolors.to_rgb(color_name)
 
@@ -423,6 +419,7 @@ def save_image_with_label_overlay(im: np.ndarray,
                                 one of 'thick', 'inner', 'outer', 'subpixel'. If None, the lab is used directly.
     :param alpha: Alpha value for the underlying image when using the overlay.
                   Setting this less than 1 can help the overlay stand out more prominently.
+    :return: The input image with corresponding label overlay if lab is a torch Tensor, or nothing otherwise.
     """
 
     import imageio
@@ -447,7 +444,6 @@ def save_image_with_label_overlay(im: np.ndarray,
             bg = (lab.sum(0) == 0)
 
             if label_boundary_mode is None:
-                from palettable import wesanderson as wes
                 from palettable.scientific import diverging as div
                 colour_cells = list((np.array(div.Berlin_12.colors[1]) /255))
                 colour_nuclei = list( np.array(div.Berlin_12.colors[11] ) /255)
@@ -528,6 +524,7 @@ def save_image_with_label_overlay(im: np.ndarray,
 def display_cells_and_nuclei(lab):
     display = save_image_with_label_overlay(torch.zeros((lab.shape[-2],lab.shape[-1],3)), lab, return_image= True, label_boundary_mode=None,alpha = 1)
     return display
+
 def display_colourized(mIF, random_seed = 0):
     from instanseg.utils.augmentations import Augmentations
     Augmenter=Augmentations()
@@ -541,7 +538,7 @@ def display_colourized(mIF, random_seed = 0):
     colour_render = _move_channel_axis(colour_render,to_back = True).detach().numpy()*255
     return colour_render.astype(np.uint8)
 
-def display_overlay(im, lab):
+def _display_overlay(im, lab):
     assert lab.ndim == 4, "lab must be 4D"
     assert im.ndim == 3, "im must be 3D"
     output_dimension = lab.shape[1]
@@ -623,7 +620,12 @@ def _choose_device(device: str = None, verbose=True) -> str:
     return device
 
 
-def count_instances(labels: Union[np.ndarray, torch.Tensor]):
+def count_instances(labels: Union[np.ndarray, torch.Tensor]) -> int:
+    """
+    Count the total number of labelled pixels in an input tensor.
+    :param labels: The input tensor.
+    :return: The total number of non-zero labels.
+    """
     import fastremap
     if isinstance(labels, torch.Tensor):
         num_labels = len(torch.unique(labels[labels > 0]))
@@ -885,9 +887,4 @@ def download_model(model_str: str, verbose : bool = True, headers: Optional[str]
         if os.path.exists(path_to_torchscript_model):
             return torch.jit.load(path_to_torchscript_model)
         else:
-            raise Exception(f"Model {path_to_torchscript_model} not found in the release data or locally. Please check the model name and try again.")
-
-
-
-
-
+            raise Exception(f"Model {model_str} not found in the release data or locally. Please check the model name and try again.")
diff --git a/pyproject.toml b/pyproject.toml
index 164ffed..38c1544 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,24 +17,26 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 
-# Dependencies (Specify your requirements here)
 dependencies = [
-    "numpy>=1.26.4",
-    "torch>=2.1.1",
-    "tqdm>=4.66.1",
-    "matplotlib>=3.8.2",
-    "einops>=0.7.0",
+    "aicsimageio>=4.11.0",
+    "bioio>=1.1.0",
     "colorcet>=3.0.1",
+    "einops>=0.7.0",
     "fastremap>=1.14.0",
+    "geojson>=3.1.0",
+    "matplotlib>=3.8.2",
+    "numpy>=1.26.4",
+    "palettable>=3.3.3",
     "rasterio>=1.3.9",
-    "tifffile>=2023.12.9",
     "scikit-image>=0.22.0",
-    "aicsimageio>=4.11.0",
-    "bioio>=1.1.0",
-    "tiffslide>=2.4.0",
     "slideio>=2.6.2",
+    "tifffile>=2023.12.9",
+    "tiffslide>=2.4.0",
+    "torch>=2.1.1",
+    "tqdm>=4.66.1"
 ]
 
+
 [project.optional-dependencies]
 test = [
     "pytest~=8.2.2",
@@ -42,36 +44,39 @@ test = [
     "genbadge~=1.1.1"
 ]
 full = [
+    "aicsimageio>=4.10.0",
+    "bioimageio.core==0.5.11",
+    "bioimageio.spec==0.4.9",
+    "colorcet>=3.0.1",
+    "cucim>=23.10.00; platform_system=='Linux'",
+    "edt",
+    "einops>=0.7.0",
+    "fastremap>=1.14.0",
+    "hyperopt>=0.2.7",
+    "imagecodecs>=2024.1.1",
     "ipykernel>=6.28.0",
     "ipywidgets>=8.1.1",
     "jupyterlab>=4.0.10",
-    "tqdm>=4.66.1",
-    "imagecodecs>=2024.1.1",
-    "seaborn>=0.13.1",
+    "kornia>=0.7.0",
+    "line_profiler>=4.1.3",
     "matplotlib>=3.8.2",
+    "monai>=1.3.0",
     "numba>=0.58.1",
-    "colorcet>=3.0.1",
-    "fastremap>=1.14.0",
-    "scipy>=1.11.4",
-    "scikit-learn>=1.3.2",
-    "rasterio>=1.3.9",
-    "tifffile>=2023.12.9",
-    "scikit-image>=0.22.0",
+    "onnx>=1.15.0",
     "opencv-python>=4.9.0",
-    "hyperopt>=0.2.7",
     "pip>=23.3.2",
-    "aicsimageio>=4.10.0",
-    "bioimageio.core==0.5.11",
-    "bioimageio.spec==0.4.9",
-    "onnx>=1.15.0",
+    "rasterio>=1.3.9",
+    "scikit-image>=0.22.0",
+    "scikit-learn>=1.3.2",
+    "scipy>=1.11.4",
+    "seaborn>=0.13.1",
+    "stardist>=0.8.5",
     "tensorboard>=2.17.0",
+    "tifffile>=2023.12.9",
+    "tkinterdnd2>=0.4.2",
     "torchstain>=1.3.0",
-    "cucim>=23.10.00; platform_system=='Linux'",
-    "stardist>=0.8.5",
     "torchvision>=0.16.1",
-    "kornia>=0.7.0",
-    "monai>=1.3.0",
-    "edt"
+    "tqdm>=4.66.1"
 ]
 
 [tool.hatch.build]


Test patch:
diff --git a/instanseg/scripts/test.py b/instanseg/scripts/test.py
index 0a31d33..425dd31 100755
--- a/instanseg/scripts/test.py
+++ b/instanseg/scripts/test.py
@@ -41,7 +41,7 @@ def instanseg_inference(val_images, val_labels, model, postprocessing_fn, device
             tta.Rotate90(angles=[0, 180, 90, 270]),
         ])
 
-    from instanseg.utils.tiling import instanseg_padding, recover_padding
+    from instanseg.utils.tiling import _instanseg_padding, _recover_padding
     count = 0
     time_dict = {'preprocessing': 0, 'model': 0, 'postprocessing': 0, 'torchscript': 0, 'combined': []}
 
@@ -57,11 +57,11 @@ def instanseg_inference(val_images, val_labels, model, postprocessing_fn, device
     imgs = imgs.to(device)
     imgs, _ = Augmenter.normalize(imgs)
     with torch.no_grad():
-        imgs, pad = instanseg_padding(imgs, extra_pad=0, min_dim=32)
+        imgs, pad = _instanseg_padding(imgs, extra_pad=0, min_dim=32)
         with torch.amp.autocast("cuda"):
             pred = model(imgs[None,])
         pred = pred.float()
-        pred = recover_padding(pred, pad).squeeze(0)
+        pred = _recover_padding(pred, pad).squeeze(0)
         if params is not None:
             with torch.amp.autocast("cuda"):
                 lab = postprocessing_fn(pred, **params, window_size=parser_args.window_size)
@@ -86,12 +86,12 @@ def instanseg_inference(val_images, val_labels, model, postprocessing_fn, device
             start = time.time()
 
             if not tta:
-                imgs, pad = instanseg_padding(imgs, extra_pad=0, min_dim=32, ensure_square=False)
+                imgs, pad = _instanseg_padding(imgs, extra_pad=0, min_dim=32, ensure_square=False)
                 with torch.amp.autocast("cuda"):
                     pred = model(imgs[None,])
 
-                pred = recover_padding(pred, pad).squeeze(0)
-                imgs = recover_padding(imgs, pad).squeeze(0)
+                pred = _recover_padding(pred, pad).squeeze(0)
+                imgs = _recover_padding(imgs, pad).squeeze(0)
                 torch.cuda.synchronize()
 
                 model_time = time.time() - start
diff --git a/tests/utils/test_tiling.py b/tests/utils/test_tiling.py
index d1174eb..81f4de1 100644
--- a/tests/utils/test_tiling.py
+++ b/tests/utils/test_tiling.py
@@ -2,7 +2,7 @@
 def test_tiling():
 
     import torch
-    from instanseg.utils.tiling import chops, tiles_from_chops, stitch
+    from instanseg.utils.tiling import _chops, _tiles_from_chops, _stitch
     from instanseg.utils.pytorch_utils import torch_sparse_onehot, fast_sparse_dual_iou, connected_components
 
     torch.random.manual_seed(0)
@@ -13,11 +13,11 @@ def test_tiling():
     max_cell_size = 20
     window_size = (70,70)
 
-    tuple_index = chops(input_tensor.shape, shape=window_size, overlap=2 * (overlap + max_cell_size))
-    tile_list = tiles_from_chops(input_tensor, shape=window_size, tuple_index=tuple_index)
+    tuple_index = _chops(input_tensor.shape, shape=window_size, overlap=2 * (overlap + max_cell_size))
+    tile_list = _tiles_from_chops(input_tensor, shape=window_size, tuple_index=tuple_index)
     assert len(tile_list) == len(tuple_index[0]) * len(tuple_index[1])
     labels_list = [lab for lab in tile_list]
-    output = stitch([lab[0,0] for lab in labels_list],
+    output = _stitch([lab[0,0] for lab in labels_list],
                                 shape=window_size,
                                 chop_list=tuple_index,
                                 offset = overlap,