Command to fetch MaterialX assets from PolyHaven and download them.
13def PolyHavenLoaderCmd():
14 '''
15 @brief Command to fetch MaterialX assets from PolyHaven and download them.
16 '''
17 parser = argparse.ArgumentParser(description="Fetch MaterialX assets from PolyHaven")
18 parser.add_argument("-id", "--download_id", type=str, default="", help="Filter ID to fetch MaterialX assets (e.g. 'polystyrene')")
19 parser.add_argument("-res", "--download_resolution", type=str, default="1k", help="Resolution of the MaterialX assets to download (e.g. '1k', '2k', '4k', '8k') ")
20 parser.add_argument("-fe", "--fetch", action='store_true', help="Fetch and save the MaterialX assets to a file")
21 parser.add_argument("-l", "--load", action='store_true', help="Load the MaterialX assets")
22 parser.add_argument("-df", "--data_folder", type=str, default="", help="Data folder to save / load MaterialX assets")
23 parser.add_argument("-c", "--count", type=int, default=None, help="Number of assets to fetch (default: 1)")
24 parser.add_argument('-exr', '--keep_exr', action='store_true', help="Keep EXR textures instead of converting to PNG (requires OpenImageIO)")
25 parser.add_argument('-x', '--extract_zip', action='store_true', help="Extract downloaded ZIP files")
26
27 args = parser.parse_args()
28 data_file = "polyhaven_materialx_assets.json"
29 fetch = args.fetch
30 download_id = args.download_id
31 resolution = args.download_resolution
32 load = args.load or args.download_id != ""
33 data_folder = args.data_folder
34
35 loader = polyHavenLoader.PolyHavenLoader()
36
37 materialx_assets = None
38
39 if fetch:
40 fetch_count = args.count
41 if fetch_count and fetch_count < 1:
42 fetch_count = 1
43 fetch_location = Path(data_folder) / data_file
44 logger.info(f"Fetching MaterialX assets to {fetch_location}...")
45 materialx_assets, all_assets, filtered_polyhaven_assets = loader.fetch_materialx_assets(download_id=download_id, max_items=fetch_count)
46 with open(fetch_location, "w") as f:
47 logger.info(f"- Saving MaterialX assets to {fetch_location}...")
48 json.dump(materialx_assets, f, indent=4)
49
50
51 all_location = Path(data_folder) / "polyhaven_assets.json"
52 with open(all_location, "w") as f:
53 json.dump(all_assets, f, indent=4)
54 logger.info(f"Saved all assets to {all_location}")
55
56 filtered_location = Path(data_folder) / "filtered_polyhaven_assets.json"
57 with open(filtered_location, "w") as f:
58 json.dump(filtered_polyhaven_assets, f, indent=4)
59 logger.info(f"Saved MaterialX assets to {filtered_location}")
60
61 elif load:
62 load_location = Path(data_folder) / data_file
63 if not load_location.exists():
64 load_location = Path(__file__).parent / "data" / "PolyHavenMaterialX" / data_file
65 if not load_location.exists():
66 logger.info(f"No MaterialX assets found at {load_location}. Please run with --fetch to fetch assets first.")
67 return
68
69 with open(load_location, "r") as f:
70 logger.info(f"Loaded MaterialX assets from {load_location}")
71 materialx_assets = json.load(f)
72
73
74
75 keep_exr = args.keep_exr if args.keep_exr else False
76 convert_exr_to_png = not keep_exr
77 extract_zip = args.extract_zip if args.extract_zip else False
78 if materialx_assets and download_id:
79
80 entry_id = download_id + '_' + resolution
81 entry = materialx_assets.get(entry_id)
82 if entry:
83 logger.info(f"Downloading asset with ID '{download_id}', resolution '{resolution}'")
84 asset_list = {entry_id: entry, resolution: resolution}
85 id, mtlx_string, texture_binaries = loader.download_asset(asset_list, convert_exr_to_png)
86
87 loader.save_materialx_with_textures(id, mtlx_string, texture_binaries, data_folder, extract_zip)
88 else:
89 logger.info(f"No asset found with ID '{entry_id}' in the MaterialX assets.")
90
91
92
