senoquant 1.0.0b2__py3-none-any.whl → 1.0.0b4__py3-none-any.whl

ufish/api.py

@@ -0,0 +1,778 @@
+import os
+import os.path as osp
+import time
+import typing as T
+from pathlib import Path
+from functools import partial
+
+import numpy as np
+import pandas as pd
+
+from .utils.log import logger
+
+if T.TYPE_CHECKING:
+    from torch import nn
+    from matplotlib.figure import Figure
+    import onnxruntime
+
+
+BASE_STORE_URL = 'https://huggingface.co/GangCaoLab/U-FISH/resolve/main/'
+DEFAULT_WEIGHTS_FILE = 'v1.0-alldata-ufish_c32.onnx'
+STATC_STORE_PATH = osp.abspath(
+    osp.join(osp.dirname(__file__), "model/weights/"))
+
+
+class UFish():
+    def __init__(
+            self, cuda: bool = True,
+            default_weights_file: T.Optional[str] = None,
+            local_store_path: str = '~/.ufish/'
+            ) -> None:
+        """
+        Args:
+            cuda: Whether to use GPU.
+            default_weight_file: The default weight file to use.
+            local_store_path: The local path to store the weights.
+        """
+        self._cuda = cuda
+        self._infer_mode = False
+        self.model: T.Optional["nn.Module"] = None
+        self.ort_session: T.Optional["onnxruntime.InferenceSession"] = None
+        if default_weights_file is None:
+            default_weights_file = DEFAULT_WEIGHTS_FILE
+        self.default_weights_file = default_weights_file
+        self.store_base_url = BASE_STORE_URL
+        self.local_store_path = Path(
+            os.path.expanduser(local_store_path))
+        self.weight_path: T.Optional[str] = None
+
+    def init_model(
+            self,
+            model_type: str = 'ufish',
+            **kwargs) -> None:
+        """Initialize the model.
+
+        Args:
+            model_type: The type of the model. For example,
+                'ufish', 'spot_learn', ...
+            kwargs: Other arguments for the model.
+        """
+        import torch
+        if model_type == 'ufish':
+            from .model.network.ufish_net import UFishNet
+            self.model = UFishNet(**kwargs)
+        elif model_type == 'spot_learn':
+            from .model.network.spot_learn import SpotLearn
+            self.model = SpotLearn(**kwargs)
+        else:
+            raise ValueError(f'Unknown model type: {model_type}')
+        params = sum(p.numel() for p in self.model.parameters())
+        logger.info(
+            f'Initializing {model_type} model with kwargs: {kwargs}')
+        logger.info(f'Number of parameters: {params}')
+        self.cuda = False
+        if self._cuda:
+            if torch.cuda.is_available():
+                self.model = self.model.cuda()
+                self.cuda = True
+                logger.info('CUDA is available, using GPU.')
+            else:
+                logger.warning('CUDA is not available, using CPU.')
+        else:
+            logger.info('CUDA is not used, using CPU.')
+
+    def convert_to_onnx(
+            self,
+            output_path: T.Union[Path, str],) -> None:
+        """Convert the model to ONNX format.
+
+        Args:
+            output_path: The path to the output ONNX file.
+        """
+        if self.model is None:
+            raise RuntimeError('Model is not initialized.')
+        self._turn_on_infer_mode(trace_model=True)
+        import torch
+        import torch.onnx
+        output_path = str(output_path)
+        logger.info(
+            f'Converting model to ONNX format, saving to {output_path}')
+        device = torch.device('cuda' if self.cuda else 'cpu')
+        inp = torch.rand(1, 1, 512, 512).to(device)
+        dyn_axes = {0: 'batch_size', 2: 'y', 3: 'x'}
+        torch.onnx.export(
+            self.model, inp, output_path,
+            input_names=['input'],
+            output_names=['output'],
+            opset_version=11,
+            do_constant_folding=True,
+            dynamic_axes={
+                'input': dyn_axes,
+                'output': dyn_axes,
+            },
+        )
+
+    def _turn_on_infer_mode(self, trace_model: bool = False) -> None:
+        """Turn on the infer mode."""
+        if self._infer_mode:
+            return
+        self._infer_mode = True
+        self.model.eval()
+        if trace_model:
+            import torch
+            device = next(self.model.parameters()).device
+            inp = torch.rand(1, 1, 512, 512).to(device)
+            self.model = torch.jit.trace(self.model, inp)
+
+    def load_weights_from_internet(
+            self,
+            weights_file: T.Optional[str] = None,
+            max_retry: int = 8,
+            force_download: bool = False,
+            ) -> None:
+        """Load weights from the huggingface repo.
+
+        Args:
+            weights_file: The name of the weights file on the internet.
+                See https://huggingface.co/GangCaoLab/U-FISH/tree/main
+                for available weights files.
+            max_retry: The maximum number of retries.
+            force_download: Whether to force download the weights.
+        """
+        import torch
+        weights_file = weights_file or self.default_weights_file
+        weight_url = self.store_base_url + weights_file
+        local_weight_path = self.local_store_path / weights_file
+        if local_weight_path.exists() and (not force_download):
+            logger.info(
+                f'Local weights {local_weight_path} exists, '
+                'skip downloading.'
+            )
+        else:
+            logger.info(
+                f'Downloading weights from {weight_url}, '
+                f'storing to {local_weight_path}')
+            local_weight_path.parent.mkdir(parents=True, exist_ok=True)
+            try_count = 0
+            while try_count < max_retry:
+                try:
+                    torch.hub.download_url_to_file(
+                        weight_url, local_weight_path)
+                    break
+                except Exception as e:
+                    logger.warning(f'Error downloading weights: {e}')
+                    try_count += 1
+                    time.sleep(0.5)
+            else:
+                raise RuntimeError(
+                    f'Error downloading weights from {weight_url}.')
+        self.load_weights_from_path(local_weight_path)
+
+    def load_weights_from_path(
+            self,
+            path: T.Union[Path, str],
+            ) -> None:
+        """Load weights from a local file.
+        The file can be a .pth file or an .onnx file.
+
+        Args:
+            path: The path to the weights file.
+        """
+        path = str(path)
+        self.weight_path = path
+        if path.endswith('.pth'):
+            self._load_pth_file(path)
+        elif path.endswith('.onnx'):
+            self._load_onnx(path)
+        else:
+            raise ValueError(
+                'Weights file must be a pth file or an onnx file.')
+
+    def load_weights(
+            self,
+            weights_path: T.Optional[str] = None,
+            weights_file: T.Optional[str] = None,
+            max_retry: int = 8,
+            force_download: bool = False,
+            ):
+        """Load weights from a local file or the internet.
+
+        Args:
+            weights_path: The path to the weights file.
+            weights_file: The name of the weights file on the internet.
+                See https://huggingface.co/GangCaoLab/U-FISH/tree/main
+                for available weights files.
+            max_retry: The maximum number of retries to download the weights.
+            force_download: Whether to force download the weights.
+        """
+        if weights_path is not None:
+            self.load_weights_from_path(weights_path)
+        else:
+            if weights_file is not None:
+                self.load_weights_from_internet(
+                    weights_file=weights_file,
+                    max_retry=max_retry,
+                    force_download=force_download,
+                )
+            else:
+                weights_path = osp.join(STATC_STORE_PATH, DEFAULT_WEIGHTS_FILE)
+                self.load_weights_from_path(weights_path)
+        return self
+
+    def _load_pth_file(self, path: T.Union[Path, str]) -> None:
+        """Load weights from a local file.
+
+        Args:
+            path: The path to the pth weights file."""
+        import torch
+        if self.model is None:
+            self.init_model()
+        assert self.model is not None
+        path = str(path)
+        logger.info(f'Loading weights from {path}')
+        device = torch.device('cuda' if self.cuda else 'cpu')
+        state_dict = torch.load(path, map_location=device)
+        self.model.load_state_dict(state_dict)
+
+    def _load_onnx(
+            self,
+            onnx_path: T.Union[Path, str],
+            providers: T.Optional[T.List[str]] = None,
+            ) -> None:
+        """Load weights from a local ONNX file,
+        and create an onnxruntime session.
+
+        Args:
+            onnx_path: The path to the ONNX file.
+            providers: The providers to use.
+        """
+        import onnxruntime
+        onnx_path = str(onnx_path)
+        logger.info(f'Loading ONNX from {onnx_path}')
+        if self._cuda:
+            providers = providers or ['CUDAExecutionProvider']
+        else:
+            providers = providers or ['CPUExecutionProvider']
+        self.ort_session = onnxruntime.InferenceSession(
+            onnx_path, providers=providers)
+
+    def infer(self, img: np.ndarray) -> np.ndarray:
+        """Infer the image using the U-Net model."""
+        if self.ort_session is not None:
+            output = self._infer_onnx(img)
+        elif self.model is not None:
+            output = self._infer_torch(img)
+        else:
+            raise RuntimeError(
+                'Both torch model and ONNX model are not initialized.')
+        return output
+
+    def _infer_torch(self, img: np.ndarray) -> np.ndarray:
+        """Infer the image using the torch model."""
+        self._turn_on_infer_mode()
+        import torch
+        tensor = torch.from_numpy(img).float()
+        if self.cuda:
+            tensor = tensor.cuda()
+        with torch.no_grad():
+            output = self.model(tensor)
+        output = output.detach().cpu().numpy()
+        return output
+
+    def _infer_onnx(self, img: np.ndarray) -> np.ndarray:
+        """Infer the image using the ONNX model."""
+        ort_inputs = {self.ort_session.get_inputs()[0].name: img}
+        ort_outs = self.ort_session.run(None, ort_inputs)
+        output = ort_outs[0]
+        return output
+
+    def _enhance_img2d(self, img: np.ndarray) -> np.ndarray:
+        """Enhance a 2D image."""
+        output = self.infer(img[np.newaxis, np.newaxis])[0, 0]
+        return output
+
+    def _enhance_img3d(
+            self, img: np.ndarray, batch_size: int = 4) -> np.ndarray:
+        """Enhance a 3D image."""
+        logger.info(
+            f'Enhancing 3D image in shape {img.shape}, '
+            f'batch size: {batch_size}')
+        output = np.zeros_like(img, dtype=np.float32)
+        for i in range(0, output.shape[0], batch_size):
+            logger.info(
+                f'Enhancing slice {i}-{i+batch_size}/{output.shape[0]}')
+            _slice = img[i:i+batch_size][:, np.newaxis]
+            output[i:i+batch_size] = self.infer(_slice)[:, 0]
+        return output
+
+    def _enhance_2d_or_3d(
+            self,
+            img: np.ndarray,
+            axes: str,
+            batch_size: int = 4,
+            blend_3d: bool = False,
+            ) -> np.ndarray:
+        """Enhance a 2D or 3D image."""
+        from .utils.img import scale_image
+        img = scale_image(img, warning=True)
+        if img.ndim == 2:
+            output = self._enhance_img2d(img)
+        elif img.ndim == 3:
+            if blend_3d:
+                if 'z' not in axes:
+                    logger.warning(
+                        'Image does not have a z axis, ' +
+                        'cannot blend along z axis.')
+                from .utils.img import enhance_blend_3d
+                logger.info(
+                    "Blending 3D image from 3 directions: z, y, x.")
+                output = enhance_blend_3d(
+                    img, self._enhance_img3d, axes=axes,
+                    batch_size=batch_size)
+            else:
+                output = self._enhance_img3d(img, batch_size=batch_size)
+        else:
+            raise ValueError('Image must be 2D or 3D.')
+        return output
+
+    def call_spots(
+            self,
+            enhanced_img: np.ndarray,
+            method: str = 'local_maxima',
+            **kwargs,
+            ) -> pd.DataFrame:
+        """Call spots from enhanced image.
+
+        Args:
+            enhanced_img: The enhanced image.
+            method: The method to use for spot calling.
+            kwargs: Other arguments for the spot calling function.
+        """
+        assert enhanced_img.ndim in (2, 3), 'Image must be 2D or 3D.'
+        if method == 'cc_center':
+            from .utils.spot_calling import call_spots_cc_center as call_func
+        else:
+            from .utils.spot_calling import call_spots_local_maxima as call_func  # noqa
+        df = call_func(enhanced_img, **kwargs)
+        return df
+
+    def _pred_2d_or_3d(
+            self, img: np.ndarray, axes: str,
+            blend_3d: bool = False,
+            batch_size: int = 4,
+            spots_calling_method: str = 'local_maxima',
+            **kwargs,
+            ) -> T.Tuple[pd.DataFrame, np.ndarray]:
+        """Predict the spots in a 2D or 3D image. """
+        assert img.ndim in (2, 3), 'Image must be 2D or 3D.'
+        assert len(axes) == img.ndim, \
+            "axes and image dimension must have the same length"
+        enhanced_img = self._enhance_2d_or_3d(
+            img, axes,
+            batch_size=batch_size,
+            blend_3d=(blend_3d and ('z' in axes))
+        )
+        df = self.call_spots(
+            enhanced_img,
+            method=spots_calling_method,
+            **kwargs)
+        return df, enhanced_img
+
+    def predict(
+            self, img: np.ndarray,
+            enh_img: T.Optional[np.ndarray] = None,
+            axes: T.Optional[str] = None,
+            blend_3d: bool = True,
+            batch_size: int = 4,
+            spots_calling_method: str = 'local_maxima',
+            **kwargs,
+            ) -> T.Tuple[pd.DataFrame, np.ndarray]:
+        """Predict the spots in an image.
+
+        Args:
+            img: The image to predict, it should be a multi dimensional array.
+                For example, shape (c, z, y, x) for a 4D image,
+                shape (z, y, x) for a 3D image,
+                shape (y, x) for a 2D image.
+            enh_img: The enhanced image, if None, will be created.
+                It can be a multi dimensional array or a zarr array.
+            axes: The axes of the image.
+                For example, 'czxy' for a 4D image,
+                'yx' for a 2D image.
+                If None, will try to infer the axes from the shape.
+            blend_3d: Whether to blend the 3D image.
+                Used only when the image contains a z axis.
+                If True, will blend the 3D enhanced images along
+                the z, y, x axes.
+            batch_size: The batch size for inference.
+                Used only when the image dimension is 3 or higher.
+            spots_calling_method: The method to use for spot calling.
+            kwargs: Other arguments for the spot calling function.
+        """
+        from .utils.img import (
+            infer_img_axes, check_img_axes,
+            map_predfunc_to_img
+        )
+        if axes is None:
+            logger.info("Axes not specified, infering from image shape.")
+            axes = infer_img_axes(img.shape)
+            logger.info(f"Infered axes: {axes}, image shape: {img.shape}")
+        check_img_axes(img, axes)
+        if not isinstance(img, np.ndarray):
+            img = np.array(img)
+        predfunc = partial(
+            self._pred_2d_or_3d,
+            blend_3d=blend_3d,
+            batch_size=batch_size,
+            spots_calling_method=spots_calling_method,
+            **kwargs,
+            )
+        df, enhanced_img = map_predfunc_to_img(
+            predfunc, img, axes)
+        if enh_img is not None:
+            enh_img[:] = enhanced_img
+        return df, enhanced_img
+
+    def predict_chunks(
+            self,
+            img: np.ndarray,
+            enh_img: T.Optional[np.ndarray] = None,
+            axes: T.Optional[str] = None,
+            blend_3d: bool = True,
+            batch_size: int = 4,
+            chunk_size: T.Optional[T.Tuple[T.Union[int, str], ...]] = None,
+            spots_calling_method: str = 'local_maxima',
+            **kwargs,
+            ):
+        """Predict the spots in an image chunk by chunk.
+
+        Args:
+            img: The image to predict, it should be a multi dimensional array.
+                For example, shape (c, z, y, x) for a 4D image,
+                shape (z, y, x) for a 3D image,
+                shape (y, x) for a 2D image.
+            enh_img: The enhanced image, if None, will be created.
+                It can be a multi dimensional array or a zarr array.
+            axes: The axes of the image.
+                For example, 'czxy' for a 4D image,
+                'yx' for a 2D image.
+                If None, will try to infer the axes from the shape.
+            blend_3d: Whether to blend the 3D image.
+                Used only when the image contains a z axis.
+                If True, will blend the 3D enhanced images along
+                the z, y, x axes.
+            batch_size: The batch size for inference.
+                Used only when the image dimension is 3 or higher.
+            chunk_size: The chunk size for processing.
+                For example, (1, 512, 512) for a 3D image,
+                (512, 512) for a 2D image.
+                Using 'image' as a dimension will use the whole image
+                as a chunk. For example, (1, 'image', 'image') for a 3D image,
+                ('image', 'image', 'image', 512, 512) for a 5D image.
+                If None, will use the default chunk size.
+            spots_calling_method: The method to use for spot calling.
+            kwargs: Other arguments for the spot calling function.
+        """
+        from .utils.img import (
+            check_img_axes, chunks_iterator,
+            process_chunk_size, infer_img_axes)
+        if axes is None:
+            axes = infer_img_axes(img.shape)
+        check_img_axes(img, axes)
+        if chunk_size is None:
+            from .utils.img import get_default_chunk_size
+            chunk_size = get_default_chunk_size(axes)
+            logger.info(f"Chunk size not specified, using {chunk_size}.")
+        chunk_size = process_chunk_size(chunk_size, img.shape)
+        logger.info(f"Chunk size: {chunk_size}")
+        total_dfs = []
+        if enh_img is None:
+            enh_img = np.zeros_like(img, dtype=np.float32)
+        for c_range, chunk in chunks_iterator(img, chunk_size):
+            logger.info("Processing chunk: " + str(c_range)
+                        + ", chunk shape: " + str(chunk.shape))
+            c_df, c_enh = self.predict(
+                chunk, axes=axes, blend_3d=blend_3d,
+                batch_size=batch_size,
+                spots_calling_method=spots_calling_method,
+                **kwargs)
+            dim_start = [c_range[i][0] for i in range(len(axes))]
+            c_df += dim_start
+            total_dfs.append(c_df)
+            c_enh = c_enh[
+                tuple(slice(0, (r[1]-r[0])) for r in c_range)]
+            enh_img[tuple(slice(*r) for r in c_range)] = c_enh
+        df = pd.concat(total_dfs, ignore_index=True)
+        return df, enh_img
+
+    def evaluate_result_dp(
+            self,
+            pred: pd.DataFrame,
+            true: pd.DataFrame,
+            mdist: float = 3.0,
+            ) -> pd.DataFrame:
+        """Evaluate the prediction result using deepblink metrics.
+
+        Args:
+            pred: The predicted spots.
+            true: The true spots.
+            mdist: The maximum distance to consider a spot as a true positive.
+
+        Returns:
+            A pandas dataframe containing the evaluation metrics."""
+        from .utils.metrics_deepblink import compute_metrics
+        axis_names = list(pred.columns)
+        axis_cols = [n for n in axis_names if n.startswith('axis')]
+        pred = pred[axis_cols].values
+        true = true[axis_cols].values
+        metrics = compute_metrics(
+            pred, true, mdist=mdist)
+        return metrics
+
+    def evaluate_result(
+            self,
+            pred: pd.DataFrame,
+            true: pd.DataFrame,
+            cutoff: float = 3.0,
+            ) -> float:
+        """Calculate the F1 score of the prediction result.
+
+        Args:
+            pred: The predicted spots.
+            true: The true spots.
+            cutoff: The maximum distance to consider a spot as a true positive.
+        """
+        from .utils.metrics import compute_metrics
+        res = compute_metrics(pred.values, true.values, cutoff=cutoff)
+        return res
+
+    def plot_result(
+            self,
+            img: np.ndarray,
+            pred: pd.DataFrame,
+            fig_size: T.Tuple[int, int] = (10, 10),
+            image_cmap: str = 'gray',
+            marker_size: int = 20,
+            marker_color: str = 'red',
+            marker_style: str = 'x',
+            ) -> "Figure":
+        """Plot the prediction result.
+
+        Args:
+            img: The image to plot.
+            pred: The predicted spots.
+            fig_size: The figure size.
+            image_cmap: The colormap for the image.
+            marker_size: The marker size.
+            marker_color: The marker color.
+            marker_style: The marker style.
+        """
+        from .utils.plot import Plot2d
+        plt2d = Plot2d()
+        plt2d.default_figsize = fig_size
+        plt2d.default_marker_size = marker_size
+        plt2d.default_marker_color = marker_color
+        plt2d.default_marker_style = marker_style
+        plt2d.default_imshow_cmap = image_cmap
+        plt2d.new_fig()
+        plt2d.image(img)
+        plt2d.spots(pred.values)
+        return plt2d.fig
+
+    def plot_evaluate(
+            self,
+            img: np.ndarray,
+            pred: pd.DataFrame,
+            true: pd.DataFrame,
+            cutoff: float = 3.0,
+            fig_size: T.Tuple[int, int] = (10, 10),
+            image_cmap: str = 'gray',
+            marker_size: int = 20,
+            tp_color: str = 'green',
+            fp_color: str = 'red',
+            fn_color: str = 'yellow',
+            tp_marker: str = 'x',
+            fp_marker: str = 'x',
+            fn_marker: str = 'x',
+            ) -> "Figure":
+        """Plot the prediction result.
+
+        Args:
+            img: The image to plot.
+            pred: The predicted spots.
+            true: The true spots.
+            cutoff: The maximum distance to consider a spot as a true positive.
+            fig_size: The figure size.
+            image_cmap: The colormap for the image.
+            marker_size: The marker size.
+            tp_color: The color for true positive.
+            fp_color: The color for false positive.
+            fn_color: The color for false negative.
+            tp_marker_style: The marker style for true positive.
+            fp_marker_style: The marker style for false positive.
+            fn_marker_style: The marker style for false negative.
+        """
+        from .utils.plot import Plot2d
+        plt2d = Plot2d()
+        plt2d.default_figsize = fig_size
+        plt2d.default_marker_size = marker_size
+        plt2d.default_imshow_cmap = image_cmap
+        plt2d.new_fig()
+        plt2d.image(img)
+        plt2d.evaluate_result(
+            pred.values, true.values,
+            cutoff=cutoff,
+            tp_color=tp_color,
+            fp_color=fp_color,
+            fn_color=fn_color,
+            tp_marker=tp_marker,
+            fp_marker=fp_marker,
+            fn_marker=fn_marker,
+        )
+        return plt2d.fig
+
+    def _load_dataset(
+            self,
+            path: str,
+            root_dir_path: T.Optional[str] = None,
+            img_glob: str = '*.tif',
+            coord_glob: str = '*.csv',
+            process_func=None,
+            transform=None,
+            ):
+        """Load a dataset from a path."""
+        from .data import FISHSpotsDataset
+        _path = Path(path)
+        if _path.is_dir():
+            if root_dir_path is not None:
+                logger.info(f"Dataset's root directory: {root_dir_path}")
+                _path = Path(root_dir_path) / _path
+            logger.info(f"Loading dataset from dir: {_path}")
+            logger.info(
+                f'Image glob: {img_glob}, Coordinate glob: {coord_glob}')
+            _path_str = str(_path)
+            dataset = FISHSpotsDataset.from_dir(
+                _path_str, _path_str,
+                img_glob=img_glob, coord_glob=coord_glob,
+                process_func=process_func, transform=transform)
+        else:
+            logger.info(f"Loading dataset using meta csv: {_path}")
+            assert _path.suffix == '.csv', \
+                "Meta file must be a csv file."
+            root_dir = root_dir_path or _path.parent
+            logger.info(f'Data root directory: {root_dir}')
+            dataset = FISHSpotsDataset.from_meta_csv(
+                root_dir=root_dir, meta_csv_path=str(_path),
+                process_func=process_func, transform=transform)
+        return dataset
+
+    def train(
+            self,
+            train_path: str,
+            valid_path: str,
+            root_dir: T.Optional[str] = None,
+            img_glob: str = '*.tif',
+            coord_glob: str = '*.csv',
+            target_process: T.Optional[str] = 'gaussian',
+            loss_type: str = 'DiceRMSELoss',
+            loader_workers: int = 4,
+            data_argu: bool = False,
+            argu_prob: float = 0.5,
+            num_epochs: int = 50,
+            batch_size: int = 8,
+            lr: float = 1e-3,
+            summary_dir: str = "runs/unet",
+            model_save_dir: str = "./models",
+            save_period: int = 5,
+            ):
+        """Train the U-Net model.
+
+        Args:
+            train_path: The path to the training dataset.
+                Path to a directory containing images and coordinates,
+                or a meta csv file.
+            valid_path: The path to the validation dataset.
+                Path to a directory containing images and coordinates,
+                or a meta csv file.
+            root_dir: The root directory of the dataset.
+                If using meta csv, the root directory of the dataset.
+            img_glob: The glob pattern for the image files.
+            coord_glob: The glob pattern for the coordinate files.
+            target_process: The target image processing method.
+                'gaussian' or 'dialation' or None.
+                If None, no processing will be applied.
+                default 'gaussian'.
+            loss_type: The loss function type.
+            loader_workers: The number of workers to use for the data loader.
+            data_argu: Whether to use data augmentation.
+            argu_prob: The probability to use data augmentation.
+            num_epochs: The number of epochs to train.
+            batch_size: The batch size.
+            lr: The learning rate.
+            summary_dir: The directory to save the TensorBoard summary to.
+            model_save_dir: The directory to save the model to.
+            save_period: Save the model every `save_period` epochs.
+        """
+        from .model.train import train_on_dataset
+        from .data import FISHSpotsDataset
+        if self.model is None:
+            logger.info('Model is not initialized. Will initialize a new one.')
+            self.init_model()
+        assert self.model is not None
+
+        if data_argu:
+            logger.info(
+                'Using data augmentation. ' +
+                f'Probability: {argu_prob}'
+            )
+            from .data import DataAugmentation
+            transform = DataAugmentation(p=argu_prob)
+        else:
+            transform = None
+
+        logger.info(f'Using {target_process} as target process.')
+        if target_process == 'gaussian':
+            process_func = FISHSpotsDataset.gaussian_filter
+        elif target_process == 'dialation':
+            process_func = FISHSpotsDataset.dialate_mask
+        elif isinstance(target_process, str):
+            from functools import partial
+            process_func = partial(
+                FISHSpotsDataset.dialate_mask,
+                footprint=target_process)
+        else:
+            process_func = None
+
+        logger.info(f"Loading training dataset from {train_path}")
+        train_dataset = self._load_dataset(
+            train_path, root_dir_path=root_dir,
+            img_glob=img_glob, coord_glob=coord_glob,
+            process_func=process_func, transform=transform,
+        )
+        logger.info(f"Loading validation dataset from {valid_path}")
+        valid_dataset = self._load_dataset(
+            valid_path, root_dir_path=root_dir,
+            img_glob=img_glob, coord_glob=coord_glob,
+            process_func=process_func,
+        )
+        logger.info(
+            f"Training dataset size: {len(train_dataset)}, "
+            f"Validation dataset size: {len(valid_dataset)}"
+        )
+        logger.info(
+            f"Number of epochs: {num_epochs}, "
+            f"Batch size: {batch_size}, "
+            f"Learning rate: {lr}"
+        )
+        train_on_dataset(
+            self.model,
+            train_dataset, valid_dataset,
+            loss_type=loss_type,
+            loader_workers=loader_workers,
+            num_epochs=num_epochs,
+            batch_size=batch_size,
+            lr=lr,
+            summary_dir=summary_dir,
+            model_save_dir=model_save_dir,
+            save_period=save_period,
+        )