Spaces:
Runtime error
Runtime error
File size: 5,650 Bytes
67fb03c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
import os
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import pyvista as pv
class Data_loader(Dataset):
def __init__(self, cfg, split, epoch_seed=None, mode='train'):
"""
data_dir: parent directory
split: list of int, e.g. [0,1,2,3,4] for train, [5] for val, [6] for test
num_points: number of points to sample per geometry
epoch_seed: seed for random sampling (for training)
mode: 'train', 'val', or 'test'
"""
self.data_dir = cfg.data_dir
self.split = split
self.num_points = cfg.num_points
self.epoch_seed = epoch_seed
self.mode = mode
self.cfg = cfg
self.meshes = []
self.mesh_names = []
for idx in split:
folder = f"{cfg.data_folder}_{idx}"
vtp_file = os.path.join(self.data_dir,folder, f"{folder}.vtp")
if not os.path.exists(vtp_file):
raise FileNotFoundError(f"{vtp_file} not found.")
mesh = pv.read(vtp_file)
self.meshes.append(mesh)
self.mesh_names.append(folder)
# For validation chunking
self.val_indices = None
self.val_chunk_ptr = 0
def set_epoch(self, epoch):
self.epoch_seed = epoch
self.val_indices = None
self.val_chunk_ptr = 0
def __len__(self):
if self.mode == 'train':
return len(self.meshes)
elif self.mode == 'val':
return len(self.meshes)
elif self.mode == 'test':
# Number of chunks = total points in all val meshes // num_points + remainder chunk
total = 0
for mesh in self.meshes:
return len(self.meshes)
else:
raise ValueError(f"Unknown mode: {self.mode}")
def __getitem__(self, idx):
if self.mode == 'train' or self.mode == 'val':
# Each item is a geometry, sample num_points randomly
mesh = self.meshes[idx]
n_pts = mesh.points.shape[0]
rng = np.random.default_rng(self.epoch_seed+idx)
indices = rng.choice(n_pts, self.num_points, replace=False)
pos = mesh.points[indices]
target = mesh["pressure"][indices]
pos = torch.tensor(pos, dtype=torch.float32)
target = torch.tensor(target, dtype=torch.float32).unsqueeze(-1)
if self.cfg.normalization == "std_norm":
target = (target - self.cfg.press_mean) / self.cfg.press_std
if self.cfg.pos_embed_sincos:
input_pos_mins = torch.tensor(self.cfg.input_pos_mins)
input_pos_maxs = torch.tensor(self.cfg.input_pos_maxs)
pos = 1000*(pos - input_pos_mins) / (input_pos_maxs - input_pos_mins)
return {"input_pos": pos, "output_feat": target ,"data_id": self.mesh_names[idx]}
elif self.mode == 'test':
# For each mesh in test, scramble all points and return the full mesh
mesh = self.meshes[idx]
n_pts = mesh.points.shape[0]
rng = np.random.default_rng(self.epoch_seed+idx)
indices = rng.permutation(n_pts)
pos = mesh.points[indices]
target = mesh["pressure"][indices]
pos = torch.tensor(pos, dtype=torch.float32)
target = torch.tensor(target, dtype=torch.float32).unsqueeze(-1)
if self.cfg.normalization == "std_norm":
target = (target - self.cfg.press_mean) / self.cfg.press_std
if self.cfg.pos_embed_sincos:
input_pos_mins = torch.tensor(self.cfg.input_pos_mins)
input_pos_maxs = torch.tensor(self.cfg.input_pos_maxs)
pos = 1000*(pos - input_pos_mins) / (input_pos_maxs - input_pos_mins)
return {"input_pos": pos, "output_feat": target ,"data_id": self.mesh_names[idx],"physical_coordinates":mesh.points[indices]}
else:
raise ValueError(f"Unknown mode: {self.mode}")
def get_dataloaders(cfg):
#with open(os.path.join(cfg.splits_file, "train_design_ids.txt")) as f:
# train_split = [int(line.strip().split('_')[-1]) for line in f if line.strip()]
# with open(os.path.join(cfg.splits_file, "val_design_ids.txt")) as f:
# val_split = [int(line.strip().split('_')[-1]) for line in f if line.strip()]
# with open(os.path.join(cfg.splits_file, "test_design_ids.txt")) as f:
# test_split = [int(line.strip().split('_')[-1]) for line in f if line.strip()]
# Create a list from 0 to 399
all_indices = set(range(400))
# Remove numbers present in test_split
test_split = [396, 397, 398, 399, 4, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 5, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 6, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 7, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 8, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 9, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99]
train_split = sorted(list(all_indices - set(test_split)))
print("Indices from 0 to 399 not in train_split:", test_split)
train_dataset = Data_loader(cfg, train_split, mode='train')
val_dataset = Data_loader(cfg, test_split, mode='val') ##??
test_dataset = Data_loader(cfg, test_split, mode='test')
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
return train_loader, val_loader, test_loader
|