160 lines
5.5 KiB
Python
160 lines
5.5 KiB
Python
from ultralytics import YOLO
|
||
import cv2
|
||
import numpy as np
|
||
import os
|
||
|
||
# ------------------ 模型与路径配置 ------------------
|
||
MODEL_PATH = 'ultralytics_yolo11-main/runs/train/exp4/weights/best.pt'
|
||
OUTPUT_DIR = 'test_image'
|
||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||
|
||
# 加载模型（全局一次）
|
||
model = YOLO(MODEL_PATH)
|
||
model.to('cuda') # 如无 GPU，可改为 'cpu'
|
||
|
||
|
||
def detect_jaw_angle(image_path, mode='show'):
|
||
"""
|
||
检测图像中两个夹具的开合角度
|
||
|
||
参数:
|
||
image_path (str): 输入图像路径
|
||
mode (str): 'show' -> 保存可视化结果；'silent' -> 只返回角度
|
||
|
||
返回:
|
||
float: 夹具开合角度 (0 ~ 90°)
|
||
"""
|
||
img = cv2.imread(image_path)
|
||
if img is None:
|
||
raise FileNotFoundError(f"无法读取图像: {image_path}")
|
||
|
||
h, w = img.shape[:2]
|
||
filename = os.path.basename(image_path)
|
||
name_only = os.path.splitext(filename)[0]
|
||
|
||
# 创建掩码并检测
|
||
composite_mask = np.zeros((h, w), dtype=np.uint8)
|
||
results = model(image_path, imgsz=1280, conf=0.5)
|
||
|
||
jaws = []
|
||
for r in results:
|
||
if r.masks is not None:
|
||
masks = r.masks.data.cpu().numpy()
|
||
boxes = r.boxes.xyxy.cpu().numpy()
|
||
|
||
for i, mask in enumerate(masks):
|
||
x1, y1, x2, y2 = map(int, boxes[i])
|
||
x1, y1 = max(0, x1), max(0, y1)
|
||
x2, y2 = min(w, x2), min(h, y2)
|
||
|
||
obj_mask = np.zeros((h, w), dtype=np.uint8)
|
||
mask_resized = cv2.resize(mask, (w, h))
|
||
obj_mask[y1:y2, x1:x2] = (mask_resized[y1:y2, x1:x2] * 255).astype(np.uint8)
|
||
|
||
contours, _ = cv2.findContours(obj_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
||
if len(contours) == 0:
|
||
continue
|
||
|
||
largest_contour = max(contours, key=cv2.contourArea)
|
||
area = cv2.contourArea(largest_contour)
|
||
if area < 100:
|
||
continue
|
||
|
||
rect = cv2.minAreaRect(largest_contour)
|
||
jaws.append({
|
||
'contour': largest_contour,
|
||
'rect': rect,
|
||
'area': area
|
||
})
|
||
composite_mask = np.maximum(composite_mask, obj_mask)
|
||
|
||
# 必须检测到至少两个夹具
|
||
if len(jaws) < 2:
|
||
print(f"检测到的夹具少于2个（共{len(jaws)}个）")
|
||
if mode == 'show':
|
||
cv2.imwrite(os.path.join(OUTPUT_DIR, f'mask_{name_only}.png'), composite_mask)
|
||
return None
|
||
|
||
# 取面积最大的两个
|
||
jaws.sort(key=lambda x: x['area'], reverse=True)
|
||
jaw1, jaw2 = jaws[0], jaws[1]
|
||
|
||
# ------------------ 计算角度 ------------------
|
||
def get_long_edge_vector(rect):
|
||
center, (w, h), angle = rect
|
||
rad = np.radians(angle + (0 if w >= h else 90))
|
||
return np.array([np.cos(rad), np.sin(rad)])
|
||
|
||
def get_center(contour):
|
||
M = cv2.moments(contour)
|
||
return np.array([M['m10']/M['m00'], M['m01']/M['m00']]) if M['m00'] != 0 else np.array([0,0])
|
||
|
||
dir1 = get_long_edge_vector(jaw1['rect'])
|
||
dir2 = get_long_edge_vector(jaw2['rect'])
|
||
|
||
# 校正方向：指向夹具中心
|
||
center1 = get_center(jaw1['contour'])
|
||
center2 = get_center(jaw2['contour'])
|
||
fixture_center = (center1 + center2) / 2
|
||
|
||
to_center1 = fixture_center - center1
|
||
to_center2 = fixture_center - center2
|
||
|
||
if np.linalg.norm(to_center1) > 1e-6 and np.dot(dir1, to_center1) < 0:
|
||
dir1 = -dir1
|
||
if np.linalg.norm(to_center2) > 1e-6 and np.dot(dir2, to_center2) < 0:
|
||
dir2 = -dir2
|
||
|
||
# 计算最小内角
|
||
cos_angle = np.clip(np.dot(dir1, dir2), -1.0, 1.0)
|
||
angle = np.degrees(np.arccos(cos_angle))
|
||
opening_angle = min(angle, 180 - angle)
|
||
|
||
# ------------------ 可视化（仅 mode='show'）------------------
|
||
if mode == 'show':
|
||
vis_img = np.stack([composite_mask]*3, axis=-1)
|
||
vis_img[composite_mask > 0] = [255, 255, 255]
|
||
|
||
# 绘制矩形框
|
||
box1 = np.int32(cv2.boxPoints(jaw1['rect']))
|
||
box2 = np.int32(cv2.boxPoints(jaw2['rect']))
|
||
cv2.drawContours(vis_img, [box1], 0, (0, 0, 255), 2)
|
||
cv2.drawContours(vis_img, [box2], 0, (255, 0, 0), 2)
|
||
|
||
# 绘制方向箭头
|
||
scale = 60
|
||
c1 = tuple(np.int32(center1))
|
||
c2 = tuple(np.int32(center2))
|
||
end1 = tuple(np.int32(center1 + scale * dir1))
|
||
end2 = tuple(np.int32(center2 + scale * dir2))
|
||
cv2.arrowedLine(vis_img, c1, end1, (0, 255, 0), 2, tipLength=0.3)
|
||
cv2.arrowedLine(vis_img, c2, end2, (0, 255, 0), 2, tipLength=0.3)
|
||
|
||
# 标注角度
|
||
cv2.putText(vis_img, f"Angle: {opening_angle:.2f}°", (20, 50),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
|
||
|
||
# 保存
|
||
save_path = os.path.join(OUTPUT_DIR, f'mask_with_angle_{name_only}.png')
|
||
cv2.imwrite(save_path, vis_img)
|
||
print(f"✅ 可视化结果已保存: {save_path}")
|
||
|
||
return round(opening_angle, 2)
|
||
|
||
# ------------------ 主函数 ------------------
|
||
if __name__ == '__main__':
|
||
# ✅ 设置输入图像路径
|
||
image_path = '/home/hx/yolo/test_image/1.png' # ← 修改为你自己的图片路径
|
||
|
||
# ✅ 模式选择：
|
||
# mode='show': 保存可视化图像
|
||
# mode='silent': 只返回角度
|
||
mode = 'silent' # 或 'silent'
|
||
|
||
print(f"🔍 正在处理图像: {image_path}")
|
||
angle = detect_jaw_angle(image_path, mode=mode)
|
||
|
||
if angle is not None:
|
||
print(f"🎉 检测到的夹具开合角度: {angle}°")
|
||
else:
|
||
print("❌ 未能检测到足够的夹具") |