Visualize the "cell" each coordinate lives in, and highlight its edges.
(coords, acc, percentile=99.)
| 131 | |
| 132 | |
| 133 | def visualize_coord_fix(coords, acc, percentile=99.): |
| 134 | """Visualize the "cell" each coordinate lives in, and highlight its edges.""" |
| 135 | |
| 136 | # Round towards zero. |
| 137 | coords_fix = jnp.int32(jnp.fix(coords)) |
| 138 | |
| 139 | # A very hacky plus-shaped edge detector. |
| 140 | coords_fix_pad = jnp.pad(coords_fix, [(1, 1), (1, 1), (0, 0)], 'edge') |
| 141 | mask = ((coords_fix == coords_fix_pad[2:, 1:-1, :]) & |
| 142 | (coords_fix == coords_fix_pad[:-2, 1:-1, :]) |
| 143 | & (coords_fix == coords_fix_pad[1:-1, 2:, :]) |
| 144 | & (coords_fix == coords_fix_pad[1:-1, :-2, :])) |
| 145 | |
| 146 | # Scale according to `acc` and clip to lie in [-1, 1]. |
| 147 | max_val = jnp.maximum( |
| 148 | 1, |
| 149 | math.weighted_percentile( |
| 150 | jnp.max(jnp.abs(coords_fix), axis=2), acc, percentile)) |
| 151 | coords_fix_unit = jnp.clip(coords_fix / max_val, -1, 1) |
| 152 | |
| 153 | # The [-1, 1] center cube is gray, and every other integer boundary gets |
| 154 | # colored with xyz \propto rgb - gray. Edge pixels are highlighted. |
| 155 | return matte( |
| 156 | jnp.where(mask, (coords_fix_unit + 1) / 2, 1 - jnp.abs(coords_fix_unit)), |
| 157 | acc) |
| 158 | |
| 159 | |
| 160 | def visualize_coord_norm(coords, acc, percentile=99.): |