|
|
from Neko.RealSense import RealSenseController
|
|
|
import cv2
|
|
|
from ultralytics import YOLO
|
|
|
import numpy as np
|
|
|
|
|
|
controller = RealSenseController()
|
|
|
|
|
|
def inline_detection(controller):
|
|
|
depth_matrix = controller.acquisition.get_depth_image()
|
|
|
color_matrix = controller.acquisition.get_color_image()
|
|
|
|
|
|
if depth_matrix is None or color_matrix is None:
|
|
|
print("Не удалось получить изображения. Проверьте подключение камеры.")
|
|
|
return
|
|
|
|
|
|
model = YOLO('./models/yolov9e_object_classification.pt')
|
|
|
results = model(color_matrix)
|
|
|
|
|
|
annotated_image = results[0].plot()
|
|
|
|
|
|
detections_info = []
|
|
|
|
|
|
for detection in results[0].boxes:
|
|
|
# Извлекаем координаты ограничивающей рамки
|
|
|
x1, y1, x2, y2 = map(int, detection.xyxy[0])
|
|
|
|
|
|
# Извлекаем класс объекта
|
|
|
class_id = int(detection.cls[0])
|
|
|
class_name = model.names[class_id]
|
|
|
|
|
|
# Извлекаем соответствующую область из карты глубины
|
|
|
depth_values = depth_matrix[y1:y2, x1:x2]
|
|
|
|
|
|
# Вычисляем среднее значение глубины
|
|
|
mean_depth = np.mean(depth_values)
|
|
|
|
|
|
# Сохраняем информацию в массив
|
|
|
detections_info.append({
|
|
|
'class_id': class_id,
|
|
|
'class_name': class_name,
|
|
|
'bbox': [x1, y1, x2, y2],
|
|
|
'mean_depth': mean_depth
|
|
|
})
|
|
|
|
|
|
# Наносим информацию на изображение
|
|
|
cv2.putText(annotated_image, f'{class_name} {mean_depth:.2f}m', (x1, y1 - 10),
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
|
|
cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
|
|
# Сохраняем аннотированное изображение
|
|
|
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_info:
|
|
|
print(info)
|
|
|
|
|
|
inline_detection(controller) |
