From e45f1b33767c09a5c7c824c9f7400eae854ea4ac Mon Sep 17 00:00:00 2001
From: archibate <1931127624@qq.com>
Date: Sun, 11 Oct 2020 10:59:44 +0800
Subject: [PATCH] enhance t3.Accumate with denosing

---
 README.md                  |  4 ++-
 examples/path_tracing.py   | 17 +++++-----
 examples/shadow_mapping.py |  8 ++---
 taichi_three/geometry.py   |  3 +-
 taichi_three/light.py      |  8 ++---
 taichi_three/raycast.py    | 63 +++++++++++++++++++++++++++++++-------
 taichi_three/shading.py    | 40 +++++++++++++++---------
 7 files changed, 99 insertions(+), 44 deletions(-)

diff --git a/README.md b/README.md
index 400e0de..44ec94c 100644
--- a/README.md
+++ b/README.md
@@ -74,7 +74,8 @@ Notable changes:
 * Framebuffered texcoor - get model surface coordinate by mouse, see `examples/screen_to_texcoor.py`.
 * Separate camera control logic from `t3.Camera` to `t3.CameraCtl`.
 * Support `t3.AmbientLight` and ambient occulsion.
-* Set up the path tracing scheme.
+* Set up the basis of path tracing scheme.
+* Shadow mapping is broken.
 
 Minor fixes:
 * Fix an artifect in perspective mode due to texture coordinate interpolation.
@@ -94,6 +95,7 @@ Minor fixes:
 * Fix color artifects on edges due to interpolation.
 * Fix shadow artifects on 90-deg faces.
 * Make camera buffer update less ad-hoc.
+* Support scaling in ORTHO mode.
 
 Major steps:
 * Standardize affine system - L2W, W2C, C2D.
diff --git a/examples/path_tracing.py b/examples/path_tracing.py
index dc255fc..8fc924a 100644
--- a/examples/path_tracing.py
+++ b/examples/path_tracing.py
@@ -8,15 +8,18 @@
 scene = t3.Scene()
 obj = t3.readobj('assets/cube.obj', scale=0.8)
 model = t3.Model.from_obj(obj)
-model.add_uniform('color', 1.0)
+model.shading_type = t3.IdealRT
+model.add_uniform('diffuse', 1.0)
 model.add_uniform('emission', 0.0)
 scene.add_model(model)
 light = t3.Model.from_obj(obj)
-light.add_uniform('color', 0.0)
-light.add_uniform('emission', 64.0)
+light.shading_type = t3.IdealRT
+light.add_uniform('diffuse', 0.0)
+light.add_uniform('emission', 1.0)
+light.add_uniform('emission_color', 16.0)
 scene.add_model(light)
 camera = t3.RTCamera(res=res)
-camera.ctl = t3.CameraCtl(pos=[1, 0, -1])
+camera.ctl = t3.CameraCtl(pos=[0.4, 0, -3.7])
 scene.add_camera(camera)
 accumator = t3.Accumator(camera.res)
 
@@ -27,10 +30,6 @@
     gui.running = not gui.is_pressed(ti.GUI.ESCAPE)
     if camera.from_mouse(gui):
         accumator.reset()
-    camera.loadrays()
-    for s in range(2):
-        camera.steprays()
-    camera.applyrays()
-    accumator.accumate(camera.img)
+    accumator.render(camera, 2)
     gui.set_image(accumator.buf)
     gui.show()
diff --git a/examples/shadow_mapping.py b/examples/shadow_mapping.py
index e902373..d605ecd 100644
--- a/examples/shadow_mapping.py
+++ b/examples/shadow_mapping.py
@@ -19,8 +19,8 @@
 scene.add_light(light)
 
 gui = ti.GUI('Model', camera.res)
-#gui2 = ti.GUI('Depth map', light.shadow.res)
-#gui2.fps_limit = None
+gui2 = ti.GUI('Depth map', light.shadow.res)
+gui2.fps_limit = None
 while gui.running:
     gui.get_event(None)
     gui.running = not gui.is_pressed(ti.GUI.ESCAPE)
@@ -29,6 +29,6 @@
     scene.render_shadows()
     scene.render()
     gui.set_image(camera.img)
-    #gui2.set_image(light.shadow.fb['idepth'])
+    gui2.set_image(light.shadow.img)
     gui.show()
