|
|
import numpy as np
|
|
|
import cv2
|
|
|
import matplotlib.pyplot as plt
|
|
|
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
|
|
|
from Neko.RealSense import RealSenseController
|
|
|
from Neko.CameraProcessor import CameraProcessor
|
|
|
from ultralytics import YOLO
|
|
|
from datetime import datetime
|
|
|
|
|
|
class TargetFollower:
|
|
|
def __init__(self, detections, annotated_image, camera_position=np.array([0, 0, 0])):
|
|
|
self.detections = detections
|
|
|
self.camera_position = camera_position
|
|
|
self.target = None
|
|
|
|
|
|
# Сохраняем аннотированное изображение
|
|
|
output_path = './annotated_image.jpg'
|
|
|
cv2.imwrite(output_path, annotated_image)
|
|
|
print(f"Annotated image saved to {output_path}")
|
|
|
|
|
|
print("Detections info:")
|
|
|
for info in detections:
|
|
|
print(info)
|
|
|
|
|
|
def select_target(self, class_name, nearest=False):
|
|
|
"""
|
|
|
Выбирает объект по указанному имени класса.
|
|
|
Если nearest=True, выбирается ближайший объект по значению mean_depth.
|
|
|
|
|
|
:param class_name: Имя класса, который нужно выбрать.
|
|
|
:param nearest: Если True, выбирает ближайший объект с данным классом.
|
|
|
:return: Найденный объект или None, если объект не найден.
|
|
|
"""
|
|
|
if nearest:
|
|
|
closest_detection = None
|
|
|
min_depth = float('inf')
|
|
|
|
|
|
for detection in self.detections:
|
|
|
if detection['class_name'] == class_name:
|
|
|
if detection['mean_depth'] < min_depth:
|
|
|
min_depth = detection['mean_depth']
|
|
|
closest_detection = detection
|
|
|
|
|
|
self.target = closest_detection
|
|
|
if self.target is not None:
|
|
|
print(f"Closest target selected: {self.target}")
|
|
|
else:
|
|
|
print(f"Target with class name '{class_name}' not found.")
|
|
|
return self.target
|
|
|
else:
|
|
|
for detection in self.detections:
|
|
|
if detection['class_name'] == class_name:
|
|
|
self.target = detection
|
|
|
print(f"Target selected: {self.target}")
|
|
|
return self.target
|
|
|
|
|
|
print(f"Target with class name '{class_name}' not found.")
|
|
|
return None
|
|
|
|
|
|
def calculate_follow_vector(self):
|
|
|
if self.target is None:
|
|
|
print("Target not selected.")
|
|
|
return None
|
|
|
|
|
|
bbox = self.target['bbox']
|
|
|
target_center_x = (bbox[0] + bbox[2]) / 2
|
|
|
target_center_y = (bbox[1] + bbox[3]) / 2
|
|
|
target_depth = self.target['mean_depth']
|
|
|
|
|
|
target_position = np.array([target_center_x, target_center_y, target_depth])
|
|
|
|
|
|
follow_vector = target_position - self.camera_position
|
|
|
return follow_vector
|
|
|
|
|
|
def get_target_position(self):
|
|
|
if self.target is None:
|
|
|
print("Target not selected.")
|
|
|
return None
|
|
|
|
|
|
bbox = self.target['bbox']
|
|
|
target_center_x = (bbox[0] + bbox[2]) / 2
|
|
|
target_center_y = (bbox[1] + bbox[3]) / 2
|
|
|
target_depth = self.target['mean_depth']
|
|
|
|
|
|
target_position = np.array([target_center_x, target_center_y, target_depth])
|
|
|
return target_position |
