Feeding_control_system/vision/detector.py

# vision/detector.py
import os
import cv2
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
        #model路径在对应的模型里面
        # self.alignment_model = os.path.join(current_dir, "align_model/yolov11_cls_640v6.rknn")
        
    def load_models(self):
        """
        加载所有视觉检测模型
        """
        from ultralytics import YOLO
        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 = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image=cv2.flip(image, 0)
        return get_current_door_angle(
            model=self.settings.angle_model_path,
            image=image
        )

    def detect_overflow(self, image_array):
        """
        通过图像检测是否溢料
        """
        # image_array=cv2.flip(image_array, 0)
        # cv2.imwrite('test.jpg', image_array)
        cv2.namedWindow("Alignment", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("Alignment", 640, 480)
        cv2.imshow("Alignment", image_array)
        cv2.waitKey(1)
        print('path:', self.settings.overflow_model_path)
        print('roi:', self.settings.roi_file_path)
        return detect_overflow_from_image(
            self.settings.overflow_model_path,
            self.settings.roi_file_path,
            image_array
        )

    def detect_vehicle_alignment(self, image_array):
        """
        通过图像检测模具车是否对齐
        """
        image_array = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)
        image_array=cv2.flip(image_array, 0)
        # cv2.namedWindow("Alignment", cv2.WINDOW_NORMAL)
        # cv2.resizeWindow("Alignment", 640, 480)
        # cv2.imshow("Alignment", image_array)
        # cv2.waitKey(1)
        
        return detect_vehicle_alignment(image_array)