from ultralytics import YOLO import cv2 import numpy as np import os # ------------------ 配置 ------------------ model_path = '../ultralytics_yolo11-main/runs/train/exp4/weights/best.pt' img_folder = '/home/hx/yolo/ultralytics_yolo11-main/dataset1/test' output_mask_dir = 'output_masks1' os.makedirs(output_mask_dir, exist_ok=True) SUPPORTED_FORMATS = ('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff') # ------------------ 加载模型 ------------------ model = YOLO(model_path) model.to('cuda') # 使用 GPU（如有） def get_long_edge_vector(rect): """ 从 minAreaRect 中提取长边的方向向量（单位向量） rect: cv2.minAreaRect 返回的 (center, (w, h), angle) """ center, (width, height), angle = rect # OpenCV 的 angle 范围是 [-90, 0) # 我们要的是长边的方向 if width >= height: rad = np.radians(angle) # 长边方向 else: rad = np.radians(angle + 90) # 短边变长边 dx = np.cos(rad) dy = np.sin(rad) direction = np.array([dx, dy]) norm = np.linalg.norm(direction) return direction if norm < 1e-8 else direction / norm def get_contour_center(contour): """计算轮廓质心""" M = cv2.moments(contour) if M["m00"] == 0: return np.array([0, 0]) return np.array([int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"])]) def calculate_jaw_angle(jaw1, jaw2): """ 计算两个夹具之间的开合角度（取最小内角） 返回: (angle, dir1, dir2) """ # 获取长边方向 dir1_orig = get_long_edge_vector(jaw1['rect']) dir2_orig = get_long_edge_vector(jaw2['rect']) # 计算夹角 cos_angle = np.clip(np.dot(dir1_orig, dir2_orig), -1.0, 1.0) angle = np.degrees(np.arccos(cos_angle)) # 取最小内角（0 ~ 90°） opening_angle = min(angle, 180 - angle) # --- 可选：让方向指向夹具中心（箭头朝内）--- center1 = get_contour_center(jaw1['contour']) center2 = get_contour_center(jaw2['contour']) fixture_center = (center1 + center2) / 2.0 to_center1 = fixture_center - center1 if np.linalg.norm(to_center1) > 1e-6: to_center1 = to_center1 / np.linalg.norm(to_center1) if np.dot(dir1_orig, to_center1) < 0: dir1_orig = -dir1_orig to_center2 = fixture_center - center2 if np.linalg.norm(to_center2) > 1e-6: to_center2 = to_center2 / np.linalg.norm(to_center2) if np.dot(dir2_orig, to_center2) < 0: dir2_orig = -dir2_orig return opening_angle, dir1_orig, dir2_orig def process_image(img_path, output_dir): img = cv2.imread(img_path) if img is None: print(f"❌ 无法读取图像: {img_path}") return h, w = img.shape[:2] filename = os.path.basename(img_path) name_only = os.path.splitext(filename)[0] print(f"\n🔄 正在处理: {filename}") # 创建单通道掩码 composite_mask = np.zeros((h, w), dtype=np.uint8) results = model(img_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) # 创建三通道可视化图 vis_mask = np.stack([composite_mask] * 3, axis=-1) vis_mask[composite_mask > 0] = [255, 255, 255] if len(jaws) < 2: print(f"⚠️ 检测到的夹具少于2个（共{len(jaws)}个）") cv2.imwrite(os.path.join(output_dir, f'mask_{name_only}.png'), composite_mask) return # 按面积排序，取最大的两个 jaws.sort(key=lambda x: x['area'], reverse=True) jaw1, jaw2 = jaws[0], jaws[1] # === 计算夹角和方向 === opening_angle, dir1, dir2 = calculate_jaw_angle(jaw1, jaw2) print(f"✅ 夹具开合角度: {opening_angle:.2f}°") # === 可视化 === center1 = get_contour_center(jaw1['contour']) center2 = get_contour_center(jaw2['contour']) fixture_center = ((center1[0] + center2[0]) // 2, (center1[1] + center2[1]) // 2) # 绘制最小外接矩形 box1 = cv2.boxPoints(jaw1['rect']).astype(int) box2 = cv2.boxPoints(jaw2['rect']).astype(int) cv2.drawContours(vis_mask, [box1], 0, (0, 0, 255), 2) # jaw1: 红色 cv2.drawContours(vis_mask, [box2], 0, (255, 0, 0), 2) # jaw2: 蓝色 # 绘制主方向箭头（绿色，长度可调） scale = 60 end1 = center1 + scale * dir1 end2 = center2 + scale * dir2 cv2.arrowedLine(vis_mask, tuple(center1), tuple(end1.astype(int)), (0, 255, 0), 2, tipLength=0.3) cv2.arrowedLine(vis_mask, tuple(center2), tuple(end2.astype(int)), (0, 255, 0), 2, tipLength=0.3) # 标注中心点和角度 cv2.circle(vis_mask, fixture_center, 5, (255, 255, 0), -1) cv2.putText(vis_mask, "Center", (fixture_center[0] + 10, fixture_center[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1) cv2.putText(vis_mask, 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_mask) print(f"✅ 结果已保存: {save_path}") # ------------------ 主程序 ------------------ if __name__ == '__main__': if not os.path.exists(img_folder): print(f"❌ 图像文件夹不存在: {img_folder}") exit() image_files = [f for f in os.listdir(img_folder) if f.lower().endswith(SUPPORTED_FORMATS)] if len(image_files) == 0: print(f"⚠️ 未找到支持的图像文件") exit() print(f"✅ 发现 {len(image_files)} 张图像，开始处理...") for file in image_files: process_image(os.path.join(img_folder, file), output_mask_dir) print(f"\n🎉 所有图像处理完成！结果保存在: {output_mask_dir}")