# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

# SPDX-License-Identifier: Apache-2.0

#############################################################################

# Example Differentiable Ray Caster

#

# Shows how to use the built-in wp.Mesh data structure and wp.mesh_query_ray()

# function to implement a basic differentiable ray caster

#

##############################################################################

import math

import os

import numpy as np

from pxr import Usd, UsdGeom

import warp as wp

import warp.examples

from warp.optim import SGD

class RenderMode:

"""Rendering modes

grayscale: Lambertian shading from multiple directional lights

texture: 2D texture map

normal_map: mesh normal computed from interpolated vertex normals

"""

grayscale = 0

texture = 1

normal_map = 2

@wp.struct

class RenderMesh:

"""Mesh to be ray casted.

Assumes a triangle mesh as input.

Per-vertex normals are computed with compute_vertex_normals()

"""

id: wp.uint64

vertices: wp.array[wp.vec3]

indices: wp.array[int]

tex_coords: wp.array[wp.vec2]

tex_indices: wp.array[int]

vertex_normals: wp.array[wp.vec3]

pos: wp.array[wp.vec3]

rot: wp.array[wp.quat]

@wp.struct

class Camera:

"""Basic camera for ray casting"""

horizontal: float

vertical: float

aspect: float

e: float

tan: float

pos: wp.vec3

rot: wp.quat

@wp.struct

class DirectionalLights:

"""Stores arrays of directional light directions and intensities."""

dirs: wp.array[wp.vec3]

intensities: wp.array[float]

num_lights: int

@wp.kernel

def vertex_normal_sum_kernel(verts: wp.array[wp.vec3], indices: wp.array[int], normal_sums: wp.array[wp.vec3]):

tid = wp.tid()

i = indices[tid * 3]

j = indices[tid * 3 + 1]

k = indices[tid * 3 + 2]

a = verts[i]

b = verts[j]

c = verts[k]

ab = b - a

ac = c - a

area_normal = wp.cross(ab, ac)

wp.atomic_add(normal_sums, i, area_normal)

wp.atomic_add(normal_sums, j, area_normal)

wp.atomic_add(normal_sums, k, area_normal)

@wp.kernel

def normalize_kernel(

normal_sums: wp.array[wp.vec3],

vertex_normals: wp.array[wp.vec3],

):

tid = wp.tid()

vertex_normals[tid] = wp.normalize(normal_sums[tid])

@wp.func

def texture_interpolation(tex_interp: wp.vec2, texture: wp.array2d[wp.vec3]):

tex_width = texture.shape[1]

tex_height = texture.shape[0]

tex = wp.vec2(tex_interp[0] * float(tex_width - 1), (1.0 - tex_interp[1]) * float(tex_height - 1))

x0 = int(tex[0])

x1 = x0 + 1

alpha_x = tex[0] - float(x0)

y0 = int(tex[1])

y1 = y0 + 1

alpha_y = tex[1] - float(y0)

c00 = texture[y0, x0]

c10 = texture[y0, x1]

c01 = texture[y1, x0]

c11 = texture[y1, x1]

lower = (1.0 - alpha_x) * c00 + alpha_x * c10

upper = (1.0 - alpha_x) * c01 + alpha_x * c11

color = (1.0 - alpha_y) * lower + alpha_y * upper

return color

@wp.kernel

def draw_kernel(

mesh: RenderMesh,

camera: Camera,

texture: wp.array2d[wp.vec3],

rays_width: int,

rays_height: int,

rays: wp.array[wp.vec3],

lights: DirectionalLights,

mode: int,

):

tid = wp.tid()

x = tid % rays_width

y = rays_height - tid // rays_width

sx = 2.0 * float(x) / float(rays_width) - 1.0

sy = 2.0 * float(y) / float(rays_height) - 1.0

# compute view ray in world space

ro_world = camera.pos

rd_world = wp.normalize(wp.quat_rotate(camera.rot, wp.vec3(sx * camera.tan * camera.aspect, sy * camera.tan, -1.0)))

# compute view ray in mesh space

inv = wp.transform_inverse(wp.transform(mesh.pos[0], mesh.rot[0]))

ro = wp.transform_point(inv, ro_world)

rd = wp.transform_vector(inv, rd_world)

color = wp.vec3(0.0, 0.0, 0.0)

