106 lines
3.9 KiB
Python
106 lines
3.9 KiB
Python
# 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
|
||
from config.settings import app_set_config
|
||
|
||
|
||
class VisionDetector:
|
||
def __init__(self):
|
||
pass
|
||
#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(app_set_config.angle_model_path):
|
||
print(f"夹角检测模型不存在: {app_set_config.angle_model_path}")
|
||
success = False
|
||
else:
|
||
# 注意：angle.pt模型通过predict_obb_best_angle函数使用，不需要预加载
|
||
print(f"夹角检测模型路径: {app_set_config.angle_model_path}")
|
||
except Exception as e:
|
||
print(f"检查夹角检测模型失败: {e}")
|
||
success = False
|
||
|
||
# 加载堆料检测模型
|
||
try:
|
||
if not os.path.exists(app_set_config.overflow_model_path):
|
||
print(f"堆料检测模型不存在: {app_set_config.overflow_model_path}")
|
||
success = False
|
||
else:
|
||
self.overflow_model = YOLO(app_set_config.overflow_model_path)
|
||
print(f"成功加载堆料检测模型: {app_set_config.overflow_model_path}")
|
||
except Exception as e:
|
||
print(f"加载堆料检测模型失败: {e}")
|
||
success = False
|
||
|
||
# 加载对齐检测模型
|
||
try:
|
||
if not os.path.exists(app_set_config.alignment_model_path):
|
||
print(f"对齐检测模型不存在: {app_set_config.alignment_model_path}")
|
||
success = False
|
||
else:
|
||
self.alignment_model = YOLO(app_set_config.alignment_model_path)
|
||
print(f"成功加载对齐检测模型: {app_set_config.alignment_model_path}")
|
||
except Exception as e:
|
||
print(f"加载对齐检测模型失败: {e}")
|
||
success = False
|
||
|
||
return success
|
||
|
||
def detect_angle(self, image=None):
|
||
"""
|
||
通过视觉系统获取当前出砼门角度
|
||
"""
|
||
|
||
# if image is None:
|
||
# image=cv2.imread(os.path.join(app_set_config.project_root,'test.jpg'))
|
||
image=cv2.flip(image, 0)
|
||
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||
# cv2.imshow('test',image)
|
||
return get_current_door_angle(
|
||
model=app_set_config.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:', app_set_config.overflow_model_path)
|
||
# print('roi:', app_set_config.roi_file_path)
|
||
return detect_overflow_from_image(
|
||
app_set_config.overflow_model_path,
|
||
app_set_config.roi_file_path,
|
||
image_array
|
||
)
|
||
|
||
def detect_vehicle_alignment(self, image_array)->bool:
|
||
"""
|
||
通过图像检测模具车是否对齐
|
||
"""
|
||
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)
|