210 lines
7.6 KiB
Python
210 lines
7.6 KiB
Python
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) |