query = wp.mesh_query_ray(mesh.id, ro, rd, 1.0e6)

if query.result:

i = mesh.indices[query.face * 3]

j = mesh.indices[query.face * 3 + 1]

k = mesh.indices[query.face * 3 + 2]

a_n = mesh.vertex_normals[i]

b_n = mesh.vertex_normals[j]

c_n = mesh.vertex_normals[k]

# vertex normal interpolation

normal = query.u * a_n + query.v * b_n + (1.0 - query.u - query.v) * c_n

if mode == 0 or mode == 1:

if mode == 0: # grayscale

color = wp.vec3(1.0)

elif mode == 1: # texture interpolation

tex_a = mesh.tex_coords[mesh.tex_indices[query.face * 3]]

tex_b = mesh.tex_coords[mesh.tex_indices[query.face * 3 + 1]]

tex_c = mesh.tex_coords[mesh.tex_indices[query.face * 3 + 2]]

tex = query.u * tex_a + query.v * tex_b + (1.0 - query.u - query.v) * tex_c

color = texture_interpolation(tex, texture)

# lambertian directional lighting

lambert = float(0.0)

for i in range(lights.num_lights):

dir = wp.transform_vector(inv, lights.dirs[i])

val = lights.intensities[i] * wp.dot(normal, dir)

if val < 0.0:

val = 0.0

lambert = lambert + val

color = lambert * color

elif mode == 2: # normal map

color = normal * 0.5 + wp.vec3(0.5, 0.5, 0.5)

if color[0] > 1.0:

color = wp.vec3(1.0, color[1], color[2])

if color[1] > 1.0:

color = wp.vec3(color[0], 1.0, color[2])

if color[2] > 1.0:

color = wp.vec3(color[0], color[1], 1.0)

rays[tid] = color

@wp.kernel

def downsample_kernel(rays: wp.array[wp.vec3], pixels: wp.array[wp.vec3], rays_width: int, num_samples: int):

tid = wp.tid()

pixels_width = rays_width / num_samples

px = tid % pixels_width

py = tid // pixels_width

start_idx = py * num_samples * rays_width + px * num_samples

color = wp.vec3(0.0, 0.0, 0.0)

for i in range(0, num_samples):

for j in range(0, num_samples):

ray = rays[start_idx + i * rays_width + j]

color = wp.vec3(color[0] + ray[0], color[1] + ray[1], color[2] + ray[2])

num_samples_sq = float(num_samples * num_samples)

color = wp.vec3(color[0] / num_samples_sq, color[1] / num_samples_sq, color[2] / num_samples_sq)

pixels[tid] = color

@wp.kernel

def loss_kernel(pixels: wp.array[wp.vec3], target_pixels: wp.array[wp.vec3], loss: wp.array[float]):

tid = wp.tid()

pixel = pixels[tid]

target_pixel = target_pixels[tid]

diff = target_pixel - pixel

# pseudo Huber loss

delta = 1.0

x = delta * delta * (wp.sqrt(1.0 + (diff[0] / delta) * (diff[0] / delta)) - 1.0)

y = delta * delta * (wp.sqrt(1.0 + (diff[1] / delta) * (diff[1] / delta)) - 1.0)

z = delta * delta * (wp.sqrt(1.0 + (diff[2] / delta) * (diff[2] / delta)) - 1.0)

sum = x + y + z

wp.atomic_add(loss, 0, sum)

@wp.kernel

def normalize(x: wp.array[wp.quat]):

tid = wp.tid()

x[tid] = wp.normalize(x[tid])

class Example:

"""

Non-differentiable variables:

camera.horizontal: camera horizontal aperture size

camera.vertical: camera vertical aperture size

camera.aspect: camera aspect ratio

camera.e: focal length

camera.pos: camera displacement

camera.rot: camera rotation (quaternion)

pix_width: final image width in pixels

pix_height: final image height in pixels

num_samples: anti-aliasing. calculated as pow(2, num_samples)

directional_lights: characterized by intensity (scalar) and direction (vec3)

render_mesh.indices: mesh vertex indices

render_mesh.tex_indices: texture indices

Differentiable variables:

render_mesh.pos: parent transform displacement

render_mesh.quat: parent transform rotation (quaternion)

render_mesh.vertices: mesh vertex positions

render_mesh.vertex_normals: mesh vertex normals

render_mesh.tex_coords: 2D texture coordinates

"""

