165 lines
5.2 KiB
Python
165 lines
5.2 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 __imul__(self, scalar: float):
|
||
self.x *= scalar
|
||
self.y *= scalar
|
||
self.z *= scalar
|
||
return self
|
||
|
||
def __mul__(self, scalar: float) -> 'Position':
|
||
pos = self.copy()
|
||
pos *= scalar
|
||
return pos
|
||
|
||
def __itruediv__(self, scalar: float):
|
||
self.x /= scalar
|
||
self.y /= scalar
|
||
self.z /= scalar
|
||
return self
|
||
|
||
def __truediv__(self, scalar: float) -> 'Position':
|
||
pos = self.copy()
|
||
pos /= scalar
|
||
return pos
|
||
|
||
def get_homography_matrix(self, K_in: np.ndarray = constants.K, K_out: np.ndarray | None = None, sliding: bool = True) -> np.ndarray:
|
||
""" Возвращает матрицу гомографии """
|
||
R = self.get_rotation_matrix()
|
||
T = self.get_translation_matrix(K_in)
|
||
if not sliding:
|
||
T[0, 2] = T[1, 2] = 0
|
||
if K_out is None: K_out = K_in
|
||
return K_out @ R @ T @ np.linalg.inv(K_in)
|
||
|
||
def copy(self) -> 'Position':
|
||
"""Создает полную копию объекта"""
|
||
return Position(self.x, self.y, self.z, self.yaw, self.pitch, self.roll)
|
||
|
||
def get_translation_matrix(self, K: np.ndarray = constants.K) -> np.ndarray:
|
||
return np.array([
|
||
[1, 0, self.x / K[0][0]],
|
||
[0, 1, self.y / K[0][0]],
|
||
[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)
|
||
|
||
# print(cv2.decomposeHomographyMat(inv(H), K))
|
||
# _, _, z, _ = cv2.decomposeHomographyMat(inv(H), K)
|
||
# print(z)
|
||
# z = z.copy()
|
||
# z *= constants._K_FOCUS_DISTANCE
|
||
# print(z)
|
||
|
||
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 = min(np.radians(5), max(np.radians(-5), rot[0]))
|
||
self.pitch = min(np.radians(5), max(np.radians(-5), rot[1]))
|
||
self.yaw = rot[2]
|
||
|
||
t = t[best_id].flatten()
|
||
self.x -= T[0] * constants._K_FOCUS_DISTANCE
|
||
self.y += T[1] * constants._K_FOCUS_DISTANCE
|
||
self.z = max(0.7, min(1.3, 1 + T[2]))
|
||
T[0] *= constants._K_FOCUS_DISTANCE
|
||
T[1] *= constants._K_FOCUS_DISTANCE
|
||
|
||
def apply(self, homography_matrix: np.ndarray, K = constants.K) -> 'Position':
|
||
"""Применяет матрицу трансформации для вычисления новой позиции и ориентации."""
|
||
pos = self.copy()
|
||
pos.iapply(homography_matrix, K)
|
||
return pos
|