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("-dt", "--download_type", type=str, default="mtlx", help="Type of asset to download (e.g. 'mtlx', 'blend', 'gltf')")
20 parser.add_argument("-res", "--download_resolution", type=str, default="1k", help="Resolution of the MaterialX assets to download (e.g. '1k', '2k', '4k', '8k') ")
21 parser.add_argument("-fe", "--fetch", action='store_true', help="Fetch and save the MaterialX assets to a file")
22 parser.add_argument("-l", "--load", action='store_true', help="Load the MaterialX assets")
23 parser.add_argument("-df", "--data_folder", type=str, default="", help="Data folder to save / load MaterialX assets")
24 parser.add_argument("-c", "--count", type=int, default=None, help="Number of assets to fetch (default: 1)")
25 parser.add_argument('-exr', '--keep_exr', action='store_true', help="Keep EXR textures instead of converting to PNG (requires OpenImageIO)")
26 parser.add_argument('-x', '--extract_zip', action='store_true', help="Extract downloaded ZIP files")
27 parser.add_argument('-sf', '--save_filtered', action='store_true', help="Save filtered MaterialX assets to a file")
28
29 args = parser.parse_args()
30 download_type = args.download_type.lower()
31 all_data_file = 'polyhaven_assets.json'
32 filtered_data_file = ''
33 if args.save_filtered:
34 filtered_data_file = 'filtered_polyhaven_assets.json'
35 blend_data_file = "polyhaven_blender_assets.json"
36 gltf_data_file = "polyhaven_gltf_assets.json"
37 mtlx_data_file = "polyhaven_materialx_assets.json"
38 fetch = args.fetch
39 download_id = args.download_id
40 resolution = args.download_resolution
41 load = args.load or args.download_id != ""
42 data_folder = args.data_folder
43
44
45 if data_folder:
46 Path(data_folder).mkdir(parents=True, exist_ok=True)
47
48 loader = polyHavenLoader.PolyHavenLoader()
49
50 materialx_assets = None
51 blend_assets = None
52 gltf_assets = None
53
54 if fetch:
55 fetch_count = args.count
56 if fetch_count and fetch_count < 1:
57 fetch_count = 1
58 materialx_assets, all_assets, \
59 filtered_polyhaven_assets, blender_assets, gltf_assets = \
60 loader.fetch_materialx_assets(max_items=fetch_count, download_id=download_id, download_type=args.download_type)
61
62
63 all_location = Path(data_folder) / all_data_file
64 with open(all_location, "w") as f:
65 json.dump(all_assets, f, indent=4)
66 logger.info(f"Saved all assets to {all_location}")
67
68
69 if filtered_data_file:
70 filtered_location = Path(data_folder) / filtered_data_file
71 with open(filtered_location, "w") as f:
72 json.dump(filtered_polyhaven_assets, f, indent=4)
73 logger.info(f"Saved MaterialX assets to {filtered_location}")
74
75
76 mtlx_fetch_location = Path(data_folder) / mtlx_data_file
77 with open(mtlx_fetch_location, "w") as f:
78 logger.info(f"- Saving MaterialX assets to {mtlx_fetch_location}...")
79 json.dump(materialx_assets, f, indent=4)
80
81 blender_location = Path(data_folder) / blend_data_file
82 with open(blender_location, "w") as f:
83 json.dump(blender_assets, f, indent=4)
84 logger.info(f"Saved Blender assets to {blender_location}")
85
86 gltf_location = Path(data_folder) / gltf_data_file
87 with open(gltf_location, "w") as f:
88 json.dump(gltf_assets, f, indent=4)
89 logger.info(f"Saved glTF assets to {gltf_location}")
90
91 elif load:
92 file_list = [mtlx_data_file, blend_data_file, gltf_data_file]
93 for data_file in file_list:
94 load_location = Path(data_folder) / data_file
95 if not load_location.exists():
96 load_location = Path(__file__).parent / "data" / "PolyHavenMaterialX" / data_file
97 if not load_location.exists():
98 logger.info(f"Assets not found at {load_location}. Please run with --fetch to fetch assets first.")
99 return
100
101 with open(load_location, "r") as f:
102 logger.info(f"Loaded assets from {load_location}")
103 if data_file == mtlx_data_file:
104 materialx_assets = json.load(f)
105 elif data_file == blend_data_file:
106 blend_assets = json.load(f)
107 elif data_file == gltf_data_file:
108 gltf_assets = json.load(f)
109
110
111
112 keep_exr = args.keep_exr if args.keep_exr else False
113 if args.download_type.lower() in ['blend']:
114 keep_exr = True
115
116 convert_exr_to_png = not keep_exr
117 extract_zip = args.extract_zip if args.extract_zip else False
118
119 if args.download_type.lower() == 'mtlx':
120
121 if materialx_assets and download_id:
122
123 entry_id = download_id + '_' + resolution
124 entry = materialx_assets.get(entry_id)
125 if entry:
126 logger.info(f"Downloading asset with ID '{download_id}', resolution '{resolution}'")
127 asset_list = {entry_id: entry, resolution: resolution}
128 id, mtlx_string, texture_binaries = loader.download_mtlx_asset(asset_list, convert_exr_to_png)
129
130 loader.save_materialx_with_textures(id, mtlx_string, texture_binaries, data_folder, extract_zip)
131 else:
132 logger.info(f"No asset found with ID '{entry_id}' in the MaterialX assets.")
133
134
135
136 elif args.download_type.lower() == 'blend':
137 if blend_assets and download_id:
138
139 entry_id = download_id + '_' + resolution
140 entry = blend_assets.get(entry_id)
141 if entry:
142 logger.info(f"Downloading Blender asset with ID '{download_id}', resolution '{resolution}'")
143 asset_list = {entry_id: entry, resolution: resolution}
144 id, blend_binary, texture_binaries = loader.download_blender_asset(asset_list)
145 loader.save_blender_with_textures(id, blend_binary, texture_binaries, data_folder, extract_zip)
146 else:
147 logger.info(f"No asset found with ID '{entry_id}' in the MaterialX assets.")
148
149
150
151 elif args.download_type.lower() == 'gltf':
152 if gltf_assets and download_id:
153
154 entry_id = download_id + '_' + resolution
155 entry = gltf_assets.get(entry_id)
156 if entry:
157 logger.info(f"Downloading glTF asset with ID '{download_id}', resolution '{resolution}'")
158 asset_list = {entry_id: entry, resolution: resolution}
159 id, gltf_ascii, texture_binaries = loader.download_gltf_asset(asset_list)
160 loader.save_gltf_with_textures(id, gltf_ascii, texture_binaries, data_folder, extract_zip)
161 else:
162 logger.info(f"No asset found with ID '{entry_id}' in the MaterialX assets.")
163
164
165
