| 38 | points = world_to_pixel(points) |
| 39 | |
| 40 | def color_function( |
| 41 | xy: Float[Tensor, "point 2"], |
| 42 | ) -> Float[Tensor, "point 4"]: |
| 43 | # Define a vector between the start and end points. |
| 44 | delta = xy[:, None] - points[None] |
| 45 | delta_norm = delta.norm(dim=-1) |
| 46 | mask = (delta_norm >= inner_radius[None]) & (delta_norm <= radius[None]) |
| 47 | |
| 48 | # Determine the sample's color. |
| 49 | selectable_color = color.broadcast_to((num_points, 3)) |
| 50 | arrangement = mask * torch.arange(num_points, device=device) |
| 51 | top_color = selectable_color.gather( |
| 52 | dim=0, |
| 53 | index=repeat(arrangement.argmax(dim=1), "s -> s c", c=3), |
| 54 | ) |
| 55 | rgba = torch.cat((top_color, mask.any(dim=1).float()[:, None]), dim=-1) |
| 56 | |
| 57 | return rgba |
| 58 | |
| 59 | return render_over_image(image, color_function, device, num_passes=num_msaa_passes) |