Получить угол прямоугольника рядом с началом координат в партии тензорных данных, учитывая любые две диагональные координаты в pytorch

Question

Допустим, у меня есть тензор pytorch пакета координат недиагональных элементов, и я хочу получить координату угла, который находится близко к началу координат.координаты в форме (x1, y1, x2, y2).

a = torch.tensor([[3,2,2,3], [1,1,2,2])
# expected output
[[2,2], [1,1]]

Answer 1

Вы можете просто перебрать все тензоры и для каждого из них вычислить расстояние до четырех углов и взять угол с минимальным расстоянием.

import torch
a = torch.tensor([[3,2,2,3], [1,1,2,2]])
c = torch.zeros(a.shape[0], 2)
for idx, x in enumerate(a):
    d1 = x[0] ** 2 + x[1] ** 2
    d2 = x[2] ** 2 + x[3] ** 2
    d3 = x[0] ** 2 + x[3] ** 2
    d4 = x[2] ** 2 + x[1] ** 2
    dmin = min(d1, d2, d3, d4)
    if d1 == dmin:
        c[idx] = torch.tensor([x[0], x[1]])
    elif d2 == dmin:
        c[idx] = torch.tensor([x[2], x[3]])
    elif d3 == dmin:
        c[idx] = torch.tensor([x[0], x[3]])
    elif d4 == dmin:
        c[idx] = torch.tensor([x[2], x[1]])
print(c)  # tensor([[2., 2.], [1., 1.]])