diff options
Diffstat (limited to 'Scripts/make_dfg_lut.py')
| -rw-r--r-- | Scripts/make_dfg_lut.py | 271 |
1 files changed, 271 insertions, 0 deletions
diff --git a/Scripts/make_dfg_lut.py b/Scripts/make_dfg_lut.py new file mode 100644 index 0000000..b4faf0f --- /dev/null +++ b/Scripts/make_dfg_lut.py @@ -0,0 +1,271 @@ +#!/usr/bin/env python3 + +import argparse +import math +import numpy as np +import OpenEXR +import Imath +import numba +import random +import concurrent.futures +import os +from functools import partial + + +@numba.njit(cache=True) +def rcp(a): + return 1.0 / a + + +@numba.njit(cache=True) +def lerp(a, b, t): + return a + (b - a) * t + + +@numba.njit(cache=True) +def saturate(a): + if a < 0.0: return 0.0 + if a > 1.0: return 1.0 + return a + + +# Standard BRDF components. +@numba.njit(cache=True) +def F_Schlick(LoH, f0, f90=1.0): + term = 1.0 - LoH + term2 = term * term + term5 = term2 * term2 * term + return f0 + (f90 - f0) * term5 + + +@numba.njit(cache=True) +def D_GGX(roughness, NoH): + r2 = roughness * roughness + NoH2 = NoH * NoH + NoH4 = NoH2 * NoH2 + k = rcp(NoH2) - 1.0 + r2_plus_k = r2 + k + denom = NoH4 * r2_plus_k * r2_plus_k + return r2 / (denom + 1e-6) + + +@numba.njit(cache=True) +def G_GGXSmith(roughness, NoL, NoV): + denom = 2.0 * lerp(2.0 * NoL * NoV, NoL + NoV, roughness) + return rcp(denom + 1e-6) + + +# Cloth BRDF components. +@numba.njit(cache=True) +def D_Cloth(roughness, NoH): + if roughness < 1e-4: return 0.0 + r_rcp = rcp(roughness) + sin2H = 1.0 - NoH * NoH + return (2.0 + r_rcp) * pow(sin2H, r_rcp * 0.5) / (2.0 * math.pi) + + +@numba.njit(cache=True) +def G_Cloth_L(x, a, b, c, d, e): + return a / (1.0 + b * pow(x, c)) + d * x + e + + +@numba.njit(cache=True) +def G_Cloth(roughness, LoH): + a0, a1 = 25.3245, 21.5473 + b0, b1 = 3.32435, 3.82987 + c0, c1 = 0.16801, 0.19823 + d0, d1 = -1.27393, -1.97760 + e0, e1 = -4.85967, -4.32054 + + one_minus_r = 1.0 - roughness + one_minus_r_sq = one_minus_r * one_minus_r + one_minus_LoH = 1.0 - LoH + + lambda_val = 0.0 + if LoH < 0.5: + L0 = G_Cloth_L(LoH, a0, b0, c0, d0, e0) + L1 = G_Cloth_L(LoH, a1, b1, c1, d1, e1) + L = lerp(L0, L1, one_minus_r_sq) + lambda_val = math.exp(L) + else: + L_05_0 = G_Cloth_L(0.5, a0, b0, c0, d0, e0) + L_05_1 = G_Cloth_L(0.5, a1, b1, c1, d1, e1) + L_05 = lerp(L_05_0, L_05_1, one_minus_r_sq) + + L_LoH_0 = G_Cloth_L(one_minus_LoH, a0, b0, c0, d0, e0) + L_LoH_1 = G_Cloth_L(one_minus_LoH, a1, b1, c1, d1, e1) + L_LoH = lerp(L_LoH_0, L_LoH_1, one_minus_r_sq) + + lambda_val = math.exp(2.0 * L_05 - L_LoH) + + # Apply terminator softening (equation 4) + return pow(lambda_val, 1.0 + 2.0 * pow(one_minus_LoH, 8.0)) + + +@numba.njit(cache=True) +def integrate_brdf_jitted(roughness, NoV, brdf_type, num_samples): + V_x = math.sqrt(1.0 - NoV * NoV) + V_y = 0.0 + V_z = NoV + + A, B = 0.0, 0.0 + + for i in range(num_samples): + e1, e2 = random.random(), random.random() + + # Importance sample GGX + a = roughness + a2 = a * a + + phi = 2.0 * math.pi * e1 + cos_theta = math.sqrt((1.0 - e2) / (1.0 + (a2 - 1.0) * e2)) + sin_theta = math.sqrt(1.0 - cos_theta * cos_theta) + + H_x = math.cos(phi) * sin_theta + H_y = math.sin(phi) * sin_theta + H_z = cos_theta + + VoH = H_x * V_x + H_y * V_y + H_z * V_z + if VoH <= 0: continue + + L_x = 2.0 * VoH * H_x - V_x + L_y = 2.0 * VoH * H_y - V_y + L_z = 2.0 * VoH * H_z - V_z + + NoL = saturate(L_z) + NoH = saturate(H_z) + NoV_proxy = saturate(V_z) # NoV is V_z + + if NoL > 0: + scale, bias = 0.0, 0.0 + # --- Standard BRDF --- + if brdf_type == 1: + # Note that the D term is present in the numerator and the denominator, so it cancels out. + #D = D_GGX(roughness, NoH) + G = G_GGXSmith(roughness, NoL, NoV_proxy) + Fc_term = pow(1.0 - VoH, 5.0) + + # PDF of GGX Importance Sampling is D * NoH / (4 * VoH). + # The full term is (D * G * NoL) / PDF, which simplifies to: + # G * NoL * (4 * VoH / NoH). + # This can be unstable when NoH is close to zero, so we clamp the denominator. + common_term = (G * NoL * 4.0 * VoH) / max(NoH, 1e-5) + + # We are baking the two components of the split-sum approximation for IBL: + # reflectance = f0 * scale + bias + scale = common_term * (1.0 - Fc_term) + bias = common_term * Fc_term + # --- Cloth BRDF --- + elif brdf_type == 2: + # We are importance sampling GGX, so must account for its PDF. + D_c = D_Cloth(roughness, NoH) + G_c = G_Cloth(roughness, VoH) + + # PDF = D_GGX(r, NoH) * NoH / (4 * VoH) + pdf_ggx = D_GGX(roughness, NoH) * NoH / (4.0 * VoH + 1e-6) + + # We must divide by the PDF and multiply by our target distribution and the cosine term. + scale = (D_c * G_c * NoL) / (pdf_ggx + 1e-6) + bias = 0.0 + + A += scale + B += bias + + return A / num_samples, B / num_samples + + +def calculate_pixel(coords, resolution, brdf_type, num_samples): + x, y = coords + u = (x + 0.5) / resolution + v = (y + 0.5) / resolution + + roughness = saturate(u) + NoV = saturate(v) + if NoV < 1e-4: return x, y, 0.0, 0.0, 0.0 + + r, g = 0.0, 0.0 + if brdf_type == 1: # standard + r, g = integrate_brdf_jitted(roughness, NoV, 1, num_samples) + elif brdf_type == 2: # cloth + if roughness < 1e-4: return x, y, 0.0, 0.0, 0.0 + r, g = integrate_brdf_jitted(roughness, NoV, 2, num_samples) + + return x, y, r, g, 0.0 + + +def generate_exr(resolution, output_filename, brdf_type, num_samples, num_workers): + print(f"Generating {resolution}x{resolution} EXR '{output_filename}' ({num_samples} samples/pixel) using {num_workers} workers.") + header = OpenEXR.Header(resolution, resolution) + pt = Imath.PixelType(Imath.PixelType.FLOAT) + header['channels'] = { 'R': Imath.Channel(pt), 'G': Imath.Channel(pt), 'B': Imath.Channel(pt) } + + pixel_data = np.zeros((resolution, resolution, 3), dtype=np.float32) + + coords_to_process = [(x, y) for y in range(resolution) for x in range(resolution)] + worker_func = partial(calculate_pixel, resolution=resolution, brdf_type=brdf_type, num_samples=num_samples) + + processed_count = 0 + total_pixels = len(coords_to_process) + print(f"Starting pixel processing...") + + with concurrent.futures.ProcessPoolExecutor(max_workers=num_workers) as executor: + futures = {executor.submit(worker_func, coord): coord for coord in coords_to_process} + + for future in concurrent.futures.as_completed(futures): + try: + x, y, r, g, b = future.result() + pixel_data[y, x] = (r, g, b) + except Exception as exc: + coord = futures[future] + print(f'\nPixel at {coord} generated an exception: {exc}') + + processed_count += 1 + print(f" ...processed {processed_count}/{total_pixels} pixels ({processed_count/total_pixels:.1%})", end='\r') + + print(f"\nProcessing complete. Writing to {output_filename}...") + try: + # Vertically flip to match UV coordinates (0,0 at bottom-left). + pixel_data = np.flipud(pixel_data) + + exr_file = OpenEXR.OutputFile(output_filename, header) + r_data = pixel_data[:, :, 0].ravel().tobytes() + g_data = pixel_data[:, :, 1].ravel().tobytes() + b_data = pixel_data[:, :, 2].ravel().tobytes() + exr_file.writePixels({'R': r_data, 'G': g_data, 'B': b_data}) + exr_file.close() + print(f"Successfully generated {output_filename}") + except Exception as e: + raise RuntimeError(f"Failed to write EXR file '{output_filename}': {e}") + +def main(): + parser = argparse.ArgumentParser(description='Generate DFG LUT EXR images for PBR.') + parser.add_argument('-t', '--type', choices=['standard', 'cloth'], default='standard', + help='Type of DFG texture to generate (default: standard)') + parser.add_argument('-r', '--resolution', type=int, default=128, + help='Resolution of the square EXR image (default: 128)') + parser.add_argument('-s', '--samples', type=int, default=1024, + help='Number of samples per pixel for integration (default: 1024)') + parser.add_argument('-o', '--output', + help='Output filename (default: dfg_standard.exr or dfg_cloth.exr)') + parser.add_argument('-j', '--workers', type=int, default=os.cpu_count(), + help=f'Number of worker processes (default: {os.cpu_count()})') + + args = parser.parse_args() + + if args.resolution <= 0: + print("Error: Resolution must be a positive integer") + return 1 + + brdf_type = 1 if args.type == 'standard' else 2 + output_filename = args.output if args.output else f'dfg_{args.type}.exr' + + try: + generate_exr(args.resolution, output_filename, brdf_type, args.samples, args.workers) + except Exception as e: + print(f"Error: {e}") + return 1 + + return 0 + +if __name__ == '__main__': + exit(main()) |