From 5ab1f0870aa18fe10500ddad4ea4d8653919d8c7 Mon Sep 17 00:00:00 2001 From: oarriaga Date: Wed, 29 Jul 2026 16:11:13 +0200 Subject: [PATCH 1/3] Add realistic primitives example --- .../realistic_primitives.py | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 examples/differentiable_shape_rendering/realistic_primitives.py diff --git a/examples/differentiable_shape_rendering/realistic_primitives.py b/examples/differentiable_shape_rendering/realistic_primitives.py new file mode 100644 index 000000000..42d1223fe --- /dev/null +++ b/examples/differentiable_shape_rendering/realistic_primitives.py @@ -0,0 +1,112 @@ +import jax +import jax.numpy as jp +import paz +import paz.graphics.renderer as paz_renderer + +soft_occlusion = paz_renderer.compute_soft_occlusion +paz_renderer.compute_soft_occlusion = paz.partial(soft_occlusion, slope=1.0) + +BLUE = jp.array([0.324, 0.692, 0.863]) # MyBlue +GREEN = jp.array([154 / 255, 213 / 255, 135 / 255]) # YlGnI +YELLOW = jp.array([1.0, 0.65, 0.0]) # MyYellow +RED = jp.array([171 / 255, 62 / 255, 66 / 255]) # cherry + +H, W = 280, 960 +Y_FOV = 0.37 +SCALE = 4 +BOUNCES = 4 +FLOOR_REFLECTIVE = 0.3 +SHAPE_REFLECTIVE = 0.12 + + +def main(): + path = paz.logger.make_directory("realistic_primitives") + scene = build_scene() + light = build_area_light() + render = build_render(light, shadows=True, num_bounces=BOUNCES, scale=SCALE) + image, _ = render(scene=scene) + save_image(path, "realistic_primitives.png", image, SCALE) + + +def build_render(lights, shadows, num_bounces, scale): + shape = (H * scale, W * scale) + args = shape, Y_FOV, CAMERA_POSE + kwargs = dict(lights=lights, tiles=(1, 1), chunk_size=2**10) + kwargs.update(shadows=shadows, num_bounces=num_bounces, mask=None) + return jax.jit(paz.partial(paz.graphics.render, *args, **kwargs)) + + +def save_image(path, name, image, scale): + image = paz.image.denormalize(image) + if scale > 1: + image = paz.image.resize(image, (H, W), "bilinear") + paz.image.write(f"{path}/{name}", image) + + +def build_scene(): + floor = build_floor() + sphere = build_sphere(-4.5) + cylinder = build_cylinder(-1.5) + cone = build_cone(1.5) + cube = build_cube(4.5) + return paz.graphics.Scene([floor, sphere, cylinder, cone, cube]) + + +def build_sphere(x): + return paz.graphics.Sphere(rest_pose(x, 1.15), shape_material(GREEN)) + + +def build_cylinder(x): + return paz.graphics.Cylinder(rest_pose(x, 1.0), shape_material(BLUE)) + + +def build_cone(x): + scale = paz.SE3.scaling(jp.array([1.25, 1.12, 1.25])) + pose = paz.SE3.translation(jp.array([x, 1.12, 0.0])) @ scale + return paz.graphics.Cone(pose, shape_material(YELLOW)) + + +def build_cube(x): + return paz.graphics.Cube(rest_pose(x, 1.0), shape_material(RED)) + + +def build_floor(): + return paz.graphics.Plane(material=floor_material()) + + +def rest_pose(x, size): + shift = paz.SE3.translation(jp.array([x, size, 0.0])) + return shift @ paz.SE3.scaling(jp.full(3, size)) + + +def shape_material(color): + args = color, 0.14, 0.75, 0.6, 200.0, SHAPE_REFLECTIVE + return paz.graphics.Material(*args) + + +def floor_material(): + args = jp.ones(3), 0.025, 0.67, 0.0, 100.0, FLOOR_REFLECTIVE + return paz.graphics.Material(*args) + + +def build_area_light(): + intensity = jp.array([1.5, 1.5, 1.5]) + corner = jp.array([4.75, 9.0, 4.75]) + edge1 = jp.array([2.5, 0.0, 0.0]) + edge2 = jp.array([0.0, 0.0, 2.5]) + args = intensity, corner, edge1, edge2, 10, 10, jax.random.key(0) + return paz.graphics.AreaLight(*args) + + +def build_camera_pose(): + position = jp.array([-1.2, 5.6, 8.6]) + target = jp.array([-0.54, 0.85, -0.04]) + up = jp.array([0.0, 1.0, 0.0]) + return paz.SE3.view_transform(position, target, up) + + +CAMERA_POSE = build_camera_pose() + + +if __name__ == "__main__": + main() From 6f31cf8c0c107b9dc947c6c6973b0f30a1c6b6a5 Mon Sep 17 00:00:00 2001 From: oarriaga Date: Wed, 29 Jul 2026 16:11:39 +0200 Subject: [PATCH 2/3] Add ignore of markdown files across the rpo --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 7dea0a634..dd694c0a7 100644 --- a/.gitignore +++ b/.gitignore @@ -7,7 +7,7 @@ !*.sh !requirements.txt !README.md -!*.md +# !*.md # render goldens: tracked so snapshot assertions compare instead of create !paz/graphics/snapshots/*.npy From 481f4f37a63d9b162a1d3018eb8c00cd3e79f5f3 Mon Sep 17 00:00:00 2001 From: oarriaga Date: Sun, 2 Aug 2026 16:00:48 +0200 Subject: [PATCH 3/3] refactor: clean up renderer.py internals Drop the two Args namedtuples and the `def f(*args)` signatures that unpacked positionally, in favour of explicit parameters. Single returns, nested defs lifted out, `transparancies` typo fixed, over-length lines wrapped. Also removes dead code found on the way: `rays` was threaded through four call levels unread, `compute_new_rays` took a `reflectance` it never used, the closest-hit argmin ran twice per bounce, and `same_shape` was computed twice per light. Adds a `Surfaces` namedtuple so points/normals/eyes stop travelling as a positional triple. No behaviour change. pytest paz/graphics/: 252 passed, 1 skipped. --- paz/graphics/renderer.py | 692 ++++++++++++++++------------------ paz/graphics/renderer_test.py | 34 +- 2 files changed, 328 insertions(+), 398 deletions(-) diff --git a/paz/graphics/renderer.py b/paz/graphics/renderer.py index 9c16bc326..d6dd9adc1 100644 --- a/paz/graphics/renderer.py +++ b/paz/graphics/renderer.py @@ -17,120 +17,96 @@ FACE_CHUNK_SIZE = 128 SHADOW_ORIGIN_EPSILON = 1e-5 SHADOW_SELF_HIT_EPSILON = 1e-5 -# BOUNCE_ORIGIN_EPSILON = 3e-3 BOUNCE_ORIGIN_EPSILON = 1e-2 -RENDER_NAMES = "shape y_FOV pose scene tiles chunk_size shadows " -RENDER_NAMES += "num_bounces face_chunk_size" + TRIANGLE_HIT_NAMES = "hit_mask depth points normals eyes albedo primitive" +TriangleHit = namedtuple("TriangleHit", TRIANGLE_HIT_NAMES.split()) + STATE_NAMES = "color depth hit_mask throughput active_mask " STATE_NAMES += "refractive_index rays" -SHADOW_COLOR_NAMES = "rays shapes lights indices mask shadow_mask " -SHADOW_COLOR_NAMES += "point normal points normals eyes" - -RenderArgs = namedtuple("RenderArgs", RENDER_NAMES.split()) -TriangleHit = namedtuple("TriangleHit", TRIANGLE_HIT_NAMES.split()) RenderState = namedtuple("RenderState", STATE_NAMES.split()) -ShadowColorArgs = namedtuple("ShadowColorArgs", SHADOW_COLOR_NAMES.split()) + +MATERIAL_NAMES = "reflectivities transparencies refractive_indices" +HitMaterial = namedtuple("HitMaterial", MATERIAL_NAMES.split()) + +Surfaces = namedtuple("Surfaces", ["points", "normals", "eyes"]) def render( - shape, - y_FOV, - pose, - scene, - mask, - lights, - tiles, - chunk_size, - shadows=False, - shadow_mask=None, - num_bounces=1, + shape, y_FOV, pose, scene, mask, lights, tiles, chunk_size, + shadows=False, shadow_mask=None, num_bounces=1, face_chunk_size=FACE_CHUNK_SIZE, ): - scene_args = scene, lights, mask, shadow_mask - compiled = paz.graphics.scene.compile(*scene_args) - args = shape, y_FOV, pose, compiled, tiles, chunk_size - args = RenderArgs(*args, shadows, num_bounces, face_chunk_size) - image, depth = scan_tiles(args, render_tile_step) - return assemble_image(args, image), assemble_depth(args, depth) + compiled = paz.graphics.scene.compile(scene, lights, mask, shadow_mask) + trace_args = compiled, shadows, num_bounces, face_chunk_size, chunk_size + trace = paz.lock(trace_chunks, *trace_args) + tile_step = paz.lock(render_tile_step, shape, y_FOV, pose, tiles, trace) + images, depths = scan_tiles(shape, tiles, tile_step) + image = assemble_tiles(shape, tiles, images) + depth = assemble_tiles(shape, tiles, depths)[..., 0] + return image, depth def render_masks( - shape, - y_FOV, - pose, - scene, - lights, - depth, - tiles, - chunk_size, - num_objects=None, - shadows=False, - shadow_mask=None, - num_bounces=1, + shape, y_FOV, pose, scene, lights, depth, tiles, chunk_size, + num_objects=None, shadows=False, shadow_mask=None, num_bounces=1, face_chunk_size=FACE_CHUNK_SIZE, ): if num_objects is None: num_objects = len(scene.nodes) min_depth, max_depth = depth - num_nodes = len(scene.nodes) masks = [] for object_arg in range(num_objects): - mask = jp.zeros((num_nodes,), dtype=bool).at[object_arg].set(True) + mask = build_object_mask(len(scene.nodes), object_arg) args = shape, y_FOV, pose, scene, mask, lights, tiles, chunk_size - render_args = args + (shadows, shadow_mask, num_bounces) - _, depth_image = render(*render_args, face_chunk_size) + args += shadows, shadow_mask, num_bounces, face_chunk_size + _, depth_image = render(*args) soft = paz.depth.to_soft_mask(depth_image, min_depth, max_depth) masks.append(jp.expand_dims(soft, axis=-1)) return jp.stack(masks) -def scan_tiles(args, render_step): - H, W = args.shape - H_tiles, W_tiles = args.tiles +def build_object_mask(num_nodes, object_arg): + return jp.zeros((num_nodes,), dtype=bool).at[object_arg].set(True) + + +def scan_tiles(shape, tiles, tile_step): + H, W = shape + H_tiles, W_tiles = tiles paz.graphics.mesh.assert_exact_tile_side(H, H_tiles) paz.graphics.mesh.assert_exact_tile_side(W, W_tiles) coordinates = paz.graphics.mesh.make_tile_coordinates(H_tiles, W_tiles) - render_step = paz.lock(render_step, args) - return jax.lax.scan(render_step, None, coordinates)[1] + return jax.lax.scan(tile_step, None, coordinates)[1] -def render_tile_step(carry, tile_arg, args): - H, W = args.shape - H_tiles, W_tiles = args.tiles - camera_to_world = jp.linalg.inv(args.pose) - tile_args = H, W, H_tiles, W_tiles, args.y_FOV, camera_to_world +def render_tile_step(carry, tile_arg, shape, y_FOV, pose, tiles, trace): + H, W = shape + H_tiles, W_tiles = tiles + camera_to_world = jp.linalg.inv(pose) + tile_args = H, W, H_tiles, W_tiles, y_FOV, camera_to_world rays = paz.graphics.mesh.build_tile_rays(*tile_args, tile_arg) + hit_mask, depth, color = trace(rays) tile_H, tile_W = H // H_tiles, W // W_tiles - trace_args = args.scene, args.shadows, args.num_bounces - trace_args = trace_args + (args.face_chunk_size,) - hit_mask, depth, color = trace_chunks(rays, trace_args, args.chunk_size) - post_args = hit_mask, depth, color, args.pose, rays, tile_H, tile_W + post_args = hit_mask, depth, color, pose, rays, tile_H, tile_W image, depth = postprocess(*post_args) return carry, (image, jp.expand_dims(depth, -1)) -def assemble_image(args, image): - H, W = args.shape - H_tiles, W_tiles = args.tiles - return paz.graphics.mesh.assemble(H, W, H_tiles, W_tiles, image) - - -def assemble_depth(args, depth): - H, W = args.shape - H_tiles, W_tiles = args.tiles - return paz.graphics.mesh.assemble(H, W, H_tiles, W_tiles, depth)[..., 0] +def assemble_tiles(shape, tiles, images): + H, W = shape + H_tiles, W_tiles = tiles + return paz.graphics.mesh.assemble(H, W, H_tiles, W_tiles, images) -def trace_chunks(rays, config, chunk_size): +def trace_chunks(rays, compiled, shadows, num_bounces, face_chunk, chunk_size): ray_chunks = split_ray_chunks(rays, chunk_size) - trace_step = paz.lock(trace_chunk_step, config) + step_args = compiled, shadows, num_bounces, face_chunk + trace_step = paz.lock(trace_chunk_step, *step_args) hit_mask, depth, color = jax.lax.scan(trace_step, None, ray_chunks)[1] return flatten_chunk_results(hit_mask, depth, color, len(rays[0])) -def trace_chunk_step(carry, rays, config): - compiled, shadows, num_bounces, face_chunk = config +def trace_chunk_step(carry, rays, compiled, shadows, num_bounces, face_chunk): args = rays, compiled, shadows, num_bounces, face_chunk return carry, trace_bounces(*args) @@ -147,10 +123,11 @@ def split_ray_chunks(rays, chunk_size): def pad_to_chunks(array, chunk_size): remainder = array.shape[0] % chunk_size if remainder == 0: - return array - pad_size = chunk_size - remainder - padding = jp.repeat(array[-1:], pad_size, axis=0) - return jp.concatenate([array, padding], axis=0) + padded = array + else: + padding = jp.repeat(array[-1:], chunk_size - remainder, axis=0) + padded = jp.concatenate([array, padding], axis=0) + return padded def flatten_chunk_results(hit_mask, depth, color, num_rays): @@ -181,25 +158,34 @@ def initialize_state(rays): throughput = jp.ones((num_rays, 3)) active_mask = jp.ones((num_rays,), dtype=bool) refractive_index = jp.ones((num_rays,)) - args = color, depth, hit_mask, throughput, active_mask - args += refractive_index, rays - return RenderState(*args) + fields = color, depth, hit_mask, throughput, active_mask + return RenderState(*fields, refractive_index, rays) def bounce_step(state, bounce, compiled, shadows, face_chunk): triangle_hit = compute_triangle_hit(compiled, state.rays, face_chunk) - intersections = intersect(compiled, state.rays, triangle_hit) - hit_masks, depths, points, normals, indices, eyes = intersections - hit_shape_args = find_closest_intersection_args(hit_masks, depths) - closest = gather_closest(*intersections) + hit_masks, depths, surfaces = intersect(compiled, state.rays, triangle_hit) + closest = gather_closest(hit_masks, depths, surfaces) state = update_first_hit(state, closest, bounce) state = update_active_mask(state, closest) - args = state.rays, compiled, hit_shape_args, closest, points - args += normals, eyes, shadows, triangle_hit - colors = compute_hit_colors(*args) + color_args = compiled, closest, surfaces, shadows, triangle_hit + colors = compute_hit_colors(*color_args) return update_state(state, compiled, closest, colors) +def update_first_hit(state, closest, bounce): + if bounce == 0: + depth, hit_mask = closest.depth, closest.hit_mask + else: + depth, hit_mask = state.depth, state.hit_mask + return state._replace(depth=depth, hit_mask=hit_mask) + + +def update_active_mask(state, closest): + active_mask = state.active_mask & closest.hit_mask + return state._replace(active_mask=active_mask) + + def compute_triangle_hit(compiled, rays, face_chunk): if compiled.triangles is None: triangle_hit = None @@ -222,53 +208,69 @@ def build_triangle_hit(compiled, rays, face_chunk): return TriangleHit(*args) -def update_first_hit(state, closest, bounce): - if bounce != 0: - return state - return state._replace(depth=closest.depth, hit_mask=closest.hit_mask) +def intersect(compiled, rays, triangle_hit): + rows = [] + if len(compiled.shapes) > 0: + rows.append(intersect_shapes(compiled.shapes, rays, compiled.mask)) + if triangle_hit is not None: + rows.append(build_triangle_row(triangle_hit)) + joined = tuple(jp.concatenate(fields, axis=0) for fields in zip(*rows)) + hit_masks, depths, points, normals, eyes = joined + return hit_masks, depths, Surfaces(points, normals, eyes) -def update_active_mask(state, closest): - active_mask = state.active_mask & closest.hit_mask - return state._replace(active_mask=active_mask) +def intersect_shapes(shapes, rays, mask): + merged = paz.graphics.shapes.field_merge(shapes, ["transform", "type"]) + intersect_fun = paz.lock(paz.graphics.shapes.intersect, *rays) + hit_masks, depths, points, normals, eyes = jax.vmap(intersect_fun)(merged) + hit_masks = jp.where(jp.expand_dims(mask, 1), hit_masks, False) + return hit_masks, depths, points, normals, eyes -def compute_hit_colors( - rays, compiled, indices, closest, points, normals, eyes, shadows, - triangle_hit, -): - shape_args = rays, compiled, indices, closest, points, normals, eyes - num_shapes = len(compiled.shapes) - if num_shapes == 0: +def build_triangle_row(triangle_hit): + depth = jp.expand_dims(triangle_hit.depth, -1) + fields = triangle_hit.hit_mask, depth, triangle_hit.points + fields += triangle_hit.normals, triangle_hit.eyes + return tuple(jp.expand_dims(field, 0) for field in fields) + + +def gather_closest(hit_masks, depths, surfaces): + indices = find_closest_intersection_args(hit_masks, depths) + fields = hit_masks, depths, *surfaces + closest = tuple(take_closest(field, indices) for field in fields) + return paz.graphics.Hit(*closest, indices) + + +def compute_hit_colors(compiled, closest, surfaces, shadows, triangle_hit): + if len(compiled.shapes) == 0: colors = compute_triangle_colors(compiled, triangle_hit) elif triangle_hit is None: - colors = compute_shape_colors(*shape_args, shadows) + colors = compute_shape_colors(compiled, closest, surfaces, shadows) else: - # TODO a mixed scene shades both paths for every ray and throws one - # away. Joining them before selection needs color_with_shadows to - # return rows instead of selecting inside its per-light scan. - shape_colors = compute_shape_colors(*shape_args, shadows) - triangle_colors = compute_triangle_colors(compiled, triangle_hit) - is_triangle = jp.expand_dims(indices == num_shapes, -1) - colors = jp.where(is_triangle, triangle_colors, shape_colors) + args = compiled, closest, surfaces, shadows, triangle_hit + colors = blend_hit_colors(*args) return colors -def compute_shape_colors( - rays, compiled, indices, closest, points, normals, eyes, shadows -): +def blend_hit_colors(compiled, closest, surfaces, shadows, triangle_hit): + # TODO a mixed scene shades both paths for every ray and throws one + # away. Joining them before selection needs color_with_shadows to + # return rows instead of selecting inside its per-light scan. + shape_colors = compute_shape_colors(compiled, closest, surfaces, shadows) + triangle_colors = compute_triangle_colors(compiled, triangle_hit) + is_triangle = closest.primitive_index == len(compiled.shapes) + is_triangle = jp.expand_dims(is_triangle, -1) + return jp.where(is_triangle, triangle_colors, shape_colors) + + +def compute_shape_colors(compiled, closest, surfaces, shadows): num_shapes = len(compiled.shapes) - shape_args = jp.minimum(indices, num_shapes - 1) - points, normals = points[:num_shapes], normals[:num_shapes] - eyes = eyes[:num_shapes] + indices = jp.minimum(closest.primitive_index, num_shapes - 1) + surfaces = slice_surfaces(surfaces, 0, num_shapes) if shadows: - color_args = rays, compiled.shapes, compiled.lights, shape_args - color_args += compiled.mask, compiled.shadow_mask - color_args += closest.point, closest.normal, points, normals, eyes - colors = color_with_shadows(ShadowColorArgs(*color_args)) + colors = color_with_shadows(compiled, closest, surfaces, indices) else: - color_args = compiled.lights, compiled.shapes, points, normals - colors = color_without_shadow(*color_args, eyes, shape_args) + colors = color_without_shadow(compiled, surfaces, indices) return colors @@ -295,149 +297,156 @@ def expand_scalar_field(field): return field -def intersect(compiled, rays, triangle_hit): - rows = [] - if len(compiled.shapes) > 0: - rows.append(intersect_shapes(compiled.shapes, rays, compiled.mask)) - if triangle_hit is not None: - rows.append(build_triangle_row(triangle_hit)) - return stack_intersection_rows(rows) - +def select_shader(material): + if isinstance(material, paz.graphics.CookTorranceMaterial): + shader = paz.graphics.cook_torrance + else: + shader = paz.graphics.phong + return shader -def intersect_shapes(shapes, rays, mask): - def hide_shapes(mask, hit_masks): - return jp.where(jp.expand_dims(mask, 1), hit_masks, False) +def color_without_shadow(compiled, surfaces, indices): + colors = [] + lights = paz.graphics.shapes.merge(*compiled.lights) + for group, start_arg, final_arg in iterate_shape_groups(compiled.shapes): + group_surfaces = slice_surfaces(surfaces, start_arg, final_arg) + albedo = compute_group_albedo(group, group_surfaces.points) + shader = select_shader(group.material) + axes = 0, 0, 0, 0, 0, None + color_per_light = jax.vmap(shader.compute_colors, axes) + color = jax.vmap(color_per_light, (None, None, None, None, None, 0)) + args = albedo, group.material, *group_surfaces, lights + colors.append(jp.sum(color(*args), axis=0)) + return take_closest(jp.concatenate(colors, axis=0), indices) - merge = paz.graphics.shapes.field_merge(shapes, ["transform", "type"]) - intersect_fun = paz.lock(paz.graphics.shapes.intersect, *rays) - hit_masks, depths, points, normals, eyes = jax.vmap(intersect_fun)(merge) - return hide_shapes(mask, hit_masks), depths, points, normals, eyes +def iterate_shape_groups(shapes): + start_arg = 0 + for group in paz.graphics.shapes.group_by_pattern_size(shapes).values(): + final_arg = start_arg + len(group) + yield paz.graphics.shapes.merge(*group), start_arg, final_arg + start_arg = final_arg -def build_triangle_row(triangle_hit): - depth = jp.expand_dims(triangle_hit.depth, -1) - rows = jp.expand_dims(triangle_hit.hit_mask, 0) - rows = rows, jp.expand_dims(depth, 0) - rows += (jp.expand_dims(triangle_hit.points, 0),) - rows += (jp.expand_dims(triangle_hit.normals, 0),) - return rows + (jp.expand_dims(triangle_hit.eyes, 0),) +def slice_surfaces(surfaces, start_arg, final_arg): + points = surfaces.points[start_arg:final_arg] + normals = surfaces.normals[start_arg:final_arg] + eyes = surfaces.eyes[start_arg:final_arg] + return Surfaces(points, normals, eyes) -def stack_intersection_rows(rows): - joined = tuple(jp.concatenate(fields, axis=0) for fields in zip(*rows)) - indices = jp.arange(joined[0].shape[0]) - hit_masks, depths, points, normals, eyes = joined - return hit_masks, depths, points, normals, indices, eyes +def compute_group_albedo(group, points): + compute_albedo = paz.graphics.albedo.compute_shape_albedo + return jax.vmap(compute_albedo)(group, group.material, points) -def gather_closest(hit_masks, depths, points, normals, indices, eyes): - closest_args = find_closest_intersection_args(hit_masks, depths) - args = take_closest(hit_masks, closest_args) - args = args, take_closest(depths, closest_args) - args += (take_closest(points, closest_args),) - args += (take_closest(normals, closest_args),) - args += (take_closest(eyes, closest_args),) - args += (indices[closest_args],) - return paz.graphics.Hit(*args) +def color_with_shadows(compiled, closest, surfaces, indices): + colors = jp.zeros((len(surfaces.points[0]), 3)) + lights = paz.graphics.shapes.merge(*compiled.lights) + body = paz.lock(scan_light_step, compiled, closest, surfaces, indices) + return jax.lax.scan(body, colors, lights)[0] -def select_shader(material): - if isinstance(material, paz.graphics.CookTorranceMaterial): - return paz.graphics.cook_torrance - return paz.graphics.phong +def scan_light_step(colors, light, compiled, closest, surfaces, indices): + args = compiled, closest, surfaces, indices, light + return colors + compute_light_colors(*args), None -def color_without_shadow(lights, shapes, points, normals, eyes, hit_shape_args): - colors, start_arg, merged_lights = [], 0, paz.graphics.shapes.merge(*lights) - for group in paz.graphics.shapes.group_by_pattern_size(shapes).values(): - final_arg = start_arg + len(group) - group = paz.graphics.shapes.merge(*group) - data = split_shape_data(points, normals, eyes, start_arg, final_arg) - albedo = compute_group_albedo(group, data[0]) - args, axes = (albedo, group.material, *data), (0, 0, 0, 0, 0, None) - shader = select_shader(group.material) - color_per_light = jax.vmap(shader.compute_colors, axes) - color = jax.vmap(color_per_light, (None, None, None, None, None, 0)) - colors.append(jp.sum(color(*args, merged_lights), axis=0)) - start_arg = final_arg - return take_closest(jp.concatenate(colors, axis=0), hit_shape_args) +def compute_light_colors(compiled, closest, surfaces, indices, light): + directions, distance = compute_light_directions(light, closest.point) + occlusion_args = compiled, closest, indices, directions, distance + is_shadow = compute_light_occlusion(*occlusion_args) + color_args = compiled.shapes, light, surfaces, is_shadow + return take_closest(compute_shadowed_colors(*color_args), indices) -def compute_group_albedo(group, points): - compute_albedo = paz.graphics.albedo.compute_shape_albedo - return jax.vmap(compute_albedo)(group, group.material, points) +def compute_light_directions(light, points): + vector = light.position - points + norm = paz.algebra.compute_norms(vector, 1) + return vector / norm, jp.squeeze(norm, axis=1) -def split_shape_data(points, normals, eyes, start_arg, final_arg): - points = points[start_arg:final_arg] - normals = normals[start_arg:final_arg] - eyes = eyes[start_arg:final_arg] - return points, normals, eyes +def compute_light_occlusion(compiled, closest, indices, directions, distance): + origins = compute_shadow_ray_origins(closest.point, closest.normal) + shadow_args = compiled.shapes, origins, directions + hit_masks, depths, _, _, _, casters = intersect_shadow_groups(*shadow_args) + masks = resolve_shadow_masks(compiled, hit_masks) + depth_args = masks, depths, casters, indices, closest.normal, directions + masks, depths = select_shadow_depths(*depth_args) + return compute_soft_occlusion(masks, depths, distance) -def intersect_shape_groups(shapes, origins, directions, intersect_shape): - def process_group(group, rays, start_arg): - indices = jp.arange(start_arg, start_arg + len(group)) - merged_group = paz.graphics.shapes.merge(*group) - intersect = paz.lock(intersect_shape, *rays) - intersections = jax.vmap(intersect)(merged_group) - return (*intersections, indices) +def compute_shadow_ray_origins(points, normals): + over_point, _ = compute_surface_points(points, normals) + return over_point - def concatenate(x): - return tuple(jp.concatenate(items, axis=0) for items in zip(*x)) - intersections, start_arg, rays = [], 0, (origins, directions) - for group in paz.graphics.shapes.group_by_pattern_size(shapes).values(): - intersections.append(process_group(group, rays, start_arg)) - start_arg = start_arg + len(group) - return concatenate(intersections) +def compute_surface_points(point, normal): + over_point = point + normal * SHADOW_ORIGIN_EPSILON + under_point = point - normal * SHADOW_ORIGIN_EPSILON + return over_point, under_point def intersect_shadow_groups(shapes, origins, directions): - args = (shapes, origins, directions, paz.graphics.shapes.intersect_all) - return intersect_shape_groups(*args) + intersect_all = paz.graphics.shapes.intersect_all + intersect_group = jax.vmap(paz.lock(intersect_all, origins, directions)) + rows = [] + for group, start_arg, final_arg in iterate_shape_groups(shapes): + indices = jp.arange(start_arg, final_arg) + rows.append((*intersect_group(group), indices)) + return tuple(jp.concatenate(fields, axis=0) for fields in zip(*rows)) + + +def resolve_shadow_masks(compiled, hit_masks): + transparencies = compute_transparencies(compiled.shapes) + masks = jp.where(jp.expand_dims(compiled.mask, 1), hit_masks, False) + masks = hide_transparent_shapes(masks, transparencies > 0.0) + if compiled.shadow_mask is None: + resolved = masks + else: + resolved = hide_non_casting_shapes(masks, compiled.shadow_mask) + return resolved -def compute_surface_points(point, normal, epsilon=SHADOW_ORIGIN_EPSILON): - over_point = point + normal * epsilon - under_point = point - normal * epsilon - return over_point, under_point +def compute_transparencies(shapes): + return jp.array([shape.material.transparency for shape in shapes]) -def compute_shadow_ray_origins(points, normals): - over_point, _ = compute_surface_points(points, normals) - return over_point +def hide_transparent_shapes(shadow_masks, is_transparent): + return jp.where(jp.expand_dims(is_transparent, 1), False, shadow_masks) -def compute_shadow_depth_thresholds(shape_indices, receiver_indices): - same_shape = shape_indices[:, None] == receiver_indices[None, :] - return jp.where(same_shape, SHADOW_SELF_HIT_EPSILON, paz.graphics.EPSILON) +def hide_non_casting_shapes(shadow_masks, shadow_mask): + return jp.where(jp.expand_dims(shadow_mask, 1), shadow_masks, False) -def compute_front_side_shadow_mask(*args): - shape_indices, receiver_indices, receiver_normals, directions = args - same_shape = shape_indices[:, None] == receiver_indices[None, :] - front_side = paz.algebra.dot(receiver_normals, directions) >= 0.0 +def select_shadow_depths( + hit_masks, depths, casters, receivers, normals, directions +): + same_shape = casters[:, None] == receivers[None, :] + front_args = same_shape, normals, directions + front_side_hits = compute_front_side_shadow_mask(*front_args) + root_args = hit_masks, depths, same_shape, front_side_hits + valid_roots = compute_valid_roots(*root_args) + depths = jp.where(valid_roots, depths, paz.graphics.FARAWAY) + return jp.any(valid_roots, axis=1), jp.min(depths, axis=1) + + +def compute_front_side_shadow_mask(same_shape, normals, directions): + front_side = paz.algebra.dot(normals, directions) >= 0.0 return jp.logical_and(same_shape, front_side[None, :]) -def select_shadow_depths(*args): - hit_masks, depths, shape_indices = args[:3] - receiver_indices, receiver_normals, directions = args[3:] - threshold_args = shape_indices, receiver_indices - thresholds = compute_shadow_depth_thresholds(*threshold_args) - front_args = shape_indices, receiver_indices, receiver_normals, directions - front_side_hits = compute_front_side_shadow_mask(*front_args) +def compute_valid_roots(hit_masks, depths, same_shape, front_side_hits): + thresholds = compute_shadow_depth_thresholds(same_shape) valid_roots = depths > thresholds[:, None, :] valid_roots = jp.logical_and(valid_roots, depths < paz.graphics.FARAWAY) valid_roots = jp.logical_and(jp.expand_dims(hit_masks, 1), valid_roots) - valid_roots = jp.logical_and(valid_roots, ~front_side_hits[:, None, :]) - depths = jp.where(valid_roots, depths, paz.graphics.FARAWAY) - hit_masks = jp.any(valid_roots, axis=1) - depths = jp.min(depths, axis=1) - return hit_masks, depths + return jp.logical_and(valid_roots, ~front_side_hits[:, None, :]) + + +def compute_shadow_depth_thresholds(same_shape): + return jp.where(same_shape, SHADOW_SELF_HIT_EPSILON, paz.graphics.EPSILON) def compute_soft_occlusion(hit_masks, depths, light_lengths, slope=0.01): @@ -451,114 +460,72 @@ def compute_soft_occlusion(hit_masks, depths, light_lengths, slope=0.01): return jp.where(blocker_mask, occlusion, 0.0) -def color_with_shadows(args): - transparencies = compute_transparencies(args.shapes) - colors = jp.zeros((len(args.points[0]), 3)) - lights = paz.graphics.shapes.merge(*args.lights) - body = paz.lock(scan_light_step, args, transparencies) - return jax.lax.scan(body, colors, lights)[0] - - -def scan_light_step(colors, light, args, transparencies): - return colors + compute_light_colors(args, light, transparencies), None - - -def compute_transparencies(shapes): - return jp.array([shape.material.transparency for shape in shapes]) - - -def compute_light_colors(args, light, transparencies): - directions, distance = compute_light_directions(light, args.point) - origins = compute_shadow_ray_origins(args.point, args.normal) - intersections = intersect_shadow_groups(args.shapes, origins, directions) - hit_masks, depths, _, _, _, shape_indices = intersections - masks = resolve_shadow_masks(args, hit_masks, transparencies) - select_args = masks, depths, shape_indices, args.indices - select_args += args.normal, directions - masks, depths = select_shadow_depths(*select_args) - is_shadow = compute_soft_occlusion(masks, depths, distance) - color_args = args.shapes, light, args.points, args.normals, args.eyes - color_args += (is_shadow,) - colors = compute_shadowed_colors(*color_args) - return take_closest(colors, args.indices) - - -def compute_light_directions(light, points): - vector = light.position - points - norm = paz.algebra.compute_norms(vector, 1) - return vector / norm, jp.squeeze(norm, axis=1) - - -def resolve_shadow_masks(args, hit_masks, transparencies): - shadow_masks = jp.where(jp.expand_dims(args.mask, 1), hit_masks, False) - is_transparent = transparencies > 0.0 - shadow_masks = hide_transparent_shapes(shadow_masks, is_transparent) - if args.shadow_mask is not None: - cast_mask = jp.expand_dims(args.shadow_mask, 1) - shadow_masks = jp.where(cast_mask, shadow_masks, False) - return shadow_masks - - -def hide_transparent_shapes(shadow_masks, is_transparent): - return jp.where(jp.expand_dims(is_transparent, 1), False, shadow_masks) - - -def compute_shadowed_colors(*args): - shapes, light, points, normals, eyes, is_shadow = args - colors, start_arg = [], 0 - for group in paz.graphics.shapes.group_by_pattern_size(shapes).values(): - final_arg = start_arg + len(group) - group = paz.graphics.shapes.merge(*group) - data = split_shape_data(points, normals, eyes, start_arg, final_arg) - albedo = compute_group_albedo(group, data[0]) - axes = 0, 0, 0, 0, 0, None, None +def compute_shadowed_colors(shapes, light, surfaces, is_shadow): + colors = [] + for group, start_arg, final_arg in iterate_shape_groups(shapes): + group_surfaces = slice_surfaces(surfaces, start_arg, final_arg) + albedo = compute_group_albedo(group, group_surfaces.points) shader = select_shader(group.material) + axes = 0, 0, 0, 0, 0, None, None color = jax.vmap(shader.compute_colors_with_shadow, axes) - color_args = albedo, group.material, *data, light, is_shadow - colors.append(color(*color_args)) - start_arg = final_arg + args = albedo, group.material, *group_surfaces, light, is_shadow + colors.append(color(*args)) return jp.concatenate(colors, axis=0) -def update_state(state, compiled, closest, intersected_colors): - material = get_material_properties(compiled, closest.primitive_index) - reflectivities, transparencies, refractivities = material - color_args = state.color, state.throughput, state.active_mask - color_args += intersected_colors, reflectivities, transparencies +def update_state(state, compiled, closest, hit_colors): + material = compute_material_properties(compiled, closest.primitive_index) + color_args = state.color, state.throughput, state.active_mask, hit_colors + color_args += material.reflectivities, material.transparencies color = accumulate_color(*color_args) - args = state.rays[1], state.refractive_index, closest.normal - args += (refractivities,) - normal, eye, n1, n2, n_ratio = prepare_computations(*args) - reflectance = schlick(normal, eye, n1, n2) - ray_args = normal, eye, n_ratio, closest.point, transparencies, reflectance - new_rays = compute_new_rays(*ray_args) - update_args = new_rays, n2, reflectivities, transparencies, reflectance - return apply_bounce_update(state._replace(color=color), *update_args) - - -def get_material_properties(compiled, hit_shape_args): - reflectivities, transparencies, refractivities = [], [], [] + new_rays, n_2, reflectance = compute_bounce(state, closest, material) + args = new_rays, n_2, material, reflectance + return apply_bounce_update(state._replace(color=color), *args) + + +def compute_material_properties(compiled, hit_shape_args): + values = collect_material_values(compiled) + gathered = tuple(jp.array(row)[hit_shape_args] for row in values) + return HitMaterial(*gathered) + + +def collect_material_values(compiled): + reflectivities, transparencies, refractive_indices = [], [], [] for shape in compiled.shapes: reflectivities.append(shape.material.reflective) transparencies.append(shape.material.transparency) - refractivities.append(shape.material.refractive_index) + refractive_indices.append(shape.material.refractive_index) if compiled.triangles is not None: reflectivities.append(0.0) transparencies.append(0.0) - refractivities.append(1.0) - reflectivities = jp.array(reflectivities)[hit_shape_args] - transparencies = jp.array(transparencies)[hit_shape_args] - refractivities = jp.array(refractivities)[hit_shape_args] - return reflectivities, transparencies, refractivities + refractive_indices.append(1.0) + return reflectivities, transparencies, refractive_indices -def accumulate_color(*args): - colors, throughput, active_mask = args[:3] - intersected_colors, reflectivities, transparencies = args[3:] +def accumulate_color( + colors, throughput, active_mask, hit_colors, reflectivities, transparencies +): weights = jp.maximum(1.0 - reflectivities - transparencies, 0.0) weights = jp.expand_dims(weights, -1) active_mask = jp.expand_dims(active_mask, -1) - return colors + (throughput * active_mask * weights * intersected_colors) + return colors + (throughput * active_mask * weights * hit_colors) + + +def compute_bounce(state, closest, material): + args = state.rays[1], closest.normal, state.refractive_index + terms = compute_refraction_terms(*args, material.refractive_indices) + normal, eye, n_1, n_2, n_ratio = terms + reflectance = compute_reflectance(normal, eye, n_1, n_2) + ray_args = normal, eye, n_ratio, closest.point, material.transparencies + return compute_new_rays(*ray_args), n_2, reflectance + + +def compute_refraction_terms(directions, normal, n_1, refractive_indices): + eye = -directions + normal, is_inside = flip_normal_if_inside(eye, normal) + # TODO why 1.0 hardcoded + n_2 = jp.where(is_inside, 1.0, refractive_indices) + return normal, eye, n_1, n_2, n_1 / n_2 def flip_normal_if_inside(eye, normal): @@ -566,93 +533,72 @@ def flip_normal_if_inside(eye, normal): return jp.where(jp.expand_dims(is_inside, -1), -normal, normal), is_inside -def displace_by_normal(point, normal): - upper_point = point + normal * BOUNCE_ORIGIN_EPSILON - lower_point = point - normal * BOUNCE_ORIGIN_EPSILON - return lower_point, upper_point - - -def prepare_computations(*args): - current_directions, refractive_index, normal, refractive_indices = args - eye = -current_directions - normal, is_inside = flip_normal_if_inside(eye, normal) - n1 = refractive_index - n2 = jp.where(is_inside, 1.0, refractive_indices) # TODO why 1.0 hardcoded - n_ratio = n1 / (n2) - return normal, eye, n1, n2, n_ratio +def compute_reflectance(normal, eye, n_1, n_2): + cosines = compute_transmission_cosines(eye, normal, n_1 / n_2) + cos_incident, sin_transmit_squared, cos_transmit = cosines + cos = jp.where(n_1 > n_2, cos_transmit, cos_incident) + base_reflectance = ((n_1 - n_2) / (n_1 + n_2)) ** 2 + grazing = (1.0 - cos) ** 5 + reflectance = base_reflectance + (1.0 - base_reflectance) * grazing + return jp.where(sin_transmit_squared > 1.0, 1.0, reflectance) -def schlick(normal, eye, n1, n2): - n_ratio = n1 / n2 +def compute_transmission_cosines(eye, normal, n_ratio): cos_incident = jp.sum(eye * normal, axis=-1) sin_transmit_squared = (n_ratio**2) * (1.0 - (cos_incident**2)) cos_transmit = jp.sqrt(jp.maximum(0.0, 1.0 - sin_transmit_squared)) + return cos_incident, sin_transmit_squared, cos_transmit - is_total_internal_reflection = sin_transmit_squared > 1.0 - cos = jp.where(n1 > n2, cos_transmit, cos_incident) - - r0 = ((n1 - n2) / (n1 + n2)) ** 2 - reflectance = r0 + (1.0 - r0) * (1.0 - cos) ** 5 - return jp.where(is_total_internal_reflection, 1.0, reflectance) - -def reflect_or_refract(transparancies, reflectance): - is_transparent = transparancies > 0.0 - do_reflect = ~is_transparent - return jp.expand_dims(do_reflect, -1) +def compute_new_rays(normal, eye, n_ratio, point, transparencies): + is_transparent = transparencies > 0.0 + do_reflect = jp.expand_dims(~is_transparent, -1) + reflection = paz.graphics.geometry.reflect(-eye, normal) + refraction = compute_refractive_direction(eye, normal, n_ratio) + direction = jp.where(do_reflect, reflection, refraction) + lower_point, upper_point = displace_by_normal(point, normal) + origin = jp.where(do_reflect, upper_point, lower_point) + return origin, paz.algebra.normalize(direction) def compute_refractive_direction(eye, normal, n_ratio): - cos_incident = jp.sum(eye * normal, axis=-1) - sin_transmit_squared = (n_ratio**2) * (1.0 - (cos_incident**2)) - cos_transmit = jp.sqrt(jp.maximum(0.0, 1.0 - sin_transmit_squared)) + args = eye, normal, n_ratio + cos_incident, _, cos_transmit = compute_transmission_cosines(*args) inside_vector = -eye * jp.expand_dims(n_ratio, -1) - up_weight = jp.expand_dims((n_ratio * cos_incident - cos_transmit), -1) - return up_weight * normal + inside_vector - + up_weight = n_ratio * cos_incident - cos_transmit + return jp.expand_dims(up_weight, -1) * normal + inside_vector -def compute_reflection_direction(eye, normal): - return paz.graphics.geometry.reflect(-eye, normal) - -def compute_new_rays(normal, eye, n_ratio, point, transparancies, reflectance): - do_reflect = reflect_or_refract(transparancies, reflectance) - reflection_direction = compute_reflection_direction(eye, normal) - refractive_direction = compute_refractive_direction(eye, normal, n_ratio) - direction = jp.where(do_reflect, reflection_direction, refractive_direction) - direction = paz.algebra.normalize(direction) - lower_point, upper_point = displace_by_normal(point, normal) - origin = jp.where(do_reflect, upper_point, lower_point) - return origin, direction +def displace_by_normal(point, normal): + upper_point = point + normal * BOUNCE_ORIGIN_EPSILON + lower_point = point - normal * BOUNCE_ORIGIN_EPSILON + return lower_point, upper_point -def apply_bounce_update(*args): - state, new_rays, n2, reflectivities, transparencies, reflectance = args - is_transparent = transparencies > 0.0 - is_reflective = reflectivities > 0.0 - factor_args = is_transparent, is_reflective, transparencies - factor_args += reflectivities, reflectance - factor = compute_bounce_factor(*factor_args) +def apply_bounce_update(state, new_rays, n_2, material, reflectance): + is_transparent = material.transparencies > 0.0 + is_reflective = material.reflectivities > 0.0 + factor = compute_bounce_factor(material, reflectance) factor = jp.where(is_transparent & (reflectance >= 1.0), 1.0, factor) throughput = state.throughput * jp.expand_dims(factor, -1) active_mask = state.active_mask & (is_transparent | is_reflective) - index_args = state, n2, is_transparent, reflectance + index_args = state, n_2, is_transparent, reflectance refractive_index = update_refractive_index(*index_args) args = throughput, active_mask, refractive_index, new_rays return replace_bounce_state(state, *args) -def compute_bounce_factor(*args): - is_transparent, is_reflective, transparencies = args[:3] - reflectivities, reflectance = args[3:] - transparent_factor = transparencies * (1.0 - reflectance) - reflective_factor = jp.where(is_reflective, reflectivities, 0.0) +def compute_bounce_factor(material, reflectance): + is_transparent = material.transparencies > 0.0 + is_reflective = material.reflectivities > 0.0 + transparent_factor = material.transparencies * (1.0 - reflectance) + reflective_factor = jp.where(is_reflective, material.reflectivities, 0.0) return jp.where(is_transparent, transparent_factor, reflective_factor) -def update_refractive_index(state, n2, is_transparent, reflectance): +def update_refractive_index(state, n_2, is_transparent, reflectance): update_mask = is_transparent & (reflectance < 1.0) - return jp.where(update_mask, n2, state.refractive_index) + return jp.where(update_mask, n_2, state.refractive_index) def replace_bounce_state(state, throughput, active_mask, index, rays): diff --git a/paz/graphics/renderer_test.py b/paz/graphics/renderer_test.py index b510f27cb..35422341c 100644 --- a/paz/graphics/renderer_test.py +++ b/paz/graphics/renderer_test.py @@ -8,7 +8,6 @@ from paz.graphics import composite, renderer from paz.graphics.shapes.sphere import intersect_canonical_sphere from paz.graphics.types import ( - Shape, Material, PointLight, Sphere, @@ -21,7 +20,6 @@ CylindricalPattern, Scene, ) -from paz.graphics import constants OLD_CAMERA_POSE = jp.array( @@ -184,15 +182,9 @@ def test_compute_new_rays_reflection_is_normalized(): normal = jp.array([[0.0, 1.0, 0.0]]) eye = jp.array([[0.0, 1.0, -1.0]]) point = jp.array([[0.0, 0.0, 0.0]]) - transparancies = jp.array([0.0]) - reflectance = jp.array([1.0]) + transparencies = jp.array([0.0]) _, direction = renderer.compute_new_rays( - normal, - eye, - jp.array([1.0]), - point, - transparancies, - reflectance, + normal, eye, jp.array([1.0]), point, transparencies ) norm = jp.linalg.norm(direction, axis=-1) assert jp.allclose(norm, 1.0, atol=1e-5) @@ -202,15 +194,9 @@ def test_compute_new_rays_refraction_is_normalized(): normal = jp.array([[0.0, 0.0, -1.0]]) eye = jp.array([[0.0, 0.0, -1.0]]) point = jp.array([[0.0, 0.0, 0.0]]) - transparancies = jp.array([1.0]) - reflectance = jp.array([0.0]) + transparencies = jp.array([1.0]) _, direction = renderer.compute_new_rays( - normal, - eye, - jp.array([1.0 / 1.5]), - point, - transparancies, - reflectance, + normal, eye, jp.array([1.0 / 1.5]), point, transparencies ) norm = jp.linalg.norm(direction, axis=-1) assert jp.allclose(norm, 1.0, atol=1e-5) @@ -324,7 +310,7 @@ def test_find_closest_intersection_args(): assert jp.array_equal(indices, jp.array([0, 1])) -def test_get_material_properties(): +def test_compute_material_properties(): mat1 = Material(reflective=0.5) mat2 = Material(transparency=0.8) shape1 = Sphere(material=mat1) @@ -332,12 +318,10 @@ def test_get_material_properties(): scene = Scene([shape1, shape2]) compiled = paz.graphics.scene.compile(scene, [], None) indices = jp.array([0, 1]) - reflectivities, transparencies, refractivities = ( - renderer.get_material_properties(compiled, indices) - ) - assert reflectivities[0] == 0.5 - assert transparencies[1] == 0.8 - assert refractivities[0] == 1.0 + material = renderer.compute_material_properties(compiled, indices) + assert material.reflectivities[0] == 0.5 + assert material.transparencies[1] == 0.8 + assert material.refractive_indices[0] == 1.0 def test_accumulate_color():