14def main():
15 parser = argparse.ArgumentParser(description="Fetch MaterialX assets from PolyHaven")
16 parser.add_argument("-id", "--download_id", type=str, default="polystyrene", help="Filter ID to fetch MaterialX assets (e.g. 'polystyrene')")
17 parser.add_argument("-res", "--download_resolution", type=str, default="1k", help="Resolution of the MaterialX assets to download (e.g. '1k', '2k', '4k', '8k') ")
18 parser.add_argument("-fe", "--fetch", action='store_true', help="Fetch and save the MaterialX assets to a file")
19 parser.add_argument("-l", "--load", action='store_true', help="Load the MaterialX assets")
20 parser.add_argument("-df", "--data_folder", type=str, default="data/PolyHavenMaterialX", help="Data folder to save / load MaterialX assets")
21 parser.add_argument("-c", "--count", type=int, default=None, help="Number of assets to fetch (default: 1)")
22
23 args = parser.parse_args()
24 data_file = "polyhaven_materialx_assets.json"
25 fetch = args.fetch
26 download_id = args.download_id
27 resolution = args.download_resolution
28 load = args.load
29 data_folder = args.data_folder
30
31 loader = polyHavenLoader.PolyHavenLoader()
32
33 materialx_assets = None
34
35 if fetch:
36 fetch_count = args.count
37 if fetch_count and fetch_count < 1:
38 fetch_count = 1
39 fetch_location = Path(data_folder) / data_file
40 logger.info(f"Fetching MaterialX assets to {fetch_location}...")
41 materialx_assets, all_assets, filtered_polyhaven_assets = loader.fetch_materialx_assets(download_id=download_id, max_items=fetch_count)
42 with open(fetch_location, "w") as f:
43 json.dump(materialx_assets, f, indent=4)
44
45
46 all_location = Path(data_folder) / "polyhaven_assets.json"
47 with open(all_location, "w") as f:
48 json.dump(all_assets, f, indent=4)
49 logger.info(f"Saved all assets to {all_location}")
50
51 filtered_location = Path(data_folder) / "filtered_polyhaven_assets.json"
52 with open(filtered_location, "w") as f:
53 json.dump(filtered_polyhaven_assets, f, indent=4)
54 logger.info(f"Saved MaterialX assets to {filtered_location}")
55
56 elif load:
57 load_location = Path(data_folder) / data_file
58 with open(load_location, "r") as f:
59 logger.info(f"Loaded MaterialX assets from {load_location}")
60 materialx_assets = json.load(f)
61
62
63
64 if materialx_assets and download_id:
65
66 entry_id = download_id + '___' + resolution
67 entry = materialx_assets.get(entry_id)
68 if entry:
69 logger.info(f"Downloading asset with ID '{download_id}', resolution '{resolution}'")
70 asset_list = {entry_id: entry, resolution: resolution}
71 id, mtlx_string, texture_binaries = loader.download_asset(asset_list)
72 loader.save_materialx_with_textures(id, mtlx_string, texture_binaries, data_folder)
73 else:
74 logger.info(f"No asset found with ID '{entry_id}' in the MaterialX assets.")
75 else:
76 logger.info("No operation specified.")
77
1.12.0