import random from typing import Optional, Dict, Any, List import h5py import numpy as np import torch from torch.utils.data import Dataset, DataLoader class NCFDispersionDataset(Dataset): """ 基于 HDF5 读取： - waveform - dispersion - mask 自动跳过以下无效样本： 1. 没有 ncf group 2. ncf group 下面没有任何 item 3. 缺少 disp_periods / disp_velocity / disp_mask 4. （可选）disp_mask 全 0 """ def __init__( self, h5_path: str, split: str = "train", waveform_length: int = 1536, random_ncf: bool = True, pad_value: float = 0.0, return_time: bool = True, seed: Optional[int] = None, skip_zero_mask: bool = True, # 新增：mask 全 0 的样本也跳过 verbose: bool = True, ): super().__init__() if split not in ("train", "test"): raise ValueError(f"split must be 'train' or 'test', got {split}") self.h5_path = h5_path self.split = split self.waveform_length = waveform_length self.random_ncf = random_ncf self.pad_value = pad_value self.return_time = return_time self.skip_zero_mask = skip_zero_mask self.verbose = verbose self.rng = random.Random(seed) self._h5_file = None with h5py.File(self.h5_path, "r") as f: key_name = f"{split}_keys" if key_name not in f: raise KeyError(f"HDF5 中不存在 {key_name}") raw_keys = f[key_name][()] raw_keys = [self._decode_if_needed(x) for x in raw_keys] valid_keys = [] dropped = [] for key in raw_keys: ok, reason = self._is_valid_key(f, key) if ok: valid_keys.append(key) else: dropped.append((key, reason)) self.keys = valid_keys if self.verbose: print(f"[NCFDispersionDataset] split={split}") print(f"[NCFDispersionDataset] 原始样本数: {len(raw_keys)}") print(f"[NCFDispersionDataset] 有效样本数: {len(self.keys)}") print(f"[NCFDispersionDataset] 跳过样本数: {len(dropped)}") if len(dropped) > 0: print("[NCFDispersionDataset] 前10个被跳过样本示例:") for k, r in dropped[:10]: print(f" - {k}: {r}") if len(self.keys) == 0: raise RuntimeError("过滤后没有可用样本，请检查 HDF5 数据集。") @staticmethod def _decode_if_needed(x): if isinstance(x, bytes): return x.decode("utf-8") return str(x) def _get_h5(self): if self._h5_file is None: self._h5_file = h5py.File(self.h5_path, "r") return self._h5_file def _is_valid_key(self, f: h5py.File, key: str): """ 检查一个 key 是否有效。 """ if "paths" not in f: return False, "HDF5 缺少 /paths" if key not in f["paths"]: return False, "key 不存在于 /paths" grp = f["paths"][key] # dispersion 必须存在 for name in ["disp_periods", "disp_velocity", "disp_mask"]: if name not in grp: return False, f"缺少 {name}" try: mask = np.asarray(grp["disp_mask"][()]) except Exception as e: return False, f"读取 disp_mask 失败: {e}" if self.skip_zero_mask and mask.sum() <= 0: return False, "disp_mask 全 0" # ncf 必须存在且至少一个 item if "ncf" not in grp: return False, "缺少 ncf group" ncf_group = grp["ncf"] item_names = sorted(list(ncf_group.keys())) if len(item_names) == 0: return False, "ncf group 为空" # 至少有一个 item 同时有 time 和 waveform found_valid_ncf = False for item_name in item_names: item = ncf_group[item_name] if ("time" in item) and ("waveform" in item): found_valid_ncf = True break if not found_valid_ncf: return False, "ncf 中没有合法 item（缺少 time/waveform）" return True, "ok" def __len__(self): return len(self.keys) def _pad_or_truncate_1d(self, arr: np.ndarray, target_len: int, pad_value: float = 0.0): arr = np.asarray(arr, dtype=np.float32).reshape(-1) n = arr.shape[0] if n == target_len: return arr elif n > target_len: return arr[:target_len] else: out = np.full((target_len,), pad_value, dtype=np.float32) out[:n] = arr return out def _get_valid_ncf_item_names(self, ncf_group) -> List[str]: """ 返回当前 ncf_group 下所有合法 item 名称。 """ valid_names = [] for item_name in sorted(list(ncf_group.keys())): item = ncf_group[item_name] if ("time" in item) and ("waveform" in item): valid_names.append(item_name) return valid_names def _choose_ncf_item_name(self, ncf_group) -> str: valid_names = self._get_valid_ncf_item_names(ncf_group) if len(valid_names) == 0: raise RuntimeError("当前样本没有可用的 NCF 子项") if self.random_ncf and len(valid_names) > 1: return self.rng.choice(valid_names) return valid_names[0] def __getitem__(self, index: int) -> Dict[str, Any]: f = self._get_h5() key = self.keys[index] try: path_group = f["paths"][key] periods = np.asarray(path_group["disp_periods"][()], dtype=np.float32) disp = np.asarray(path_group["disp_velocity"][()], dtype=np.float32) mask = np.asarray(path_group["disp_mask"][()], dtype=np.float32) ncf_group = path_group["ncf"] chosen_item_name = self._choose_ncf_item_name(ncf_group) chosen_ncf = ncf_group[chosen_item_name] time = np.asarray(chosen_ncf["time"][()], dtype=np.float32) waveform = np.asarray(chosen_ncf["waveform"][()], dtype=np.float32) waveform = self._pad_or_truncate_1d( waveform, target_len=self.waveform_length, pad_value=self.pad_value, ) time = self._pad_or_truncate_1d( time, target_len=self.waveform_length, pad_value=0.0, ) sample = { "key": key, "waveform": torch.from_numpy(waveform), # [T] "disp": torch.from_numpy(disp), # [49] "mask": torch.from_numpy(mask), # [49] "periods": torch.from_numpy(periods), # [49] } if self.return_time: sample["ncf_time"] = torch.from_numpy(time) return sample except Exception as e: # 极端情况下，单个样本还是坏了，就随机回退到另一个样本 new_index = self.rng.randint(0, len(self.keys) - 1) if new_index == index and len(self.keys) > 1: new_index = (index + 1) % len(self.keys) if self.verbose: print(f"[WARN] __getitem__ 跳过坏样本 key={key}, error={e}, 改为读取 index={new_index}") return self.__getitem__(new_index) def close(self): if self._h5_file is not None: try: self._h5_file.close() except Exception: pass self._h5_file = None def __del__(self): self.close() def ncf_disp_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]: out = { "key": [item["key"] for item in batch], "waveform": torch.stack([item["waveform"] for item in batch], dim=0), # [B, T] "disp": torch.stack([item["disp"] for item in batch], dim=0), # [B, 49] "mask": torch.stack([item["mask"] for item in batch], dim=0), # [B, 49] "periods": torch.stack([item["periods"] for item in batch], dim=0), # [B, 49] } if "ncf_time" in batch[0]: out["ncf_time"] = torch.stack([item["ncf_time"] for item in batch], dim=0) return out def build_dataloader( h5_path: str, split: str = "train", batch_size: int = 16, num_workers: int = 0, shuffle: Optional[bool] = None, waveform_length: int = 1536, random_ncf: bool = True, pin_memory: bool = False, drop_last: bool = False, seed: Optional[int] = None, skip_zero_mask: bool = True, verbose: bool = True, ) -> DataLoader: if shuffle is None: shuffle = (split == "train") dataset = NCFDispersionDataset( h5_path=h5_path, split=split, waveform_length=waveform_length, random_ncf=random_ncf, return_time=True, seed=seed, skip_zero_mask=skip_zero_mask, verbose=verbose, ) loader = DataLoader( dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory, drop_last=drop_last, collate_fn=ncf_disp_collate_fn, ) return loader