ref: create vision_chunk

This commit is contained in:
2025-12-21 20:13:51 +03:00
parent 92b6b0b0fe
commit 9d4839aa2a
6 changed files with 221 additions and 152 deletions

View File

@@ -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