dataset.py

"""
data generator for feeding data into pytorch models

Augmentations
-------------
    - label smoothing (label, on the fly)
    - flip (signal, on the fly)
    - (re-)normalize to random mean and std (signal, on the fly)
    - baseline wander (signal, on the fly, combination of sinusoidal noise of several different frequencies, together with an optional Gaussian noise)
    - sinusoidal noise (signal, on the fly, done in baseline wander)
    - Gaussian noise (signal, on the fly, done in baseline wander)
    - stretch and compress (signal, offline)

Issues
------
1. flat segments are found in the original CPSC dataset, e.g. 170*4000 to 185*4000,
which could be checked for example via:
>>> raw_data = ds.reader.load_data("A02", data_format="flat")
>>> flat_segs = []
>>> for idx in range(len(raw_data)//ds.seglen):
>>>     seg_data = raw_data[idx*ds.seglen:(idx+1)*ds.seglen]
>>>     if ds._get_seg_ampl(seg_data) < 0.1:
>>>         flat_segs.append(idx)
>>>     print(f"{idx+1}/{len(raw_data)//ds.seglen}", end="\r")
>>> ds.reader.plot(rec="A02", sampfrom=169*ds.seglen, sampto=186*ds.seglen)
also for sliced segments via:
>>> flat_segs = {rec:[] for rec in ds.reader.all_records}
>>> valid_segs = {rec:[] for rec in ds.reader.all_records}
>>> for i, rec in enumerate(ds.reader.all_records):
>>>     for idx, seg in enumerate(ds.all_segments[rec]):
>>>         seg_data = ds._load_seg_data(seg)
>>>         if ds._get_seg_ampl(seg_data) < 0.1:
>>>             flat_segs[rec].append(seg)
>>>         else:
>>>             valid_segs[rec].append(seg)
>>>         print(f"{idx+1}/{len(ds.all_segments[rec])} @ {i+1}/{len(ds.reader.all_records)}", end="\r")
UPDATE: using `ecg_denoise` as follows:
>>> valid_segs = {rec:[] for rec in ds.all_segments.keys()}
>>> invalid_segs = []
>>> for i, (rec, l_seg) in enumerate(ds.all_segments.items()):
>>>     for idx, seg in enumerate(l_seg):
>>>         if ecg_denoise(ds._load_seg_data(seg), ds.reader.fs ,{"ampl_min":0.15}) == [[0,ds.seglen]]:
>>>             valid_segs[rec].append(seg)
>>>         else:
>>>             invalid_segs.append(seg)
>>>         print(f"{idx+1}/{len(l_seg)} @ {i+1}/{len(ds.all_segments)}", end="\r")

References
----------
[1] Cai, Wenjie, and Danqin Hu. "QRS complex detection using novel deep learning neural networks." IEEE Access (2020).
[2] Tan, Jen Hong, et al. "Application of stacked convolutional and long short-term memory network for accurate identification of CAD ECG signals." Computers in biology and medicine 94 (2018): 19-26.
[3] Yao, Qihang, et al. "Multi-class Arrhythmia detection from 12-lead varied-length ECG using Attention-based Time-Incremental Convolutional Neural Network." Information Fusion 53 (2020): 174-182.
"""

import json
from copy import deepcopy
from itertools import repeat
from pathlib import Path
from random import randint, sample, shuffle, uniform
from typing import Dict, List, Optional, Sequence, Tuple, Union

import numpy as np
import torch
from scipy import signal as SS
from scipy.io import loadmat, savemat
from torch.utils.data.dataset import Dataset

try:
    import torch_ecg  # noqa: F401
except ModuleNotFoundError:
    import sys

    sys.path.insert(0, str(Path(__file__).absolute().parents[2]))

import signal_processing as SP
from cfg import ModelCfg, PreprocCfg

# from torch_ecg._preprocessors import PreprocManager
from torch_ecg.augmenters.baseline_wander import gen_baseline_wander
from torch_ecg.cfg import CFG
from torch_ecg.databases import CPSC2020 as CR
from torch_ecg.utils.misc import ReprMixin, get_record_list_recursive3, list_sum
from torch_ecg.utils.utils_data import mask_to_intervals

