|
|
|
|
|
import pyrealsense2 as rs
|
|
|
|
|
|
import numpy as np
|
|
|
|
|
|
import cv2
|
|
|
|
|
|
import torch
|
|
|
|
|
|
|
|
|
|
|
|
# Настройка потоков RealSense
|
|
|
|
|
|
pipeline = rs.pipeline()
|
|
|
|
|
|
config = rs.config()
|
|
|
|
|
|
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
|
|
|
|
|
|
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
|
|
|
|
|
|
|
|
|
|
|
|
# Запуск потока
|
|
|
|
|
|
pipeline.start(config)
|
|
|
|
|
|
|
|
|
|
|
|
# Создание объекта для выравнивания глубины и цвета
|
|
|
|
|
|
align_to = rs.stream.color
|
|
|
|
|
|
align = rs.align(align_to)
|
|
|
|
|
|
|
|
|
|
|
|
# Загрузка модели YOLOv5
|
|
|
|
|
|
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
|
|
|
|
|
|
|
|
|
|
|
|
# Функция для детекции объектов с помощью YOLO
|
|
|
|
|
|
def detect_objects(image):
|
|
|
|
|
|
results = model(image)
|
|
|
|
|
|
detections = results.xyxy[0].numpy()
|
|
|
|
|
|
boxes = detections[:, :4] # x1, y1, x2, y2
|
|
|
|
|
|
confidences = detections[:, 4]
|
|
|
|
|
|
class_ids = detections[:, 5].astype(int)
|
|
|
|
|
|
return boxes, confidences, class_ids
|
|
|
|
|
|
|
|
|
|
|
|
recognized_objects = []
|
|
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
|
while True:
|
|
|
|
|
|
# Получение и выравнивание кадров
|
|
|
|
|
|
frames = pipeline.wait_for_frames()
|
|
|
|
|
|
aligned_frames = align.process(frames)
|
|
|
|
|
|
depth_frame = aligned_frames.get_depth_frame()
|
|
|
|
|
|
color_frame = aligned_frames.get_color_frame()
|
|
|
|
|
|
if not depth_frame or not color_frame:
|
|
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
|
|
# Преобразование в numpy массивы
|
|
|
|
|
|
depth_image = np.asanyarray(depth_frame.get_data())
|
|
|
|
|
|
color_image = np.asanyarray(color_frame.get_data())
|
|
|
|
|
|
|
|
|
|
|
|
# Детекция объектов
|
|
|
|
|
|
boxes, confidences, class_ids = detect_objects(color_image)
|
|
|
|
|
|
|
|
|
|
|
|
# Для каждого обнаруженного объекта
|
|
|
|
|
|
for i, box in enumerate(boxes):
|
|
|
|
|
|
x1, y1, x2, y2 = box.astype(int)
|
|
|
|
|
|
w = x2 - x1
|
|
|
|
|
|
h = y2 - y1
|
|
|
|
|
|
|
|
|
|
|
|
# Проверка границ изображения
|
|
|
|
|
|
x1, y1 = max(0, x1), max(0, y1)
|
|
|
|
|
|
x2, y2 = min(color_image.shape[1], x2), min(color_image.shape[0], y2)
|
|
|
|
|
|
|
|
|
|
|
|
# Получение средней глубины в области бокса
|
|
|
|
|
|
depth_roi = depth_image[y1:y2, x1:x2]
|
|
|
|
|
|
depth_roi = depth_roi.astype(float)
|
|
|
|
|
|
depth_scale = depth_frame.get_units()
|
|
|
|
|
|
depth_value = depth_roi.mean() * depth_scale # В метрах
|
|
|
|
|
|
|
|
|
|
|
|
# Обновление массива распознанных объектов
|
|
|
|
|
|
recognized_objects.append({
|
|
|
|
|
|
'box': (x1, y1, x2, y2),
|
|
|
|
|
|
'confidence': confidences[i],
|
|
|
|
|
|
'class_id': class_ids[i],
|
|
|
|
|
|
'depth': depth_value,
|
|
|
|
|
|
})
|
|
|
|
|
|
|
|
|
|
|
|
# Отображение результатов
|
|
|
|
|
|
cv2.rectangle(color_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
|
|
label = f"{model.names[class_ids[i]]}: {depth_value:.2f}m"
|
|
|
|
|
|
cv2.putText(color_image, label, (x1, y1 - 10),
|
|
|
|
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (36, 255, 12), 2)
|
|
|
|
|
|
|
|
|
|
|
|
# Отображение изображения
|
|
|
|
|
|
cv2.imshow('RealSense', color_image)
|
|
|
|
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
|
|
|
|
break
|
|
|
|
|
|
|
|
|
|
|
|
# Очистка списка распознанных объектов для следующего кадра
|
|
|
|
|
|
recognized_objects.clear()
|
|
|
|
|
|
|
|
|
|
|
|
finally:
|
|
|
|
|
|
# Остановка потока
|
|
|
|
|
|
pipeline.stop()
|
|
|
|
|
|
cv2.destroyAllWindows()
|
