LoadBatches1D.py

# -*- coding: utf-8 -*-
"""
Created on Sun Apr 21 13:46:44 2019

@author: Winham

LoadBatches1D.py: 迭代生成训练时的batch

实现参考：https://github.com/divamgupta/image-segmentation-keras/blob/master/LoadBatches.py

"""

import os
import itertools
import numpy as np
from sklearn import preprocessing as prep


def getSigArr(path, sigNorm='scale'):
    sig = np.load(path)
    if sigNorm == 'scale':
        sig = prep.scale(sig)
    elif sigNorm == 'minmax':
        min_max_scaler = prep.MinMaxScaler()
        sig = min_max_scaler.fit_transform(sig)
    return np.expand_dims(sig, axis=1)


def getSegmentationArr(path, nClasses=3, output_length=1800, class_value=[0, 0.5, 1]):
    # class_value是在generate_labels.py中定义的，背景0，正常0.5，早搏1，必须要保持一致
    seg_labels = np.zeros([output_length, nClasses])
    seg = np.load(path)
    for i in range(nClasses):
        seg_labels[:, i] = (seg == class_value[i]).astype(float)
    return seg_labels


def SigSegmentationGenerator(sigs_path, segs_path, batch_size, n_classes, output_length=1800):
    sigs = os.listdir(sigs_path)
    segmentations = os.listdir(segs_path)
    sigs.sort()
    segmentations.sort()
    for i in range(len(sigs)):
        sigs[i] = sigs_path + sigs[i]
        segmentations[i] = segs_path + segmentations[i]
    assert len(sigs) == len(segmentations)
    for sig, seg in zip(sigs, segmentations):
        assert (sig.split('/')[-1].split(".")[0] == seg.split('/')[-1].split(".")[0])
    zipped = itertools.cycle(zip(sigs, segmentations))
    while True:
        X = []
        Y = []
        for _ in range(batch_size):
            sig, seg = next(zipped)
            X.append(getSigArr(sig))
            Y.append(getSegmentationArr(seg, n_classes, output_length))
        yield np.array(X), np.array(Y)