-    #gui2.show()
+    gui2.show()
diff --git a/taichi_three/geometry.py b/taichi_three/geometry.py
index fbe21ae..c411c7c 100644
--- a/taichi_three/geometry.py
+++ b/taichi_three/geometry.py
@@ -154,7 +154,8 @@ def render_triangle(model, camera, face):
                     continue
 
                 clr = [a * w_A + b * w_B + c * w_C for a, b, c in zip(clra, clrb, clrc)]
-                camera.fb.update(X, model.pixel_shader(*clr))
+                #camera.fb.update(X, model.pixel_shader(*clr))
+                camera.img[X].fill(1)
 
 
 @ti.func
diff --git a/taichi_three/light.py b/taichi_three/light.py
index 9d816cc..d81447e 100644
--- a/taichi_three/light.py
+++ b/taichi_three/light.py
@@ -69,11 +69,11 @@ def get_dir(self, pos):
     def set_view(self, camera):
         self.viewdir[None] = (camera.L2W[None].inverse() @ ts.vec4(self.dir[None], 0)).xyz
 
-    def make_shadow_camera(self, dis=10, fov=60, **kwargs):
-        shadow = Camera()
+    def make_shadow_camera(self, res=(512, 512), dis=10, fov=60, **kwargs):
+        shadow = Camera(res=res)
         shadow.fov = math.radians(fov)
-        shadow.ctl = CameraCtl(pos=(self.dir[None].value * dis).entries, **kwargs)
-        @ti.materialize_callback
+        #shadow.ctl = CameraCtl(pos=(-self.dir[None].value * dis).entries, **kwargs)
+        #@ti.materialize_callback
         def init_camera():
             shadow.ctl.apply(shadow)
         shadow.type = shadow.ORTHO
diff --git a/taichi_three/raycast.py b/taichi_three/raycast.py
index 3f88a87..24a7626 100644
--- a/taichi_three/raycast.py
+++ b/taichi_three/raycast.py
@@ -14,8 +14,8 @@ def __init__(self, shape=(512, 512)):
     @ti.kernel
     def accumate(self, src: ti.template()):
         self.count[None] += 1
+        alpha = max(1 / 256, 1 / self.count[None])
         for I in ti.grouped(self.buf):
-            alpha = 1 / self.count[None]
             self.buf[I] = self.buf[I] * (1 - alpha) + src[I] * alpha
 
     @ti.kernel
@@ -24,6 +24,26 @@ def reset(self):
         for I in ti.grouped(self.buf):
             self.buf[I] *= 0
 
