Spaces:
Sleeping
Sleeping
| # SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. | |
| # SPDX-FileCopyrightText: All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import numpy as np | |
| from torch.utils.data import DataLoader, Dataset | |
| class DatasetTopoDiff(Dataset): | |
| def __init__(self, topologies, stress, strain, load_im, constraints): | |
| self.topologies = topologies | |
| self.constraints = constraints | |
| self.image_size = topologies.shape[1] | |
| self.stress = stress | |
| self.strain = strain | |
| self.load_im = load_im | |
| def __len__(self): | |
| return self.topologies.shape[0] | |
| def __getitem__(self, idx): | |
| cons = self.constraints[idx] | |
| vol_frac = cons["VOL_FRAC"] | |
| cons = np.zeros((5, self.image_size, self.image_size)) | |
| cons[0] = self.stress[idx] | |
| cons[1] = self.strain[idx] | |
| cons[2] = self.load_im[idx][:, :, 0] | |
| cons[3] = self.load_im[idx][:, :, 1] | |
| cons[4] = np.ones((self.image_size, self.image_size)) * vol_frac | |
| return np.expand_dims(self.topologies[idx], 0) * 2 - 1, cons | |
| def load_data_topodiff( | |
| topologies, constraints, stress, strain, load_img, batch_size, deterministic=False | |
| ): | |
| dataset = DatasetTopoDiff(topologies, stress, strain, load_img, constraints) | |
| if deterministic: | |
| loader = DataLoader( | |
| dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True | |
| ) | |
| else: | |
| loader = DataLoader( | |
| dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True | |
| ) | |
| while True: | |
| yield from loader | |