File size: 3,065 Bytes
2568013 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
from typing import Literal, Optional
import torch
from einops import einsum, repeat
from jaxtyping import Float
from torch import Tensor
from .coordinate_conversion import generate_conversions
from .rendering import render_over_image
from .types import Pair, Scalar, Vector, sanitize_scalar, sanitize_vector
def draw_lines(
image: Float[Tensor, "3 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,
) -> Float[Tensor, "3 height width"]:
device = image.device
start = sanitize_vector(start, 2, device)
end = sanitize_vector(end, 2, device)
color = sanitize_vector(color, 3, device)
width = sanitize_scalar(width, device)
(num_lines,) = torch.broadcast_shapes(
start.shape[0],
end.shape[0],
color.shape[0],
width.shape,
)
# Convert world-space points to pixel space.
_, h, w = image.shape
world_to_pixel, _ = generate_conversions((h, w), device, x_range, y_range)
start = world_to_pixel(start)
end = world_to_pixel(end)
def color_function(
xy: Float[Tensor, "point 2"],
) -> Float[Tensor, "point 4"]:
# Define a vector between the start and end points.
delta = end - start
delta_norm = delta.norm(dim=-1, keepdim=True)
u_delta = delta / delta_norm
# Define a vector between each sample and the start point.
indicator = xy - start[:, None]
# Determine whether each sample is inside the line in the parallel direction.
extra = 0.5 * width[:, None] if cap == "square" else 0
parallel = einsum(u_delta, indicator, "l xy, l s xy -> l s")
parallel_inside_line = (parallel <= delta_norm + extra) & (parallel > -extra)
# Determine whether each sample is inside the line perpendicularly.
perpendicular = indicator - parallel[..., None] * u_delta[:, None]
perpendicular_inside_line = perpendicular.norm(dim=-1) < 0.5 * width[:, None]
inside_line = parallel_inside_line & perpendicular_inside_line
# Compute round caps.
if cap == "round":
near_start = indicator.norm(dim=-1) < 0.5 * width[:, None]
inside_line |= near_start
end_indicator = indicator = xy - end[:, None]
near_end = end_indicator.norm(dim=-1) < 0.5 * width[:, None]
inside_line |= near_end
# Determine the sample's color.
selectable_color = color.broadcast_to((num_lines, 3))
arrangement = inside_line * torch.arange(num_lines, device=device)[:, None]
top_color = selectable_color.gather(
dim=0,
index=repeat(arrangement.argmax(dim=0), "s -> s c", c=3),
)
rgba = torch.cat((top_color, inside_line.any(dim=0).float()[:, None]), dim=-1)
return rgba
return render_over_image(image, color_function, device, num_passes=num_msaa_passes)
|