if ModelCfg.torch_dtype == torch.float64:
    torch.set_default_tensor_type(torch.DoubleTensor)
    _DTYPE = np.float64
else:
    _DTYPE = np.float32


__all__ = [
    "CPSC2020",
]


class CPSC2020(ReprMixin, Dataset):
    """Data generator for deep learning models.

    Parameters
    ----------
    config : dict
        Configurations for the Dataset, ref. `cfg.TrainCfg`
    training : bool, default True
        If True, the training set will be loaded, otherwise the test set.

    strategy
    --------
    1. slice each record into short segments of length `TrainCfg.input_len`,
       and of overlap length `TrainCfg.overlap_len` around premature beats.
    2. do augmentations for premature segments.

    """

    __DEBUG__ = False
    __name__ = "CPSC2020"

    def __init__(self, config: CFG, training: bool = True) -> None:
        super().__init__()
        self.config = deepcopy(config)
        assert self.config.db_dir is not None, "db_dir must be specified"
        self.config.db_dir = Path(self.config.db_dir)
        self.reader = CR(db_dir=self.config.db_dir)
        if ModelCfg.torch_dtype.lower() == "double":
            self.dtype = np.float64
        else:
            self.dtype = np.float32
        self.allowed_preproc = PreprocCfg.preproc

        self.all_classes = self.config.classes
        self.n_classes = len(self.config.classes)

        self.training = training
        split_res = self.reader.train_test_split_rec(test_rec_num=self.config.test_rec_num)
        self.__data_aug = self.training

        self.seglen = self.config.input_len  # alias, for simplicity

        # create directories if needed
        # preprocess_dir stores pre-processed signals
        self.preprocess_dir = self.config.db_dir / "preprocessed"
        self.preprocess_dir.mkdir(parents=True, exist_ok=True)
        # segments_dir for sliced segments
        self.segments_dir = self.config.db_dir / "segments"
        self.segments_dir.mkdir(parents=True, exist_ok=True)
        # rpeaks_dir for detected r peaks, for optional use
        self.rpeaks_dir = self.config.db_dir / "rpeaks"
        self.rpeaks_dir.mkdir(parents=True, exist_ok=True)

        if self.config.model_name.lower() in ["crnn", "seq_lab"]:
            # for classification, or for sequence labeling
            self.segments_dirs = CFG()
            self.__all_segments = CFG()
            self.segments_json = self.segments_dir / "crnn_segments.json"
            self._ls_segments()

            if self.training:
                self.segments = list_sum([self.__all_segments[rec] for rec in split_res.train])
                shuffle(self.segments)
            else:
                self.segments = list_sum([self.__all_segments[rec] for rec in split_res.test])
        # elif self.config.model_name.lower() == "od":  # object detection
        #     pass
        else:
            raise NotImplementedError(f"data generator for model \042{self.config.model_name}\042 not implemented")

        if self.config.bw:
            self._n_bw_choices = len(self.config.bw_ampl_ratio)
            self._n_gn_choices = len(self.config.bw_gaussian)

        if len(self) == 0:
            raise ValueError(
                "No segments found, please check the data directory, "
                "or call the `persistence` method to preprocess and slice the recordings into segments."
            )

    def _ls_segments(self) -> None:
        """List all segments of each record, and save the list into a json file."""
        for item in ["data", "ann"]:
            self.segments_dirs[item] = CFG()
            for rec in self.reader.all_records:
                self.segments_dirs[item][rec] = self.segments_dir / item / rec
                self.segments_dirs[item][rec].mkdir(parents=True, exist_ok=True)
        if self.segments_json.is_file():
            self.__all_segments = json.loads(self.segments_json.read_text())
            return
        print(f"please allow the reader a few minutes to collect the segments from {self.segments_dir}...")
        seg_filename_pattern = f"S\\d{{2}}_\\d{{7}}\\.{self.reader.rec_ext}"
        self.__all_segments = CFG(
            {
                rec: get_record_list_recursive3(str(self.segments_dirs.data[rec]), seg_filename_pattern)
                for rec in self.reader.all_records
            }
        )
        if all([len(self.__all_segments[rec]) > 0 for rec in self.reader.all_records]):
            self.segments_json.write_text(json.dumps(self.__all_segments, ensure_ascii=False))

    @property
    def all_segments(self):
        return self.__all_segments

    def __getitem__(self, index: int) -> Tuple[np.ndarray, np.ndarray]:
        seg_name = self.segments[index]
        seg_data = self._load_seg_data(seg_name)
        if self.config.model_name.lower() == "crnn":
            seg_label = self._load_seg_label(seg_name)
        elif self.config.model_name.lower() == "seq_lab":
            seg_label = self._load_seg_seq_lab(
                seg=seg_name,
                reduction=self.config.seq_lab_reduction,
            )
        # seg_ampl = np.max(seg_data) - np.min(seg_data)
        seg_ampl = self._get_seg_ampl(seg_data)
        # spb_indices = ann["SPB_indices"]
        # pvc_indices = ann["PVC_indices"]
        if self.__data_aug:
            if self.config.bw:
                ar = self.config.bw_ampl_ratio[randint(0, self._n_bw_choices - 1)]
                gm, gs = self.config.bw_gaussian[randint(0, self._n_gn_choices - 1)]
                bw_ampl = ar * seg_ampl
                g_ampl = gm * seg_ampl
                bw = gen_baseline_wander(
                    siglen=self.seglen,
                    fs=self.config.fs,
                    bw_fs=self.config.bw_fs,
                    amplitude=bw_ampl,
                    amplitude_mean=gm,
                    amplitude_std=gs,
                )
                seg_data = seg_data + bw
            if len(self.config.flip) > 0:
                sign = sample(self.config.flip, 1)[0]
                seg_data *= sign
            if self.config.random_normalize:
                rn_mean = uniform(
                    self.config.random_normalize_mean[0],
                    self.config.random_normalize_mean[1],
                )
                rn_std = uniform(
                    self.config.random_normalize_std[0],
                    self.config.random_normalize_std[1],
                )
                seg_data = (seg_data - np.mean(seg_data) + rn_mean) / np.std(seg_data) * rn_std
            if self.config.label_smoothing > 0:
                seg_label = (1 - self.config.label_smoothing) * seg_label + self.config.label_smoothing / self.n_classes

        if self.__DEBUG__:
            self.reader.plot(
                rec="",  # unnecessary indeed
                data=seg_data,
                ann=self._load_seg_beat_ann(seg_name),
                ticks_granularity=2,
            )

        seg_data = seg_data.reshape((self.config.n_leads, self.seglen))

        return seg_data, seg_label

    def __len__(self) -> int:
        return len(self.segments)

    def _get_seg_ampl(self, seg_data: np.ndarray, window: int = 80) -> float:
        """Get amplitude of a segment.

        Parameters
        ----------
        seg_data : numpy.ndarray
            Data of the segment.
        window : int, default 80 (corr. to 200ms)
            Window length of a window for computing amplitude, with units in number of sample points.

        Returns
        -------
        ampl : float
            Amplitude of `seg_data`.

        """
        half_window = window // 2
        ampl = 0
        for idx in range(len(seg_data) // half_window - 1):
            s = seg_data[idx * half_window : idx * half_window + window]
            ampl = max(ampl, np.max(s) - np.min(s))
        return ampl

    def _get_seg_data_path(self, seg: str) -> Path:
        """Get the path of the data file of the segment.

        Parameters
        ----------
        seg : str
            Name of the segment, of pattern like "S01_0000193".

        Returns
        -------
        fp : pathlib.Path
            Path of the data file of the segment.

        """
        rec = seg.split("_")[0].replace("S", "A")
        fp = self.segments_dir / "data" / rec / f"{seg}{self.reader.rec_ext}"
        return fp

    def _get_seg_ann_path(self, seg: str) -> Path:
        """Get the path of the annotation file of the segment.

        Parameters
        ----------
        seg : str
            Name of the segment, of pattern like "S01_0000193".

        Returns
        -------
        fp : pathlib.Path
            Path of the annotation file of the segment.

        """
        rec = seg.split("_")[0].replace("S", "A")
        fp = self.segments_dir / "ann" / rec / f"{seg}{self.reader.rec_ext}"
        return fp

    def _load_seg_data(self, seg: str) -> np.ndarray:
        """Load the data of the segment.

        Parameters
        ----------
        seg : str
            Name of the segment, of pattern like "S01_0000193".

        Returns
        -------
        seg_data : numpy.ndarray
            Data of the segment, of shape ``(self.seglen,)``.

        """
        seg_data_fp = self._get_seg_data_path(seg)
        seg_data = loadmat(str(seg_data_fp))["ecg"].squeeze()
        return seg_data

    def _load_seg_label(self, seg: str) -> np.ndarray:
        """Load the label of the segment.

        Parameters
        ----------
        seg : str
            Name of the segment, of pattern like "S01_0000193".

        Returns
        -------
        seg_label : numpy.ndarray
            Label of the segment, of shape ``(self.n_classes,)``.

        """
        seg_ann_fp = self._get_seg_ann_path(seg)
        seg_label = loadmat(str(seg_ann_fp))["label"].squeeze()
        return seg_label

    def _load_seg_beat_ann(self, seg: str) -> Dict[str, np.ndarray]:
        """Load the beat annotation of the segment.

        Parameters
        ----------
        seg : str
            Name of the segment, of pattern like "S01_0000193".

        Returns
        -------
        seg_beat_ann : dict
            "SPB_indices", "PVC_indices", each of :class:`numpy.ndarray` values.

        """
        seg_ann_fp = self._get_seg_ann_path(seg)
        seg_beat_ann = loadmat(str(seg_ann_fp))
        seg_beat_ann = {k: v.flatten() for k, v in seg_beat_ann.items() if k in ["SPB_indices", "PVC_indices"]}
        return seg_beat_ann

    def _load_seg_seq_lab(self, seg: str, reduction: int = 8) -> np.ndarray:
        """Load the sequence label of the segment.

        Parameters
        ----------
        seg : str
            Name of the segment, of pattern like "S01_0000193".
        reduction : int, default 8
            Reduction (granularity) of length of the model output,
            compared to the original signal length.

        Returns
        -------
        seq_lab : numpy.ndarray
            Label of the sequence,
            of shape ``(self.seglen//reduction, self.n_classes)``.

        """
        seg_beat_ann = {k: np.round(v / reduction).astype(int) for k, v in self._load_seg_beat_ann(seg).items()}
        bias_thr = int(round(self.config.bias_thr / reduction))
        seq_lab = np.zeros(
            shape=(self.seglen // reduction, self.n_classes),
            dtype=_DTYPE,
        )
        for p in seg_beat_ann["SPB_indices"]:
            start_idx = max(0, p - bias_thr)
            end_idx = min(seq_lab.shape[0], p + bias_thr + 1)
            seq_lab[start_idx:end_idx, self.config.classes.index("S")] = 1
        for p in seg_beat_ann["PVC_indices"]:
            start_idx = max(0, p - bias_thr)
            end_idx = min(seq_lab.shape[0], p + bias_thr + 1)
            seq_lab[start_idx:end_idx, self.config.classes.index("V")] = 1

        return seq_lab

    def disable_data_augmentation(self) -> None:
        self.__data_aug = False

    def enable_data_augmentation(self) -> None:
        self.__data_aug = True

    def persistence(self, force_recompute: bool = False, verbose: int = 0) -> None:
        """Make the dataset persistent w.r.t. the ratios in `self.config`.

        Parameters
        ----------
        force_recompute : bool, default False
            If True, recompute regardless of possible existing files.
        verbose : int, default 0
            Verbosity level for logging messages.

        """
        # if force_recompute:
        #     self.__all_segments = CFG({rec: [] for rec in self.reader.all_records})
        self._preprocess_data(
            self.allowed_preproc,
            force_recompute=force_recompute,
            verbose=verbose,
        )
        self._slice_data(
            force_recompute=force_recompute,
            verbose=verbose,
        )

    def _preprocess_data(self, preproc: List[str], force_recompute: bool = False, verbose: int = 0) -> None:
        """Preprocesses the ecg data in advance for further use,
        offline for `self.persistence`.

        Parameters
        ----------
        preproc : list of str
            Type of preprocesses to perform,
            should be sublist of `self.allowed_preproc`
        force_recompute : bool, default False
            If True, recompute regardless of possible existing files.
        verbose : int, default 0
            Verbosity level for logging messages.

        """
        preproc = self._normalize_preprocess_names(preproc, True)
        config = deepcopy(PreprocCfg)
        config.preproc = preproc
        for idx, rec in enumerate(self.reader.all_records):
            self._preprocess_one_record(
                rec=rec,
                config=config,
                force_recompute=force_recompute,
                verbose=verbose,
            )
            if verbose >= 1:
                print(f"{idx+1}/{len(self.reader.all_records)} records", end="\r")

    def _preprocess_one_record(
        self,
        rec: Union[int, str],
        config: dict,
        force_recompute: bool = False,
        verbose: int = 0,
    ) -> None:
        """Preprocesses the ecg data in advance for further use,
        offline for `self.persistence`.

        Parameters
        ----------
        rec : int or str
            Number of the record, NOTE that rec_no starts from 1,
            or the record name.
        config : dict
            Configurations of preprocessing
        force_recompute : bool, default False
            If True, recompute regardless of possible existing files
        verbose : int, default 0
            Verbosity level for logging messages.

        """
        # format save path
        save_fp = CFG()
        rec_name = self.reader._get_rec_name(rec)
        suffix = self._get_rec_suffix(config.preproc)
        save_fp.data = self.preprocess_dir / f"{rec_name}-{suffix}{self.reader.rec_ext}"
        save_fp.rpeaks = self.rpeaks_dir / f"{rec_name}-{suffix}{self.reader.rec_ext}"
        if (not force_recompute) and save_fp.data.is_file() and save_fp.rpeaks.is_file():
            return
        # perform pre-process
        pps = SP.parallel_preprocess_signal(
            self.reader.load_data(rec, data_format="flat"),
            fs=self.reader.fs,
            config=config,
            verbose=verbose,
        )
        # `rpeaks_skip_dist` useless for `seq_lab_detect`, as already set internally
        # pps['rpeaks'] = pps['rpeaks'][np.where( (pps['rpeaks']>=config.rpeaks_skip_dist) & (pps['rpeaks']<len(pps['filtered_ecg'])-config.rpeaks_skip_dist) )[0]]
        # save mat, keep in accordance with original mat files
        savemat(save_fp.data, {"ecg": np.atleast_2d(pps["filtered_ecg"]).T}, format="5")
        savemat(save_fp.rpeaks, {"rpeaks": np.atleast_2d(pps["rpeaks"]).T}, format="5")

    def _normalize_preprocess_names(self, preproc: List[str], ensure_nonempty: bool) -> List[str]:
        """Transform all preproc into lower case,
        and keep them in a specific ordering.

        Parameters
        ----------
        preproc : list of str
            List of preprocesses types,
            should be sublist of `self.allowd_features`.
        ensure_nonempty : bool,
            If True, when the passed `preproc` is empty,
            `self.allowed_preproc` will be returned.

        Returns
        -------
        list of str
            Normalized list of preprocess types.

        """
        _p = [item.lower() for item in preproc] if preproc else []
        if ensure_nonempty:
            _p = _p or self.allowed_preproc
        # ensure ordering
        _p = [item for item in self.allowed_preproc if item in _p]
        # assert all([item in self.allowed_preproc for item in _p])
        return _p

    def _get_rec_suffix(self, operations: List[str]) -> str:
        """Get the suffix of the filename of the preprocessed ECG signal.

        Parameters
        ----------
        operations : list of str
            Names of operations to perform (or has performed),
            should be sublist of `self.allowed_preproc`.

        Returns
        -------
        suffix : str
            Suffix of the filename of the preprocessed ECG signal.

        """
        suffix = "-".join(sorted([item.lower() for item in operations]))
        return suffix

    def _slice_data(self, force_recompute: bool = False, verbose: int = 0) -> None:
        """Slice all records into segments of length `self.config.input_len`, i.e. `self.seglen`,
        and perform data augmentations specified in `self.config`.

        Parameters
        ----------
        force_recompute : bool, default False
            If True, recompute regardless of possible existing files.
        verbose : int, default 0
            Verbosity level for logging messages.

        """
        for idx, rec in enumerate(self.reader.all_records):
            self._slice_one_record(
                rec=rec,
                force_recompute=force_recompute,
                update_segments_json=False,
                verbose=verbose,
            )
            if verbose >= 1:
                print(f"{idx+1}/{len(self.reader.all_records)} records", end="\r")
        if force_recompute:
            self.segments_json.write_text(json.dumps(self.__all_segments, ensure_ascii=False))

    def _slice_one_record(
        self,
        rec: Union[int, str],
        force_recompute: bool = False,
        update_segments_json: bool = False,
        verbose: int = 0,
    ) -> None:
        """Slice one record into segments of length `self.config.input_len`, i.e. `self.seglen`,
        and perform data augmentations specified in `self.config`.

        Parameters
        ----------
        rec : int or str,
            Number of the record, NOTE that rec_no starts from 1,
            or the record name.
        force_recompute : bool, default False
            If True, recompute regardless of possible existing files.
        update_segments_json : bool, default False
            If both `force_recompute` and `update_segments_json` are True,
            the file `self.segments_json` will be updated,
            useful when slicing not all records.
        verbose : int, default 0
            Verbosity level for logging messages.

        """
        rec_name = self.reader._get_rec_name(rec)
        self.segments_dirs.data[rec_name].mkdir(parents=True, exist_ok=True)
        self.segments_dirs.ann[rec_name].mkdir(parents=True, exist_ok=True)
        if (not force_recompute) and len(self.__all_segments[rec_name]) > 0:
            return
        elif force_recompute:
            self.__all_segments[rec_name] = []

        data = self.reader.load_data(rec, units="mV", data_format="flat")
        ann = self.reader.load_ann(rec)
        border_dist = int(0.5 * self.config.fs)
        forward_len = self.seglen - self.config.overlap_len

        spb_mask = np.zeros((len(data),), dtype=int)
        pvc_mask = np.zeros((len(data),), dtype=int)
        spb_mask[ann["SPB_indices"]] = 1
        pvc_mask[ann["PVC_indices"]] = 1

        # generate initial segments with no overlap for non premature beats
        n_init_seg = len(data) // self.seglen
        segments = (data[: self.seglen * n_init_seg]).reshape((n_init_seg, self.seglen))
        labels = np.zeros((n_init_seg, self.n_classes))
        labels[..., self.config.class_map["N"]] = 1
        # leave only non premature segments
        non_premature = np.logical_or(spb_mask, pvc_mask)[: self.seglen * n_init_seg]
        non_premature = non_premature.reshape((n_init_seg, self.seglen)).sum(axis=1)
        non_premature = np.where(non_premature == 0)[0]
        segments = segments[non_premature, ...]
        labels = labels[non_premature, ...]
        beat_ann = list(
            repeat(
                {
                    "SPB_indices": np.array([], dtype=int),
                    "PVC_indices": np.array([], dtype=int),
                },
                len(non_premature),
            )
        )

        if verbose >= 1:
            print(f"\nn_init_seg = {n_init_seg}")
            print(f"segments.shape = {segments.shape}")
            print(f"finish extracting non-premature segments, totally {len(non_premature)}")
            print("start doing augmentation...")

        # do data augmentation for premature beats
        # first locate all possible premature segments
        # mask for segment start indices
        premature_mask = np.zeros((len(data),), dtype=int)
        for idx in np.concatenate((ann["SPB_indices"], ann["PVC_indices"])):
            start_idx = max(0, idx - self.seglen + border_dist)
            end_idx = max(start_idx, min(idx - border_dist, len(data) - self.seglen))
            premature_mask[start_idx:end_idx] = 1
        # intervals for allowed start of augmented segments
        premature_intervals = mask_to_intervals(premature_mask, 1)
        n_original = 0
        n_added = 0
        for itv in premature_intervals:
            start_idx = itv[0]
            n_original += 1
            while start_idx < itv[1]:
                # stretch and compress the signal
                if self.config.stretch_compress != 0:
                    sign = sample([-1, 1] + [0] * 4, 1)[0]
                    if sign != 0:
                        sc_ratio = self.config.stretch_compress
                        sc_ratio = 1 + (uniform(sc_ratio / 4, sc_ratio) * sign) / 100
                        sc_len = int(round(sc_ratio * self.seglen))
                        end_idx = start_idx + sc_len
                        aug_seg = data[start_idx:end_idx]
                        aug_seg = SS.resample(x=aug_seg, num=self.seglen).reshape((1, -1))
                    else:
                        end_idx = start_idx + self.seglen
                        # the segment of original signal, with no augmentation
                        aug_seg = data[start_idx:end_idx]
                        sc_ratio = 1

                    seg_label = np.zeros((self.n_classes,))
                    seg_spb_inds = np.where(spb_mask[start_idx:end_idx] == 1)[0]
                    seg_pvc_inds = np.where(pvc_mask[start_idx:end_idx] == 1)[0]
                    seg_spb_inds = np.round(seg_spb_inds * sc_ratio).astype(int)
                    seg_pvc_inds = np.round(seg_pvc_inds * sc_ratio).astype(int)
                    seg_beat_ann = {
                        "SPB_indices": seg_spb_inds,
                        "PVC_indices": seg_pvc_inds,
                        "interval": np.array([start_idx, end_idx]),
                    }
                    if len(seg_spb_inds) > 0:
                        seg_label[self.config.class_map["S"]] = 1
                    if len(seg_pvc_inds) > 0:
                        seg_label[self.config.class_map["V"]] = 1
                    seg_label = seg_label.reshape((1, -1))
                    segments = np.append(segments, aug_seg.reshape((1, -1)), axis=0)
                    labels = np.append(labels, seg_label.copy(), axis=0)
                    beat_ann.append(seg_beat_ann.copy())
                    n_added += 1
                    if verbose >= 2:
                        print(
                            f"{n_added} aug seg generated, start_idx at {start_idx}/{len(data)}",
                            end="\r",
                        )

                # start_idx += forward_len  # should be randomly forwarded
                start_idx += randint(forward_len // 4, forward_len)

        if verbose >= 1:
            print(f"\ngenerate {n_added} premature segments out from {n_original} in total via data augmentation")

        # randomly shuffle the data and save into separate files
        seg_inds = list(range(segments.shape[0]))
        shuffle(seg_inds)
        for i, ind in enumerate(seg_inds):
            save_fp = CFG()
            seg_name = f"{rec_name.replace('A', 'S')}_{i:07d}"
            save_fp.data = self.segments_dirs.data[rec_name] / f"{seg_name}{self.reader.rec_ext}"
            save_fp.ann = self.segments_dirs.ann[rec_name] / f"{seg_name}{self.reader.rec_ext}"
            seg = segments[ind, ...]
            # if self._get_seg_ampl(seg) < 0.1:  # drop out flat segments
            #     continue
            if SP.ecg_denoise(seg, self.reader.fs, config={"ampl_min": 0.15}) != [[0, self.seglen]]:
                continue
            savemat(str(save_fp.data), {"ecg": seg}, format="5")
            seg_label = labels[ind, ...]
            seg_beat_ann = beat_ann[ind]
            save_ann_dict = seg_beat_ann.copy()
            save_ann_dict.update({"label": seg_label})
            savemat(str(save_fp.ann), save_ann_dict, format="5")
            self.__all_segments[rec_name].append(seg_name)
            if verbose >= 2:
                print(f"saving {i+1}/{len(seg_inds)}...", end="\r")
        if update_segments_json:
            self.segments_json.write_text(json.dumps(self.__all_segments, ensure_ascii=False))

    def plot_seg(
        self,
        seg: str,
        ticks_granularity: int = 0,
        rpeak_inds: Optional[Union[Sequence[int], np.ndarray]] = None,
    ) -> None:
        """Plot the segment. A helper function for debugging visualization.

        Parameters
        ----------
        seg : str
            Name of the segment, of pattern like "S01_0000193".
        ticks_granularity : int, default 0
            Granularity to plot axis ticks, the higher the more,
            0 (no ticks) --> 1 (major ticks) --> 2 (major + minor ticks).
        rpeak_inds : array_like, optional
            Indices of R peaks.

        """
        seg_data = self._load_seg_data(seg)
        seg_beat_ann = self._load_seg_beat_ann(seg)
        rec_name = seg.split("_")[0].replace("S", "A")
        self.reader.plot(
            rec=rec_name,  # unnecessary indeed
            data=seg_data,
            ann=seg_beat_ann,
            ticks_granularity=ticks_granularity,
            rpeak_inds=rpeak_inds,
        )

    def extra_repr_keys(self) -> List[str]:
        return [
            "training",
            "reader",
        ]