YOLO obb全流程

内容：xanylabeling 数据标注工具；pytorch（python）；yolo-obb 模型

一、数据集

1、数据集工具xanylabeling的安装

（详细配置与使用方法参考：X-Anylabeling自动标注软件安装使用教程含conda环境搭建安装(详细教程)_x-anylabeling安装教程-CSDN博客）

2、xanylabeling的使用

直接图像中标注矩形框（在选中该矩形框后，使用 Z\X\C\V四个按钮设置旋转角度，围绕中心点）。

打开图片目录并对图片操作保存后，在同目录下出现同名的json文件，需要使用json转化成yolo格式的txt文件。

3、数据集搭建

保持图片名称与txt名称一致；整个数据集格式如下，images 存放图片文件，labels的存放txt文件（txt内容为：0 0.530889 0.337543 0.235365 0.369421 0.224447 0.224803 0.519971 0.192925，九个分别表示类别第一顶点x和y 第二顶点x和y 第三顶点x和y 第四顶点x和y）；

二、模型

1、下载yolo-obb

ultralytics/ultralytics: Ultralytics YOLO11 🚀

2、修改配置文件

cfg/datasets文件夹下设置yaml文件：（有些时候相对路径出现问题，所以本图直接使用绝对路径）

cfg/models文件夹下找到对应版本，修改obb的yaml文件，修改nc（也可能不用修改，程序运行中自动获取datasets的yaml设置类别），若是需要修改模型结构，也在这个文件

3、终端运行

下载代码后，出现两个层级的ultralytics文件夹，第一层表示代码根目录，包括docker等部署方法的内容，第二层ultralytics才是我们所需要的代码部分。

将数据集放在第一层ultralytics内后，打开终端运行，直接使用 yolo obb train data=cfg\datasets\dota8.yaml model=cfg\models\v8\yolov8s-obb.yaml epochs=100 batch=4 device=0 imgsz=640 这句命令运行，data表示数据文件，model表示加载模型；

提示：1、运行过程中，出现提示未查询到数据集，自动下载dota8，表示数据文件没有配置好。

2、出现下载yolo11n.pt表示正常运行，需要借助yolo11n进行提高能力

三、模型使用

from ultralytics import YOLO
# 加载训练好的模型
model = YOLO("cfg/models/v8/yolov8s-obb.yaml")
model.load("best.pt")   # 训练后的模型
# 图片地址
image_path = r"..\dataset\test\images"
import os
for img in os.listdir(image_path):# if img == "35.jpg":img_path = os.path.join(image_path,img)# 在图片上运行推理results = model(image_path) # 推理图片r =  results[0]image = r.plot()  # 获取图片以及检测框等数据# 显示结果图像cv2.imshow('Image with Center Points', image )cv2.waitKey(0)cv2.destroyAllWindows()

四、修改模型（可选）

混合注意力为例

1、修改cfg/models/v8/yolov8-obb.yaml

在对应版本的yaml的文件中添加模型（也可以删除或变更）

添加混合注意力层次（-1表示上层输出通道作为本层输入通道、1表示模块数量、CBAM表示模块名称、[1024，7]表示输入的参数，1024是本层输出通道数，7表示卷积核尺寸）。

2、修改nn文件夹

第一处：一般而言在尝试用的模块（CBAM\DWCONV等卷积）在nn/conv.py内存在，若在nn文件夹下的其他py均未定义需要使用模块，那么需要自行添加（根据其功能在对应文件内添加，或者直接新增文件）

添加模块：使用class定义该模块，参数与yaml文件保持一致

第二处：在 nn/_init_.py 文件内添加模块导入导出（图表示该模块位于conv文件，以此类推）

