286 lines
9.7 KiB
Python
286 lines
9.7 KiB
Python
|
import os
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
from torch.utils.data import Dataset, DataLoader
|
|||
|
|
from sklearn.model_selection import train_test_split
|
|||
|
|
import h5py
|
|||
|
|
|
|||
|
|
def custom_collate_fn(batch):
|
|||
|
|
X, y_SBP, y_DBP = zip(*batch)
|
|||
|
|
|
|||
|
|
X = torch.tensor(np.array(X), dtype=torch.float32)
|
|||
|
|
y_SBP = torch.tensor(y_SBP, dtype=torch.float32)
|
|||
|
|
y_DBP = torch.tensor(y_DBP, dtype=torch.float32)
|
|||
|
|
|
|||
|
|
return X, y_SBP, y_DBP
|
|||
|
|
|
|||
|
|
|
|||
|
|
class BPDataset(Dataset):
|
|||
|
|
def __init__(self, X_data, y_SBP, y_DBP):
|
|||
|
|
self.X_data = X_data
|
|||
|
|
self.y_SBP = y_SBP
|
|||
|
|
self.y_DBP = y_DBP
|
|||
|
|
|
|||
|
|
def __len__(self):
|
|||
|
|
return len(self.y_SBP)
|
|||
|
|
|
|||
|
|
def __getitem__(self, idx):
|
|||
|
|
# X_sample = self.X_data[idx * 250:(idx + 1) * 250]
|
|||
|
|
X_sample = self.X_data[idx]
|
|||
|
|
y_SBP_sample = self.y_SBP[idx]
|
|||
|
|
y_DBP_sample = self.y_DBP[idx]
|
|||
|
|
|
|||
|
|
return X_sample, y_SBP_sample, y_DBP_sample
|
|||
|
|
|
|||
|
|
|
|||
|
|
class BPDataLoader:
|
|||
|
|
def __init__(self, data_dir, val_split=0.2, batch_size=32, shuffle=True, data_type='npy'):
|
|||
|
|
self.data_dir = data_dir
|
|||
|
|
self.val_split = val_split
|
|||
|
|
self.batch_size = batch_size
|
|||
|
|
self.shuffle = shuffle
|
|||
|
|
self.train_dataloader = None
|
|||
|
|
self.val_dataloader = None
|
|||
|
|
self.data_type = data_type
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_data(self):
|
|||
|
|
X_BP_path = os.path.join(self.data_dir, 'X_BP.npy')
|
|||
|
|
y_DBP_path = os.path.join(self.data_dir, 'Y_DBP.npy')
|
|||
|
|
y_SBP_path = os.path.join(self.data_dir, 'Y_SBP.npy')
|
|||
|
|
|
|||
|
|
X_BP = np.load(X_BP_path)
|
|||
|
|
# 将数据reshape成(batch_size, 250,1)的形状
|
|||
|
|
X_BP = X_BP.reshape(-1, 250, 1)
|
|||
|
|
|
|||
|
|
y_DBP = np.load(y_DBP_path)
|
|||
|
|
y_SBP = np.load(y_SBP_path)
|
|||
|
|
|
|||
|
|
return X_BP, y_DBP, y_SBP
|
|||
|
|
|
|||
|
|
def load_data_UKL_h5(self):
|
|||
|
|
|
|||
|
|
X_BP_path = os.path.join(self.data_dir, 'rPPG-BP-UKL_rppg_7s.h5')
|
|||
|
|
with h5py.File(X_BP_path, 'r') as f:
|
|||
|
|
rppg = f.get('rppg')
|
|||
|
|
BP = f.get('label')
|
|||
|
|
rppg = np.array(rppg)
|
|||
|
|
BP = np.array(BP)
|
|||
|
|
|
|||
|
|
# 将数据从(875, 7851)reshape成(7851, 875, 1)的形状
|
|||
|
|
rppg = rppg.transpose(1, 0)
|
|||
|
|
rppg = rppg.reshape(-1, 875, 1)
|
|||
|
|
|
|||
|
|
X_BP = rppg
|
|||
|
|
y_DBP = BP[1]
|
|||
|
|
y_SBP = BP[0]
|
|||
|
|
|
|||
|
|
return X_BP, y_DBP, y_SBP
|
|||
|
|
|
|||
|
|
def load_data_MIMIC_h5(self):
|
|||
|
|
|
|||
|
|
X_BP_path = os.path.join(self.data_dir, 'MIMIC-III_ppg_dataset.h5')
|
|||
|
|
|
|||
|
|
#
|
|||
|
|
# 获取data_dir下文件列表
|
|||
|
|
files = os.listdir(self.data_dir)
|
|||
|
|
|
|||
|
|
# 检查是否存在已经处理好的数据
|
|||
|
|
if 'X_MIMIC_BP.npy' in files and 'Y_MIMIC_DBP.npy' in files and 'Y_MIMIC_SBP.npy' in files:
|
|||
|
|
print('loading preprocessed data.....')
|
|||
|
|
|
|||
|
|
X_BP = np.load(os.path.join(self.data_dir, 'X_MIMIC_BP.npy'))
|
|||
|
|
y_DBP = np.load(os.path.join(self.data_dir, 'Y_MIMIC_DBP.npy'))
|
|||
|
|
y_SBP = np.load(os.path.join(self.data_dir, 'Y_MIMIC_SBP.npy'))
|
|||
|
|
|
|||
|
|
return X_BP, y_DBP, y_SBP
|
|||
|
|
|
|||
|
|
with h5py.File(X_BP_path, 'r') as f:
|
|||
|
|
ppg = f.get('ppg')
|
|||
|
|
BP = f.get('label')
|
|||
|
|
ppg = np.array(ppg)
|
|||
|
|
BP = np.array(BP)
|
|||
|
|
|
|||
|
|
# 统计BP中SBP的最大值和最小值
|
|||
|
|
max_sbp = np.max(BP[:, 0])
|
|||
|
|
min_sbp = np.min(BP[:, 0])
|
|||
|
|
|
|||
|
|
max_sbp = 10 - max_sbp % 10 + max_sbp
|
|||
|
|
min_sbp = min_sbp - min_sbp % 10
|
|||
|
|
|
|||
|
|
# 划分区间
|
|||
|
|
bins = np.arange(min_sbp, max_sbp, 10)
|
|||
|
|
|
|||
|
|
print(bins)
|
|||
|
|
|
|||
|
|
sampled_ppg_data = []
|
|||
|
|
sampled_bp_data = []
|
|||
|
|
|
|||
|
|
for i in range(len(bins) - 1):
|
|||
|
|
# 获取当前区间的数据
|
|||
|
|
bin_data_sbp_dbp = BP[(BP[:, 0] >= bins[i]) & (BP[:, 0] < bins[i + 1])]
|
|||
|
|
bin_data_ppg = ppg[(BP[:, 0] >= bins[i]) & (BP[:, 0] < bins[i + 1])]
|
|||
|
|
|
|||
|
|
# 如果当前区间有数据
|
|||
|
|
if len(bin_data_sbp_dbp) > 0:
|
|||
|
|
# 从当前区间中随机抽取20%的数据
|
|||
|
|
num_samples = int(len(bin_data_sbp_dbp) * 0.1)
|
|||
|
|
indices = np.random.choice(len(bin_data_sbp_dbp), num_samples, replace=False)
|
|||
|
|
sampled_bin_data_sbp_dbp = bin_data_sbp_dbp[indices]
|
|||
|
|
sampled_bin_data_ppg = bin_data_ppg[indices]
|
|||
|
|
|
|||
|
|
# 将抽取的数据添加到最终的列表中
|
|||
|
|
sampled_bp_data.append(sampled_bin_data_sbp_dbp)
|
|||
|
|
sampled_ppg_data.append(sampled_bin_data_ppg)
|
|||
|
|
|
|||
|
|
# 将列表中的数据合并成NumPy数组
|
|||
|
|
ppg = np.concatenate(sampled_ppg_data, axis=0)
|
|||
|
|
BP = np.concatenate(sampled_bp_data, axis=0)
|
|||
|
|
|
|||
|
|
print(ppg.shape, BP.shape)
|
|||
|
|
|
|||
|
|
# 将数据从(9054000, 875)reshape成(9054000, 875, 1)的形状
|
|||
|
|
ppg = ppg.reshape(-1, 875, 1)
|
|||
|
|
|
|||
|
|
X_BP = ppg
|
|||
|
|
|
|||
|
|
# 取出第一列赋值给y_DBP，第0列赋值给y_SBP
|
|||
|
|
y_DBP = BP[:, 1]
|
|||
|
|
y_SBP = BP[:, 0]
|
|||
|
|
|
|||
|
|
# 将数据保存到文件中
|
|||
|
|
np.save('data/X_MIMIC_BP.npy', X_BP)
|
|||
|
|
np.save('data/Y_MIMIC_DBP.npy', y_DBP)
|
|||
|
|
np.save('data/Y_MIMIC_SBP.npy', y_SBP)
|
|||
|
|
|
|||
|
|
return X_BP, y_DBP, y_SBP
|
|||
|
|
|
|||
|
|
def load_data_MIMIC_h5_full(self):
|
|||
|
|
|
|||
|
|
X_BP_path = os.path.join(self.data_dir, 'MIMIC-III_ppg_dataset.h5')
|
|||
|
|
|
|||
|
|
# 获取data_dir下文件列表
|
|||
|
|
files = os.listdir(self.data_dir)
|
|||
|
|
|
|||
|
|
# 检查是否存在已经处理好的数据
|
|||
|
|
if 'X_MIMIC_BP_full.npy' in files and 'Y_MIMIC_DBP_full.npy' in files and 'Y_MIMIC_SBP_full.npy' in files:
|
|||
|
|
print('loading preprocessed data.....')
|
|||
|
|
|
|||
|
|
X_BP = np.load(os.path.join(self.data_dir, 'X_MIMIC_BP_full.npy'))
|
|||
|
|
y_DBP = np.load(os.path.join(self.data_dir, 'Y_MIMIC_DBP_full.npy'))
|
|||
|
|
y_SBP = np.load(os.path.join(self.data_dir, 'Y_MIMIC_SBP_full.npy'))
|
|||
|
|
|
|||
|
|
return X_BP, y_DBP, y_SBP
|
|||
|
|
|
|||
|
|
with h5py.File(X_BP_path, 'r') as f:
|
|||
|
|
ppg = f.get('ppg')
|
|||
|
|
BP = f.get('label')
|
|||
|
|
ppg = np.array(ppg)
|
|||
|
|
BP = np.array(BP)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 将数据从(9054000, 875)reshape成(9054000, 875, 1)的形状
|
|||
|
|
ppg = ppg.reshape(-1, 875, 1)
|
|||
|
|
|
|||
|
|
X_BP = ppg
|
|||
|
|
|
|||
|
|
# 取出第一列赋值给y_DBP，第0列赋值给y_SBP
|
|||
|
|
y_DBP = BP[:, 1]
|
|||
|
|
y_SBP = BP[:, 0]
|
|||
|
|
|
|||
|
|
print("data shape:", X_BP.shape, y_DBP.shape, y_SBP.shape)
|
|||
|
|
|
|||
|
|
print("saving data.....")
|
|||
|
|
|
|||
|
|
# 将数据保存到文件中
|
|||
|
|
np.save('data/X_MIMIC_BP_full.npy', X_BP)
|
|||
|
|
np.save('data/Y_MIMIC_DBP_full.npy', y_DBP)
|
|||
|
|
np.save('data/Y_MIMIC_SBP_full.npy', y_SBP)
|
|||
|
|
|
|||
|
|
print("data saved.....")
|
|||
|
|
|
|||
|
|
return X_BP, y_DBP, y_SBP
|
|||
|
|
|
|||
|
|
def create_dataset(self, X_data, y_SBP, y_DBP):
|
|||
|
|
return BPDataset(X_data, y_SBP, y_DBP)
|
|||
|
|
|
|||
|
|
def split_data(self, X_data, y_SBP, y_DBP):
|
|||
|
|
X_train, X_val, y_train_SBP, y_val_SBP, y_train_DBP, y_val_DBP = train_test_split(
|
|||
|
|
X_data, y_SBP, y_DBP, test_size=self.val_split, random_state=42
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# print(X_train.shape, X_val.shape, y_train_SBP.shape, y_val_SBP.shape, y_train_DBP.shape, y_val_DBP.shape)
|
|||
|
|
|
|||
|
|
train_dataset = self.create_dataset(X_train, y_train_SBP, y_train_DBP)
|
|||
|
|
val_dataset = self.create_dataset(X_val, y_val_SBP, y_val_DBP)
|
|||
|
|
|
|||
|
|
return train_dataset, val_dataset
|
|||
|
|
|
|||
|
|
def create_dataloaders(self):
|
|||
|
|
if self.data_type == 'UKL':
|
|||
|
|
X_data, y_DBP, y_SBP = self.load_data_UKL_h5()
|
|||
|
|
elif self.data_type == 'MIMIC':
|
|||
|
|
X_data, y_DBP, y_SBP = self.load_data_MIMIC_h5()
|
|||
|
|
elif self.data_type == 'MIMIC_full':
|
|||
|
|
X_data, y_DBP, y_SBP = self.load_data_MIMIC_h5_full()
|
|||
|
|
else:
|
|||
|
|
X_data, y_DBP, y_SBP = self.load_data()
|
|||
|
|
train_dataset, val_dataset = self.split_data(X_data, y_SBP, y_DBP)
|
|||
|
|
|
|||
|
|
self.train_dataloader = DataLoader(
|
|||
|
|
train_dataset, batch_size=self.batch_size, shuffle=self.shuffle, collate_fn=custom_collate_fn
|
|||
|
|
)
|
|||
|
|
self.val_dataloader = DataLoader(
|
|||
|
|
val_dataset, batch_size=self.batch_size, shuffle=False, collate_fn=custom_collate_fn
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def get_dataloaders(self):
|
|||
|
|
if self.train_dataloader is None or self.val_dataloader is None:
|
|||
|
|
self.create_dataloaders()
|
|||
|
|
|
|||
|
|
return self.train_dataloader, self.val_dataloader
|
|||
|
|
|
|||
|
|
def get_distributed_dataloaders(self, world_size, rank):
|
|||
|
|
|
|||
|
|
if self.data_type == 'UKL':
|
|||
|
|
X_data, y_DBP, y_SBP = self.load_data_UKL_h5()
|
|||
|
|
elif self.data_type == 'MIMIC':
|
|||
|
|
X_data, y_DBP, y_SBP = self.load_data_MIMIC_h5()
|
|||
|
|
elif self.data_type == 'MIMIC_full':
|
|||
|
|
X_data, y_DBP, y_SBP = self.load_data_MIMIC_h5_full()
|
|||
|
|
else:
|
|||
|
|
X_data, y_DBP, y_SBP = self.load_data()
|
|||
|
|
train_dataset, val_dataset = self.split_data(X_data, y_SBP, y_DBP)
|
|||
|
|
|
|||
|
|
train_sampler = torch.utils.data.distributed.DistributedSampler(
|
|||
|
|
train_dataset, num_replicas=world_size, rank=rank, shuffle=True
|
|||
|
|
)
|
|||
|
|
val_sampler = torch.utils.data.distributed.DistributedSampler(
|
|||
|
|
val_dataset, num_replicas=world_size, rank=rank, shuffle=False
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
train_dataloader = DataLoader(
|
|||
|
|
train_dataset,
|
|||
|
|
batch_size=self.batch_size,
|
|||
|
|
sampler=train_sampler,
|
|||
|
|
collate_fn=custom_collate_fn,
|
|||
|
|
)
|
|||
|
|
val_dataloader = DataLoader(
|
|||
|
|
val_dataset,
|
|||
|
|
batch_size=self.batch_size,
|
|||
|
|
sampler=val_sampler,
|
|||
|
|
collate_fn=custom_collate_fn,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return train_dataloader, val_dataloader, train_sampler, val_sampler
|
|||
|
|
|
|||
|
|
# 使用示例
|
|||
|
|
#
|
|||
|
|
# data_loader = BPDataLoader(data_dir='data', val_split=0.2, batch_size=32,data_type='MIMIC')
|
|||
|
|
# train_dataloader, val_dataloader = data_loader.get_dataloaders()
|
|||
|
|
#
|
|||
|
|
# for i, (X, y_SBP, y_DBP) in enumerate(train_dataloader):
|
|||
|
|
# print(f"Batch {i+1}: X.shape={X.shape }, y_SBP.shape={y_SBP.shape}, y_DBP.shape={y_DBP.shape}")
|
|||
|
|
# if i == 2:
|
|||
|
|
# break
|