| import sys |
| from src.models.geometry.render import renderutils as ru |
| import torch |
| from src.models.geometry.render import util |
| import nvdiffrast.torch as dr |
| import os |
|
|
| from PIL import Image |
| import torchvision.transforms.functional as TF |
| import torchvision.utils as vutils |
| import imageio |
| os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" |
| LIGHT_MIN_RES = 16 |
|
|
| MIN_ROUGHNESS = 0.04 |
| MAX_ROUGHNESS = 1.00 |
|
|
| class cubemap_mip(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx, cubemap): |
| return util.avg_pool_nhwc(cubemap, (2,2)) |
|
|
| @staticmethod |
| def backward(ctx, dout): |
| res = dout.shape[1] * 2 |
| out = torch.zeros(6, res, res, dout.shape[-1], dtype=torch.float32, device="cuda") |
| for s in range(6): |
| gy, gx = torch.meshgrid(torch.linspace(-1.0 + 1.0 / res, 1.0 - 1.0 / res, res, device="cuda"), |
| torch.linspace(-1.0 + 1.0 / res, 1.0 - 1.0 / res, res, device="cuda"), |
| indexing='ij') |
| v = util.safe_normalize(util.cube_to_dir(s, gx, gy)) |
| out[s, ...] = dr.texture(dout[None, ...] * 0.25, v[None, ...].contiguous(), filter_mode='linear', boundary_mode='cube') |
| return out |
| |
| def build_mips(base, cutoff=0.99): |
| specular = [base] |
| while specular[-1].shape[1] > LIGHT_MIN_RES: |
| specular.append(cubemap_mip.apply(specular[-1])) |
| |
|
|
| diffuse = ru.diffuse_cubemap(specular[-1]) |
|
|
| for idx in range(len(specular) - 1): |
| roughness = (idx / (len(specular) - 2)) * (MAX_ROUGHNESS - MIN_ROUGHNESS) + MIN_ROUGHNESS |
| specular[idx] = ru.specular_cubemap(specular[idx], roughness, cutoff) |
| specular[-1] = ru.specular_cubemap(specular[-1], 1.0, cutoff) |
|
|
| return specular, diffuse |
|
|
| |
| |
| def _load_env_hdr(fn, scale=1.0): |
| latlong_img = torch.tensor(util.load_image(fn), dtype=torch.float32, device='cuda')*scale |
| cubemap = util.latlong_to_cubemap(latlong_img, [512, 512]) |
|
|
| specular, diffuse = build_mips(cubemap) |
|
|
| return specular, diffuse |
|
|
| def main(path_hdr, save_path_map): |
| all_envs = os.listdir(path_hdr) |
|
|
| for env in all_envs: |
| env_path = os.path.join(path_hdr, env) |
| base_n = os.path.basename(env_path).split('.')[0] |
|
|
| try: |
| if not os.path.exists(os.path.join(save_path_map, base_n)): |
| os.makedirs(os.path.join(save_path_map, base_n)) |
| specular, diffuse = _load_env_hdr(env_path) |
| for i in range(len(specular)): |
| tensor = specular[i] |
| torch.save(tensor, os.path.join(save_path_map, base_n, f'specular_{i}.pth')) |
| |
| torch.save(diffuse, os.path.join(save_path_map, base_n, 'diffuse.pth')) |
| except Exception as e: |
| print(f"Error processing {env}: {e}") |
| continue |
|
|
| if __name__ == "__main__": |
| if len(sys.argv) != 3: |
| print("Usage: python script.py <path_hdr> <save_path_map>") |
| sys.exit(1) |
| |
| path_hdr = sys.argv[1] |
| save_path_map = sys.argv[2] |
|
|
| if not os.path.exists(path_hdr): |
| print(f"Error: path_hdr '{path_hdr}' does not exist.") |
| sys.exit(1) |
|
|
| if not os.path.exists(save_path_map): |
| os.makedirs(save_path_map) |
|
|
| main(path_hdr, save_path_map) |