def __init__(self, height=1024, train_iters=150, rot_array=None):

cam_pos = wp.vec3(0.0, 0.75, 7.0)

cam_rot = wp.quat(0.0, 0.0, 0.0, 1.0)

horizontal_aperture = 36.0

vertical_aperture = 20.25

aspect = horizontal_aperture / vertical_aperture

focal_length = 50.0

self.height = height

self.width = int(aspect * self.height)

self.num_pixels = self.width * self.height

if rot_array is None:

rot_array = [0.0, 0.0, 0.0, 1.0]

asset_stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "bunny.usd"))

mesh_geom = UsdGeom.Mesh(asset_stage.GetPrimAtPath("/root/bunny"))

points = np.array(mesh_geom.GetPointsAttr().Get())

indices = np.array(mesh_geom.GetFaceVertexIndicesAttr().Get())

num_points = points.shape[0]

num_faces = int(indices.shape[0] / 3)

# manufacture texture coordinates + indices for this asset

distance = np.linalg.norm(points, axis=1)

radius = np.max(distance)

distance = distance / radius

tex_coords = np.stack((distance, distance), axis=1)

tex_indices = indices

# manufacture texture for this asset

x = np.arange(256.0)

xx, yy = np.meshgrid(x, x)

zz = np.zeros_like(xx)

texture_host = np.stack((xx, yy, zz), axis=2) / 255.0

# set anti-aliasing

self.num_samples = 1

# set render mode

self.render_mode = RenderMode.texture

# set training iterations

self.train_rate = 5.00e-8

self.momentum = 0.5

self.dampening = 0.1

self.weight_decay = 0.0

self.train_iters = train_iters

self.period = 10 # Training iterations between render() calls

self.iter = 0

# storage for training animation

self.images = np.zeros((self.height, self.width, 3, max(int(self.train_iters / self.period), 1)))

self.image_counter = 0

# construct RenderMesh

self.render_mesh = RenderMesh()

self.mesh = wp.Mesh(

points=wp.array(points, dtype=wp.vec3, requires_grad=True),

indices=wp.array(indices, dtype=int),

)

self.render_mesh.id = self.mesh.id

self.render_mesh.vertices = self.mesh.points

self.render_mesh.indices = self.mesh.indices

self.render_mesh.tex_coords = wp.array(tex_coords, dtype=wp.vec2, requires_grad=True)

self.render_mesh.tex_indices = wp.array(tex_indices, dtype=int)

self.normal_sums = wp.zeros(num_points, dtype=wp.vec3, requires_grad=True)

self.render_mesh.vertex_normals = wp.zeros(num_points, dtype=wp.vec3, requires_grad=True)

self.render_mesh.pos = wp.zeros(1, dtype=wp.vec3, requires_grad=True)

self.render_mesh.rot = wp.array(np.array(rot_array), dtype=wp.quat, requires_grad=True)

# compute vertex normals

wp.launch(

kernel=vertex_normal_sum_kernel,

dim=num_faces,

inputs=[self.render_mesh.vertices, self.render_mesh.indices, self.normal_sums],

)

wp.launch(

kernel=normalize_kernel,

dim=num_points,

inputs=[self.normal_sums, self.render_mesh.vertex_normals],

)

# construct camera

self.camera = Camera()

self.camera.horizontal = horizontal_aperture

self.camera.vertical = vertical_aperture

self.camera.aspect = aspect

self.camera.e = focal_length

self.camera.tan = vertical_aperture / (2.0 * focal_length)

self.camera.pos = cam_pos

self.camera.rot = cam_rot

# construct texture

self.texture = wp.array2d(texture_host, dtype=wp.vec3, requires_grad=True)

# construct lights

self.lights = DirectionalLights()

self.lights.dirs = wp.array(np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), dtype=wp.vec3, requires_grad=True)

self.lights.intensities = wp.array(np.array([2.0, 0.2]), dtype=float, requires_grad=True)

self.lights.num_lights = 2

# construct rays

self.rays_width = self.width * pow(2, self.num_samples)

self.rays_height = self.height * pow(2, self.num_samples)

self.num_rays = self.rays_width * self.rays_height

self.rays = wp.zeros(self.num_rays, dtype=wp.vec3, requires_grad=True)

# construct pixels

