87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
|
# 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)
|