Add script for fast rotation figure.

.gitignore

@@ -1,4 +1,7 @@
 **/__pycache__/*
+*.zip
 data
+scenes
 envmaps
+renders
 textures

figures/fast_rotation/gt.py

@@ -0,0 +1,24 @@
+import mitsuba as mi
+mi.set_variant("cuda_rgb")
+mi.set_log_level(mi.LogLevel.Info)
+
+import os
+import plugins as _
+from utils.render import render_multi_pass
+from figures.fast_rotation.scene import config
+
+scene_path = os.path.join("scenes", "teapot", "scene.xml")
+
+for i in range(2):
+ scene = mi.load_file(
+ scene_path,
+ resx=config["resx"], resy=config["resy"],
+ spp=config["spp"], alpha=config[f"alpha_{i}"], glint=""
+ )
+
+ render = render_multi_pass(
+ mi.render,
+ config["resx"], config["resy"],
+ scene, config["spp"],
+ os.path.join("renders", "fast_rotation", f"gt_{i}.exr")
+ )

figures/fast_rotation/onfly.py

@@ -0,0 +1,44 @@
+import mitsuba as mi
+mi.set_variant("cuda_rgb", "cuda_ad_rgb")
+mi.set_log_level(mi.LogLevel.Info)
+
+import os
+import numpy as np
+import plugins as _
+from figures.fast_rotation.scene import config
+from plugins.ratio_estimator import ratio_estimator
+from plugins.sh_half_integrator import render_half_fastdot
+
+sh_order = config["sh_order"]
+scene_path = os.path.join("scenes", "teapot", "scene.xml")
+envmap_path = os.path.join("scenes", "teapot", "envmaps", "indoor_1.hdr")
+
+for i in range(2):
+ scene = mi.load_file(
+ scene_path,
+ resx=config["resx"], resy=config["resy"],
+ spp=config["spp"], alpha=config[f"alpha_{i}"],
+ glint="-glint"
+ )
+ render, sh_pixels = render_half_fastdot(
+ scene,
+ envmap_path,
+ resx=config["resx"], resy=config["resy"],
+ fast_rotation=True
+ )
+ sh_pixels = sh_pixels.astype(bool)
+
+ scene = mi.load_file(
+ scene_path,
+ resx=config["resx"], resy=config["resy"],
+ spp=config["spp"], alpha=config[f"alpha_{i}"],
+ glint=""
+ )
+ shadowed, unshadowed, _ = ratio_estimator(scene)
+ ratio = np.where(unshadowed > 1e-6, shadowed / unshadowed, 0)
+ bg = mi.render(scene, spp=64).numpy()
+
+ render *= ratio
+ render[~sh_pixels] = bg[~sh_pixels]
+
+ mi.Bitmap(render).write(os.path.join("renders", "fast_rotation", f"onfly_{i}.exr"))

figures/fast_rotation/precomp.py

@@ -0,0 +1,46 @@
+import mitsuba as mi
+mi.set_variant("cuda_rgb", "cuda_ad_rgb")
+mi.set_log_level(mi.LogLevel.Info)
+
+import os
+import numpy as np
+import plugins as _
+from figures.fast_rotation.scene import config
+from plugins.ratio_estimator import ratio_estimator
+from plugins.sh_half_integrator import render_half_fastdot
+
+sh_order = config["sh_order"]
+scene_path = os.path.join("scenes", "teapot", "scene.xml")
+envmap_path = os.path.join("scenes", "teapot", "envmaps", "indoor_1.hdr")
+
+for i in range(2):
+ scene = mi.load_file(
+ scene_path,
+ resx=config["resx"], resy=config["resy"],
+ spp=config["spp"], alpha=config[f"alpha_{i}"],
+ glint="-glint"
+ )
+
+ render, sh_pixels = render_half_fastdot(
+ scene,
+ envmap_path,
+ resx=config["resx"], resy=config["resy"],
+ fast_rotation=False
+ )
+
+ sh_pixels = sh_pixels.astype(bool)
+
+ scene = mi.load_file(
+ scene_path,
+ resx=config["resx"], resy=config["resy"],
+ spp=config["spp"], alpha=config[f"alpha_{i}"],
+ glint=""
+ )
+ shadowed, unshadowed, _ = ratio_estimator(scene)
+ ratio = np.where(unshadowed > 1e-6, shadowed / unshadowed, 0)
+ bg = mi.render(scene, spp=64).numpy()
+
+ render *= ratio
+ render[~sh_pixels] = bg[~sh_pixels]
+
+ mi.Bitmap(render).write(os.path.join("renders", "fast_rotation", f"precomp_{i}.exr"))

figures/fast_rotation/run_all.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+mkdir -p renders/fast_rotation
+python figures/fast_rotation/gt.py
+python figures/fast_rotation/onfly.py
+python figures/fast_rotation/precomp.py

figures/fast_rotation/scene.py

@@ -0,0 +1,10 @@
+import mitsuba as mi
+
+config = {
+ "resx": 1600,
+ "resy": 1080,
+ "alpha_0": 0.05,
+ "alpha_1": 0.1,
+ "sh_order": 100,
+ "spp": 10000
+}

plugins/__init__.py

@@ -0,0 +1,5 @@
+import mitsuba as mi
+from plugins.glint import GlintDummy
+
+# BSDFs
+mi.register_bsdf("glint_dummy", lambda props: GlintDummy(props))

plugins/glint.py

@@ -0,0 +1,78 @@
+import mitsuba as mi
+import drjit as dr
+from utils.ior import complex_ior_from_file
+
+class GlintDummy(mi.BSDF):
+ def __init__(self, props: mi.Properties):
+ super().__init__(props)
+
+ self.glint_idx = props.get("glint_idx", mi.Texture.D65(1))
+ self.alpha: mi.Texture = props.get("alpha", mi.Texture.D65(0.01))
+ self.material = props.get("material", "none")
+ if self.material == "none" or props.has_property("eta"):
+ self.eta: mi.Texture = props.get("eta", mi.Texture.D65(0.01))
+ self.k: mi.Texture = props.get("k", mi.Texture.D65(1))
+ else:
+ self.eta, self.k = complex_ior_from_file(self.material)
+
+ self.alpha_mul: mi.Texture = mi.Texture.D65(props.get("alpha_mul", 1.0))
+ self.alpha_add: mi.Texture = mi.Texture.D65(props.get("alpha_add", 0.01))
+ self.clearcoat_alpha: mi.Texture = mi.Texture.D65(props.get("clearcoat_alpha", 0.05))
+ self.specular_reflectance: mi.Texture = props.get("specular_reflectance", mi.Texture.D65(1))
+
+ def traverse(self, callback: mi.TraversalCallback):
+ callback.put_object("alpha", self.alpha, mi.ParamFlags.Differentiable)
+ callback.put_object("alpha_mul", self.alpha_mul, mi.ParamFlags.Differentiable)
+ callback.put_object("alpha_add", self.alpha_add, mi.ParamFlags.Differentiable)
+ callback.put_object("glint_idx", self.glint_idx, mi.ParamFlags.Differentiable)
+ callback.put_object("clearcoat_alpha", self.clearcoat_alpha, mi.ParamFlags.Differentiable)
+
+ # This function returns the FG term for the base layer
+ def eval(self, ctx: mi.BSDFContext, si: mi.SurfaceInteraction3f, wo: mi.Vector3f, active: bool = True):
+ eta = self.eta.eval(si, active)
+ k = self.k.eval(si, active)
+ alpha = self.alpha.eval_1(si, active)
+ alpha_mul = self.alpha_mul.eval_1(si, active)
+ alpha_add = self.alpha_add.eval_1(si, active)
+ # alpha = dr.fma(alpha, alpha_mul, alpha_add)
+ alpha *= alpha_mul
+ alpha = dr.clamp(alpha, 0.01, 0.99)
+
+ # Find half vector
+ m = dr.normalize(si.wi + wo)
+
+ # Find fresnel term
+ f = mi.Color3f(0)
+ for i in range(3):
+ f[i] = mi.fresnel_conductor(dr.dot(si.wi, m), mi.Complex2f(eta[i], k[i]))
+
+ # Find shadowing
+ ndf = mi.MicrofacetDistribution(mi.MicrofacetType.Beckmann, alpha, sample_visible=False)
+ g = ndf.G(si.wi, wo, m)
+
+ # Specular reflectance
+ spec = self.specular_reflectance.eval(si, active)
+ return spec * f * g / (4 * si.wi.z)
+
+ # This function returns the FG term for clearcoat layer
+ def eval_pdf(self, ctx: mi.BSDFContext, si: mi.SurfaceInteraction3f, wo: mi.Vector3f, active: bool = True):
+ clearcoat_alpha = self.clearcoat_alpha.eval_1(si, active)
+ alpha = dr.clamp(clearcoat_alpha, 0.01, 0.99)
+
+ # Find half vector
+ m = dr.normalize(si.wi + wo)
+
+ # Find fresnel term
+ f = mi.Color3f(0)
+ for i in range(3):
+ f[i] = mi.fresnel_conductor(dr.dot(si.wi, m), mi.Complex2f(1.5, 1))
+
+ # Find shadowing
+ ndf = mi.MicrofacetDistribution(mi.MicrofacetType.Beckmann, alpha, sample_visible=False)
+ g = ndf.G(si.wi, wo, m)
+
+ return f * g / (4 * si.wi.z), mi.Float(1)
+
+ def sample(self, ctx: mi.BSDFContext, si: mi.SurfaceInteraction3f, sample1: float, sample2: mi.Point2f, active: bool = True):
+ bs = dr.zeros(mi.BSDFSample3f)
+ return bs, mi.Color3f(1)

plugins/ratio_estimator.py

@@ -0,0 +1,94 @@
+import mitsuba as mi
+import drjit as dr
+import numpy as np
+import time
+from utils.render import generate_rays
+
+def ratio_estimator(scene: mi.Scene, spp: int=1024):
+ start_time = time.time()
+ rays, _, pos = generate_rays(scene, spp, random_offset=True)
+ si: mi.SurfaceInteraction3f = scene.ray_intersect(rays)
+
+ ctx = mi.BSDFContext() 
+
+ active = si.is_valid()
+
+ rng = mi.PCG32(dr.width(active))
+
+ unshadowed = dr.select(active, mi.Color3f(0), 1)
+ shadowed = dr.select(active, mi.Color3f(0), 1)
+
+ bsdf: mi.BSDF = si.bsdf(rays)
+
+ # Emitter Sampling
+ ds, emitter_val = scene.sample_emitter_direction(si, mi.Point2f(rng.next_float32(active), rng.next_float32(active)), False, active)
+ active_e = active & dr.neq(ds.pdf, 0.0)
+ wo = si.to_local(ds.d)
+ bsdf_val, bsdf_pdf = bsdf.eval_pdf(ctx, si, wo, active_e)
+ # Note that we are give 0 weight to delta lights
+ # This is because it doesn't make sense to use ratio estimator
+ # for delta lights.
+ mis = dr.select(ds.delta, 0, ds.pdf / (ds.pdf + bsdf_pdf))
+ occluded = scene.ray_test(si.spawn_ray(ds.d), active_e)
+ unshadowed += dr.select(active_e, emitter_val * bsdf_val * mis, 0)
+ shadowed += dr.select(active_e & ~occluded, bsdf_val * emitter_val * mis, 0)
+
+ # BSDF Sampling
+ # WARNING: Delta BSDFs are not handled
+ bs, bsdf_val = bsdf.sample(ctx, si, rng.next_float32(active), mi.Point2f(rng.next_float32(active), rng.next_float32(active)), active)
+ active_b = active & dr.any(dr.neq(bsdf_val, 0))
+ wo = si.to_world(bs.wo)
+ occluded = scene.ray_test(si.spawn_ray(wo), active_b)
+ envmap: mi.Emitter = scene.environment()
+ si_bsdf = dr.zeros(mi.SurfaceInteraction3f, dr.width(active_b))
+ si_bsdf.wi = -wo
+ emitter_val = envmap.eval(si_bsdf, active_b)
+ emitter_pdf = scene.pdf_emitter_direction(si, mi.DirectionSample3f(scene, si_bsdf, si), active_b)
+ mis = bs.pdf / (bs.pdf + emitter_pdf)
+ unshadowed += dr.select(active_b, bsdf_val * emitter_val * mis, 0)
+ shadowed += dr.select(~occluded & active_b, bsdf_val * emitter_val * mis, 0)
+
+ film: mi.Film = scene.sensors()[0].film()
+
+ # Develop the film for unshadowed
+ result = [unshadowed.x, unshadowed.y, unshadowed.z, mi.Float(1)]
+ film.clear()
+ # Image block
+ block = film.create_block()
+ # Offset is the currect location of the block
+ # In case of GPU, the block covers the entire image, hence offset is 0
+ block.set_offset(film.crop_offset())
+
+ ################################
+ # Save image
+ ################################
+ block.put(pos, result)
+ film.put_block(block)
+ unshadowed = film.develop().numpy()
+
+ # Develop the film for unshadowed
+ result = [shadowed.x, shadowed.y, shadowed.z, mi.Float(1)]
+ film.clear()
+ # Image block
+ block = film.create_block()
+ # Offset is the currect location of the block
+ # In case of GPU, the block covers the entire image, hence offset is 0
+ block.set_offset(film.crop_offset())
+
+ ################################
+ # Save image
+ ################################
+ block.put(pos, result)
+ film.put_block(block)
+ shadowed = film.develop().numpy()
+
+ if np.any(np.isnan(shadowed)):
+ count = np.count_nonzero(np.isnan(shadowed).any(-1))
+ mi.Log(mi.LogLevel.Warn, f"{count} NaN found in 'shadowed'")
+ if np.any(np.isnan(unshadowed)):
+ count = np.count_nonzero(np.isnan(unshadowed).any(-1))
+ mi.Log(mi.LogLevel.Warn, f"{count} NaN found in 'unshadowed'")
+
+ end_time = time.time()
+
+ return shadowed, unshadowed, end_time - start_time