diff --git a/examples/differentiable_shape_rendering/phong_realism_levels.py b/examples/differentiable_shape_rendering/phong_realism_levels.py index 41b08d4be..125e65c13 100644 --- a/examples/differentiable_shape_rendering/phong_realism_levels.py +++ b/examples/differentiable_shape_rendering/phong_realism_levels.py @@ -1,10 +1,10 @@ import jax import jax.numpy as jp import paz -import paz.graphics.renderer as paz_renderer +from paz.graphics.renderer import shadow -soft_occlusion = paz_renderer.compute_soft_occlusion -paz_renderer.compute_soft_occlusion = paz.partial(soft_occlusion, slope=1.0) +soft_occlusion = shadow.compute_soft_occlusion +shadow.compute_soft_occlusion = paz.partial(soft_occlusion, slope=1.0) H, W = 480, 640 diff --git a/examples/differentiable_shape_rendering/realistic_primitives.py b/examples/differentiable_shape_rendering/realistic_primitives.py index 42d1223fe..85eca7f23 100644 --- a/examples/differentiable_shape_rendering/realistic_primitives.py +++ b/examples/differentiable_shape_rendering/realistic_primitives.py @@ -1,10 +1,10 @@ import jax import jax.numpy as jp import paz -import paz.graphics.renderer as paz_renderer +from paz.graphics.renderer import shadow -soft_occlusion = paz_renderer.compute_soft_occlusion -paz_renderer.compute_soft_occlusion = paz.partial(soft_occlusion, slope=1.0) +soft_occlusion = shadow.compute_soft_occlusion +shadow.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 diff --git a/examples/differentiable_shape_rendering/visualize_scene_gradients.py b/examples/differentiable_shape_rendering/visualize_scene_gradients.py index 0b5db7024..f89314dab 100644 --- a/examples/differentiable_shape_rendering/visualize_scene_gradients.py +++ b/examples/differentiable_shape_rendering/visualize_scene_gradients.py @@ -8,7 +8,7 @@ import matplotlib.pyplot as plt import numpy as np import paz -import paz.graphics.renderer as paz_renderer +from paz.graphics.renderer import shadow import paz.utils.plot as plot from mpl_toolkits.axes_grid1 import make_axes_locatable @@ -23,7 +23,7 @@ CURVATURE_RATIO = 1e-3 LUMINANCE_WEIGHTS = jp.array([0.2126, 0.7152, 0.0722]) SIGNED_CMAP = "RdBu_r" -SOFT_OCCLUSION = paz_renderer.compute_soft_occlusion +SOFT_OCCLUSION = shadow.compute_soft_occlusion CAMERA_ARGS = ( jp.array([0.0, 2.0, 2.0]), jp.array([0.0, 0.0, 0.0]), @@ -77,7 +77,7 @@ def build_scene(shape_transform=jp.eye(4)): def configure_soft_occlusion(): - paz_renderer.compute_soft_occlusion = paz.partial(SOFT_OCCLUSION, slope=1.0) + shadow.compute_soft_occlusion = paz.partial(SOFT_OCCLUSION, slope=1.0) def compute_autodiff_gradient(function, args, basis): diff --git a/paz/graphics/mesh_test.py b/paz/graphics/mesh_test.py index d2ebbccf2..50bfa9a80 100644 --- a/paz/graphics/mesh_test.py +++ b/paz/graphics/mesh_test.py @@ -854,6 +854,101 @@ def test_scene_mixes_meshes_and_shapes(): assert jp.any(depth > 0) +def build_cube_mesh_with_material(material): + vertices, faces, edges = build_cube(1.0) + colors = build_vertex_colors(vertices, [0.7, 0.3, 0.1]) + transform = SE3.translation(jp.array([-0.45, 0.0, 0.0])) + args = vertices, colors, transform, material, faces, edges + return Mesh(*args) + + +def render_mesh_with_material(material): + shape, y_FOV, pose, _, _, lights = make_multi_mesh_scene() + mesh = build_cube_mesh_with_material(material) + scene = paz.graphics.Scene([mesh, build_sphere_mesh()]) + args = shape, y_FOV, pose, scene, None, lights, (1, 1), 1024 + return paz.graphics.render(*args, False, None, 2) + + +def test_mesh_reflective_material_changes_render(): + matte = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100, 0.0) + mirror = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100, 0.9) + matte_image, _ = render_mesh_with_material(matte) + mirror_image, _ = render_mesh_with_material(mirror) + assert compute_max_abs_difference(matte_image, mirror_image) > 1e-3 + + +def test_mesh_transparent_material_changes_render(): + opaque = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100, 0.0, 0.0) + glass = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100, 0.0, 0.8, 1.5) + opaque_image, _ = render_mesh_with_material(opaque) + glass_image, _ = render_mesh_with_material(glass) + assert compute_max_abs_difference(opaque_image, glass_image) > 1e-3 + + +def build_mesh_over_floor_scene(): + vertices, faces, edges = build_cube(1.0) + colors = build_vertex_colors(vertices, [0.7, 0.3, 0.1]) + material = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100) + transform = SE3.translation(jp.array([0.0, 1.0, 0.0])) + args = vertices, colors, transform, material, faces, edges + floor = paz.graphics.Plane(SE3.translation(jp.array([0.0, -0.5, 0.0]))) + return paz.graphics.Scene([Mesh(*args), floor]) + + +def render_mesh_over_floor(shadows, shadow_mask=None): + camera_pose = SE3.view_transform( + jp.array([0.0, 2.0, -4.0]), jp.zeros(3), jp.array([0.0, 1.0, 0.0]) + ) + lights = [PointLight(jp.ones(3), jp.array([0.0, 5.0, -1.0]))] + scene = build_mesh_over_floor_scene() + args = (32, 32), jp.pi / 3.0, camera_pose, scene, None, lights + return paz.graphics.render(*args, (1, 1), 1024, shadows, shadow_mask) + + +def test_mesh_casts_shadow_on_shape(): + lit, _ = render_mesh_over_floor(False) + shadowed, _ = render_mesh_over_floor(True) + assert compute_max_abs_difference(lit, shadowed) > 1e-2 + assert jp.all(shadowed <= lit + 1e-4) + + +def test_mesh_shadow_mask_stops_mesh_casting(): + casting, _ = render_mesh_over_floor(True) + blocked, _ = render_mesh_over_floor(True, jp.array([False, True])) + assert compute_max_abs_difference(casting, blocked) > 1e-2 + assert jp.all(casting <= blocked + 1e-4) + + +def build_mesh_only_shadow_scene(): + vertices, faces, edges = build_cube(1.0) + colors = build_vertex_colors(vertices, [0.7, 0.4, 0.2]) + material = Material(jp.zeros(3), 0.1, 0.9, 0.1, 100) + blocker_pose = SE3.translation(jp.array([0.0, 1.4, 0.0])) + blocker = Mesh(vertices, colors, blocker_pose, material, faces, edges) + slab_pose = SE3.translation(jp.array([0.0, -0.6, 0.0])) + slab_pose = slab_pose @ SE3.scaling(jp.array([6.0, 0.2, 6.0])) + slab = Mesh(vertices, colors, slab_pose, material, faces, edges) + return paz.graphics.Scene([blocker, slab]) + + +def render_mesh_only_shadow(shadows): + camera_pose = SE3.view_transform( + jp.array([0.0, 2.4, -5.0]), jp.zeros(3), jp.array([0.0, 1.0, 0.0]) + ) + lights = [PointLight(jp.ones(3), jp.array([0.0, 6.0, -1.5]))] + scene = build_mesh_only_shadow_scene() + args = (32, 32), jp.pi / 3.0, camera_pose, scene, None, lights + return paz.graphics.render(*args, (1, 1), 1024, shadows) + + +def test_mesh_receives_shadow_from_mesh(): + lit, _ = render_mesh_only_shadow(False) + shadowed, _ = render_mesh_only_shadow(True) + assert compute_max_abs_difference(lit, shadowed) > 1e-2 + assert jp.all(shadowed <= lit + 1e-4) + + def test_scene_rejects_meshes_with_mixed_pattern_sizes(): plain = build_cube_mesh() textured = build_textured_quad_mesh() diff --git a/paz/graphics/renderer.py b/paz/graphics/renderer.py deleted file mode 100644 index d6dd9adc1..000000000 --- a/paz/graphics/renderer.py +++ /dev/null @@ -1,606 +0,0 @@ -from collections import namedtuple - -import jax -import jax.numpy as jp - -import paz - -from paz.graphics.composite import ( - compute_scene_hit_mask, - find_closest_intersection_args, - postprocess, - take_closest, -) - -# TODO retune: measured best under 6k faces, but 2048 is 1.75x faster -# at 82k, so this is wrong for the room-scale meshes ahead. -FACE_CHUNK_SIZE = 128 -SHADOW_ORIGIN_EPSILON = 1e-5 -SHADOW_SELF_HIT_EPSILON = 1e-5 -BOUNCE_ORIGIN_EPSILON = 1e-2 - -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" -RenderState = namedtuple("RenderState", STATE_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, - face_chunk_size=FACE_CHUNK_SIZE, -): - 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, - face_chunk_size=FACE_CHUNK_SIZE, -): - if num_objects is None: - num_objects = len(scene.nodes) - min_depth, max_depth = depth - masks = [] - for object_arg in range(num_objects): - mask = build_object_mask(len(scene.nodes), object_arg) - args = shape, y_FOV, pose, scene, mask, lights, tiles, 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 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) - return jax.lax.scan(tile_step, None, coordinates)[1] - - -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 - 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_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, compiled, shadows, num_bounces, face_chunk, chunk_size): - ray_chunks = split_ray_chunks(rays, chunk_size) - 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, compiled, shadows, num_bounces, face_chunk): - args = rays, compiled, shadows, num_bounces, face_chunk - return carry, trace_bounces(*args) - - -def split_ray_chunks(rays, chunk_size): - origins, directions = rays - origins = pad_to_chunks(origins, chunk_size) - directions = pad_to_chunks(directions, chunk_size) - num_chunks = origins.shape[0] // chunk_size - shape = num_chunks, chunk_size, 3 - return origins.reshape(shape), directions.reshape(shape) - - -def pad_to_chunks(array, chunk_size): - remainder = array.shape[0] % chunk_size - if remainder == 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): - hit_mask = flatten_chunk_array(hit_mask, num_rays) - depth = flatten_chunk_array(depth, num_rays) - color = flatten_chunk_array(color, num_rays) - return hit_mask, depth, color - - -def flatten_chunk_array(array, num_rays): - shape = (-1,) + array.shape[2:] - return array.reshape(shape)[:num_rays] - - -def trace_bounces(rays, compiled, shadows, bounces, face_chunk): - state = initialize_state(rays) - bounce = paz.lock(bounce_step, compiled, shadows, face_chunk) - for step_arg in range(bounces): - state = bounce(state, step_arg) - return state.hit_mask, state.depth, state.color - - -def initialize_state(rays): - num_rays = rays[0].shape[0] - color = jp.zeros((num_rays, 3)) - depth = jp.full((num_rays,), paz.graphics.FARAWAY) - hit_mask = jp.zeros((num_rays,), dtype=bool) - throughput = jp.ones((num_rays, 3)) - active_mask = jp.ones((num_rays,), dtype=bool) - refractive_index = jp.ones((num_rays,)) - 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) - 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) - 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 - else: - triangle_hit = build_triangle_hit(compiled, rays, face_chunk) - return triangle_hit - - -def build_triangle_hit(compiled, rays, face_chunk): - triangles = compiled.triangles - args = triangles, rays, face_chunk - result = paz.graphics.mesh.intersect_triangles(*args) - hit_mask, depth, points, normals, eyes, face_index, u, v = result - primitive = triangles.primitive_index[face_index] - hit_mask = jp.logical_and(hit_mask, compiled.triangle_mask[primitive]) - depth = jp.where(hit_mask, depth, paz.graphics.FARAWAY) - albedo_args = triangles, face_index, u, v - albedo = paz.graphics.albedo.compute_triangle_albedo(*albedo_args) - args = hit_mask, depth, points, normals, eyes, albedo, primitive - return TriangleHit(*args) - - -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 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 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(compiled, closest, surfaces, shadows) - else: - args = compiled, closest, surfaces, shadows, triangle_hit - colors = blend_hit_colors(*args) - return colors - - -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) - indices = jp.minimum(closest.primitive_index, num_shapes - 1) - surfaces = slice_surfaces(surfaces, 0, num_shapes) - if shadows: - colors = color_with_shadows(compiled, closest, surfaces, indices) - else: - colors = color_without_shadow(compiled, surfaces, indices) - return colors - - -def compute_triangle_colors(compiled, triangle_hit): - materials = compiled.triangles.materials - material = gather_triangle_material(materials, triangle_hit.primitive) - shader = select_shader(materials) - colors = jp.zeros_like(triangle_hit.albedo) - for light in compiled.lights: - args = triangle_hit.albedo, material, triangle_hit.points - args += triangle_hit.normals, triangle_hit.eyes, light - colors = colors + shader.compute_colors(*args) - return colors - - -def gather_triangle_material(materials, primitive): - material = jax.tree.map(lambda field: field[primitive], materials) - return jax.tree.map(expand_scalar_field, material) - - -def expand_scalar_field(field): - if field.ndim == 1: - field = jp.expand_dims(field, -1) - return field - - -def select_shader(material): - if isinstance(material, paz.graphics.CookTorranceMaterial): - shader = paz.graphics.cook_torrance - else: - shader = paz.graphics.phong - return shader - - -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) - - -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 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 compute_group_albedo(group, points): - compute_albedo = paz.graphics.albedo.compute_shape_albedo - return jax.vmap(compute_albedo)(group, group.material, points) - - -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 scan_light_step(colors, light, compiled, closest, surfaces, indices): - args = compiled, closest, surfaces, indices, light - return colors + compute_light_colors(*args), None - - -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_light_directions(light, points): - vector = light.position - points - norm = paz.algebra.compute_norms(vector, 1) - return vector / norm, jp.squeeze(norm, axis=1) - - -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 compute_shadow_ray_origins(points, normals): - over_point, _ = compute_surface_points(points, normals) - return over_point - - -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): - 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_transparencies(shapes): - return jp.array([shape.material.transparency for shape in shapes]) - - -def hide_transparent_shapes(shadow_masks, is_transparent): - return jp.where(jp.expand_dims(is_transparent, 1), False, shadow_masks) - - -def hide_non_casting_shapes(shadow_masks, shadow_mask): - return jp.where(jp.expand_dims(shadow_mask, 1), shadow_masks, False) - - -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 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) - 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): - closest_depths = jp.where(hit_masks, depths, paz.graphics.FARAWAY) - closest_depths = jp.min(closest_depths, axis=0) - scene_hit_mask = compute_scene_hit_mask(hit_masks) - blockers = closest_depths <= light_lengths - blocker_mask = jp.logical_and(scene_hit_mask, blockers) - difference = light_lengths - closest_depths - occlusion = jax.nn.sigmoid(slope * difference) - return jp.where(blocker_mask, occlusion, 0.0) - - -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) - 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, 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) - 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) - refractive_indices.append(shape.material.refractive_index) - if compiled.triangles is not None: - reflectivities.append(0.0) - transparencies.append(0.0) - refractive_indices.append(1.0) - return reflectivities, transparencies, refractive_indices - - -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 * 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): - is_inside = jp.sum(normal * eye, axis=-1) < 0.0 - return jp.where(jp.expand_dims(is_inside, -1), -normal, normal), is_inside - - -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 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 - - -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): - args = eye, normal, n_ratio - cos_incident, _, cos_transmit = compute_transmission_cosines(*args) - inside_vector = -eye * jp.expand_dims(n_ratio, -1) - up_weight = n_ratio * cos_incident - cos_transmit - return jp.expand_dims(up_weight, -1) * normal + inside_vector - - -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(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, 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(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, n_2, is_transparent, reflectance): - update_mask = is_transparent & (reflectance < 1.0) - return jp.where(update_mask, n_2, state.refractive_index) - - -def replace_bounce_state(state, throughput, active_mask, index, rays): - state = state._replace(throughput=throughput, active_mask=active_mask) - return state._replace(refractive_index=index, rays=rays) diff --git a/paz/graphics/renderer/__init__.py b/paz/graphics/renderer/__init__.py new file mode 100644 index 000000000..65efde72d --- /dev/null +++ b/paz/graphics/renderer/__init__.py @@ -0,0 +1,48 @@ +import jax.numpy as jp + +import paz + +from paz.graphics.renderer import rays, tiling + +# TODO retune: measured best under 6k faces, but 2048 is 1.75x faster +# at 82k, so this is wrong for the room-scale meshes ahead. +FACE_CHUNK_SIZE = 128 + + +def render( + shape, y_FOV, pose, scene, mask, lights, tiles, chunk_size, + shadows=False, shadow_mask=None, num_bounces=1, + face_chunk_size=FACE_CHUNK_SIZE, +): + compiled = paz.graphics.scene.compile(scene, lights, mask, shadow_mask) + ray_args = compiled, shadows, num_bounces, face_chunk_size, chunk_size + render_rays = paz.lock(rays.render_chunks, *ray_args) + step_args = shape, y_FOV, pose, tiles, render_rays + tile_step = paz.lock(tiling.render_step, *step_args) + images, depths = tiling.scan(shape, tiles, tile_step) + image = tiling.assemble(shape, tiles, images) + depth = tiling.assemble(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, + face_chunk_size=FACE_CHUNK_SIZE, +): + if num_objects is None: + num_objects = len(scene.nodes) + min_depth, max_depth = depth + masks = [] + for object_arg in range(num_objects): + mask = build_object_mask(len(scene.nodes), object_arg) + args = shape, y_FOV, pose, scene, mask, lights, tiles, 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 build_object_mask(num_nodes, object_arg): + return jp.zeros((num_nodes,), dtype=bool).at[object_arg].set(True) diff --git a/paz/graphics/renderer/intersect.py b/paz/graphics/renderer/intersect.py new file mode 100644 index 000000000..333b11287 --- /dev/null +++ b/paz/graphics/renderer/intersect.py @@ -0,0 +1,92 @@ +from collections import namedtuple + +import jax +import jax.numpy as jp + +import paz + +from paz.graphics.composite import find_closest_intersection_args, take_closest +from paz.graphics.types import Surfaces + +TRIANGLE_HIT_NAMES = "hit_mask depth points normals eyes albedo primitive" +TriangleHit = namedtuple("TriangleHit", TRIANGLE_HIT_NAMES.split()) + + +def build_candidates(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 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 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 find_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_triangle_hit(compiled, rays, face_chunk): + if compiled.triangles is None: + triangle_hit = None + else: + triangle_hit = build_triangle_hit(compiled, rays, face_chunk) + return triangle_hit + + +def build_triangle_hit(compiled, rays, face_chunk): + triangles = compiled.triangles + args = triangles, rays, face_chunk + result = paz.graphics.mesh.intersect_triangles(*args) + hit_mask, depth, points, normals, eyes, face_index, u, v = result + primitive = triangles.primitive_index[face_index] + hit_mask = jp.logical_and(hit_mask, compiled.triangle_mask[primitive]) + depth = jp.where(hit_mask, depth, paz.graphics.FARAWAY) + albedo_args = triangles, face_index, u, v + albedo = paz.graphics.albedo.compute_triangle_albedo(*albedo_args) + args = hit_mask, depth, points, normals, eyes, albedo, primitive + return TriangleHit(*args) + + +def intersect_shadow_groups(shapes, origins, directions): + 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 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 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) diff --git a/paz/graphics/renderer/material.py b/paz/graphics/renderer/material.py new file mode 100644 index 000000000..c426a63e0 --- /dev/null +++ b/paz/graphics/renderer/material.py @@ -0,0 +1,51 @@ +from collections import namedtuple + +import jax.numpy as jp + +MATERIAL_NAMES = "reflectivities transparencies refractive_indices" +HitMaterial = namedtuple("HitMaterial", MATERIAL_NAMES.split()) + + +def compute_material_properties(compiled, closest, triangle_hit): + shape_material = gather_shape_material(compiled, closest.primitive_index) + if triangle_hit is None: + resolved = shape_material + else: + materials = compiled.triangles.materials + triangle = gather_primitive_material(materials, triangle_hit.primitive) + is_triangle = closest.primitive_index == len(compiled.shapes) + resolved = select_material(is_triangle, triangle, shape_material) + return resolved + + +def gather_shape_material(compiled, hit_shape_args): + values = collect_shape_values(compiled) + gathered = tuple(jp.array(row)[hit_shape_args] for row in values) + return HitMaterial(*gathered) + + +def collect_shape_values(compiled): + reflectivities, transparencies, refractive_indices = [], [], [] + for shape in compiled.shapes: + reflectivities.append(shape.material.reflective) + transparencies.append(shape.material.transparency) + refractive_indices.append(shape.material.refractive_index) + if compiled.triangles is not None: + reflectivities.append(0.0) + transparencies.append(0.0) + refractive_indices.append(1.0) + return reflectivities, transparencies, refractive_indices + + +def gather_primitive_material(materials, primitive): + reflectivities = materials.reflective[primitive] + transparencies = materials.transparency[primitive] + refractive_indices = materials.refractive_index[primitive] + return HitMaterial(reflectivities, transparencies, refractive_indices) + + +def select_material(is_triangle, triangle, shape): + fields = [] + for triangle_field, shape_field in zip(triangle, shape): + fields.append(jp.where(is_triangle, triangle_field, shape_field)) + return HitMaterial(*fields) diff --git a/paz/graphics/renderer/optics.py b/paz/graphics/renderer/optics.py new file mode 100644 index 000000000..b4ccf90ec --- /dev/null +++ b/paz/graphics/renderer/optics.py @@ -0,0 +1,77 @@ +import jax.numpy as jp + +import paz + +BOUNCE_ORIGIN_EPSILON = 1e-2 + + +def compute_bounce(directions, closest, refractive_index, material): + args = directions, closest.normal, 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): + is_inside = jp.sum(normal * eye, axis=-1) < 0.0 + return jp.where(jp.expand_dims(is_inside, -1), -normal, normal), is_inside + + +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 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 + + +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): + args = eye, normal, n_ratio + cos_incident, _, cos_transmit = compute_transmission_cosines(*args) + inside_vector = -eye * jp.expand_dims(n_ratio, -1) + up_weight = n_ratio * cos_incident - cos_transmit + return jp.expand_dims(up_weight, -1) * normal + inside_vector + + +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 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) diff --git a/paz/graphics/renderer/rays.py b/paz/graphics/renderer/rays.py new file mode 100644 index 000000000..136ef6e33 --- /dev/null +++ b/paz/graphics/renderer/rays.py @@ -0,0 +1,146 @@ +from collections import namedtuple + +import jax +import jax.numpy as jp + +import paz + +from paz.graphics.renderer import intersect, material, optics, shade + +STATE_NAMES = "color depth hit_mask throughput active_mask " +STATE_NAMES += "refractive_index rays" +RenderState = namedtuple("RenderState", STATE_NAMES.split()) + + +def render(rays, compiled, shadows, num_bounces, face_chunk): + state = initialize_state(rays) + step = paz.lock(bounce, compiled, shadows, face_chunk) + for step_arg in range(num_bounces): + state = step(state, step_arg) + return state.hit_mask, state.depth, state.color + + +def bounce(state, step_arg, compiled, shadows, face_chunk): + hit_args = compiled, state.rays, face_chunk + triangle_hit = intersect.compute_triangle_hit(*hit_args) + candidates = intersect.build_candidates(compiled, state.rays, triangle_hit) + hit_masks, depths, surfaces = candidates + closest = intersect.find_closest(hit_masks, depths, surfaces) + state = update_first_hit(state, closest, step_arg) + state = update_active_mask(state, closest) + color_args = compiled, closest, surfaces, shadows, triangle_hit + colors = shade.compute_hit_colors(*color_args, face_chunk) + return advance(state, compiled, closest, colors, triangle_hit) + + +def advance(state, compiled, closest, hit_colors, triangle_hit): + material_args = compiled, closest, triangle_hit + properties = material.compute_material_properties(*material_args) + color_args = state.color, state.throughput, state.active_mask, hit_colors + color_args += properties.reflectivities, properties.transparencies + color = accumulate_color(*color_args) + bounce_args = state.rays[1], closest, state.refractive_index, properties + new_rays, n_2, reflectance = optics.compute_bounce(*bounce_args) + args = new_rays, n_2, properties, reflectance + return apply_bounce_update(state._replace(color=color), *args) + + +def render_chunks(rays, compiled, shadows, num_bounces, face_chunk, chunk): + ray_chunks = split_ray_chunks(rays, chunk) + step_args = compiled, shadows, num_bounces, face_chunk + chunk_step = paz.lock(render_chunk_step, *step_args) + hit_mask, depth, color = jax.lax.scan(chunk_step, None, ray_chunks)[1] + return flatten_chunk_results(hit_mask, depth, color, len(rays[0])) + + +def render_chunk_step(carry, rays, compiled, shadows, num_bounces, face_chunk): + args = rays, compiled, shadows, num_bounces, face_chunk + return carry, render(*args) + + +def initialize_state(rays): + num_rays = rays[0].shape[0] + color = jp.zeros((num_rays, 3)) + depth = jp.full((num_rays,), paz.graphics.FARAWAY) + hit_mask = jp.zeros((num_rays,), dtype=bool) + throughput = jp.ones((num_rays, 3)) + active_mask = jp.ones((num_rays,), dtype=bool) + refractive_index = jp.ones((num_rays,)) + fields = color, depth, hit_mask, throughput, active_mask + return RenderState(*fields, refractive_index, rays) + + +def update_first_hit(state, closest, step_arg): + if step_arg == 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 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 * hit_colors) + + +def apply_bounce_update(state, new_rays, n_2, properties, reflectance): + is_transparent = properties.transparencies > 0.0 + is_reflective = properties.reflectivities > 0.0 + factor = optics.compute_bounce_factor(properties, 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, 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 update_refractive_index(state, n_2, is_transparent, reflectance): + update_mask = is_transparent & (reflectance < 1.0) + return jp.where(update_mask, n_2, state.refractive_index) + + +def replace_bounce_state(state, throughput, active_mask, index, new_rays): + state = state._replace(throughput=throughput, active_mask=active_mask) + return state._replace(refractive_index=index, rays=new_rays) + + +def split_ray_chunks(rays, chunk_size): + origins, directions = rays + origins = pad_to_chunks(origins, chunk_size) + directions = pad_to_chunks(directions, chunk_size) + num_chunks = origins.shape[0] // chunk_size + shape = num_chunks, chunk_size, 3 + return origins.reshape(shape), directions.reshape(shape) + + +def pad_to_chunks(array, chunk_size): + remainder = array.shape[0] % chunk_size + if remainder == 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): + hit_mask = flatten_chunk_array(hit_mask, num_rays) + depth = flatten_chunk_array(depth, num_rays) + color = flatten_chunk_array(color, num_rays) + return hit_mask, depth, color + + +def flatten_chunk_array(array, num_rays): + shape = (-1,) + array.shape[2:] + return array.reshape(shape)[:num_rays] diff --git a/paz/graphics/renderer/shade.py b/paz/graphics/renderer/shade.py new file mode 100644 index 000000000..f266b6177 --- /dev/null +++ b/paz/graphics/renderer/shade.py @@ -0,0 +1,170 @@ +import jax +import jax.numpy as jp + +import paz + +from paz.graphics.composite import take_closest +from paz.graphics.renderer import shadow +from paz.graphics.renderer.intersect import ( + iterate_shape_groups, + slice_surfaces, +) + + +def compute_hit_colors( + compiled, closest, surfaces, shadows, triangle_hit, face_chunk +): + if len(compiled.shapes) == 0: + triangle_args = compiled, triangle_hit, shadows, face_chunk + colors = compute_triangle_colors(*triangle_args) + elif triangle_hit is None: + shape_args = compiled, closest, surfaces, shadows, face_chunk + colors = compute_shape_colors(*shape_args) + else: + args = compiled, closest, surfaces, shadows, triangle_hit + colors = blend_hit_colors(*args, face_chunk) + return colors + + +def blend_hit_colors( + compiled, closest, surfaces, shadows, triangle_hit, face_chunk +): + # 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_args = compiled, closest, surfaces, shadows, face_chunk + shape_colors = compute_shape_colors(*shape_args) + triangle_args = compiled, triangle_hit, shadows, face_chunk + triangle_colors = compute_triangle_colors(*triangle_args) + 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, face_chunk): + num_shapes = len(compiled.shapes) + indices = jp.minimum(closest.primitive_index, num_shapes - 1) + surfaces = slice_surfaces(surfaces, 0, num_shapes) + if shadows: + args = compiled, closest, surfaces, indices, face_chunk + colors = color_with_shadows(*args) + else: + colors = color_without_shadow(compiled, surfaces, indices) + return colors + + +def compute_triangle_colors(compiled, triangle_hit, shadows, face_chunk): + if shadows: + args = compiled, triangle_hit, face_chunk + colors = color_triangles_with_shadows(*args) + else: + colors = color_triangles(compiled, triangle_hit) + return colors + + +def color_triangles(compiled, triangle_hit): + shader, material = select_triangle_shader(compiled, triangle_hit) + colors = jp.zeros_like(triangle_hit.albedo) + for light in compiled.lights: + args = build_triangle_shader_args(triangle_hit, material, light) + colors = colors + shader.compute_colors(*args) + return colors + + +def color_triangles_with_shadows(compiled, triangle_hit, face_chunk): + shader, material = select_triangle_shader(compiled, triangle_hit) + receiver = shadow.build_triangle_receiver(triangle_hit) + colors = jp.zeros_like(triangle_hit.albedo) + for light in compiled.lights: + occlusion_args = compiled, receiver, light, face_chunk + is_shadow = shadow.compute_occlusion(*occlusion_args) + args = build_triangle_shader_args(triangle_hit, material, light) + colors = colors + shader.compute_colors_with_shadow(*args, is_shadow) + return colors + + +def select_triangle_shader(compiled, triangle_hit): + materials = compiled.triangles.materials + material = gather_triangle_material(materials, triangle_hit.primitive) + return select_shader(materials), material + + +def build_triangle_shader_args(triangle_hit, material, light): + args = triangle_hit.albedo, material, triangle_hit.points + return args + (triangle_hit.normals, triangle_hit.eyes, light) + + +def gather_triangle_material(materials, primitive): + material = jax.tree.map(lambda field: field[primitive], materials) + return jax.tree.map(expand_scalar_field, material) + + +def expand_scalar_field(field): + if field.ndim == 1: + field = jp.expand_dims(field, -1) + return field + + +def select_shader(material): + if isinstance(material, paz.graphics.CookTorranceMaterial): + shader = paz.graphics.cook_torrance + else: + shader = paz.graphics.phong + return shader + + +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) + + +def compute_group_albedo(group, points): + compute_albedo = paz.graphics.albedo.compute_shape_albedo + return jax.vmap(compute_albedo)(group, group.material, points) + + +def color_with_shadows(compiled, closest, surfaces, indices, face_chunk): + colors = jp.zeros((len(surfaces.points[0]), 3)) + lights = paz.graphics.shapes.merge(*compiled.lights) + step_args = compiled, closest, surfaces, indices, face_chunk + body = paz.lock(scan_light_step, *step_args) + return jax.lax.scan(body, colors, lights)[0] + + +def scan_light_step( + colors, light, compiled, closest, surfaces, indices, face_chunk +): + args = compiled, closest, surfaces, indices, light, face_chunk + return colors + compute_light_colors(*args), None + + +def compute_light_colors( + compiled, closest, surfaces, indices, light, face_chunk +): + receiver = shadow.build_shape_receiver(closest, indices) + is_shadow = shadow.compute_occlusion(compiled, receiver, light, face_chunk) + color_args = compiled.shapes, light, surfaces, is_shadow + return take_closest(compute_shadowed_colors(*color_args), indices) + + +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) + args = albedo, group.material, *group_surfaces, light, is_shadow + colors.append(color(*args)) + return jp.concatenate(colors, axis=0) diff --git a/paz/graphics/renderer/shadow.py b/paz/graphics/renderer/shadow.py new file mode 100644 index 000000000..4e0fe24c5 --- /dev/null +++ b/paz/graphics/renderer/shadow.py @@ -0,0 +1,155 @@ +from collections import namedtuple + +import jax +import jax.numpy as jp + +import paz + +from paz.graphics.composite import compute_scene_hit_mask +from paz.graphics.renderer.intersect import intersect_shadow_groups + +SHADOW_ORIGIN_EPSILON = 1e-5 +SHADOW_SELF_HIT_EPSILON = 1e-5 +# A shadow ray leaving a mesh surface can re-hit its own triangle at a +# grazing angle. Shapes reject that by identity; triangles have no such +# index, so they need a distance that clears float error at scene scale. +TRIANGLE_SELF_HIT_EPSILON = 1e-3 +NO_SHAPE = -1 + +Receiver = namedtuple("Receiver", ["points", "normals", "indices"]) + + +def build_shape_receiver(closest, indices): + return Receiver(closest.point, closest.normal, indices) + + +def build_triangle_receiver(triangle_hit): + indices = jp.full(len(triangle_hit.points), NO_SHAPE) + return Receiver(triangle_hit.points, triangle_hit.normals, indices) + + +def compute_occlusion(compiled, receiver, light, face_chunk): + directions, distance = compute_light_directions(light, receiver.points) + origins = compute_shadow_ray_origins(receiver.points, receiver.normals) + rows = [] + if len(compiled.shapes) > 0: + shape_args = compiled, receiver, origins, directions + rows.append(compute_shape_blockers(*shape_args)) + if compiled.triangles is not None: + blocker_args = compiled, origins, directions, face_chunk + rows.append(compute_triangle_blockers(*blocker_args)) + masks = jp.concatenate([row[0] for row in rows], axis=0) + depths = jp.concatenate([row[1] for row in rows], axis=0) + return compute_soft_occlusion(masks, depths, distance) + + +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 compute_shape_blockers(compiled, receiver, origins, directions): + shadow_args = compiled.shapes, origins, directions + intersections = intersect_shadow_groups(*shadow_args) + hit_masks, depths, _, _, _, casters = intersections + masks = resolve_shadow_masks(compiled, hit_masks) + depth_args = masks, depths, casters, receiver.indices + depth_args += receiver.normals, directions + return select_shadow_depths(*depth_args) + + +def compute_triangle_blockers(compiled, origins, directions, face_chunk): + triangles = compiled.triangles + args = triangles.vertices, triangles.faces, (origins, directions) + result = paz.graphics.mesh.intersect_chunked(*args, face_chunk) + hit_mask, depth, _, _, face_index = result + primitive = triangles.primitive_index[face_index] + hit_mask = jp.logical_and(hit_mask, compiled.triangle_mask[primitive]) + hit_mask = jp.logical_and(hit_mask, depth > TRIANGLE_SELF_HIT_EPSILON) + hit_mask = hide_non_casting_triangles(compiled, hit_mask, primitive) + depth = jp.where(hit_mask, depth, paz.graphics.FARAWAY) + return jp.expand_dims(hit_mask, 0), jp.expand_dims(depth, 0) + + +def hide_non_casting_triangles(compiled, hit_mask, primitive): + if compiled.triangle_shadow_mask is None: + casting = hit_mask + else: + casting = compiled.triangle_shadow_mask[primitive] + casting = jp.logical_and(hit_mask, casting) + return casting + + +def compute_shadow_ray_origins(points, normals): + over_point, _ = compute_surface_points(points, normals) + return over_point + + +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 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_transparencies(shapes): + return jp.array([shape.material.transparency for shape in shapes]) + + +def hide_transparent_shapes(shadow_masks, is_transparent): + return jp.where(jp.expand_dims(is_transparent, 1), False, shadow_masks) + + +def hide_non_casting_shapes(shadow_masks, shadow_mask): + return jp.where(jp.expand_dims(shadow_mask, 1), shadow_masks, False) + + +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 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) + 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): + closest_depths = jp.where(hit_masks, depths, paz.graphics.FARAWAY) + closest_depths = jp.min(closest_depths, axis=0) + scene_hit_mask = compute_scene_hit_mask(hit_masks) + blockers = closest_depths <= light_lengths + blocker_mask = jp.logical_and(scene_hit_mask, blockers) + difference = light_lengths - closest_depths + occlusion = jax.nn.sigmoid(slope * difference) + return jp.where(blocker_mask, occlusion, 0.0) diff --git a/paz/graphics/renderer/tiling.py b/paz/graphics/renderer/tiling.py new file mode 100644 index 000000000..54ed1b127 --- /dev/null +++ b/paz/graphics/renderer/tiling.py @@ -0,0 +1,34 @@ +import jax +import jax.numpy as jp + +import paz + +from paz.graphics.composite import postprocess + + +def scan(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) + return jax.lax.scan(tile_step, None, coordinates)[1] + + +def render_step(carry, tile_arg, shape, y_FOV, pose, tiles, render_rays): + 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 = render_rays(rays) + tile_H, tile_W = H // H_tiles, W // W_tiles + 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(shape, tiles, images): + H, W = shape + H_tiles, W_tiles = tiles + return paz.graphics.mesh.assemble(H, W, H_tiles, W_tiles, images) diff --git a/paz/graphics/renderer_test.py b/paz/graphics/renderer_test.py index 35422341c..9e36923a3 100644 --- a/paz/graphics/renderer_test.py +++ b/paz/graphics/renderer_test.py @@ -115,15 +115,15 @@ def compute_selected_shadow_depths(camera_pose, image_shape=(120, 160)): rays = paz.graphics.camera.build_rays(image_shape, jp.pi / 3.0, camera_pose) compiled = paz.graphics.scene.compile(scene, lights, None) shapes, mask, lights = compiled.shapes, compiled.mask, compiled.lights - intersections = renderer.intersect(compiled, rays, None) - closest = renderer.gather_closest(*intersections) + intersections = renderer.intersect.build_candidates(compiled, rays, None) + closest = renderer.intersect.find_closest(*intersections) vector = lights[0].position - closest.point distance = jp.squeeze(paz.algebra.compute_norms(vector, 1), axis=1) light_directions = vector / jp.expand_dims(distance, 1) - shadow_ray_origins = renderer.compute_shadow_ray_origins( + shadow_ray_origins = renderer.shadow.compute_shadow_ray_origins( closest.point, closest.normal ) - intersections = renderer.intersect_shadow_groups( + intersections = renderer.intersect.intersect_shadow_groups( shapes, shadow_ray_origins, light_directions ) hit_masks, depths, _, _, _, shape_indices = intersections @@ -132,7 +132,7 @@ def compute_selected_shadow_depths(camera_pose, image_shape=(120, 160)): shadow_masks = jp.where( jp.expand_dims(transparencies > 0.0, 1), False, shadow_masks ) - shadow_masks, depths = renderer.select_shadow_depths( + shadow_masks, depths = renderer.shadow.select_shadow_depths( shadow_masks, depths, shape_indices, @@ -169,7 +169,7 @@ def test_compute_soft_occlusion(): ) rows = [[True, True, True, True], [False, False, True, False]] hit_masks = jp.array(rows) - result = renderer.compute_soft_occlusion( + result = renderer.shadow.compute_soft_occlusion( hit_masks, depths, light_lengths, slope=10.0 ) assert float(result[1]) > 0.9 @@ -183,7 +183,7 @@ def test_compute_new_rays_reflection_is_normalized(): eye = jp.array([[0.0, 1.0, -1.0]]) point = jp.array([[0.0, 0.0, 0.0]]) transparencies = jp.array([0.0]) - _, direction = renderer.compute_new_rays( + _, direction = renderer.optics.compute_new_rays( normal, eye, jp.array([1.0]), point, transparencies ) norm = jp.linalg.norm(direction, axis=-1) @@ -195,7 +195,7 @@ def test_compute_new_rays_refraction_is_normalized(): eye = jp.array([[0.0, 0.0, -1.0]]) point = jp.array([[0.0, 0.0, 0.0]]) transparencies = jp.array([1.0]) - _, direction = renderer.compute_new_rays( + _, direction = renderer.optics.compute_new_rays( normal, eye, jp.array([1.0 / 1.5]), point, transparencies ) norm = jp.linalg.norm(direction, axis=-1) @@ -205,8 +205,9 @@ def test_compute_new_rays_refraction_is_normalized(): def test_compute_surface_points_offset_hit(): point = jp.array([[0.0, 0.0, 0.0]]) normal = jp.array([[0.0, 0.0, -1.0]]) - over_point, under_point = renderer.compute_surface_points(point, normal) - assert over_point[0, 2] < -(renderer.SHADOW_ORIGIN_EPSILON / 2.0) + surface_points = renderer.shadow.compute_surface_points(point, normal) + over_point, under_point = surface_points + assert over_point[0, 2] < -(renderer.shadow.SHADOW_ORIGIN_EPSILON / 2.0) assert point[0, 2] > over_point[0, 2] assert under_point[0, 2] > 0.0 @@ -223,7 +224,7 @@ def test_select_shadow_depths_discard_front_side_same_shape_hits(): receiver_indices = jp.array([0]) receiver_normals = jp.array([[0.0, 1.0, 0.0]]) light_directions = jp.array([[0.0, 1.0, 0.0]]) - hit_masks, depths = renderer.select_shadow_depths( + hit_masks, depths = renderer.shadow.select_shadow_depths( hit_masks, depths, shape_indices, @@ -249,7 +250,7 @@ def test_select_shadow_depths_keep_back_side_second_root(): receiver_indices = jp.array([0]) receiver_normals = jp.array([[0.0, 1.0, 0.0]]) light_directions = jp.array([[0.0, -1.0, 0.0]]) - hit_masks, depths = renderer.select_shadow_depths( + hit_masks, depths = renderer.shadow.select_shadow_depths( hit_masks, depths, shape_indices, @@ -258,7 +259,7 @@ def test_select_shadow_depths_keep_back_side_second_root(): light_directions, ) args = hit_masks, depths, jp.array([0.01]) - result = renderer.compute_soft_occlusion(*args) + result = renderer.shadow.compute_soft_occlusion(*args) assert bool(hit_masks[0, 0]) assert float(depths[0, 0]) == pytest.approx(0.2) assert bool(hit_masks[1, 0]) @@ -295,7 +296,7 @@ def test_select_colors(): def test_initialize_render_state(): num_rays = 10 rays = (jp.zeros((num_rays, 3)), jp.ones((num_rays, 3))) - state = renderer.initialize_state(rays) + state = renderer.rays.initialize_state(rays) assert state.color.shape == (num_rays, 3) assert state.throughput.shape == (num_rays, 3) assert jp.all(state.refractive_index == 1.0) @@ -310,7 +311,7 @@ def test_find_closest_intersection_args(): assert jp.array_equal(indices, jp.array([0, 1])) -def test_compute_material_properties(): +def test_gather_shape_material(): mat1 = Material(reflective=0.5) mat2 = Material(transparency=0.8) shape1 = Sphere(material=mat1) @@ -318,12 +319,26 @@ def test_compute_material_properties(): scene = Scene([shape1, shape2]) compiled = paz.graphics.scene.compile(scene, [], None) indices = jp.array([0, 1]) - material = renderer.compute_material_properties(compiled, indices) + material = renderer.material.gather_shape_material(compiled, indices) assert material.reflectivities[0] == 0.5 assert material.transparencies[1] == 0.8 assert material.refractive_indices[0] == 1.0 +def test_gather_primitive_material_reads_mesh_fields(): + materials = Material( + reflective=jp.array([0.0, 0.6]), + transparency=jp.array([0.0, 0.3]), + refractive_index=jp.array([1.0, 1.5]), + ) + primitive = jp.array([1, 0, 1]) + gather = renderer.material.gather_primitive_material + material = gather(materials, primitive) + expected_indices = jp.array([1.5, 1.0, 1.5]) + assert jp.array_equal(material.reflectivities, jp.array([0.6, 0.0, 0.6])) + assert jp.array_equal(material.refractive_indices, expected_indices) + + def test_accumulate_color(): num_rays = 2 colors = jp.zeros((num_rays, 3)) @@ -332,7 +347,7 @@ def test_accumulate_color(): intersected_colors = jp.ones((num_rays, 3)) reflectivities = jp.zeros((num_rays,)) transparencies = jp.zeros((num_rays,)) - result = renderer.accumulate_color( + result = renderer.rays.accumulate_color( colors, throughput, active_mask, @@ -349,7 +364,7 @@ def test_compute_shadow_ray_origins_avoid_lit_side_self_hit(): normals = jp.array([[0.0, 1.0, 0.0]]) light_position = jp.array([[0.0, 3.0, -3.0]]) directions = paz.algebra.normalize(light_position - points) - origins = renderer.compute_shadow_ray_origins(points, normals) + origins = renderer.shadow.compute_shadow_ray_origins(points, normals) hit_mask, _, _ = intersect_canonical_sphere(origins, directions) assert not bool(hit_mask[0]) @@ -721,8 +736,9 @@ def test_saved_pose_sphere_self_shadow_keeps_later_roots(): ) args = depths, shape_indices, receiver_indices, 0 sphere_depths = take_shape_depths(*args) + epsilon = renderer.shadow.SHADOW_SELF_HIT_EPSILON assert int(jp.sum(sphere_depths < 1e-2)) > 0 - assert float(jp.min(sphere_depths)) > renderer.SHADOW_SELF_HIT_EPSILON + assert float(jp.min(sphere_depths)) > epsilon def test_saved_pose_floor_self_hits_stay_filtered(): diff --git a/paz/graphics/scene.py b/paz/graphics/scene.py index 167f5d81a..8c76763c4 100644 --- a/paz/graphics/scene.py +++ b/paz/graphics/scene.py @@ -189,11 +189,13 @@ def compile(scene, lights, mask, shadow_mask=None): if shadow_mask is not None: shadow_mask = prepare_mask(shadow_mask, len(flat_scene), scene) + mesh_args = select_meshes(flat_scene, mask, shadow_mask) + meshes, triangle_mask, triangle_shadow_mask = mesh_args shape_args = select_shapes(flat_scene, mask, shadow_mask) shapes, shape_mask, shadow_mask = sort_by_group(*shape_args) - meshes, triangle_mask = select_meshes(flat_scene, mask) args = shapes, build_triangles(meshes), lights, shape_mask - return paz.graphics.CompiledScene(*args, shadow_mask, triangle_mask) + args += shadow_mask, triangle_mask, triangle_shadow_mask + return paz.graphics.CompiledScene(*args) def select_shapes(flat_scene, mask, shadow_mask): @@ -205,10 +207,15 @@ def select_shapes(flat_scene, mask, shadow_mask): return shapes, mask[indices], shadow_mask -def select_meshes(flat_scene, mask): +def select_meshes(flat_scene, mask, shadow_mask): args = [arg for arg, node in enumerate(flat_scene) if not is_shape(node)] + indices = jp.array(args, dtype=jp.int32) meshes = [flat_scene[arg] for arg in args] - return meshes, mask[jp.array(args, dtype=jp.int32)] + if shadow_mask is None: + mesh_shadow_mask = None + else: + mesh_shadow_mask = shadow_mask[indices] + return meshes, mask[indices], mesh_shadow_mask def is_shape(node): diff --git a/paz/graphics/types.py b/paz/graphics/types.py index e7fc4da51..bb2ab18e3 100644 --- a/paz/graphics/types.py +++ b/paz/graphics/types.py @@ -17,6 +17,8 @@ HIT_NAMES = "hit_mask depth point normal eye primitive_index" Hit = namedtuple("Hit", HIT_NAMES.split()) +Surfaces = namedtuple("Surfaces", ["points", "normals", "eyes"]) + MESH_NAMES = "vertices vertex_colors transform material faces edges " MESH_NAMES += "pattern vertex_uvs" Mesh = namedtuple("Mesh", MESH_NAMES.split(), defaults=(None, None)) @@ -25,7 +27,8 @@ TRIANGLE_NAMES += "materials patterns" Triangles = namedtuple("Triangles", TRIANGLE_NAMES.split()) -COMPILED_NAMES = "shapes triangles lights mask shadow_mask triangle_mask" +COMPILED_NAMES = "shapes triangles lights mask shadow_mask triangle_mask " +COMPILED_NAMES += "triangle_shadow_mask" CompiledScene = namedtuple("CompiledScene", COMPILED_NAMES.split())