YOLO作为单阶段检测方法，可以直接端到端的输出目标对象位置和类别，而在一些大型无人机遥感等目标检测任务中，使用单阶段检测往往会产生类别预测错误的问题。 正好，YOLOv5-6.2版本提供了一个图像分类的网络，那么就可以借此将YOLOv5进行改造，分离检测和识别的过程。 一阶段识别目标，并将目标框裁剪出来得到图片，然后输入到图像分类网络进行筛选，最后进行显示。

这个思路的核心是设定一套编码规则，来让两阶段能够平滑地进行过渡。 我的思路是将一个裁剪出的对象直接通过文件名来和对应的label标签进行绑定，具体规则如下： 以下面这幅图片命名为例：

DJI_0001_03700__0.367188__0.738889__0.009375__0.025926.jpg

注：上述四值均为百分比数据，取值0-1

想要实现这个操作，就要修改​​detect.py​​中的这个部分：

​​save_one_box​​​位于​​utils/general.py​​，修改如下：

首先和通常一样训练一样，准备数据集，训练好模型。

然后运行​​detect.py​​​，注意​​save-crop​​参数设为True，检测完之后，可以得到输出结果：

下面进入到二阶段的图像分类的训练，在开始之前，需要拉取YOLOv5仓库的最新版本，注意不要拉取Tag6.2版本，该版本在图像分类识别中仅支持单文件识别，而最新版本已经支持文件夹的批量识别。

然后需要人工对数据进行一个校正，因为单阶段输出的很多类别是存在错误的，需要手工处理，将其划分到正确的文件夹，同时对一些虚检的对象进行剔除。这一步可能比较费劲，特别是处理小物体时，有时候比较难判断。

然后， 需要对数据进行一个划分，可以通过手动或程序的方式划分成train和test，结构如下图所示(图中数据类别做了脱敏处理)：

准备好数据之后，运行​​classify/train.py​​​： 注意这里图片尺寸的设置可以设置为图片的长边尺寸

训练完之后，得到模型，然后将所有数据放到同一个文件夹里进行检测，即抛弃各种类别，混到同一个文件夹。

然后运行​​classify/predict.py​​进行识别，识别完成的结果如下：

这里每一张图片会对应一个label，label中包含了最大5个类别的概率和名称。

下面就需要把识别出来的结果转回到YOLO检测格式。

编写脚本​​cls2corp.py​​

classes = ['class1', 'class2', 'class3'] # 这里修改为自己的类别cls_path = 'runs/predict-cls/exp2'labels_path = Path('cls_result')value_threshold = 0.5 # 阈值，概率大于该阈值才会进行显示

if __name__ == '__main__': if labels_path.exists(): shutil.rmtree(labels_path) os.mkdir(labels_path) for i in classes: dir = labels_path.joinpath(i) if dir.exists(): shutil.rmtree(dir) os.mkdir(dir) for img in tqdm(os.listdir(cls_path)): suffix = img[-4:] if suffix == ".jpg" or suffix == ".png": img_name = img[:-4] for label in os.listdir(Path(f"{cls_path}/{'labels'}")): label_name = label[:-4] if img_name == label_name: with open(f"{cls_path}/{'labels'}/{label}", 'r', encoding='utf-8') as f: lines = f.readline() value = lines.split(" ")[0] cls = lines.split(" ")[1].strip() break if float(value) >= value_threshold: shutil.copy(Path(f"{cls_path}/{img}"), Path(f"{labels_path}/{cls}"))

这里设定一个阈值​​value_threshold​​，主要考虑到虚检情况，对于一些模糊或者没有对象的图片，它的每一个类别概率都不超过0.5，此时就把这个对象给抛弃。

运行完之后，可以在​​cls_result​​中得到每个类别的对应图片，这里取概率值最大的图片作为一个对象的识别类别。

然后要把这个文件夹的内容还原成labels，编写脚本​​crop_merge.py​​

import osfrom pathlib import Pathimport shutilfrom tqdm import tqdm

classes = ['class1', 'class2', 'class3'] # 这里修改为自己的类别crops_path = 'cls_result'labels_path = Path('cls_labels')

if __name__ == '__main__': # 由于后续是追加写入txt，因此先要删除labels，再进行创建 if labels_path.exists(): shutil.rmtree(labels_path) os.mkdir(labels_path) for dir_class in tqdm(os.listdir(crops_path)): index = classes.index(dir_class) # 根据类别顺序分配0-5 for img in os.listdir(Path(f"{crops_path}/{dir_class}")): img = img[:-4] # 去除文件名后缀 parts = img.split('__') img_name = parts[0] x = float(parts[1]) y = float(parts[2]) w = float(parts[3]) h = float(parts[4]) line = (index, x, y, w, h) with open('cls_labels' + '/' + img_name + '.txt', mode='a') as f: f.write(('%g ' * len(line)).rstrip() % line + '\n')

运行完后，得到如下结果：

最后就是把生成的labels在原图上标记出来，编写可视化脚本​​visualize.py​​

