如何在PyTorch Tensor上绘制线条？仿numpy数组操作实现

你习惯用OpenCV在numpy数组上绘制线条，示例代码如下：

pts = [[160, 160], [280, 280]]
mask = np.zeros([640,640]).astype(np.uint8)
mask = cv2.line(mask, *pts, (255, 255), 15)
plt.imshow(mask)
plt.show()

对应的绘制结果：

现在要对如下PyTorch Tensor完成相同的线条绘制操作：

mask_t = torch.zeros([640,640],dtype=torch.uint8).to(device)

可以通过以下两种方式实现：

方法一：借助OpenCV（简单高效） #

利用OpenCV对numpy数组的良好支持，将Tensor转换为numpy数组完成绘制后再转回Tensor，步骤如下：

import torch
import cv2

pts = [[160, 160], [280, 280]]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
mask_t = torch.zeros([640,640], dtype=torch.uint8).to(device)

# 将Tensor转回CPU并转为numpy数组
mask_np = mask_t.cpu().numpy()
# 用OpenCV绘制线条
mask_np = cv2.line(mask_np, *pts, (255, 255), 15)
# 将numpy数组转回Tensor并放回原设备
mask_t = torch.from_numpy(mask_np).to(device)

这种方法直接复用OpenCV优化成熟的绘制逻辑，代码简洁，适合大多数场景。

方法二：纯PyTorch实现（避免设备切换开销） #

如果需要全程在GPU上操作、避免数据在CPU与GPU间传输的开销，可以手动实现画线算法（比如Bresenham算法），示例代码如下：

import torch

def draw_line_tensor(tensor, pt1, pt2, value, thickness=1):
    assert len(tensor.shape) == 2, "输入必须是2D Tensor"
    y0, x0 = pt1
    y1, x1 = pt2

    dx = abs(x1 - x0)
    dy = abs(y1 - y0)
    sx = 1 if x0 < x1 else -1
    sy = 1 if y0 < y1 else -1
    err = dx - dy

    while True:
        # 绘制当前点及厚度范围内的区域
        y_min = max(0, y0 - thickness // 2)
        y_max = min(tensor.shape[0] - 1, y0 + thickness // 2)
        x_min = max(0, x0 - thickness // 2)
        x_max = min(tensor.shape[1] - 1, x0 + thickness // 2)
        tensor[y_min:y_max+1, x_min:x_max+1] = value

        if x0 == x1 and y0 == y1:
            break
        e2 = 2 * err
        if e2 > -dy:
            err -= dy
            x0 += sx
        if e2 < dx:
            err += dx
            y0 += sy
    return tensor

# 使用示例
pts = [[160, 160], [280, 280]]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
mask_t = torch.zeros([640,640], dtype=torch.uint8).to(device)
mask_t = draw_line_tensor(mask_t, pts[0], pts[1], 255, thickness=15)

这种方法适合对性能要求较高、需要批量GPU操作的场景，无需依赖OpenCV。

内容的提问来源于stack exchange，提问作者ShinNShirley