pokouqiege/models/line_detect/te.py

import numpy as np
import cv2
import torch

def arc_to_mask(xc, yc, a, b, theta, phi1, phi2, H, W, line_width=1):
    """
    Generate a binary mask of an elliptical arc.

    Args:
        xc, yc (float): 椭圆中心
        a, b (float): 长半轴、短半轴 (a >= b)
        theta (float): 椭圆旋转角度（**弧度**，逆时针，相对于 x 轴）
        phi1, phi2 (float): 起始和终止参数角（**弧度**，在 [0, 2π) 内）
        H, W (int): 输出 mask 的高度和宽度
        line_width (int): 弧线宽度（像素）

    Returns:
        mask (Tensor): [H, W], dtype=torch.uint8, 0/255
    """
    # 确保 phi1 -> phi2 是正向（可处理跨 2π 的情况）
    if phi2 < phi1:
        phi2 += 2 * np.pi

    # 生成参数角（足够密集，避免断线）
    num_points = max(int(200 * abs(phi2 - phi1) / (2 * np.pi)), 10)
    phi = np.linspace(phi1, phi2, num_points)

    # 椭圆参数方程（先在未旋转坐标系下计算）
    x_local = a * np.cos(phi)
    y_local = b * np.sin(phi)

    # 应用旋转和平移
    cos_t = np.cos(theta)
    sin_t = np.sin(theta)
    x_rot = x_local * cos_t - y_local * sin_t + xc
    y_rot = x_local * cos_t + y_local * sin_t + yc

    # 转为整数坐标（OpenCV 需要 int32）
    points = np.stack([x_rot, y_rot], axis=1).astype(np.int32)

    # 创建空白图像
    img = np.zeros((H, W), dtype=np.uint8)

    # 绘制折线（antialias=False 更适合 mask）
    cv2.polylines(img, [points], isClosed=False, color=255, thickness=line_width, lineType=cv2.LINE_AA)

    return torch.from_numpy(img).byte()  # [H, W], values: 0 or 255


# 椭圆参数
xc, yc = 100.0, 100.0
a, b = 80.0, 40.0
theta = np.radians(30)      # 30度 → 弧度
phi1 = np.radians(0)        # 从右侧开始
phi2 = np.radians(180)      # 到左侧结束（上半椭圆）

H, W = 200, 200

mask = arc_to_mask(xc, yc, a, b, theta, phi1, phi2, H, W, line_width=2)
# print(mask.shape)  # torch.Size([200, 200])
# print(mask.dtype)  # torch.uint8

# 可视化（调试用）
import matplotlib.pyplot as plt
plt.imshow(mask.numpy(), cmap='gray')
plt.title("Elliptical Arc Mask")
plt.show()