self.pixels = wp.zeros(self.num_pixels, dtype=wp.vec3, requires_grad=True)

self.target_pixels = wp.zeros(self.num_pixels, dtype=wp.vec3)

# loss array

self.loss = wp.zeros(1, dtype=float, requires_grad=True)

# capture graph

self.use_cuda_graph = wp.get_device().is_cuda

if self.use_cuda_graph:

with wp.ScopedCapture() as capture:

self.tape = wp.Tape()

with self.tape:

self.forward()

self.tape.backward(self.loss)

self.graph = capture.graph

self.optimizer = SGD(

[self.render_mesh.rot],

self.train_rate,

momentum=self.momentum,

dampening=self.dampening,

weight_decay=self.weight_decay,

)

def ray_cast(self):

# raycast

wp.launch(

kernel=draw_kernel,

dim=self.num_rays,

inputs=[

self.render_mesh,

self.camera,

self.texture,

self.rays_width,

self.rays_height,

self.rays,

self.lights,

self.render_mode,

],

)

# downsample

wp.launch(

kernel=downsample_kernel,

dim=self.num_pixels,

inputs=[self.rays, self.pixels, self.rays_width, pow(2, self.num_samples)],

)

def forward(self):

self.ray_cast()

# compute pixel loss

wp.launch(loss_kernel, dim=self.num_pixels, inputs=[self.pixels, self.target_pixels, self.loss])

def step(self):

with wp.ScopedTimer("step"):

if self.use_cuda_graph:

wp.capture_launch(self.graph)

else:

self.tape = wp.Tape()

with self.tape:

self.forward()

self.tape.backward(self.loss)

rot_grad = self.tape.gradients[self.render_mesh.rot]

self.optimizer.step([rot_grad])

wp.launch(normalize, dim=1, inputs=[self.render_mesh.rot])

if self.iter % self.period == 0:

print(f"Iter: {self.iter} Loss: {self.loss}")

self.tape.zero()

self.loss.zero_()

self.iter = self.iter + 1

def render(self):

with wp.ScopedTimer("render"):

self.images[:, :, :, self.image_counter] = self.get_image()

self.image_counter += 1

def get_image(self):

return self.pixels.numpy().reshape((self.height, self.width, 3))

def get_animation(self):

fig, ax = plt.subplots()

plt.axis("off")

plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)

plt.margins(0, 0)

frames = []

for i in range(self.images.shape[3]):

frame = ax.imshow(self.images[:, :, :, i], animated=True)

frames.append([frame])

ani = animation.ArtistAnimation(fig, frames, interval=50, blit=True, repeat_delay=1000)

return ani

if __name__ == "__main__":

import argparse

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")

parser.add_argument("--train-iters", type=int, default=150, help="Total number of training iterations.")

parser.add_argument("--height", type=int, default=1024, help="Height of rendered image in pixels.")

parser.add_argument(

"--headless",

action="store_true",

help="Run in headless mode, suppressing the opening of any graphical windows.",

)

args = parser.parse_known_args()[0]

with wp.ScopedDevice(args.device):

reference_example = Example(height=args.height)

# render target rotation

reference_example.ray_cast()

# offset mesh rotation

example = Example(

train_iters=args.train_iters,

height=args.height,

rot_array=[

0.0,

(math.sqrt(3) - 1) / (2.0 * math.sqrt(2.0)),

0.0,

(math.sqrt(3) + 1) / (2.0 * math.sqrt(2.0)),

],

)

wp.copy(example.target_pixels, reference_example.pixels)

# recover target rotation

for i in range(example.train_iters):

example.step()

if i % example.period == 0:

example.render()

if not args.headless:

import matplotlib.animation as animation

import matplotlib.image as img

import matplotlib.pyplot as plt

target_image = reference_example.get_image()

target_image_filename = "example_diffray_target_image.png"

img.imsave(target_image_filename, target_image)

print(f"Saved the target image at `{target_image_filename}`")

final_image = example.get_image()

final_image_filename = "example_diffray_final_image.png"

img.imsave(final_image_filename, final_image)

print(f"Saved the final image at `{final_image_filename}`")

anim = example.get_animation()

anim_filename = "example_diffray_animation.gif"

anim.save(anim_filename, dpi=300, writer=animation.PillowWriter(fps=5))

print(f"Saved the animation at `{anim_filename}`")