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)