feat: add calculate_optical_flow
This commit is contained in:
80
autopilot.py
80
autopilot.py
@@ -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),
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user