| import io |
| import os |
| import tarfile |
| import pickle |
| import zstandard |
| import torch |
| from torch.utils.data import IterableDataset, DataLoader |
|
|
| class MultiFileWrapper(io.BufferedIOBase): |
| """将多个文件对象包装成一个连续的缓冲流，提高读取性能""" |
| def __init__(self, paths): |
| self.files = [open(p, 'rb') for p in paths] |
| self.current_idx = 0 |
|
|
| def read(self, size=-1): |
| if self.current_idx >= len(self.files): |
| return b'' |
| |
| data = self.files[self.current_idx].read(size) |
| while (not data or len(data) == 0) and self.current_idx < len(self.files) - 1: |
| self.current_idx += 1 |
| data = self.files[self.current_idx].read(size) |
| return data |
|
|
| def readable(self): |
| return True |
|
|
| def close(self): |
| for f in self.files: |
| f.close() |
|
|
| class DyMeshIterableDataset(IterableDataset): |
| def __init__(self, archive_paths=None): |
| super(DyMeshIterableDataset, self).__init__() |
| if archive_paths is None: |
| base = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/" |
| self.archive_paths = [os.path.join(base, f"DyMesh_50000v_16f_0000_part_{str(i).zfill(2)}") for i in range(4)] |
| else: |
| self.archive_paths = archive_paths |
|
|
| def __iter__(self): |
| wrapper = MultiFileWrapper(self.archive_paths) |
| try: |
| dctx = zstandard.ZstdDecompressor() |
| with dctx.stream_reader(wrapper) as reader: |
| with tarfile.open(fileobj=reader, mode='r|') as tar: |
| for member in tar: |
| if member.isfile(): |
| f = tar.extractfile(member) |
| if f: |
| data = pickle.load(f) |
| if isinstance(data, dict) and 'vertices' in data: |
| yield { |
| 'vertices': data['vertices'], |
| 'faces': data['faces'], |
| 'caption': data.get('caption', '无文本描述') |
| } |
| except Exception as e: |
| print(f"❌ 跨文件流读取异常: {e}") |
| finally: |
| wrapper.close() |
|
|
| def collate_fn(batch): |
| """仅负责将数据组织成 Tensor，不再进行 VAE 推理""" |
| vertices_list = [torch.tensor(item['vertices'], dtype=torch.float32) for item in batch] |
| faces_list = [torch.tensor(item['faces'], dtype=torch.long) for item in batch] |
| caption_list = [item['caption'] for item in batch] |
| |
| return { |
| 'vertices': [item['vertices'] for item in batch], |
| 'faces': [item['faces'] for item in batch], |
| 'caption': [item['caption'] for item in batch] |
| } |
|
|
| def get_mesh_dataloader(batch_size=4, num_workers=0, archive_paths=None): |
| """统一的 DataLoader 获取接口""" |
| dataset = DyMeshIterableDataset(archive_paths=archive_paths) |
| return DataLoader( |
| dataset, |
| batch_size=batch_size, |
| collate_fn=collate_fn, |
| num_workers=num_workers, |
| pin_memory=True if torch.cuda.is_available() else False |
| ) |
