from ultralytics import YOLO import cv2 import numpy as np import os # ================================ # 配置区（所有固定参数放在这里） # ================================ # --- 模型与类别 --- CLASS_NAMES = ['ban', 'bag'] # 类别名称 WEIGHT_PATH = r"/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb_ailai2/weights/best.pt" IMAGE_PATH = r"/home/hx/ailai_image_obb/ailai_obb/camera01/1.jpg" # --- 结果保存路径 --- RESULT_DIR = "./inference_results" os.makedirs(RESULT_DIR, exist_ok=True) # 确保目录存在 RESULT_IMAGE_PATH = os.path.join(RESULT_DIR, os.path.basename(IMAGE_PATH)) # 保存路径 FIXED_BAG_CENTER_FILE = os.path.join(RESULT_DIR, "bag_center.txt") # 固定参考点，只读' # --- ban 的实际物理尺寸（单位：毫米 mm）--- BAN_REAL_WIDTH_MM = 890 # 实际宽度（例如：800 mm） BAN_REAL_HEIGHT_MM = 1990 # 实际高度（例如：1990 mm） # --- ban 在图像中的平均像素尺寸（由模型检测得到，稳定值）--- BAN_PIXEL_WIDTH_PX = 536.35 # 图像中检测到的平均宽度（像素） BAN_PIXEL_HEIGHT_PX = 242.83 # 图像中检测到的平均高度（像素） # --- 计算缩放因子（mm/px）--- SCALE_X = BAN_REAL_WIDTH_MM / BAN_PIXEL_WIDTH_PX # mm/px SCALE_Y = BAN_REAL_HEIGHT_MM / BAN_PIXEL_HEIGHT_PX # mm/px print(f"标定信息 1 像素 = {SCALE_X:.3f} mm (水平) {SCALE_Y:.3f} mm (垂直)") # --- 固定 bag 角度（单位：度）--- FIXED_BAG_ANGLE_DEG = 0.0 # 固定参考角度 # ================================ # 工具函数 # ================================ def pixels_to_physical(dx_px, dy_px): """ 将像素偏移量转换为物理偏移量（毫米） """ dx_mm = dx_px * SCALE_X dy_mm = dy_px * SCALE_Y return dx_mm, dy_mm def load_center_from_txt(txt_path): """加载文本中的中心点 (cx, cy)""" if not os.path.exists(txt_path): print(f"文件不存在 {txt_path}") return None try: with open(txt_path, "r") as f: line = f.readline().strip() if not line: return None x, y = map(float, line.split()) return (x, y) except Exception as e: print(f"读取文件失败 {txt_path} {e}") return None def get_angles_and_centers(image_path, weight_path): """获取每个类别的角度和中心点""" img = cv2.imread(image_path) if img is None: print(f"无法读取图像 {image_path}") return {cls: None for cls in CLASS_NAMES}, None model = YOLO(weight_path) results = model(img, imgsz=640, conf=0.15, task='obb') result = results[0] boxes = result.obb if boxes is None or len(boxes) == 0: print("未检测到任何目标") return {cls: None for cls in CLASS_NAMES}, img xywhr = boxes.xywhr.cpu().numpy() confs = boxes.conf.cpu().numpy() class_ids = boxes.cls.cpu().numpy().astype(int) angles_and_centers = {} for class_id, class_name in enumerate(CLASS_NAMES): mask = (class_ids == class_id) if not np.any(mask): print(f"未检测到类别 {class_name}") angles_and_centers[class_name] = None continue idx_in_class = np.argmax(confs[mask]) global_idx = np.where(mask)[0][idx_in_class] cx, cy, w, h, r = xywhr[global_idx] angle_deg = np.degrees(r) angles_and_centers[class_name] = { 'angle': angle_deg, 'center': (float(cx), float(cy)), 'width': float(w), 'height': float(h), 'box_points': cv2.boxPoints(((cx, cy), (w, h), angle_deg)) # 获取四个角点 } return angles_and_centers, img def draw_results_and_save(image, angles_and_centers, fixed_bag_center, save_path): """绘制旋转框、中心点、标签、偏移向量，并保存图像""" vis_img = image.copy() # 颜色定义 colors = { 'ban': (0, 255, 0), # 绿色 'bag': (255, 0, 0) # 蓝色 } # 绘制每个检测结果 for cls_name, info in angles_and_centers.items(): if info is None: continue # 绘制旋转框 box_pts = np.int32(info['box_points']) cv2.drawContours(vis_img, [box_pts], 0, colors[cls_name], 2) # 绘制中心点 cx, cy = int(info['center'][0]), int(info['center'][1]) cv2.circle(vis_img, (cx, cy), 5, (0, 0, 255), -1) # 红色实心圆 # 添加标签 label = f"{cls_name} {info['angle']:.1f}°" cv2.putText(vis_img, label, (cx - 30, cy - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, colors[cls_name], 2) # 绘制 fixed_bag_center 和偏移向量 current_bag = angles_and_centers.get('bag') if fixed_bag_center and current_bag: fx, fy = int(fixed_bag_center[0]), int(fixed_bag_center[1]) cx, cy = int(current_bag['center'][0]), int(current_bag['center'][1]) # 绘制 fixed center（空心圆） cv2.circle(vis_img, (fx, fy), 7, (255, 255, 0), 2) # 青色 # 绘制从 fixed 到 current 的箭头（偏移向量） cv2.arrowedLine(vis_img, (fx, fy), (cx, cy), (0, 255, 255), 2, tipLength=0.05) # 添加偏移文字 dx_px = cx - fx dy_px = cy - fy dx_mm, dy_mm = pixels_to_physical(dx_px, dy_px) offset_text = f"Δx={dx_mm:.1f}mm Δy={dy_mm:.1f}mm" cv2.putText(vis_img, offset_text, (fx + 10, fy + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2) # 保存图像 cv2.imwrite(save_path, vis_img) print(f"\n✅ 可视化结果已保存至: {save_path}") # =============================== # 主程序 # =============================== if __name__ == "__main__": # 1. 获取当前检测结果和图像 angles_and_centers, img = get_angles_and_centers(IMAGE_PATH, WEIGHT_PATH) if img is None: exit(1) print("\n各类别旋转角度 中心点") for cls_name, info in angles_and_centers.items(): if info is not None: print(f" {cls_name} 角度 {info['angle']:.2f} 中心点 {info['center'][0]:.2f} {info['center'][1]:.2f} 尺寸 {info['width']:.2f} {info['height']:.2f} px") else: print(f" {cls_name} 未检测到") # 2. 获取固定参考 bag 中心点 fixed_bag_center = load_center_from_txt(FIXED_BAG_CENTER_FILE) print(f"\n固定参考 bag 中心点 {fixed_bag_center}") # 3. 计算偏移量（像素） current_bag = angles_and_centers.get('bag') if current_bag and current_bag['center'] and fixed_bag_center: dx_px = current_bag['center'][0] - fixed_bag_center[0] dy_px = current_bag['center'][1] - fixed_bag_center[1] print(f"\n像素偏移量") print(f" Δx {dx_px:.2f} px Δy {dy_px:.2f} px") # 转换为物理偏移量 dx_mm, dy_mm = pixels_to_physical(dx_px, dy_px) print(f"\n物理偏移量 基于 ban 尺寸标定") print(f" ΔX {dx_mm:.2f} mm ΔY {dy_mm:.2f} mm") # 计算角度偏移量 current_bag_angle = current_bag['angle'] angle_diff = current_bag_angle - FIXED_BAG_ANGLE_DEG print(f"\n角度偏移量") print(f" 当前 bag 角度 {current_bag_angle:.2f}") print(f" 固定 bag 角度 {FIXED_BAG_ANGLE_DEG}") print(f" 角度差 {angle_diff:.2f}") else: print(f"\n偏移量计算失败 数据缺失") if not current_bag: print(" 未检测到 bag") if not fixed_bag_center: print(" 固定参考点文件为空或不存在") # 4. 绘制并保存可视化结果 draw_results_and_save(img, angles_and_centers, fixed_bag_center, RESULT_IMAGE_PATH)