+    @ti.kernel
+    def denoise(self, alpha: ti.template(), throttle: ti.template()):
+        ti.static_print('denoise', alpha)
+        for I in ti.grouped(self.buf):
+            center = ts.clamp(self.buf[I])
+            around = ts.clamp((self.buf[I + ts.D.x_] + self.buf[I + ts.D.X_] + self.buf[I + ts.D._x] + self.buf[I + ts.D._X]) / 4)
+            diff = (center - around).norm_sqr()
+            if diff < throttle**2:
+                self.buf[I] = center * (1 - alpha) + around * alpha
+
+    def render(self, camera, depth, baseres=3, regrate=128):
+        rate = max(0, baseres - self.count[None] // regrate)
+        region = camera.res[0] // 2**rate, camera.res[1] // 2**rate
+        camera.loadrays((0, 0), region, 2**rate)
+        for step in range(depth):
+            camera.steprays()
+        camera.applyrays()
+        self.accumate(camera.img)
+        self.denoise(0.5 * (rate / baseres), 0.5)
+
 
 class RTCamera(Camera):
     def __init__(self, res=None):
@@ -35,9 +55,13 @@ def __init__(self, res=None):
         self.rc = ti.Vector.field(3, float, self.N)
         self.rI = ti.Vector.field(2, int, self.N)
 
-    @ti.kernel
     def steprays(self):
-        for i in self.ro:
+        nrays = self.region[0] * self.region[1]
+        self._steprays(nrays)
+
+    @ti.kernel
+    def _steprays(self, nrays: ti.template()):
+        for i in range(nrays):
             hit = 1e6
             orig, dir = self.ro[i], self.rd[i]
             if self.rd[i].norm_sqr() >= 1e-3:
@@ -49,22 +73,39 @@ def steprays(self):
                 self.ro[i], self.rd[i] = orig, dir
                 self.rc[i] *= clr
 
+    def loadrays(self, topleft=None, region=None, skipstep=None):
+        self.topleft = topleft or (0, 0)
+        self.region = region or self.res
+        self.skipstep = skipstep or 1
+        self._loadrays(self.topleft, self.region, self.skipstep)
+
     @ti.kernel
-    def loadrays(self):
-        self.fb.clear_buffer()
-        for I in ti.grouped(ti.ndrange(*self.res)):
-            i = I.dot(ts.vec(1, self.res[0]))
+    def _loadrays(self, topleft: ti.template(), region: ti.template(), skipstep: ti.template()):
+        ti.static_print('loadrays:', topleft, region, skipstep)
+        for II in ti.grouped(ti.ndrange(*region)):
+            I = II * skipstep + topleft
+            for J in ti.static(ti.grouped(ti.ndrange(skipstep, skipstep))):
+                self.img[I + J] *= 0
+        for II in ti.grouped(ti.ndrange(*region)):
+            i = II.dot(ts.vec(1, region[0]))
+            I = II * skipstep + topleft + skipstep / 2
             coor = ts.vec2((I.x - self.cx) / self.fx, (I.y - self.cy) / self.fy)
             orig, dir = self.generate(coor)
             self.ro[i] = orig
             self.rd[i] = dir
             self.rc[i] = ts.vec3(1.0)
-            self.rI[i] = I
+            self.rI[i] = II
 
-    @ti.kernel
     def applyrays(self):
-        for i in self.ro:
-            self.img[self.rI[i]] = self.rc[i]
+        nrays = self.region[0] * self.region[1]
+        self._applyrays(nrays, self.topleft, self.skipstep)
+
+    @ti.kernel
+    def _applyrays(self, nrays: ti.template(), topleft: ti.template(), skipstep: ti.template()):
+        for i in range(nrays):
+            I = self.rI[i] * skipstep + topleft
+            for J in ti.static(ti.grouped(ti.ndrange(skipstep, skipstep))):
+                self.img[I + J] = self.rc[i]
 
     @ti.func
     def generate(self, coor):
diff --git a/taichi_three/shading.py b/taichi_three/shading.py
index 9c013a5..5999c6e 100644
--- a/taichi_three/shading.py
+++ b/taichi_three/shading.py
@@ -120,6 +120,32 @@ def brdf(self, normal, lightdir, viewdir):
         return strength
 
 
+class IdealRT(Shading):
+    emission = 0.0
+    diffuse = 1.0
+    specular = 0.0
+    emission_color = 1.0
+    diffuse_color = 1.0
+    specular_color = 1.0
+    parameters = ['emission', 'diffuse', 'specular', 'emission_color', 'diffuse_color', 'specular_color']
+
+    @ti.func
+    def radiance(self, pos, indir, normal):
+        outdir = ts.vec3(0.0)
+        clr = ts.vec3(0.0)
+        if ti.random() < self.emission:
+            clr = ts.vec3(self.emission_color)
+        elif ti.random() < self.specular:
+            clr = ts.vec3(self.specular_color)
+            outdir = ts.reflect(indir, normal)
+        elif ti.random() < self.diffuse:
+            clr = ts.vec3(self.diffuse_color)
+            outdir = ts.randUnit3D()
+            if outdir.dot(normal) < 0:
+                outdir = -outdir
+        return pos, outdir, clr
+
+
 # https://zhuanlan.zhihu.com/p/37639418
 class CookTorrance(Shading):
     color = 1.0
@@ -136,20 +162,6 @@ class CookTorrance(Shading):
     def __init__(self, **kwargs):
         self.__dict__.update(kwargs)
 
-    @ti.func
-    def radiance(self, pos, indir, normal):
-        outdir = ts.vec3(0.0)
-        clr = ts.vec3(0.0)
-        if ti.random() < self.emission:
-            clr = ts.vec3(1.0)
-        elif ti.random() < self.color:
-            clr = ts.vec3(1.0)
-            outdir = ts.reflect(indir, normal)
-            outdir = ts.randUnit3D()
-            if outdir.dot(normal) < 0:
-                outdir = -outdir
-        return pos, outdir, clr
-
     @ti.func
     def ischlick(self, cost):
         k = (self.roughness + 1)**2 / 8