()
| 266 | |
| 267 | |
| 268 | def main(): |
| 269 | |
| 270 | #argparse |
| 271 | parser = argparse.ArgumentParser(description='Train sdf and color') |
| 272 | parser.add_argument('--dataset_path', required=True, help='Path to the hair_synth dataset to train on') |
| 273 | parser.add_argument('--dataset_processed_path', required=True, help='Path to the hair_synth processed dataset to train on') |
| 274 | parser.add_argument('--exp_info', default="", help='Experiment info string useful for distinguishing one experiment for another') |
| 275 | parser.add_argument('--device', default="cuda") |
| 276 | args = parser.parse_args() |
| 277 | |
| 278 | #get the output path which will be at the root of the package |
| 279 | hair_forge_root=os.path.dirname(os.path.abspath(__file__)) |
| 280 | output_training_path=os.path.join(hair_forge_root, "out_training") |
| 281 | os.makedirs(output_training_path, exist_ok=True) |
| 282 | |
| 283 | |
| 284 | experiment_name="hair_forge" |
| 285 | if args.exp_info: |
| 286 | experiment_name+="_"+args.exp_info |
| 287 | |
| 288 | |
| 289 | loader_train, loader_test= create_dataloaders(args.dataset_path, args.dataset_processed_path) |
| 290 | |
| 291 | |
| 292 | hyperparams=HyperParamsRGB2Mat() |
| 293 | train(args, hyperparams, loader_train, loader_test, experiment_name, output_training_path) |
| 294 | |
| 295 | #finished training |
| 296 | return |
| 297 | |
| 298 | |
| 299 | if __name__ == '__main__': |
no test coverage detected