(
image: Float[Tensor, "3 height width"] | Float[Tensor, "4 height width"],
points: Vector,
color: Vector = [1, 1, 1],
radius: Scalar = 1,
inner_radius: Scalar = 0,
num_msaa_passes: int = 1,
x_range: Optional[Pair] = None,
y_range: Optional[Pair] = None,
)
| 11 | |
| 12 | |
| 13 | def draw_points( |
| 14 | image: Float[Tensor, "3 height width"] | Float[Tensor, "4 height width"], |
| 15 | points: Vector, |
| 16 | color: Vector = [1, 1, 1], |
| 17 | radius: Scalar = 1, |
| 18 | inner_radius: Scalar = 0, |
| 19 | num_msaa_passes: int = 1, |
| 20 | x_range: Optional[Pair] = None, |
| 21 | y_range: Optional[Pair] = None, |
| 22 | ) -> Float[Tensor, "3 height width"] | Float[Tensor, "4 height width"]: |
| 23 | device = image.device |
| 24 | points = sanitize_vector(points, 2, device) |
| 25 | color = sanitize_vector(color, 3, device) |
| 26 | radius = sanitize_scalar(radius, device) |
| 27 | inner_radius = sanitize_scalar(inner_radius, device) |
| 28 | (num_points,) = torch.broadcast_shapes( |
| 29 | points.shape[0], |
| 30 | color.shape[0], |
| 31 | radius.shape, |
| 32 | inner_radius.shape, |
| 33 | ) |
| 34 | |
| 35 | # Convert world-space points to pixel space. |
| 36 | _, h, w = image.shape |
| 37 | world_to_pixel, _ = generate_conversions((h, w), device, x_range, y_range) |
| 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) |
no test coverage detected