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 = sorted(os.listdir(data_path)) 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 = [] # What datasets are used? phy, kiss, shhs1,... 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 # self.chn_grp_map[f'{dataset}/{m}'] = CHNGRP_MAP[f'{dataset}/{m}'] 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] # Assume signal shape (sum(num_epochs), sig_len_epoch) 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] # if num_epochs <= EPOCH_LEN_MIN: # print(f"Warning: num_epochs less than EPOCH_LEN_MIN for {dataset}, {subject_name}. Num_epochs: {num_epochs}, skip this row") # continue 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): # Return frames, label, (subject_name, epoch) 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) # (in seconds) elif shift_max == 0: shift = 0 else: shift = np.random.randint(shift_max, 0) # (in seconds) # inputs["label"] = torch.tensor(label, dtype=torch.long) 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 == 'sleep_staging' and self.num_classes == 4: # self.label_change = {0:0, 1:1, 2:1, 3:2, 4:3} 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) # (num_epochs, 1) 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): # Return frames, label, (subject_name, epoch) 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 an empty boolean tensor for mask so default_collate can batch samples. # Downstream training code treats an empty tensor as 'no mask'. return inputs, torch.tensor(label, dtype=torch.long), torch.tensor([], dtype=torch.bool), torch.tensor([], dtype=torch.bool), subject_name # No mask 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) # Validate age 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 # Validate sex if sex is None: print(f"Warning: Sex is missing for {dataset}, {subject_name}. Skipping subject") return False, None, None # Validate BMI 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: # For finetuning, just collect subjects index_map.append((dataset, subject_name)) continue # For pretraining with macro demographics, validate demographics 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 # For pretraining: reorder by demographics and sequence length 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)) # Alternate sorting by BMI (even epochs) and age (odd epochs) if self.epoch_accum % 2 == 0: sorted_pairs = sorted(zipped_pairs, key=lambda x: x[3]) # Sort by BMI else: sorted_pairs = sorted(zipped_pairs, key=lambda x: x[2]) # Sort by age index_map = [pair[0] for pair in sorted_pairs] len_map = [pair[1] for pair in sorted_pairs] self.total_subjects = len(index_map) # Group by age and apply secondary sorting within groups 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] # Sort by length: ascending for epochs 1,2; descending for 0,3 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 )] # Group by 4 subjects 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:]) # Interleave groups and add residuals 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() # Datasets to skip 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 # Get demographics file path 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': # Disease prediction 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'] # 1. 데이터 병합 (censdate 포함) # 날짜 보정을 위해 'censdate' 열을 함께 가져옵니다. df_filtered = pd.concat([ event_df[['nsrrid', 'censdate']], event_df[is_event_cols], event_df[event_time_cols] ], axis=1).set_index('nsrrid') # 2. 정합성 검사 및 보정 for event_col, time_col in zip(is_event_cols, event_time_cols): # 조건 A: 이벤트가 1 이상인데 날짜가 NaN이거나 0인 경우 # 이 경우 이벤트를 0으로 변경하고 날짜는 censdate를 참조하게 함 invalid_mask = (df_filtered[event_col] >= 1) & (df_filtered[time_col].isna() | (df_filtered[time_col] == 0)) # invalid_mask = (df_filtered[event_col] >= 1) & ((df_filtered[time_col] == 0)) df_filtered.loc[invalid_mask, event_col] = 0 # 조건 B: 날짜열의 값이 없거나(NaN) 0인 모든 경우에 대해 'censdate' 값으로 채움 # fillna를 사용하여 NaN을 채우고, 0인 값은 replace를 통해 처리할 수 있습니다. 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'] # 3. 최종 결측치 처리 (censdate마저 없는 경우를 대비한 방어 코드) df_filtered[is_event_cols] = df_filtered[is_event_cols].fillna(0) # 만약 censdate도 NaN이라서 여전히 결측치라면 1e6으로 최종 처리 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] # 중복 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) # (num_epochs, 1) 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: # Finetuning: include optional demographics 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) # Convert to float and handle NaN 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 # Build demographic dictionary demo = self._build_demo_dict(dataset, subject_name) # Initialize label and mask label = None if not self.pretrain_macro: label = self.all_labels[idx] # Get signal info from first modality to create mask m_first = modalities_to_use[0] _, _, _, sub_epochs = self._get_signal_for_modality(dataset, subject_name, m_first) # Update label and create mask based on truncation/padding 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 # Process signals for dual or single views 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) # Return based on task type 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. """ # Generate random shift parameters for two views 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 # Random modality dropout 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) # Create two shifted views signals_v1 = self._shift_signals(signals, shift_epochs_1, shift_fwd_1) signals_v2 = self._shift_signals(signals, shift_epochs_2, shift_fwd_2) # Apply random modality dropout 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) # Pad and process 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): # Randomly select two modalities for this batch all_modalities = dataset.modalities selected_modalities = sorted(random.sample(all_modalities, 2)) # selected_modalities = all_modalities # Use all modalities for now # Fetch data for each index with the selected modalities batch = [dataset.__getitem__(idx, modalities_override=selected_modalities) for idx in batch_indices] # inputs = zip(*batch) 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 # Just return the index # build loader with sample 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: # Create dataset 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, # collate_fn=lambda batch_indices: collate_fn_random_modalities(batch_indices, dataset), 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': # Balanced Sampler labels = dataset.get_all_labels() # class_count = Counter(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) # Length : min(class_count) sample_weights = [class_weights[i] for i in labels] sample_weights = torch.tensor(sample_weights, dtype=torch.float) print("Class weights:", sample_weights) # Length : min(class_count) 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': # print("Random sampler, num_samples : ", int(len(dataset)*proportion), "Length dataset", len(dataset)) # Proportion does not look effective.... # sampler = RandomSampler(dataset, replacement=True, num_samples=int(len(dataset)*proportion)) 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.tensor(class_count, dtype=torch.float) 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, prefetch_factor=2) 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 # 2. Manually create the DistributedSampler with shuffle=False sampler = DistributedSampler( dataset, shuffle=shuffle # <--- The key override ) 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