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)
|