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()