# vision/detector.py
import os
from ultralytics import YOLO
from vision.angle_detector import get_current_door_angle
from vision.overflow_detector import detect_overflow_from_image
from vision.alignment_detector import detect_vehicle_alignment


class VisionDetector:
    def __init__(self, settings):
        self.settings = settings

        # 模型实例
        self.angle_model = None
        self.overflow_model = None
        self.alignment_model = None

    def load_models(self):
        """
        加载所有视觉检测模型
        """
        success = True

        # 加载夹角检测模型
        try:
            if not os.path.exists(self.settings.angle_model_path):
                print(f"夹角检测模型不存在: {self.settings.angle_model_path}")
                success = False
            else:
                # 注意：angle.pt模型通过predict_obb_best_angle函数使用，不需要预加载
                print(f"夹角检测模型路径: {self.settings.angle_model_path}")
        except Exception as e:
            print(f"检查夹角检测模型失败: {e}")
            success = False

        # 加载堆料检测模型
        try:
            if not os.path.exists(self.settings.overflow_model_path):
                print(f"堆料检测模型不存在: {self.settings.overflow_model_path}")
                success = False
            else:
                self.overflow_model = YOLO(self.settings.overflow_model_path)
                print(f"成功加载堆料检测模型: {self.settings.overflow_model_path}")
        except Exception as e:
            print(f"加载堆料检测模型失败: {e}")
            success = False

        # 加载对齐检测模型
        try:
            if not os.path.exists(self.settings.alignment_model_path):
                print(f"对齐检测模型不存在: {self.settings.alignment_model_path}")
                success = False
            else:
                self.alignment_model = YOLO(self.settings.alignment_model_path)
                print(f"成功加载对齐检测模型: {self.settings.alignment_model_path}")
        except Exception as e:
            print(f"加载对齐检测模型失败: {e}")
            success = False

        return success

    def detect_angle(self, image=None, image_path=None):
        """
        通过视觉系统获取当前出砼门角度
        """
        return get_current_door_angle(
            model=self.angle_model,
            image=image,
            image_path=image_path
        )

    def detect_overflow(self, image_array):
        """
        通过图像检测是否溢料
        """
        return detect_overflow_from_image(
            image_array,
            self.overflow_model,
            self.settings.roi_file_path
        )

    def detect_vehicle_alignment(self, image_array):
        """
        通过图像检测模具车是否对齐
        """
        return detect_vehicle_alignment(image_array, self.alignment_model)