|
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) |