import osimport shutilfrom pathlib import Pathimport numpy as npimport cv2from tqdm import tqdm

# 修改输入图片文件夹img_folder = "D:/Desktop/Work/XDUAV-dataset/images/"img_list = os.listdir(img_folder)img_list.sort()# 修改输入标签文件夹label_folder = "cls_labels/"label_list = os.listdir(label_folder)label_list.sort()# 输出图片文件夹位置path = os.getcwd()output_folder = path + '/' + str("visualize_output")

labels = ['class1', 'class2', 'class3'] # 这里修改为自己的类别colormap = [(0, 255, 0), (132, 112, 255), (255, 255, 255)] # 色盘，可根据类别添加新颜色

# 坐标转换def xywh2xyxy(x, w1, h1, img): label, x, y, w, h = x label = int(label) # print("原图宽高:\nw1={}\nh1={}".format(w1, h1)) # 边界框反归一化 x_t = x * w1 y_t = y * h1 w_t = w * w1 h_t = h * h1 # print("反归一化后输出：\n第一个:{}\t第二个:{}\t第三个:{}\t第四个:{}\t\n\n".format(x_t, y_t, w_t, h_t)) # 计算坐标 top_left_x = x_t - w_t / 2 top_left_y = y_t - h_t / 2 bottom_right_x = x_t + w_t / 2 bottom_right_y = y_t + h_t / 2 # print('标签:{}'.format(labels[int(label)])) # print("左上x坐标:{}".format(top_left_x)) # print("左上y坐标:{}".format(top_left_y)) # print("右下x坐标:{}".format(bottom_right_x)) # print("右下y坐标:{}".format(bottom_right_y)) p1, p2 = (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)) # 绘制矩形框 cv2.rectangle(img, p1, p2, colormap[1], thickness=2, lineType=cv2.LINE_AA) label = labels[label] if label: w, h = cv2.getTextSize(label, 0, fontScale=2 / 3, thickness=2)[0] # text width, height outside = p1[1] - h - 3 >= 0 # label fits outside box p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3 # 绘制矩形框填充 cv2.rectangle(img, p1, p2, colormap[1], -1, cv2.LINE_AA) # 绘制标签 cv2.putText(img, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, 2 / 3, colormap[2], thickness=2, lineType=cv2.LINE_AA) """ # (可选)给不同目标绘制不同的颜色框 if int(label) == 0: cv2.rectangle(img, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), (0, 255, 0), 2) elif int(label) == 1: cv2.rectangle(img, (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y)), (255, 0, 0), 2) """ return img

if __name__ == '__main__': # 创建输出文件夹 if Path(output_folder).exists(): shutil.rmtree(output_folder) os.mkdir(output_folder) # labels和images可能存在不相等的情况，需要额外判断对应关系 img_index = 0 label_index = 0 for _ in tqdm(range(len(label_list))): image_path = img_folder + "/" + img_list[img_index] label_path = label_folder + "/" + label_list[label_index] if img_list[img_index][:-4] != label_list[label_index][:-4]: img_index += 1 continue # 读取图像文件 img = cv2.imread(str(image_path)) h, w = img.shape[:2] # 读取 labels with open(label_path, 'r') as f: lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # 绘制每一个目标 for x in lb: # 反归一化并得到左上和右下坐标，画出矩形框 img = xywh2xyxy(x, w, h, img) """ # 直接查看生成结果图 cv2.imshow('show', img) cv2.waitKey(0) """ cv2.imwrite(output_folder + '/' + '{}.png'.format(image_path.split('/')[-1][:-4]), img) img_index += 1 label_index += 1

在​​visualize_output​​文件夹下，得到最终效果。

使用二阶段目标检测带来的明显好处是：

不过存在的问题是：

在模型训练过程中，我注意到官方提供了新的预训练模型yolov5s6，那么这个预训练模型和yolov5s有什么区别呢？有些博文说yolov5s6是在更大的图像尺寸(1280x1280)得到的，不过这比较片面。

使用​​python export.py --weight yolov5s.pt --include onnx​​将其转换成onnx格式后，可以用Netron打开查看其结构：

可以看到，yolov5s6在模型最后的输出部分新增了一个检测头，当然上面也有一些小改动，不过改动不大。

在​​model/hub/yolov5s6.yaml​​中，对应了yolov5s6的模型，那么，如果用yolov5s.yaml来加载yolov5s6也是可行的。

在​​train.py​​中，这两行代码加载了预训练模型：

注意这里的​​strict=False​​表示模型如果和预训练模型不匹配，会加载能够匹配上的层权，下面的这个日志打印出了具体有多少匹配上的内容。

因此，模型构建还是由yaml文件来决定的，下面我又做了个小测试，用​​yolov5.yaml​​​分别来加载​​yolov5s.pt​​​和​​yolov5s6.pt​​​，两者得到的​​best.pt​​大小存在略微差异，而转换成onnx之后，两者大小一致，印证了结论的正确性。