ref: create vision_chunk
This commit is contained in:
141
autopilot.py
141
autopilot.py
@@ -9,6 +9,7 @@ from PIL import Image
|
||||
|
||||
from timer import Timer
|
||||
|
||||
from vision_chunk import VisionChunk
|
||||
from visualization import VisualizationManager
|
||||
from geolocation import Geolocation
|
||||
|
||||
@@ -42,28 +43,23 @@ class AutoPilot(Pilot):
|
||||
# Состояние одометрии
|
||||
geo: Geolocation
|
||||
|
||||
# Ориентиры
|
||||
points: list[tuple[float,float]] # Координаты
|
||||
keypoints: list[any] # Ключевые точки ориентиров
|
||||
target_idx: int # Текущая целевая метка
|
||||
chunks: list[VisionChunk] # Ориентиры
|
||||
target_idx: int # Текущий ориентир
|
||||
|
||||
# Всякие вспомогательные штуки
|
||||
timer: Timer
|
||||
prev_image: np.ndarray | None
|
||||
orb_detector: cv2.ORB
|
||||
bf_matcher: cv2.BFMatcher
|
||||
prev_chunk: VisionChunk | None
|
||||
frame_count: int
|
||||
image_center: tuple # Центр изображения (x, y)
|
||||
vis_manager: VisualizationManager # Менеджер визуализации (опционально)
|
||||
|
||||
# Положение на основе ориентира
|
||||
reserved_geo: Geolocation | None
|
||||
|
||||
def __init__(self, points = [], keypoints = [], viz_manager=None):
|
||||
self.prev_image = None
|
||||
def __init__(self, points = [], chunks = [], viz_manager=None):
|
||||
self.prev_chunk = None
|
||||
self.geo = Geolocation(0, 0, 1, 0)
|
||||
self.chunks = chunks
|
||||
self.frame_count = 0
|
||||
self.image_center = (0, 0) # Будет обновлено при получении первого изображения
|
||||
self.vis_manager = viz_manager # Менеджер визуализации
|
||||
self.reserved_geo = None
|
||||
|
||||
@@ -74,22 +70,7 @@ class AutoPilot(Pilot):
|
||||
self.min_scale: float = 0.7
|
||||
self.max_scale: float = 1.4
|
||||
|
||||
# Инициализация ORB детектора
|
||||
self.orb_detector = cv2.ORB_create(
|
||||
nfeatures=1000,
|
||||
scaleFactor=1.3,
|
||||
nlevels=4,
|
||||
edgeThreshold=19,
|
||||
patchSize=31,
|
||||
fastThreshold=15,
|
||||
scoreType=cv2.ORB_FAST_SCORE # FAST обычно быстрее
|
||||
)
|
||||
|
||||
# Инициализация матчера для сопоставления ключевых точек (kNN + Lowe ratio)
|
||||
self.bf_matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False)
|
||||
|
||||
self.points = points
|
||||
self.keypoints = keypoints
|
||||
self.target_idx = 0
|
||||
|
||||
self.timer = Timer()
|
||||
@@ -97,74 +78,6 @@ class AutoPilot(Pilot):
|
||||
def get_position(self) -> tuple[float, float]:
|
||||
return self.geo.x, self.geo.y
|
||||
|
||||
def image_to_numpy(self, image: Image.Image) -> np.ndarray:
|
||||
"""Конвертирует PIL Image в numpy array для OpenCV"""
|
||||
return np.array(image)
|
||||
|
||||
def detect_and_match_keypoints(self, img1: np.ndarray, img2: np.ndarray):
|
||||
"""
|
||||
Обнаруживает ключевые точки ORB и сопоставляет их между двумя изображениями
|
||||
"""
|
||||
# Предобработка (повышение контраста и устойчивости)
|
||||
def preprocess(gray_img: np.ndarray) -> np.ndarray:
|
||||
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
|
||||
return clahe.apply(gray_img)
|
||||
|
||||
# В градации серого
|
||||
g1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY) if len(img1.shape) == 3 else img1
|
||||
g2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY) if len(img2.shape) == 3 else img2
|
||||
g1 = preprocess(g1)
|
||||
g2 = preprocess(g2)
|
||||
|
||||
# Обнаружение ключевых точек и дескрипторов
|
||||
kp1, des1 = self.orb_detector.detectAndCompute(g1, None)
|
||||
kp2, des2 = self.orb_detector.detectAndCompute(g2, None)
|
||||
|
||||
if des1 is None or des2 is None or len(kp1) < 4 or len(kp2) < 4:
|
||||
return None, None, None, None, None
|
||||
|
||||
# kNN сопоставление и тест Лоу
|
||||
matches_knn = self.bf_matcher.knnMatch(des1, des2, k=2)
|
||||
good_matches: list[cv2.DMatch] = []
|
||||
for m_n in matches_knn:
|
||||
if len(m_n) < 2:
|
||||
continue
|
||||
m, n = m_n
|
||||
if m.distance < 0.75 * n.distance:
|
||||
good_matches.append(m)
|
||||
|
||||
# Дополнительная фильтрация по расстоянию (мягкий порог)
|
||||
good_matches = sorted(good_matches, key=lambda x: x.distance)
|
||||
good_matches = [m for m in good_matches if m.distance < 64]
|
||||
|
||||
if len(good_matches) < 4:
|
||||
return None, None, None, None, None
|
||||
|
||||
# Получаем центр изображения
|
||||
height1, width1 = g1.shape[:2]
|
||||
height2, width2 = g2.shape[:2]
|
||||
center_x1, center_y1 = width1 // 2, height1 // 2
|
||||
center_x2, center_y2 = width2 // 2, height2 // 2
|
||||
|
||||
# Извлекаем координаты сопоставленных точек и отцентрируем их
|
||||
src_pts = []
|
||||
dst_pts = []
|
||||
|
||||
for match in good_matches:
|
||||
# Координаты точки в первом изображении
|
||||
pt1 = kp1[match.queryIdx].pt
|
||||
src_pts.append([pt1[0] - center_x1, center_y1 - pt1[1]])
|
||||
|
||||
# Координаты точки во втором изображении
|
||||
pt2 = kp2[match.trainIdx].pt
|
||||
dst_pts.append([pt2[0] - center_x2, center_y2 - pt2[1]])
|
||||
|
||||
# Конвертируем в numpy массивы
|
||||
src_pts = np.float32(src_pts).reshape(-1, 1, 2)
|
||||
dst_pts = np.float32(dst_pts).reshape(-1, 1, 2)
|
||||
|
||||
return src_pts, dst_pts, good_matches, kp1, kp2
|
||||
|
||||
def estimate_transformation_matrix(self, src_pts: np.ndarray, dst_pts: np.ndarray):
|
||||
"""
|
||||
Оценивает матрицу трансформации на основе сопоставленных ключевых точек
|
||||
@@ -251,18 +164,16 @@ class AutoPilot(Pilot):
|
||||
'frame_count': self.frame_count
|
||||
}
|
||||
|
||||
|
||||
|
||||
def get_position_by_chunk(self) -> Geolocation | None:
|
||||
# Пытаемся найти ориентир на картинке:
|
||||
current_image = self.prev_image
|
||||
landmark_image = cv2.imread(Path('chunks') / f'chunk_{self.target_idx}.png', cv2.IMREAD_COLOR_RGB)
|
||||
src_pts, dst_pts, matches, kp1, kp2 = self.detect_and_match_keypoints(landmark_image, current_image)
|
||||
current_chunk = self.prev_chunk
|
||||
landmark_chunk = self.chunks[self.target_idx]
|
||||
src_pts, dst_pts, matches, kp1, kp2 = landmark_chunk.detect_and_match_keypoints(current_chunk)
|
||||
|
||||
if src_pts is not None and dst_pts is not None and self.vis_manager:
|
||||
was_enabled = self.timer.enabled
|
||||
if was_enabled: self.timer.stop()
|
||||
self.vis_manager.update_chunk_matches(landmark_image, current_image, kp1, kp2, matches)
|
||||
self.vis_manager.update_chunk_matches(landmark_chunk.to_numpy(), current_chunk.to_numpy(), kp1, kp2, matches)
|
||||
if was_enabled: self.timer.start()
|
||||
|
||||
if src_pts is not None and dst_pts is not None:
|
||||
@@ -293,26 +204,15 @@ class AutoPilot(Pilot):
|
||||
return None
|
||||
|
||||
|
||||
def handle(self, image: Image) -> PilotCommand:
|
||||
def handle(self, current_chunk: VisionChunk) -> PilotCommand:
|
||||
self.timer.start()
|
||||
self.frame_count += 1
|
||||
|
||||
if self.prev_image is None:
|
||||
self.prev_image = self.image_to_numpy(image)
|
||||
# Вычисляем центр изображения
|
||||
height, width = self.prev_image.shape[:2]
|
||||
self.image_center = (width // 2, height // 2)
|
||||
|
||||
if self.prev_chunk is None:
|
||||
self.prev_chunk = current_chunk
|
||||
self.timer.stop()
|
||||
return PilotCommand(processing_time=self.timer.get_elapsed())
|
||||
|
||||
# Конвертируем текущее изображение
|
||||
current_image = self.image_to_numpy(image)
|
||||
|
||||
# Обновляем центр изображения
|
||||
height, width = current_image.shape[:2]
|
||||
self.image_center = (width // 2, height // 2)
|
||||
|
||||
# Расстояние до цели
|
||||
distance_to_target = math.sqrt(
|
||||
(self.points[self.target_idx][0] - self.geo.x) ** 2 +
|
||||
@@ -320,7 +220,7 @@ class AutoPilot(Pilot):
|
||||
)
|
||||
|
||||
# Обнаруживаем и сопоставляем ключевые точки
|
||||
src_pts, dst_pts, matches, kp1, kp2 = self.detect_and_match_keypoints(self.prev_image, current_image)
|
||||
src_pts, dst_pts, matches, kp1, kp2 = self.prev_chunk.detect_and_match_keypoints(current_chunk)
|
||||
|
||||
# Оцениваем смещение БПЛА
|
||||
if src_pts is not None and dst_pts is not None:
|
||||
@@ -334,7 +234,6 @@ class AutoPilot(Pilot):
|
||||
self.timer.stop()
|
||||
|
||||
# Выводим текущее состояние БПЛА
|
||||
drone_state = self.get_drone_state()
|
||||
# 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}")
|
||||
@@ -346,14 +245,14 @@ class AutoPilot(Pilot):
|
||||
|
||||
# Обновляем сопоставление точек
|
||||
self.vis_manager.update_matches(
|
||||
self.prev_image,
|
||||
current_image,
|
||||
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_image,
|
||||
current_chunk.to_numpy(),
|
||||
kp1, kp2,
|
||||
np.array(matches)[mask.ravel().astype(bool)])
|
||||
|
||||
@@ -361,14 +260,14 @@ class AutoPilot(Pilot):
|
||||
# Используем новую визуализацию движения точек по сетке
|
||||
homography_matrix = transformation_info['homography']
|
||||
self.vis_manager.update_homography_grid(
|
||||
current_image,
|
||||
current_chunk.to_numpy(),
|
||||
homography_matrix,
|
||||
grid_step=85)
|
||||
|
||||
self.timer.start()
|
||||
|
||||
# Пытаемся найти ориентир на картинке:
|
||||
self.prev_image = current_image
|
||||
self.prev_chunk = current_chunk
|
||||
npos = self.get_position_by_chunk()
|
||||
if npos is not None:
|
||||
self.reserved_geo = npos
|
||||
|
||||
Reference in New Issue
Block a user