predict_data_module

class PredictDataModule(L.LightningDataModule):
    """
    CAREamics Lightning prediction data module.

    The data module can be used with Path, str or numpy arrays. The data can be either
    a folder containing images or a single file.

    To read custom data types, you can set `data_type` to `custom` in `data_config`
    and provide a function that returns a numpy array from a path as
    `read_source_func` parameter. The function will receive a Path object and
    an axies string as arguments, the axes being derived from the `data_config`.

    You can also provide a `fnmatch` and `Path.rglob` compatible expression (e.g.
    "*.czi") to filter the files extension using `extension_filter`.

    Parameters
    ----------
    pred_config : InferenceModel
        Pydantic model for CAREamics prediction configuration.
    pred_data : pathlib.Path or str or numpy.ndarray
        Prediction data, can be a path to a folder, a file or a numpy array.
    read_source_func : Callable, optional
        Function to read custom types, by default None.
    extension_filter : str, optional
        Filter to filter file extensions for custom types, by default "".
    dataloader_params : dict, optional
        Dataloader parameters, by default {}.
    """

    def __init__(
        self,
        pred_config: InferenceConfig,
        pred_data: Union[Path, str, NDArray],
        read_source_func: Optional[Callable] = None,
        extension_filter: str = "",
        dataloader_params: Optional[dict] = None,
    ) -> None:
        """
        Constructor.

        The data module can be used with Path, str or numpy arrays. The data can be
        either a folder containing images or a single file.

        To read custom data types, you can set `data_type` to `custom` in `data_config`
        and provide a function that returns a numpy array from a path as
        `read_source_func` parameter. The function will receive a Path object and
        an axies string as arguments, the axes being derived from the `data_config`.

        You can also provide a `fnmatch` and `Path.rglob` compatible expression (e.g.
        "*.czi") to filter the files extension using `extension_filter`.

        Parameters
        ----------
        pred_config : InferenceModel
            Pydantic model for CAREamics prediction configuration.
        pred_data : pathlib.Path or str or numpy.ndarray
            Prediction data, can be a path to a folder, a file or a numpy array.
        read_source_func : Callable, optional
            Function to read custom types, by default None.
        extension_filter : str, optional
            Filter to filter file extensions for custom types, by default "".
        dataloader_params : dict, optional
            Dataloader parameters, by default {}.

        Raises
        ------
        ValueError
            If the data type is `custom` and no `read_source_func` is provided.
        ValueError
            If the data type is `array` and the input is not a numpy array.
        ValueError
            If the data type is `tiff` and the input is neither a Path nor a str.
        """
        if dataloader_params is None:
            dataloader_params = {}
        if dataloader_params is None:
            dataloader_params = {}
        super().__init__()

        # check that a read source function is provided for custom types
        if pred_config.data_type == SupportedData.CUSTOM and read_source_func is None:
            raise ValueError(
                f"Data type {SupportedData.CUSTOM} is not allowed without "
                f"specifying a `read_source_func` and an `extension_filer`."
            )

        # check correct input type
        if (
            isinstance(pred_data, np.ndarray)
            and pred_config.data_type != SupportedData.ARRAY
        ):
            raise ValueError(
                f"Received a numpy array as input, but the data type was set to "
                f"{pred_config.data_type}. Set the data type "
                f"to {SupportedData.ARRAY} to predict on numpy arrays."
            )

        # and that Path or str are passed, if tiff file type specified
        elif (isinstance(pred_data, Path) or isinstance(pred_config, str)) and (
            pred_config.data_type != SupportedData.TIFF
            and pred_config.data_type != SupportedData.CUSTOM
        ):
            raise ValueError(
                f"Received a path as input, but the data type was neither set to "
                f"{SupportedData.TIFF} nor {SupportedData.CUSTOM}. Set the data type "
                f" to {SupportedData.TIFF} or "
                f"{SupportedData.CUSTOM} to predict on files."
            )

        # configuration data
        self.prediction_config = pred_config
        self.data_type = pred_config.data_type
        self.batch_size = pred_config.batch_size
        self.dataloader_params = dataloader_params

        self.pred_data = pred_data
        self.tile_size = pred_config.tile_size
        self.tile_overlap = pred_config.tile_overlap

        # check if it is tiled
        self.tiled = self.tile_size is not None and self.tile_overlap is not None

        # read source function
        if pred_config.data_type == SupportedData.CUSTOM:
            # mypy check
            assert read_source_func is not None

            self.read_source_func: Callable = read_source_func
        elif pred_config.data_type != SupportedData.ARRAY:
            self.read_source_func = get_read_func(pred_config.data_type)

        self.extension_filter = extension_filter

    def prepare_data(self) -> None:
        """Hook used to prepare the data before calling `setup`."""
        # if the data is a Path or a str
        if not isinstance(self.pred_data, np.ndarray):
            self.pred_files = list_files(
                self.pred_data, self.data_type, self.extension_filter
            )

    def setup(self, stage: Optional[str] = None) -> None:
        """
        Hook called at the beginning of predict.

        Parameters
        ----------
        stage : Optional[str], optional
            Stage, by default None.
        """
        # if numpy array
        if self.data_type == SupportedData.ARRAY:
            if self.tiled:
                self.predict_dataset: PredictDatasetType = InMemoryTiledPredDataset(
                    prediction_config=self.prediction_config,
                    inputs=self.pred_data,
                )
            else:
                self.predict_dataset = InMemoryPredDataset(
                    prediction_config=self.prediction_config,
                    inputs=self.pred_data,
                )
        else:
            if self.tiled:
                self.predict_dataset = IterableTiledPredDataset(
                    prediction_config=self.prediction_config,
                    src_files=self.pred_files,
                    read_source_func=self.read_source_func,
                )
            else:
                self.predict_dataset = IterablePredDataset(
                    prediction_config=self.prediction_config,
                    src_files=self.pred_files,
                    read_source_func=self.read_source_func,
                )

    def predict_dataloader(self) -> DataLoader:
        """
        Create a dataloader for prediction.

        Returns
        -------
        DataLoader
            Prediction dataloader.
        """
        return DataLoader(
            self.predict_dataset,
            batch_size=self.batch_size,
            collate_fn=collate_tiles if self.tiled else None,
            **self.dataloader_params,
        )