133 lines
4.3 KiB
Python
133 lines
4.3 KiB
Python
|
||
import cv2
|
||
import math
|
||
import numpy as np
|
||
from numpy.linalg import inv
|
||
from typing import Optional
|
||
import constants
|
||
from scipy.spatial.transform import Rotation
|
||
import utility
|
||
|
||
class Position:
|
||
"""Класс позиции с полной ориентацией БПЛА в 3D пространстве"""
|
||
|
||
x: float # Координата X
|
||
y: float # Координата Y
|
||
z: float # Масштаб
|
||
yaw: float # Рыскание (rotation around Z-axis)
|
||
pitch: float # Тангаж (rotation around Y-axis)
|
||
roll: float # Крен (rotation around X-axis)
|
||
|
||
def __init__(
|
||
self,
|
||
x: float = 0,
|
||
y: float = 0,
|
||
z: float = 1,
|
||
yaw: float = 0,
|
||
pitch: float = 0,
|
||
roll: float = 0
|
||
):
|
||
self.x = x
|
||
self.y = y
|
||
self.z = z
|
||
self.yaw = yaw
|
||
self.pitch = pitch
|
||
self.roll = roll
|
||
|
||
def __str__(self) -> str:
|
||
return (
|
||
f"Position(x={self.x:.2f}, y={self.y:.2f}, z={self.z:.2f}, "
|
||
f"yaw={math.degrees(self.yaw):.1f}°, "
|
||
f"pitch={math.degrees(self.pitch):.1f}°, "
|
||
f"roll={math.degrees(self.roll):.1f}°)"
|
||
)
|
||
|
||
def get_homography_matrix(self, K: np.ndarray = constants.K, sliding: bool = True) -> np.ndarray:
|
||
""" Возвращает матрицу гомографии """
|
||
R = self.get_rotation_matrix()
|
||
T = self.get_translation_matrix()
|
||
if not sliding:
|
||
T[0, 2] = T[1, 2] = 0
|
||
return K @ R @ T @ np.linalg.inv(K)
|
||
|
||
def copy(self) -> 'Position':
|
||
"""Создает полную копию объекта"""
|
||
return Position(self.x, self.y, self.z, self.yaw, self.pitch, self.roll)
|
||
|
||
def get_translation_matrix(self) -> np.ndarray:
|
||
return np.array([
|
||
[1, 0, self.x / constants._K_FOCUS_DISTANCE],
|
||
[0, 1, self.y / constants._K_FOCUS_DISTANCE],
|
||
[0, 0, self.z]
|
||
])
|
||
|
||
def get_rotation_matrix(self) -> np.ndarray:
|
||
"""
|
||
Матрица вращения с порядком применения: yaw → pitch → roll
|
||
Умножение: R = Rx(roll) * Ry(pitch) * Rz(yaw)
|
||
"""
|
||
cy, sy = math.cos(self.yaw), math.sin(self.yaw)
|
||
cp, sp = math.cos(self.pitch), math.sin(self.pitch)
|
||
cr, sr = math.cos(self.roll), math.sin(self.roll)
|
||
|
||
Rz = np.array([
|
||
[cy, -sy, 0],
|
||
[sy, cy, 0],
|
||
[0, 0, 1],
|
||
])
|
||
|
||
Ry = np.array([
|
||
[cp, 0, sp],
|
||
[0, 1, 0],
|
||
[-sp, 0, cp],
|
||
])
|
||
|
||
Rx = np.array([
|
||
[1, 0, 0],
|
||
[0, cr, -sr],
|
||
[0, sr, cr],
|
||
])
|
||
|
||
return Rx @ Ry @ Rz
|
||
|
||
def iapply(self, homography_matrix: np.ndarray, K = constants.K) -> 'Position':
|
||
"""Применяет матрицу трансформации для вычисления новой позиции и ориентации."""
|
||
|
||
np.set_printoptions(suppress=True)
|
||
H = homography_matrix @ self.get_homography_matrix(sliding=False)
|
||
|
||
# Decompose homography
|
||
_, R, t, _ = cv2.decomposeHomographyMat(H, K)
|
||
R = np.array(R)
|
||
t = np.array(t)
|
||
|
||
T = inv(R) @ inv(K) @ H @ K
|
||
ind = np.array([A[2][0] ** 2 + A[2][1] ** 2 for A in T])
|
||
top_k = max(1, len(T) // 2)
|
||
if (len(T) == 3): raise "len(T) == 3"
|
||
ind = np.argpartition(ind, top_k - 1)[:top_k]
|
||
T = T[ind[0]]
|
||
T = T @ np.array([0, 0, 1]) / np.mean((T[0][0], T[1][1]))
|
||
T[2] -= 1
|
||
R = R[ind]
|
||
t = t[ind]
|
||
|
||
best_id = ((t - T) ** 2).sum((1, 2)).argmin()
|
||
|
||
R = R[best_id]
|
||
rot = Rotation.from_matrix(R).as_euler('XYZ').flatten()
|
||
self.roll = rot[0]
|
||
self.pitch = rot[1]
|
||
self.yaw = rot[2]
|
||
|
||
t = t[best_id].flatten()
|
||
self.x += -T[0] * constants._K_FOCUS_DISTANCE * self.z
|
||
self.y += T[1] * constants._K_FOCUS_DISTANCE * self.z
|
||
self.z = 1 + T[2]
|
||
|
||
def apply(self, homography_matrix: np.ndarray, K = constants.K) -> 'Position':
|
||
"""Применяет матрицу трансформации для вычисления новой позиции и ориентации."""
|
||
pos = self.copy()
|
||
pos.iapply(homography_matrix, K)
|
||
return pos
|