import cv2
from ultralytics import YOLO
import numpy as np

class DataProcessor:
    def __init__(self, model_path='./../app/models/yolov9e_object_classification.pt'):
        self.model = YOLO(model_path)

    def __decode_image_from_base64(self, data, color=True):
        decoded = base64.b64decode(data)
        nparr = np.frombuffer(decoded, np.uint8)
        return cv2.imdecode(nparr, cv2.IMREAD_COLOR if color else cv2.IMREAD_GRAYSCALE)

    def inline_detection(self, data):
        color_matrix = self.__decode_image_from_base64(data['rgb'], color=True)
        depth_matrix = self.__decode_image_from_base64(data['depth'], color=False)

        if depth_matrix is None or color_matrix is None:
            print("Не удалось получить изображения. Проверьте подключение камеры.")
            return [], None

        results = self.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 = self.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': float(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)

        return detections_info, annotated_image