31 lines
949 B
Python
31 lines
949 B
Python
|
# vision/alignment_detector.py
|
||
|
|
def detect_vehicle_alignment(image_array, alignment_model):
|
||
|
|
"""
|
||
|
|
通过图像检测模具车是否对齐
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
# 检查模型是否已加载
|
||
|
|
if alignment_model is None:
|
||
|
|
print("对齐检测模型未加载")
|
||
|
|
return False
|
||
|
|
|
||
|
|
if image_array is None:
|
||
|
|
print("输入图像为空")
|
||
|
|
return False
|
||
|
|
|
||
|
|
# 直接使用模型进行推理
|
||
|
|
results = alignment_model(image_array)
|
||
|
|
pared_probs = results[0].probs.data.cpu().numpy().flatten()
|
||
|
|
|
||
|
|
# 类别0: 未对齐, 类别1: 对齐
|
||
|
|
class_id = int(pared_probs.argmax())
|
||
|
|
confidence = float(pared_probs[class_id])
|
||
|
|
|
||
|
|
# 只有当对齐且置信度>95%时才认为对齐
|
||
|
|
if class_id == 1 and confidence > 0.95:
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
except Exception as e:
|
||
|
|
print(f"对齐检测失败: {e}")
|
||
|
|
return False
|