|
|
| import h5py |
| import os, sys |
| import numpy as np |
| import json |
| import torch |
| from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler |
| import pandas as pd |
| import matplotlib.pyplot as plt |
| from collections import Counter, defaultdict |
| import cv2 |
| from multiprocessing import Manager |
| import random |
| from scipy.signal import find_peaks |
| import librosa |
| import pickle |
| from tqdm import tqdm |
| import math |
|
|
| sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| from configs.global_cfg import TARGET_RESAMPLE, DATASET_PATH, LABELS_PATH, EPOCH_LEN_MAX, EPOCH_LEN_MIN, DEMO_MAPS, PID_PREFIX_MAP |
|
|
| CHNGRP_PICK = { |
| 'EEG': ['C3-M2','C3-A2','C4-M1'], |
| 'EOG': ['E1-M2','EOG-Diff','Left-A2'], |
| 'ECG': ['ECG'], |
| 'AIRFLOW': ['AIRFLOW'], |
| 'EMG': ['EMG','Chin1-Chin2'], |
| 'SaO2': ['SaO2'], |
| 'ABD': ['ABD'], |
| 'RESP': ['AIRFLOW'], |
| } |
|
|
|
|
| class Pretrain_Dataset(Dataset): |
| def __init__(self, |
| data_root, |
| data_list, |
| transforms=None, |
| split='train', |
| mode=None, |
| num_classes=5, |
| modalities=["EEG"], |
| sig_len_target=60, |
| ): |
|
|
| super().__init__() |
| |
| self.data_root = data_root |
| self.transforms = transforms |
| self.split = split |
| if mode is None: |
| self.mode = split |
| else: |
| self.mode = mode |
| |
| self.data_list = data_list |
| self.total_num_epochs = 0 |
| self.total_num_subjects = len(self.data_list) |
| |
| self.subject_epoch_dict = {} |
| self.all_labels = [] |
| self.modalities = modalities |
| self.sig_len_target = sig_len_target |
| self.datasets = [] |
|
|
| |
| self.datasets = set([d[0] for d in self.data_list]) |
| |
| self.data_dict = {} |
| self.stats = {} |
| self.sig_info_dict = {} |
| self.chn_grp_map = dict() |
|
|
|
|
| for dataset in self.datasets: |
| for m in self.modalities: |
| chn_grp_dir = os.path.join(self.data_root, DATASET_PATH[dataset], m) |
| if not os.path.exists(chn_grp_dir): |
| chn_grp_dir = os.path.join(self.data_root, DATASET_PATH[dataset], self.split, m) |
| chns = [f.split('.')[0] for f in os.listdir(chn_grp_dir) if f.endswith('mmap')] |
| self.chn_grp_map[f'{dataset}/{m}'] = chns |
| |
| for chn in chns: |
| info_path = os.path.join(chn_grp_dir, f'{chn}.pickle') |
| with open(info_path, 'rb') as f: |
| sig_info = pickle.load(f) |
| |
| self.sig_info_dict[f'{dataset}/{m}/{chn}'] = sig_info |
| sum_mu, sum_musq, cnt = 0.0, 0.0, 0.0 |
| for k, v in sig_info['sig_info'].items(): |
| if (dataset, k) in self.data_list: |
| start, end = v['pos'] |
| mu, musq = v['stats']['sum'], v['stats']['sumsq'] |
| sum_mu += mu |
| sum_musq += musq |
| cnt += (end - start) |
|
|
| self.data_dict[f'{dataset}/{m}/{chn}'] = np.memmap(os.path.join(chn_grp_dir, f'{chn}.mmap'), dtype="float32", |
| mode="r", |
| shape=sig_info['data_shape']) |
| sig_len = sig_info['data_shape'][1] |
| self.stats[f'{dataset}/{m}/{chn}'] = {'mean': sum_mu / cnt / sig_len if cnt > 0 else 0, |
| 'std': np.sqrt(sum_musq / cnt / sig_len - (sum_mu / cnt / sig_len) ** 2) if cnt > 0 else 1.0} |
| |
| self.index_map = self.create_index() |
| |
| def __len__(self): |
| return self.total_num_epochs |
| def __repr__(self): |
| return "Video_Frames_Dataset : Total Subjects : {}, Total Epochs : {}".format(self.total_num_subjects, self.total_num_epochs) |
|
|
| def _get_channel(self, dataset, modality): |
| """Select a channel for a modality based on the current split.""" |
| chns = self.chn_grp_map[f'{dataset}/{modality}'] |
| if self.mode == 'train': |
| return random.choice(chns) |
| for chn in chns: |
| if chn in CHNGRP_PICK[modality]: |
| return chn |
| return chns[0] |
| |
| def _get_signal_for_modality(self, dataset, subject_name, m): |
| """Get signal data for a specific modality.""" |
| |
| chn = self._get_channel(dataset, m) |
| sig_info = self.sig_info_dict[f'{dataset}/{m}/{chn}'] |
| sub_start, sub_end = sig_info['sig_info'][subject_name]['pos'] |
| sub_epochs = sub_end - sub_start |
| signals = self.data_dict[f'{dataset}/{m}/{chn}'][sub_start:sub_end] |
| sig_len_epoch = signals.shape[1] |
| |
| return signals, chn, sig_len_epoch, sub_epochs |
|
|
| def _load_signal_epoch(self, dataset, subject_name, modality, epoch, shift=0): |
| """Load and slice a single modality epoch from mmap storage.""" |
| |
| signals, chn, sig_len_epoch, _ = self._get_signal_for_modality(dataset, subject_name, modality) |
|
|
| sig_len_epoch = signals.shape[1] |
| hz = sig_len_epoch // 30 |
| signals_reshape = np.reshape(signals, (-1, 1)) |
|
|
| start = max(0, sig_len_epoch * epoch - shift * hz) |
| end = start + self.sig_len_target * hz |
| if signals_reshape.shape[0] < end: |
| end = signals_reshape.shape[0] |
| start = end - self.sig_len_target * hz |
|
|
| return chn, signals_reshape[start:end], sig_len_epoch, hz |
|
|
| def _signal_to_tensor(self, signal_epoch, sig_len_epoch, dataset, modality, chn): |
| """Convert a sliced epoch to a normalized tensor.""" |
| hz = sig_len_epoch // 30 |
| if modality == 'SaO2': |
| target_signal = torch.tensor(signal_epoch, dtype=torch.float32) |
| if hz != TARGET_RESAMPLE: |
| target_signal = target_signal.repeat_interleave(TARGET_RESAMPLE, dim=0) |
| return (torch.clamp(target_signal, 0.7, 1.0) - 0.7) / 0.3 |
|
|
| std = self.stats[f'{dataset}/{modality}/{chn}']['std'] |
| signal_epoch = signal_epoch * 1 / std |
| if hz != TARGET_RESAMPLE: |
| signal_epoch = librosa.resample(signal_epoch, orig_sr=sig_len_epoch // 30, target_sr=TARGET_RESAMPLE, |
| axis=0, res_type='soxr_qq') |
| return torch.tensor(signal_epoch, dtype=torch.float32) |
|
|
| def _build_modality_tensor(self, dataset, subject_name, epoch, modality, shift=0): |
| """Shared path for reading, slicing, and normalizing a modality.""" |
| chn, signal_epoch, sig_len_epoch, _ = self._load_signal_epoch(dataset, subject_name, modality, epoch, shift=shift) |
| return self._signal_to_tensor(signal_epoch, sig_len_epoch, dataset, modality, chn) |
| |
| def create_index(self): |
| ''' |
| Create index for the dataset |
| ''' |
| index_map = [] |
|
|
| for dataset, subject_name in self.data_list: |
| m = self.modalities[0] |
| chn = self.chn_grp_map[f'{dataset}/{m}'][0] |
| sig_info = self.sig_info_dict[f'{dataset}/{m}/{chn}'] |
| num_epochs = sig_info['sig_info'][subject_name]['pos'][1] - sig_info['sig_info'][subject_name]['pos'][0] |
| |
| |
| |
| for epoch in range(num_epochs): |
| index_map.append((dataset, subject_name, epoch)) |
|
|
| self.total_num_epochs = len(index_map) |
|
|
| return index_map |
| |
| def __getitem__(self, idx, modalities_override=None): |
| |
| |
| dataset, subject_name, epoch = self.index_map[idx] |
|
|
| inputs = dict() |
| modalities_to_use = modalities_override if modalities_override is not None else self.modalities |
| |
| |
| shift_max = self.sig_len_target - 30 |
| if shift_max > 0: |
| shift = np.random.randint(0, shift_max) |
| elif shift_max == 0: |
| shift = 0 |
| else: |
| shift = np.random.randint(shift_max, 0) |
| |
| |
| for m in modalities_to_use: |
| inputs[m] = self._build_modality_tensor(dataset, subject_name, epoch, m, shift=shift) |
|
|
| return inputs |
| |
| |
| |
|
|
|
|
|
|
| class SFT_Dataset(Pretrain_Dataset): |
| def __init__(self, |
| data_root, |
| data_list, |
| ft_task='sleep_staging', |
| transforms=None, |
| split='train', |
| mode=None, |
| num_classes=5, |
| modalities=["EEG"], |
| sig_len_target=30 |
| ): |
| super().__init__(data_root, data_list, transforms, split, mode, num_classes, modalities, sig_len_target) |
| self.ft_task = ft_task |
| self.num_classes = num_classes |
| |
| |
| |
| |
| if self.ft_task == 'apnea_seg' and self.num_classes == 2: |
| self.label_change = {0:0, 1:1, 2:1} |
| |
| self.get_labels() |
|
|
| def __len__(self): |
| return self.total_num_epochs |
|
|
| |
| def get_labels(self): |
| self.all_labels = [] |
| shift = (self.sig_len_target - 30) // 2 |
| prev_subject_name = '' |
|
|
| for dataset, subject_name, epoch in tqdm(self.index_map): |
| if not subject_name == prev_subject_name: |
| suffix = 'ah' if self.ft_task == 'apnea_seg' else 'ann' |
| label_dir = os.path.join(self.data_root, LABELS_PATH[self.ft_task][dataset]) |
| label_path = os.path.join(label_dir, f'{subject_name}_{suffix}.npy') |
| if not os.path.exists(label_path): |
| label_path = os.path.join(label_dir, self.split, f'{subject_name}_{suffix}.npy') |
| labels = np.load(label_path).astype(np.int64) |
| |
| if self.ft_task == 'apnea_seg': |
| labels = labels.reshape(-1,) |
| labels_slice = labels[epoch*30 : (epoch+1)*30] |
| |
| if self.num_classes == 2: |
| labels_slice = np.array([self.label_change[lbl] for lbl in labels_slice]) |
| |
| self.all_labels.append(labels_slice) |
| |
| else: |
| self.all_labels.append(labels[epoch]) |
| prev_subject_name = subject_name |
|
|
| def get_all_labels(self): |
| return self.all_labels |
|
|
|
|
| def __getitem__(self, idx, modalities_override=None): |
| |
|
|
| |
| dataset, subject_name, epoch = self.index_map[idx] |
| label = self.all_labels[idx] |
| |
| inputs = dict() |
| modalities_to_use = modalities_override if modalities_override is not None else self.modalities |
| |
| shift = (self.sig_len_target - 30) // 2 |
| |
| for m in modalities_to_use: |
| inputs[m] = self._build_modality_tensor(dataset, subject_name, epoch, m, shift=shift) |
|
|
| |
| |
| return inputs, torch.tensor(label, dtype=torch.long), torch.tensor([], dtype=torch.bool), torch.tensor([], dtype=torch.bool), subject_name |
|
|
|
|
| class Seq_Dataset(Pretrain_Dataset): |
| def __init__(self, |
| data_root, |
| data_list, |
| transforms=None, |
| split='train', |
| mode=None, |
| num_classes=5, |
| modalities=["EEG"], |
| sig_len_target=30, |
| pretrain_macro=False, |
| seq_truncate=False, |
| ft_task='sleep_staging', |
| dual_views=False, |
| ): |
| super().__init__(data_root, data_list, transforms, split, mode, num_classes, modalities, sig_len_target) |
|
|
| self.num_classes = num_classes |
| self.get_demographics() |
| self.epoch_accum = 0 |
| self.seq_truncate = seq_truncate |
| self.max_num_epochs = EPOCH_LEN_MAX if seq_truncate else 1440 |
| self.min_num_epochs = EPOCH_LEN_MIN if seq_truncate else 0 |
| self.pretrain_macro = pretrain_macro |
| self.dual_views = dual_views |
| self.ft_task = ft_task |
| |
| self.create_index_full() |
| print("Max length:", self.max_num_epochs) |
| |
| if not self.pretrain_macro: |
| self.get_labels() |
|
|
| def reshuffle_full_night_index(self): |
| """ |
| Reshuffle the full night index. Call this at the start of each epoch |
| to get a different ordering of samples. |
| """ |
| self.epoch_accum += 1 |
| self.create_index_full() |
| if not self.pretrain_macro: |
| self.get_labels() |
| |
| def _validate_demographics(self, dataset, subject_name): |
| """Validate demographics data for a subject. Returns (is_valid, age, bmi).""" |
| demo = self.demographics.get(f'{dataset}/{subject_name}', {}) |
| |
| age = demo.get('Age', None) |
| sex = demo.get('Sex', None) |
| bmi = demo.get('BMI', None) |
| |
| |
| if age is None or (isinstance(age, float) and math.isnan(age)): |
| print(f"Warning: Age is missing for {dataset}, {subject_name}. Skipping subject") |
| return False, None, None |
| |
| |
| if sex is None: |
| print(f"Warning: Sex is missing for {dataset}, {subject_name}. Skipping subject") |
| return False, None, None |
| |
| |
| if bmi is None: |
| print(f"Warning: BMI is missing for {dataset}, {subject_name}. Skipping subject") |
| return False, None, None |
| |
| try: |
| bmi = float(bmi) |
| except (ValueError, TypeError): |
| print(f"Warning: BMI is not a number for {dataset}, {subject_name}. Skipping subject") |
| return False, None, None |
| |
| if isinstance(bmi, float) and math.isnan(bmi): |
| print(f"Warning: BMI is NaN for {dataset}, {subject_name}. Skipping subject") |
| return False, None, None |
| |
| return True, age, bmi |
| |
| def create_index_full(self): |
| '''Create index for the dataset''' |
| index_map, len_map, ages, bmis = [], [], [], [] |
|
|
| for dataset, subject_name in self.data_list: |
| m = self.modalities[0] |
| chn = self.chn_grp_map[f'{dataset}/{m}'][0] |
| sig_info = self.sig_info_dict[f'{dataset}/{m}/{chn}'] |
| num_epochs = sig_info['sig_info'][subject_name]['pos'][1] - sig_info['sig_info'][subject_name]['pos'][0] |
| |
| if not self.seq_truncate: |
| self.max_num_epochs = max(num_epochs, self.max_num_epochs) |
| |
| if num_epochs <= self.min_num_epochs: |
| print(f"Warning: Sequence length < EPOCH_LEN_MIN for {dataset}, {subject_name}. Num_epochs: {num_epochs}, skipping") |
| continue |
| |
| if not self.pretrain_macro: |
| |
| index_map.append((dataset, subject_name)) |
| continue |
| |
| |
| is_valid, age, bmi = self._validate_demographics(dataset, subject_name) |
| if not is_valid: |
| continue |
| |
| ages.append(age) |
| bmis.append(bmi) |
| len_map.append(num_epochs) |
| index_map.append((dataset, subject_name)) |
| |
| if not self.pretrain_macro: |
| self.index_map = index_map |
| self.total_subjects = len(index_map) |
| return |
|
|
| |
| self._reorder_index_by_demographics(index_map, len_map, ages, bmis) |
| |
| def _reorder_index_by_demographics(self, index_map, len_map, ages, bmis): |
| """Reorder index map by demographics and sequence length for balanced batching.""" |
| zipped_pairs = list(zip(index_map, len_map, ages, bmis)) |
| |
| |
| if self.epoch_accum % 2 == 0: |
| sorted_pairs = sorted(zipped_pairs, key=lambda x: x[3]) |
| else: |
| sorted_pairs = sorted(zipped_pairs, key=lambda x: x[2]) |
| |
| index_map = [pair[0] for pair in sorted_pairs] |
| len_map = [pair[1] for pair in sorted_pairs] |
| |
| self.total_subjects = len(index_map) |
| |
| |
| num_age_grp = 4 |
| age_grp_size = self.total_subjects // num_age_grp |
| age_groups = [] |
| residuals = [] |
| |
| for i in range(num_age_grp): |
| start_idx = i * age_grp_size |
| end_idx = (i + 1) * age_grp_size |
| index_map_grp = index_map[start_idx:end_idx] |
| len_map_grp = len_map[start_idx:end_idx] |
| |
| |
| reverse = (self.epoch_accum % 4 == 0 or self.epoch_accum % 4 == 3) |
| index_map_grp_sorted = [pair[0] for pair in sorted( |
| zip(index_map_grp, len_map_grp), |
| key=lambda x: x[1], |
| reverse=reverse |
| )] |
| |
| |
| len_sub4 = len(index_map_grp_sorted) // 4 |
| sub4 = [index_map_grp_sorted[4*k:4*(k+1)] for k in range(len_sub4)] |
| age_groups.append(sub4) |
| residuals.extend(index_map_grp_sorted[4*len_sub4:]) |
| |
| if i == num_age_grp - 1 and end_idx < self.total_subjects: |
| residuals.extend(index_map[end_idx:]) |
| |
| |
| final_index_map = [] |
| for items in zip(*age_groups): |
| flattened = [item for sublist in items for item in sublist] |
| final_index_map.extend(flattened) |
| final_index_map.extend(residuals) |
| |
| assert len(final_index_map) == self.total_subjects, \ |
| f"Index map length mismatch: {len(final_index_map)} vs {self.total_subjects}" |
| |
| self.index_map = final_index_map |
| self.total_subjects = len(self.index_map) |
| |
|
|
| def __len__(self): |
| return self.total_subjects |
| |
| def get_demographics(self): |
| """Load demographic data for all datasets.""" |
| self.demographics = dict() |
| |
|
|
| |
| |
| skip_datasets = {'isruc', 'cfs', 'sof'} |
| |
| dataset_list = list(set([d[0] for d in self.data_list])) |
| for dataset in tqdm(dataset_list): |
| if any(s in dataset for s in skip_datasets): |
| continue |
| |
| |
| demo_file = DEMO_MAPS.get(dataset, f'{dataset}_demographics.pkl') |
| demo_path = os.path.join(self.data_root, 'demographics', demo_file) |
| |
| if not os.path.exists(demo_path): |
| print(f"Warning: Demographics file not found: {demo_path}") |
| continue |
| |
| with open(demo_path, 'rb') as f: |
| demo_data = pickle.load(f) |
|
|
| for pid, demo in demo_data.items(): |
| if dataset in PID_PREFIX_MAP: |
| prefix = PID_PREFIX_MAP[dataset] |
| if dataset in ['mros1', 'mros1-val', 'mros1-te', 'mros2', 'mros2-val', 'mros2-te']: |
| pid = f'{prefix}-{pid.lower()}' |
| elif dataset in ['mesa', 'mesa-val', 'mesa-te']: |
| pid = f'{prefix}-{pid:04d}' |
| elif dataset == 'cfs': |
| pid = f'{dataset}-visit5-{pid}' |
| else: |
| pid = f'{dataset}-{pid}' |
| |
| self.demographics[f'{dataset}/{pid}'] = demo |
| |
| def get_labels(self): |
| self.all_labels = [] |
| shift = (self.sig_len_target - 30) // 2 |
| prev_subject_name = '' |
| |
| if self.ft_task == 'diagnosis': |
| |
| events_csv = 'datasets/#_SHHS/datasets/shhs-cvd-summary-dataset-0.20.0.csv' |
| event_df = pd.read_csv(events_csv) |
|
|
| is_event_cols = ['angina','cvd_death','chf','chd_death','mi','stroke'] |
| event_time_cols = ['ang_date','cvd_dthdt','chf_date','chd_dthdt','mi_date','stk_date'] |
|
|
| |
| |
| df_filtered = pd.concat([ |
| event_df[['nsrrid', 'censdate']], |
| event_df[is_event_cols], |
| event_df[event_time_cols] |
| ], axis=1).set_index('nsrrid') |
|
|
| |
| for event_col, time_col in zip(is_event_cols, event_time_cols): |
| |
| |
| invalid_mask = (df_filtered[event_col] >= 1) & (df_filtered[time_col].isna() | (df_filtered[time_col] == 0)) |
| |
| df_filtered.loc[invalid_mask, event_col] = 0 |
| |
| |
| |
| df_filtered[time_col] = df_filtered[time_col].fillna(df_filtered['censdate']) |
| df_filtered.loc[df_filtered[time_col] == 0, time_col] = df_filtered['censdate'] |
|
|
| |
| df_filtered[is_event_cols] = df_filtered[is_event_cols].fillna(0) |
| |
| df_filtered[event_time_cols] = df_filtered[event_time_cols].fillna(365*6) |
| |
| df_filtered[is_event_cols] = np.minimum(df_filtered[is_event_cols], 1) |
|
|
| study_ids = set(df_filtered.index) |
| |
|
|
| labels_dict = {} |
| for study_id in study_ids: |
| row = df_filtered.loc[study_id] |
| |
| |
| if isinstance(row, pd.DataFrame): |
| row = row.iloc[0] |
| |
| is_event_row = [int(val) for val in row[is_event_cols]] |
| event_time_row = [float(val) for val in row[event_time_cols]] |
| |
| labels_dict[str(study_id)] = { |
| "is_event": is_event_row, |
| "event_time": event_time_row |
| } |
| |
| for sub_info in tqdm(self.index_map): |
|
|
| dataset, subject = sub_info |
| |
| if self.ft_task == 'diagnosis': |
| if not dataset == 'shhs1': |
| raise NotImplementedError("Diagnosis task is only implemented for shhs1 dataset.") |
| study_id = subject.split('-')[1] |
| self.all_labels.append(labels_dict[study_id]) |
| |
| else: |
| suffix = 'ah' if self.ft_task == 'apnea_seg' else 'ann' |
| if self.ft_task in ['age','sex','ahi']: |
| ft_task = 'sleep_staging' |
| else: |
| ft_task = self.ft_task |
| label_dir = os.path.join(self.data_root, LABELS_PATH[ft_task][dataset]) |
| label_path = os.path.join(label_dir, f'{subject}_{suffix}.npy') |
| if not os.path.exists(label_path): |
| label_path = os.path.join(label_dir, self.split, f'{subject}_{suffix}.npy') |
| labels = np.load(label_path).astype(np.int64) |
| labels = labels.reshape(-1,) |
| self.all_labels.append(labels) |
| |
| |
| def get_all_labels(self): |
| return self.all_labels |
| |
| def _process_signal(self, signal_epochs, sig_len_epoch, dataset, m, chn): |
| """Convert raw signal to target tensor with appropriate normalization.""" |
| signals_reshape = np.reshape(signal_epochs, (-1, 1)) |
| Hz = sig_len_epoch // 30 |
| |
| if m == 'SaO2': |
| resampled = signals_reshape |
| target_signal = torch.tensor(resampled, dtype=torch.float32) |
| if Hz != TARGET_RESAMPLE: |
| target_signal = target_signal.repeat_interleave(TARGET_RESAMPLE, dim=0) |
| target_signal = (torch.clamp(target_signal, 0.7, 1.0) - 0.7) / 0.3 |
| else: |
| std = self.stats[f'{dataset}/{m}/{chn}']['std'] |
| signals_reshape = signals_reshape * 1 / std |
| if Hz != TARGET_RESAMPLE: |
| resampled = librosa.resample(signals_reshape, orig_sr=sig_len_epoch//30, |
| target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq') |
| else: |
| resampled = signals_reshape |
| target_signal = torch.tensor(resampled, dtype=torch.float32) |
| |
| return target_signal |
| |
| def _shift_signals(self, signals, shift_epochs, shift_fwd): |
| """Shift signals by specified number of epochs.""" |
| if shift_epochs == 0: |
| return signals |
| if shift_fwd: |
| return np.concatenate([signals[shift_epochs:], signals[-shift_epochs:]]) |
| return np.concatenate([signals[:shift_epochs], signals[:-shift_epochs]]) |
| |
| def _pad_signals(self, signals, sub_epochs): |
| """Pad signals to max_num_epochs.""" |
| if sub_epochs >= self.max_num_epochs: |
| return signals[:self.max_num_epochs] |
| pad_len = self.max_num_epochs - sub_epochs |
| padding = np.zeros_like(signals[-1:]).repeat(pad_len, axis=0) |
| return np.concatenate([signals, padding], axis=0) |
| |
| def _build_demo_dict(self, dataset, subject_name): |
| """Build demographic dictionary for a subject.""" |
| demo = {} |
| demo_key = f'{dataset}/{subject_name}' |
| |
| if self.pretrain_macro: |
| age = self.demographics[demo_key]['Age'] |
| sex = self.demographics[demo_key]['Sex'] |
| bmi = self.demographics[demo_key]['BMI'] |
| demo['age'] = float(age) |
| demo['sex'] = sex |
| demo['bmi'] = float(bmi) |
| else: |
| |
| demo_values = self.demographics.get(demo_key, {}) |
| age = demo_values.get('Age', -1) |
| sex = demo_values.get('Sex', -1) |
| bmi = demo_values.get('BMI', -1) |
| ahi = demo_values.get('AHI', -1) |
| |
| |
| demo['age'] = float(age) if age is not None and not (isinstance(age, float) and math.isnan(age)) else -1.0 |
| demo['sex'] = sex |
| demo['bmi'] = float(bmi) if bmi is not None and not (isinstance(bmi, float) and math.isnan(bmi)) else -1.0 |
| demo['ahi'] = float(ahi) if ahi is not None and not (isinstance(ahi, float) and math.isnan(ahi)) else -1.0 |
| |
| return demo |
| |
| def __getitem__(self, idx, modalities_override=None): |
| """Get a full-night sequence of signals with optional dual views.""" |
| dataset, subject_name = self.index_map[idx] |
| modalities_to_use = modalities_override or self.modalities |
| |
| |
| demo = self._build_demo_dict(dataset, subject_name) |
| |
| |
| label = None |
| if not self.pretrain_macro: |
| label = self.all_labels[idx] |
| |
| |
| m_first = modalities_to_use[0] |
| _, _, _, sub_epochs = self._get_signal_for_modality(dataset, subject_name, m_first) |
| |
| |
| if not self.pretrain_macro and self.ft_task in ['age', 'sex', 'ahi', 'sleep_staging']: |
| if sub_epochs >= self.max_num_epochs: |
| label = label[:self.max_num_epochs] |
| mask = torch.ones(self.max_num_epochs, dtype=torch.bool) |
| else: |
| label = np.concatenate([label, label[-1:].repeat(self.max_num_epochs - sub_epochs, axis=0)]) |
| mask = torch.zeros(self.max_num_epochs, dtype=torch.bool) |
| mask[:sub_epochs] = True |
| else: |
| if sub_epochs >= self.max_num_epochs: |
| mask = torch.ones(self.max_num_epochs, dtype=torch.bool) |
| else: |
| mask = torch.zeros(self.max_num_epochs, dtype=torch.bool) |
| mask[:sub_epochs] = True |
| |
| |
| if self.dual_views: |
| inputs_view1, inputs_view2 = self._process_dual_views(dataset, subject_name, modalities_to_use, sub_epochs) |
| inputs_out = (inputs_view1, inputs_view2) |
| else: |
| inputs_out = self._process_single_view(dataset, subject_name, modalities_to_use, sub_epochs) |
| |
| |
| if self.pretrain_macro: |
| return inputs_out, torch.tensor([], dtype=torch.long), mask, demo, subject_name |
| elif self.ft_task == 'diagnosis': |
| is_event = torch.tensor(label["is_event"]) |
| event_time = torch.tensor(label["event_time"], dtype=torch.float32) |
| return inputs_out, is_event, event_time, mask, demo, subject_name |
| else: |
| return inputs_out, torch.tensor(label, dtype=torch.long), mask, demo, subject_name |
| |
| def _process_single_view(self, dataset, subject_name, modalities_to_use, sub_epochs): |
| """Process signals for single view.""" |
| inputs = {} |
| for m in modalities_to_use: |
| signals, chn, sig_len_epoch, _ = self._get_signal_for_modality(dataset, subject_name, m) |
| signals = self._pad_signals(signals, sub_epochs) |
| inputs[m] = self._process_signal(signals, sig_len_epoch, dataset, m, chn) |
| return inputs |
| |
| def _process_dual_views(self, dataset, subject_name, modalities_to_use, sub_epochs): |
| """ |
| Used for Macro-Encoder alternative training methods (SimCLR and DINO) |
| Process signals for dual views with random augmentation. |
| """ |
| |
| shift_epochs_1 = np.random.randint(0, 60) |
| shift_fwd_1 = np.random.rand() > 0.5 |
| shift_epochs_2 = np.random.randint(0, 60) |
| shift_fwd_2 = np.random.rand() > 0.5 |
| |
| |
| zero_modalities = np.random.choice(len(modalities_to_use), 6, replace=False) |
| k1 = np.random.randint(6) |
| zero_modalities_1 = zero_modalities[:k1] |
| zero_modalities_2 = zero_modalities[k1:] |
| |
| inputs_view1, inputs_view2 = {}, {} |
| |
| for i, m in enumerate(modalities_to_use): |
| signals, chn, sig_len_epoch, _ = self._get_signal_for_modality(dataset, subject_name, m) |
| |
| |
| signals_v1 = self._shift_signals(signals, shift_epochs_1, shift_fwd_1) |
| signals_v2 = self._shift_signals(signals, shift_epochs_2, shift_fwd_2) |
| |
| |
| if i in zero_modalities_1: |
| signals_v1 = np.zeros_like(signals_v1) |
| if i in zero_modalities_2: |
| signals_v2 = np.zeros_like(signals_v2) |
| |
| |
| signals_v1 = self._pad_signals(signals_v1, sub_epochs) |
| signals_v2 = self._pad_signals(signals_v2, sub_epochs) |
| |
| inputs_view1[m] = self._process_signal(signals_v1, sig_len_epoch, dataset, m, chn) |
| inputs_view2[m] = self._process_signal(signals_v2, sig_len_epoch, dataset, m, chn) |
| |
| return inputs_view1, inputs_view2 |
| |
| |
| |
| |
|
|
| def collate_fn_random_modalities(batch_indices, dataset): |
| |
| all_modalities = dataset.modalities |
| selected_modalities = sorted(random.sample(all_modalities, 2)) |
| |
| |
| batch = [dataset.__getitem__(idx, modalities_override=selected_modalities) for idx in batch_indices] |
| |
| inputs = batch |
| inputs_ = {m: [] for m in selected_modalities} |
| for m in selected_modalities: |
| for item in inputs: |
| inputs_[m].append(item[m]) |
| inputs = {m: torch.stack(inputs_[m], dim=0) for m in selected_modalities} |
| return inputs |
|
|
| class Pretrain_Dataset_Wrapper(Dataset): |
| def __init__(self, dataset): |
| self.dataset = dataset |
| def __len__(self): |
| return len(self.dataset) |
| def __getitem__(self, idx): |
| return idx |
|
|
| |
| def build_loader(data_root, |
| data_list=None, |
| batch_size=1, |
| shuffle=True, |
| transforms=None, |
| sampler='weighted', |
| num_workers=8, |
| proportion = 1.0, |
| split='train', |
| mode=None, |
| num_classes=5, |
| modalities=["EEG"], |
| pretrain_micro=False, |
| pretrain_macro=False, |
| sig_len_target=30, |
| ft_task='sleep_staging', |
| full_night=False, |
| seq_truncate=False, |
| dual_views=False, |
| ): |
| """_summary_ |
| |
| Args: |
| data_path (_type_): _description_ |
| mode (str, optional): _description_. Defaults to 'epochs'. |
| batch_size (int, optional): _description_. Defaults to 1. |
| shuffle (bool, optional): _description_. Defaults to True. |
| transforms (_type_, optional): _description_. Defaults to None. |
| sampler (str, optional): _description_. Defaults to 'weighted'. |
| train (bool, optional): _description_. Defaults to True. |
| |
| Returns: |
| _type_: _description_ |
| """ |
| if pretrain_micro: |
| |
| print("Using Pretrain Dataset") |
| dataset = Pretrain_Dataset(data_root, |
| data_list, |
| transforms, |
| mode=mode, |
| split=split, |
| num_classes=num_classes, |
| modalities=modalities, |
| sig_len_target=sig_len_target |
| ) |
| |
| print("Dataset created") |
| loader = DataLoader(dataset, |
| batch_size=batch_size, |
| num_workers=num_workers, |
| shuffle=shuffle, |
| |
| prefetch_factor=2) |
|
|
| else: |
| if full_night: |
| print("Using sequence dataset") |
| dataset = Seq_Dataset(data_root, |
| data_list, |
| transforms=transforms, |
| split=split, |
| mode=mode, |
| num_classes=num_classes, |
| modalities=modalities, |
| sig_len_target=sig_len_target, |
| ft_task=ft_task, |
| pretrain_macro=pretrain_macro, |
| seq_truncate=seq_truncate, |
| dual_views=dual_views |
| ) |
| else: |
| print("Using Epoch-level Fine-tuning Dataset. Task:", ft_task) |
| dataset = SFT_Dataset(data_root, |
| data_list, |
| ft_task=ft_task, |
| transforms=transforms, |
| split=split, |
| mode=mode, |
| num_classes=num_classes, |
| modalities=modalities, |
| sig_len_target=sig_len_target |
| ) |
| |
| print("Dataset created") |
| |
| if sampler=='weighted': |
| |
| labels = dataset.get_all_labels() |
| |
| |
| class_count = np.zeros(num_classes, dtype=np.int64) |
| unique, counts = np.unique(labels, return_counts=True) |
| for i, cnt in zip(unique, counts): |
| class_count[i] += cnt |
| print("Class count:", class_count) |
|
|
| class_weights = 1.0 / torch.tensor(class_count, dtype=torch.float) |
| |
| sample_weights = [class_weights[i] for i in labels] |
| sample_weights = torch.tensor(sample_weights, dtype=torch.float) |
| print("Class weights:", sample_weights) |
| |
| sampler = torch.utils.data.sampler.WeightedRandomSampler(sample_weights, int(len(labels)*proportion), replacement=True) |
| loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, num_workers=num_workers, prefetch_factor=2) |
| |
| elif sampler == 'random': |
| |
| |
| |
| labels = dataset.get_all_labels() |
| |
| class_count = np.zeros(num_classes, dtype=np.int64) |
| unique, counts = np.unique(labels, return_counts=True) |
| for i, cnt in zip(unique, counts): |
| class_count[i] += cnt |
| print("Class count:", class_count) |
|
|
| |
| class_weights = torch.ones(num_classes, dtype=torch.float) / num_classes |
| sample_weights = [1.0 for i in labels] |
| sampler = torch.utils.data.sampler.WeightedRandomSampler(sample_weights, int(len(labels)*proportion), replacement=True) |
| |
| loader = DataLoader(dataset, |
| batch_size=batch_size, |
| sampler=sampler, |
| num_workers=num_workers, |
| collate_fn=lambda batch_indices: collate_fn_random_modalities(batch_indices, dataset), |
| prefetch_factor=2) |
| elif sampler == 'distributed': |
| from torch.utils.data import DistributedSampler |
| |
| sampler = DistributedSampler( |
| dataset, |
| shuffle=shuffle |
| ) |
| loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, num_workers=num_workers, prefetch_factor=2, drop_last=True) |
|
|
| elif sampler is None: |
| loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, prefetch_factor=2) |
| else: |
| assert False, "Not Implemented loader" |
| return loader |
|
|
|
|
|
|
|
|
|
|
