(
image: Float[Tensor, "3 height width"] | Float[Tensor, "4 height width"],
start: Vector,
end: Vector,
color: Vector,
width: Scalar,
cap: Literal["butt", "round", "square"] = "round",
num_msaa_passes: int = 1,
x_range: Optional[Pair] = None,
y_range: Optional[Pair] = None,
)
| 11 | |
| 12 | |
| 13 | def draw_lines( |
| 14 | image: Float[Tensor, "3 height width"] | Float[Tensor, "4 height width"], |
| 15 | start: Vector, |
| 16 | end: Vector, |
| 17 | color: Vector, |
| 18 | width: Scalar, |
| 19 | cap: Literal["butt", "round", "square"] = "round", |
| 20 | num_msaa_passes: int = 1, |
| 21 | x_range: Optional[Pair] = None, |
| 22 | y_range: Optional[Pair] = None, |
| 23 | ) -> Float[Tensor, "3 height width"] | Float[Tensor, "4 height width"]: |
| 24 | device = image.device |
| 25 | start = sanitize_vector(start, 2, device) |
| 26 | end = sanitize_vector(end, 2, device) |
| 27 | color = sanitize_vector(color, 3, device) |
| 28 | width = sanitize_scalar(width, device) |
| 29 | (num_lines,) = torch.broadcast_shapes( |
| 30 | start.shape[0], |
| 31 | end.shape[0], |
| 32 | color.shape[0], |
| 33 | width.shape, |
| 34 | ) |
| 35 | |
| 36 | # Convert world-space points to pixel space. |
| 37 | _, h, w = image.shape |
| 38 | world_to_pixel, _ = generate_conversions((h, w), device, x_range, y_range) |
| 39 | start = world_to_pixel(start) |
| 40 | end = world_to_pixel(end) |
| 41 | |
| 42 | def color_function( |
| 43 | xy: Float[Tensor, "point 2"], |
| 44 | ) -> Float[Tensor, "point 4"]: |
| 45 | # Define a vector between the start and end points. |
| 46 | delta = end - start |
| 47 | delta_norm = delta.norm(dim=-1, keepdim=True) |
| 48 | u_delta = delta / delta_norm |
| 49 | |
| 50 | # Define a vector between each sample and the start point. |
| 51 | indicator = xy - start[:, None] |
| 52 | |
| 53 | # Determine whether each sample is inside the line in the parallel direction. |
| 54 | extra = 0.5 * width[:, None] if cap == "square" else 0 |
| 55 | parallel = einsum(u_delta, indicator, "l xy, l s xy -> l s") |
| 56 | parallel_inside_line = (parallel <= delta_norm + extra) & (parallel > -extra) |
| 57 | |
| 58 | # Determine whether each sample is inside the line perpendicularly. |
| 59 | perpendicular = indicator - parallel[..., None] * u_delta[:, None] |
| 60 | perpendicular_inside_line = perpendicular.norm(dim=-1) < 0.5 * width[:, None] |
| 61 | |
| 62 | inside_line = parallel_inside_line & perpendicular_inside_line |
| 63 | |
| 64 | # Compute round caps. |
| 65 | if cap == "round": |
| 66 | near_start = indicator.norm(dim=-1) < 0.5 * width[:, None] |
| 67 | inside_line |= near_start |
| 68 | end_indicator = indicator = xy - end[:, None] |
| 69 | near_end = end_indicator.norm(dim=-1) < 0.5 * width[:, None] |
| 70 | inside_line |= near_end |
no test coverage detected