diff --git a/scripts/unity_dataset_processing/decimate.py b/scripts/unity_dataset_processing/decimate.py index e6ca4b62c6..a219c6f63d 100644 --- a/scripts/unity_dataset_processing/decimate.py +++ b/scripts/unity_dataset_processing/decimate.py @@ -116,16 +116,15 @@ def decimate( for i in range(importer.mesh_count): mesh = importer.mesh(i) - # Transform the mesh to its max scale in the scene. For quantized meshes - # this expands the position attribute to a floating-point Vector3. - scaled_mesh = meshtools.transform3d( - mesh, Matrix4.scaling(max_mesh_scaling.get(i, Vector3(1.0))) - ) - # Calculate total triangle area of the *transformed* mesh. You might want # to fiddle with this heuristics, another option is calculating the mesh # AABB but that won't do the right thing for planar meshes. if simplify: + # Transform the mesh to its max scale in the scene. For quantized meshes + # this expands the position attribute to a floating-point Vector3. + scaled_mesh = meshtools.transform3d( + mesh, Matrix4.scaling(max_mesh_scaling.get(i, Vector3(1.0))) + ) if not scaled_mesh.is_indexed: converter.end_file() importer.close() diff --git a/scripts/unity_dataset_processing/unity_dataset_processing.py b/scripts/unity_dataset_processing/unity_dataset_processing.py index 458759c89b..29f590850c 100644 --- a/scripts/unity_dataset_processing/unity_dataset_processing.py +++ b/scripts/unity_dataset_processing/unity_dataset_processing.py @@ -7,6 +7,7 @@ import argparse import json import os +import time from multiprocessing import Manager, Pool from pathlib import Path from typing import Callable, List, Set @@ -16,6 +17,7 @@ OUTPUT_DIR = "data/hitl_simplified/data/" OMIT_BLACK_LIST = False OMIT_GRAY_LIST = False +PROCESS_COUNT = os.cpu_count() class Job: @@ -24,6 +26,13 @@ class Job: simplify: bool +class Config: + # Increase logging verbosity. + verbose: bool = False + # Activate multiprocessing. Disable when debugging. + use_multiprocessing: bool = False + + def file_is_scene_config(filepath: str) -> bool: """ Return whether or not the file is an scene_instance.json @@ -93,8 +102,14 @@ def get_model_ids_from_scene_instance_json(filepath: str) -> List[str]: return model_ids +def validate_jobs(jobs: List[Job]): + for job in jobs: + assert Path(job.source_path).exists + assert job.dest_path != None and job.dest_path != "" + + def process_model(args): - job, counter, lock, total_models = args + job, counter, lock, total_models, verbose = args if os.path.isfile(job.dest_path): print(f"Skipping: {job.source_path}") @@ -110,19 +125,21 @@ def process_model(args): source_tris, target_tris, simplified_tris = decimate.decimate( inputFile=job.source_path, outputFile=job.dest_path, - quiet=True, + quiet=not verbose, + verbose=verbose, sloppy=False, simplify=job.simplify, ) except Exception: try: print( - f"Unable to decimate: {job.source_path}. Trying passthrough (no decimation)." + f"Unable to decimate: {job.source_path}. Trying without decimation." ) source_tris, target_tris, simplified_tris = decimate.decimate( inputFile=job.source_path, outputFile=job.dest_path, - quiet=True, + quiet=not verbose, + verbose=verbose, simplify=False, ) except Exception: @@ -135,8 +152,8 @@ def process_model(args): ) result = { - "source_tris": source_tris, - "simplified_tris": simplified_tris, + "source_tris": str(source_tris), + "simplified_tris": str(simplified_tris), "source_path": job.source_path, "status": "ok", } @@ -161,7 +178,8 @@ def process_model(args): return result -def simplify_models(jobs: List[Job]): +def simplify_models(jobs: List[Job], config: Config): + start_time = time.time() total_source_tris = 0 total_simplified_tris = 0 black_list = [] @@ -171,6 +189,9 @@ def simplify_models(jobs: List[Job]): total_skipped = 0 total_error = 0 + validate_jobs(jobs) + total_models = len(jobs) + # Initialize counter and lock manager = Manager() counter = manager.Value("i", 0) @@ -179,13 +200,14 @@ def simplify_models(jobs: List[Job]): total_models = len(jobs) # Pair up the model paths with the counter and lock - args_lists = [(job, counter, lock, total_models) for job in jobs] + args_lists = [ + (job, counter, lock, total_models, config.verbose) for job in jobs + ] results = [] - use_multiprocessing = False # total_models > 6 - if use_multiprocessing: - max_processes = 6 + if config.use_multiprocessing: + max_processes = PROCESS_COUNT with Pool(processes=min(max_processes, total_models)) as pool: results = list(pool.map(process_model, args_lists)) else: @@ -226,6 +248,11 @@ def simplify_models(jobs: List[Job]): print(" " + item + ",") print("]") + elapsed_time = time.time() - start_time + print(f"Elapsed time (s): {elapsed_time}") + if not config.use_multiprocessing: + print("Add --use-multiprocessing to speed-up processing.") + def find_model_paths_in_scenes(hssd_hab_root_dir, scene_ids) -> List[str]: model_filepaths: Set[str] = set() @@ -279,8 +306,23 @@ def main(): type=str, help="one or more scene ids", ) + parser.add_argument( + "--verbose", + action="store_true", + default=False, + help="Increase logging verbosity.", + ) + parser.add_argument( + "--use-multiprocessing", + action="store_true", + default=False, + help="Enable multiprocessing.", + ) args = parser.parse_args() + config = Config() + config.verbose = args.verbose + config.use_multiprocessing = args.use_multiprocessing # Force input paths to have a trailing slash if args.hssd_hab_root_dir[-1] != "/": @@ -352,7 +394,7 @@ def main(): job.simplify = False jobs.append(job) - simplify_models(jobs) + simplify_models(jobs, config) if __name__ == "__main__":