This function is called first in the training process. Returns: train_mesh_filenames: list of mesh filenames for training (full path) test_mesh_filenames: list of mesh filenames for testing (full path) or None
(
dataset_name: str,
dataset_root_dir: str,
cleaned_root_dir: str,
clean_mesh: bool,
train_mesh_filenames: T.List[str],
test_mesh_filenames: T.List[str] = None,
rank: int = 0,
world_size: int = 1,
printout: bool = True,
)
| 148 | |
| 149 | |
| 150 | def gather_and_clean_dataset( |
| 151 | dataset_name: str, |
| 152 | dataset_root_dir: str, |
| 153 | cleaned_root_dir: str, |
| 154 | clean_mesh: bool, |
| 155 | train_mesh_filenames: T.List[str], |
| 156 | test_mesh_filenames: T.List[str] = None, |
| 157 | rank: int = 0, |
| 158 | world_size: int = 1, |
| 159 | printout: bool = True, |
| 160 | ) -> T.Dict[str, T.List[str]]: |
| 161 | """ |
| 162 | This function is called first in the training process. |
| 163 | |
| 164 | Returns: |
| 165 | train_mesh_filenames: |
| 166 | list of mesh filenames for training (full path) |
| 167 | test_mesh_filenames: |
| 168 | list of mesh filenames for testing (full path) or None |
| 169 | """ |
| 170 | |
| 171 | if printout: |
| 172 | print(f'dataset_name = {dataset_name}') |
| 173 | print(f'dataset_root_dir = {dataset_root_dir}') |
| 174 | print(f'cleaned_root_dir = {cleaned_root_dir}') |
| 175 | print(f'clean_mesh = {clean_mesh}') |
| 176 | print(f'train_mesh_filenames:') |
| 177 | print(train_mesh_filenames) |
| 178 | print(f'test_mesh_filenames:') |
| 179 | print(test_mesh_filenames) |
| 180 | print(f'rank = {rank}', flush=True) |
| 181 | |
| 182 | if train_mesh_filenames is None: |
| 183 | train_mesh_filenames = [] |
| 184 | if test_mesh_filenames is None: |
| 185 | test_mesh_filenames = [] |
| 186 | |
| 187 | # download ShapeNet dataset |
| 188 | if rank == 0: |
| 189 | |
| 190 | if dataset_name.lower() in { |
| 191 | 'shapenet', 'shapenet-debug', 'sketchfab', |
| 192 | 'sketchfab-small', 'sketchfab-small-debug', |
| 193 | 'tex-models', |
| 194 | }: |
| 195 | assert os.path.exists(dataset_root_dir) |
| 196 | |
| 197 | # compile the mesh_filenames |
| 198 | train_mesh_filenames = [ |
| 199 | os.path.join( |
| 200 | dataset_root_dir, |
| 201 | mesh_filename, |
| 202 | ) |
| 203 | for mesh_filename in train_mesh_filenames |
| 204 | ] |
| 205 | |
| 206 | if test_mesh_filenames is not None: |
| 207 | test_mesh_filenames = [ |
no test coverage detected