feat: add calculate_optical_flow

This commit is contained in:
2025-12-30 01:37:08 +03:00
parent a02c20a8ff
commit a06442b729
2 changed files with 53 additions and 29 deletions

View File

@@ -78,9 +78,55 @@ class AutoPilot(Pilot):
def get_position(self) -> tuple[float, float]:
return self.geo.x, self.geo.y
def calculate_optical_flow(self, prev_chunk: VisionChunk, current_chunk: VisionChunk):
"""
Вычисляет оптический поток между двумя кадрами
Возвращает src_pts и dst_pts в отцентрированных координатах
"""
prev_gray = prev_chunk.to_cv2_gray()
current_gray = current_chunk.to_cv2_gray()
h, w = prev_gray.shape[:2]
cx, cy = w // 2, h // 2
# Создаем сетку точек для отслеживания (аналогично вашему step=20)
step = 20
grid_points = []
for y in range(step, h - step, step):
for x in range(step, w - step, step):
grid_points.append([x, y])
if len(grid_points) < 4:
return None, None
p0 = np.float32(grid_points).reshape(-1, 1, 2)
# Параметры для Lucas-Kanade
lk_params = dict(
winSize=(21, 21),
maxLevel=3,
criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03)
)
# Вычисляем разреженный оптический поток
p1, status, err = cv2.calcOpticalFlowPyrLK(prev_gray, current_gray, p0, None, **lk_params)
# Фильтруем только успешно отслеженные точки
good_old = p0[status == 1]
good_new = p1[status == 1]
if len(good_old) < 4:
return None, None
# Преобразуем в отцентрированные координаты (Y вверх)
src_pts = np.float32([[x - cx, cy - y] for x, y in good_old]).reshape(-1, 1, 2)
dst_pts = np.float32([[x - cx, cy - y] for x, y in good_new]).reshape(-1, 1, 2)
return src_pts, dst_pts
def estimate_transformation_matrix(self, src_pts: np.ndarray, dst_pts: np.ndarray):
"""
Оценивает матрицу трансформации на основе сопоставленных ключевых точек
Оценивает матрицу трансформации на основе сопоставленных точек
Точки уже отцентрированы относительно центра изображения
"""
# Используем RANSAC для оценки матрицы гомографии
@@ -219,8 +265,8 @@ class AutoPilot(Pilot):
(self.points[self.target_idx][1] - self.geo.y) ** 2
)
# Обнаруживаем и сопоставляем ключевые точки
src_pts, dst_pts, matches, kp1, kp2 = self.prev_chunk.detect_and_match_keypoints(current_chunk)
# Вычисляем оптический поток для покадрового сравнения
src_pts, dst_pts = self.calculate_optical_flow(self.prev_chunk, current_chunk)
# Оцениваем смещение БПЛА
if src_pts is not None and dst_pts is not None:
@@ -232,32 +278,10 @@ class AutoPilot(Pilot):
self.update_drone_position(transformation_info)
self.timer.stop()
# Выводим текущее состояние БПЛА
# print(f" [Pilot] Drone Position: ({drone_state['x']:.2f}, {drone_state['y']:.2f})")
# print(f" [Pilot] Angle: {drone_state['angle_degrees']:.1f}°")
# print(f" [Pilot] Target Index: {self.target_idx}")
# print(f" [Pilot] Target Position: {self.points[self.target_idx]}")
# print(f" [Pilot] Distance: {distance_to_target}")
# Обновляем визуализацию
if self.vis_manager:
# Обновляем сопоставление точек
self.vis_manager.update_matches(
self.prev_chunk.to_numpy(),
current_chunk.to_numpy(),
kp1, kp2, matches,
transformation_info)
mask = transformation_info['mask']
self.vis_manager.update_motion_vectors(
current_chunk.to_numpy(),
kp1, kp2,
np.array(matches)[mask.ravel().astype(bool)])
# Используем новую визуализацию движения точек по сетке
# Используем визуализацию движения точек по сетке
homography_matrix = transformation_info['homography']
self.vis_manager.update_homography_grid(
current_chunk.to_numpy(),
@@ -302,8 +326,8 @@ class AutoPilot(Pilot):
if angle_diff >= math.pi:
angle_diff -= 2 * math.pi
d_r = max(10, min(75., distance_to_target / 2))
d_a_limit = d_r / 10 * 0.07
d_r = max(10, min(35., distance_to_target / 2))
d_a_limit = d_r / 10 * 0.01
command = PilotCommand(
max(min(d_a_limit, angle_diff), -d_a_limit),

View File

@@ -7,7 +7,7 @@ from typing import Literal, Optional, Tuple
from PIL import Image
FeatureMethod = Literal["orb", "sift", "akaze", "brisk"]
DEFAULT_METHOD = "akaze"
DEFAULT_METHOD = "orb"
@dataclass
class VisionChunk: