YOLO obb全流程

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

一、数据集

1、数据集工具xanylabeling的安装

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

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 = point1
    x2, y2 = point2
    x_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_vector1
    point2 = np.array(center_point) + scaled_vector2
    
    # 获取新的坐标和中心点
    x1, y1 = point1
    x2, y2 = point2
    x_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 = point1
    x2, y2 = point2
    x_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_vector1
    point2 = np.array(center_point) + scaled_vector2
    
    # 获取新的坐标和中心点
    x1, y1 = point1
    x2, y2 = point2
    x_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

def 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)

  效果展示：