第三处：在/task.py 文件中 from ultralytics.nn.modules import (）内添加CBAM模块；在base_modules = frozenset(）中添加CBAM；在解析模块中添加参数解析，将第一个参数设置为输出通道和输入通道，第二个参数设置为卷积核大小

3、直接训练

修改完毕后继续训练，在训练输出内容中查看模型结构是否得到应用。

验证方式与之前一致。

五、具体任务（参考）

此次任务是需要进行检测旋转框，并根据检测内容分析物体方向（我使用yoloobb+opencv的方法），任务目标：检测巧克力的位置并判断其视觉方向。

思路：

索取yolo检测结果

results = model(image_path) # 推理图片
r =  results[0]
image = r.plot()  # 获取图片以及检测框等数据
boxes = r.obb.xyxyxyxy.detach() # 获取四个角点坐标

根据检测的四个点进行判断长度和中心点

points = []
for box in boxes:zreo_point = box[0]first_point = box[1]second_point = box[2]three_point = box[3]loss0 = torch.sum(torch.pow(zreo_point-first_point, 2))loss1 = torch.sum(torch.pow(first_point-second_point, 2))if loss0<=loss1:points.append((zreo_point+first_point)/2)points.append((second_point+three_point)/2)else:   points.append((first_point+second_point)/2)points.append((zreo_point+three_point)/2)

画出宽中心点连接中心的两个线段，根据图像色彩判断方向（二值化的腐蚀与膨胀），保留目标线段并画在检测结果的图像数据中。

def get_line_colors(image, point1, point2):# 获取坐标和中心点x1, y1 = point1x2, y2 = point2x_center = int((x1 + x2)/2)y_center = int((y1 + y2)/2)center_point = (x_center, y_center)# 缩放点位 让宽中心点向物体中心点靠近vector1 = np.array(point1) - np.array(center_point)vector2 = np.array(point2) - np.array(center_point)# 缩放向量scaled_vector1 = vector1 * 3/4 scaled_vector2 = vector2 * 3/4# 计算新的终点坐标point1 = np.array(center_point) + scaled_vector1point2 = np.array(center_point) + scaled_vector2# 获取新的坐标和中心点x1, y1 = point1x2, y2 = point2x_center = int((x1 + x2)/2)y_center = int((y1 + y2)/2)# 使用 OpenCV 的 line iterator 来获取沿线的像素，判断像素的变化得到视觉的方向line1_iterator = cv2.line(np.zeros_like(image), (int(x1), int(y1)), (int(x_center), int(y_center)), color=1, thickness=2)line2_iterator = cv2.line(np.zeros_like(image), (int(x2), int(y2)), (int(x_center), int(y_center)), color=1, thickness=2)colors1 = []colors2 = []# 图片处理mage = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)h, s, v = cv2.split(mage)# 增强S通道（饱和度）s_enhanced = cv2.equalizeHist(s)enhanced_image = cv2.merge((h, s_enhanced, v))BGR_image = cv2.cvtColor(enhanced_image, cv2.COLOR_HSV2BGR)GRAY_image = cv2.cvtColor(BGR_image, cv2.COLOR_BGR2GRAY)_, bin_image = cv2.threshold(GRAY_image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)kernel = np.ones((7, 7), np.uint8)  # 创建一个5x5的方形结构元素bin_image = cv2.dilate(bin_image, kernel, iterations=5)bin_image = cv2.erode(bin_image, kernel, iterations=9)for point in zip(*np.where(line1_iterator)):colors1.append(bin_image[point[0], point[1]])for point in zip(*np.where(line2_iterator)):colors2.append(bin_image[point[0], point[1]])sum_color1 = np.average(colors1)sum_color2 = np.average(colors2) if sum_color1>sum_color2:cv2.arrowedLine(image, (int(x_center), int(y_center)), (int(x1), int(y1)), (255, 0, 0), 2)  # 标记连接点的线段else:cv2.arrowedLine(image, (int(x_center), int(y_center)), (int(x2), int(y2)), (255, 0, 0), 2)  # 标记连接点的线段return image

完整代码如下所示：

import numpy as np
import cv2 
import torch
from ultralytics import YOLO# 获取boxes的框信息
def point_list(boxes):points = []for box in boxes:zreo_point = box[0]first_point = box[1]second_point = box[2]three_point = box[3]loss0 = torch.sum(torch.pow(zreo_point-first_point, 2))loss1 = torch.sum(torch.pow(first_point-second_point, 2))if loss0<=loss1:points.append((zreo_point+first_point)/2)points.append((second_point+three_point)/2)else:   points.append((first_point+second_point)/2)points.append((zreo_point+three_point)/2)return points
# 根据两个线段判断正反情况
def get_line_colors(image, point1, point2):# 获取坐标和中心点x1, y1 = point1x2, y2 = point2x_center = int((x1 + x2)/2)y_center = int((y1 + y2)/2)center_point = (x_center, y_center)# 缩放点位vector1 = np.array(point1) - np.array(center_point)vector2 = np.array(point2) - np.array(center_point)# 缩放向量scaled_vector1 = vector1 * 3/4scaled_vector2 = vector2 * 3/4# 计算新的终点坐标point1 = np.array(center_point) + scaled_vector1point2 = np.array(center_point) + scaled_vector2# 获取新的坐标和中心点x1, y1 = point1x2, y2 = point2x_center = int((x1 + x2)/2)y_center = int((y1 + y2)/2)# 使用 OpenCV 的 line iterator 来获取沿线的像素line1_iterator = cv2.line(np.zeros_like(image), (int(x1), int(y1)), (int(x_center), int(y_center)), color=1, thickness=2)line2_iterator = cv2.line(np.zeros_like(image), (int(x2), int(y2)), (int(x_center), int(y_center)), color=1, thickness=2)colors1 = []colors2 = []# 图片处理mage = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)h, s, v = cv2.split(mage)# 增强S通道（饱和度）s_enhanced = cv2.equalizeHist(s)enhanced_image = cv2.merge((h, s_enhanced, v))BGR_image = cv2.cvtColor(enhanced_image, cv2.COLOR_HSV2BGR)GRAY_image = cv2.cvtColor(BGR_image, cv2.COLOR_BGR2GRAY)_, bin_image = cv2.threshold(GRAY_image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)kernel = np.ones((7, 7), np.uint8)  # 创建一个5x5的方形结构元素bin_image = cv2.dilate(bin_image, kernel, iterations=5)bin_image = cv2.erode(bin_image, kernel, iterations=9)for point in zip(*np.where(line1_iterator)):colors1.append(bin_image[point[0], point[1]])for point in zip(*np.where(line2_iterator)):colors2.append(bin_image[point[0], point[1]])sum_color1 = np.average(colors1)sum_color2 = np.average(colors2) if sum_color1>sum_color2:cv2.arrowedLine(image, (int(x_center), int(y_center)), (int(x1), int(y1)), (255, 0, 0), 2)  # 标记连接点的线段else:cv2.arrowedLine(image, (int(x_center), int(y_center)), (int(x2), int(y2)), (255, 0, 0), 2)  # 标记连接点的线段return imagedef predict(model,image_path):# 在图片上运行推理results = model(image_path) # 推理图片r =  results[0]image = r.plot()  # 获取图片以及检测框等数据boxes = r.obb.xyxyxyxy.detach() # 获取四个角点坐标points =  point_list(boxes) # 获取检测框宽的两个中心点# print(point_list(boxes)) # 遍历每个检测框for i in range(0,len(points)-1, 2):# 将坐标转换为整数point1 = tuple(map(int, points[i])) point2 = tuple(map(int, points[i+1]))image = get_line_colors(image, point1, point2) # 根据两个端点获取线段数据并判断正反new_width = int(image.shape[1] * 0.5)new_height = int(image.shape[0] * 0.5)resized_image =  cv2.resize(image, (new_width, new_height))return resized_image# 加载训练好的模型
model = YOLO("./CBAM300best.pt")
# model.load("../CBAM300best.pt")  
# 图片地址
image_path = r"..\dataset\test\images"
import os
for img in os.listdir(image_path):# if img == "35.jpg":img_path = os.path.join(image_path,img)resized_image = predict(model,img_path)image_name = img.split(".")image_name[0] = image_name[0]+"-1"new_image_path = ".".join(image_name)save_path = os.path.join(image_path,new_image_path)print(save_path)cv2.imwrite(save_path,resized_image)

效果展示：