Overview
This module provides functions for drawing geometric primitives (points, lines, rectangles, and convex polygons) directly on image tensors using pure PyTorch operations.
Description
The draw module in the Kornia image package implements rasterization of basic geometric shapes on torch.Tensor images. All operations are performed in-place on the image tensor and are fully differentiable (where applicable). The module contains four public functions: draw_point2d for setting pixel colors at specific coordinates, draw_line for Bresenham-style line drawing, draw_rectangle for axis-aligned rectangle outlines or fills, and draw_convex_polygon for filled convex polygon rendering using scanline fill. Helper functions include _draw_pixel, _get_convex_edges for computing polygon scanline boundaries, and _batch_polygons for batching variable-length polygon lists.
Usage
Import these functions when you need to annotate or draw on image tensors programmatically, such as visualizing bounding boxes, keypoints, or regions of interest without leaving the PyTorch tensor ecosystem.
Code Reference
Source Location
Signature
def draw_point2d(image: Tensor, points: Tensor, color: Tensor) -> Tensor: ...
def draw_line(image: torch.Tensor, p1: torch.Tensor, p2: torch.Tensor, color: torch.Tensor) -> torch.Tensor: ...
def draw_rectangle(
image: torch.Tensor, rectangle: torch.Tensor,
color: Optional[torch.Tensor] = None, fill: Optional[bool] = None
) -> torch.Tensor: ...
def draw_convex_polygon(
images: Tensor, polygons: Union[Tensor, List[Tensor]], colors: Tensor
) -> Tensor: ...
Import
from kornia.image.draw import draw_point2d, draw_line, draw_rectangle, draw_convex_polygon
I/O Contract
draw_point2d
| Name |
Type |
Required |
Description
|
| image |
Tensor |
Yes |
Input image of shape (C, H, W) or (H, W).
|
| points |
Tensor |
Yes |
Points [x, y] to draw on the image.
|
| color |
Tensor |
Yes |
Color of the pixel with shape (C,).
|
draw_line
| Name |
Type |
Required |
Description
|
| image |
torch.Tensor |
Yes |
Input image of shape (C, H, W).
|
| p1 |
torch.Tensor |
Yes |
Start point [x, y] with shape (2,) or (B, 2).
|
| p2 |
torch.Tensor |
Yes |
End point [x, y] with shape (2,) or (B, 2).
|
| color |
torch.Tensor |
Yes |
Line color with shape (C,).
|
draw_rectangle
| Name |
Type |
Required |
Description
|
| image |
torch.Tensor |
Yes |
Batch of images of shape (B, C, H, W).
|
| rectangle |
torch.Tensor |
Yes |
Rectangle coordinates (B, N, 4) as [x1, y1, x2, y2].
|
| color |
torch.Tensor |
No |
Color tensor of shape (1,), (3,), (B, N, 1), or (B, N, 3). Default: zeros.
|
| fill |
bool |
No |
Whether to fill rectangles with color (default False).
|
draw_convex_polygon
| Name |
Type |
Required |
Description
|
| images |
Tensor |
Yes |
Batch of images of shape (B, C, H, W).
|
| polygons |
Tensor or List[Tensor] |
Yes |
Polygon vertices (B, N, 2) as [x, y] or list of variable-length polygons.
|
| colors |
Tensor |
Yes |
Fill colors of shape (B, 3) or (3,).
|
Outputs
| Name |
Type |
Description
|
| output |
torch.Tensor |
The modified image tensor (operations are in-place but also return the tensor).
|
Usage Examples
import torch
from kornia.image.draw import draw_line, draw_rectangle, draw_convex_polygon
# Draw a line
image = torch.zeros(1, 8, 8)
result = draw_line(image, torch.tensor([6, 4]), torch.tensor([1, 4]), torch.tensor([255]))
# Draw rectangles on a batch
img = torch.rand(2, 3, 10, 12)
rect = torch.tensor([[[0, 0, 4, 4]], [[4, 4, 10, 10]]])
out = draw_rectangle(img, rect)
# Draw a filled convex polygon
img = torch.rand(1, 3, 12, 16)
poly = torch.tensor([[[4, 4], [12, 4], [12, 8], [4, 8]]])
color = torch.tensor([[0.5, 0.5, 0.5]])
out = draw_convex_polygon(img, poly, color)
Related Pages
