| 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 |
| 71 | |
| 72 | # Determine the sample's color. |
| 73 | selectable_color = color.broadcast_to((num_lines, 3)) |
| 74 | arrangement = inside_line * torch.arange(num_lines, device=device)[:, None] |
| 75 | top_color = selectable_color.gather( |
| 76 | dim=0, |
| 77 | index=repeat(arrangement.argmax(dim=0), "s -> s c", c=3), |
| 78 | ) |
| 79 | rgba = torch.cat((top_color, inside_line.any(dim=0).float()[:, None]), dim=-1) |
| 80 | |
| 81 | return rgba |
| 82 | |
| 83 | return render_over_image(image, color_function, device, num_passes=num_msaa_passes) |