axondeepseg / axondeepseg

Axon/Myelin segmentation using Deep Learning
https://axondeepseg.rtfd.io/
MIT License
114 stars 32 forks source link

Add mandatory preprocessing steps before inference #805

Open hermancollin opened 1 month ago

hermancollin commented 1 month ago

Describe the problem Previously with the ivadomed API, images were loaded internally by ivadomed which took care of file formats and input shape. The problem is that nnunet is more rigid concerning its inputs. Namely, it will fail when:

  1. images are not in the expected file format. For instance, we trained all nnunet models on PNG images, so these models will only accept PNG inputs. We would need to manually save a temporary PNG file for inference.
  2. images don't have the expected shape (i.e. nb of channels). During my years working on ADS, I have seen countless occurences of black&white images saved in RGB format. We also need to take this into account.

Proposed solution In the segment.py CLI, before calling the core inference function (apply_model.axon_segmentation(...)), we check if the inputs meet the 2 conditions above, and otherwise save an intermediary file + edit the input list.

Related to #800

hermancollin commented 2 weeks ago

I recently added this feature in #800. The relevant code is centralized in the prepare_input() function defined in segment.py. Here is its signature and docstring:

def prepare_inputs(path_imgs: List[Path], file_format: str, n_channels: int) -> List[Path]:
    '''
    Verifies if the input images can be sent to axon_segmentation(). Otherwise, 
    converts and saves in expected format.

    Parameters
    ----------
    path_imgs : List(Path)
        List of paths to the images to prepare.
    file_format : str
        Expected file format for the images, e.g. '.png' or '.tif'.
    n_channels : int
        Number of channels expected by the model, e.g. 1 for grayscale.

    Error
    -----
    If n_channels > 1 and the image doesn't have enough channels, an error is raised.
    Note that if n_channels == 1, the image will be converted to grayscale if needed.

    Returns
    -------
    filelist : List(Path)
        List of paths to sanitized images. If an image is already in the 
        expected format, the path is the same as the one in path_imgs. Otherwise, 
        will be the path to the converted image.
    '''

Again, this is required because we are feeding nnunet a list of filenames. The only way to avoid this annoying step would be to load the images ourselves and feed a list of npy arrays instead. This would have many advantages, but one big down side is we would need to deal with memory management on our side. nnunet does it pretty graciously from what I've seen, so I don't see an immediate need for this. After all, ivadomed used to save an intermediary nifti file for ALL inputs, whereas we only save intermediary files when necessary.