利用Yolov8实现多个检测模型融合
- 1 背景
- 2 代码实现
1 背景
(1)例如,现在有Yolov8两个模型,一个模型用于识别人,一个模型用于识别树。现要求将识别人和识别树的两个模型合并成一个模型,仅利用一个模型就能同时识别人和树。
(2)思路:将同时存在人和树的数据集重新标注,使label文件里存在具有人和树的类别,然后再次训练出新的模型。
(3)重点解决问题:生成新的label文件,得到新的数据集。在下面的代码中,通过两个类的Yolov8模型对原有数据集对应的训练集和验证集的images进行检测,然后将检测结果进行处理,进而生成新的包含两个类别的数据集。
2 代码实现
from ultralytics import YOLO
import os, cv2
import numpy as np
# 存放全局变量
source_dir_path = [r"F:\0_work\数据集\datasets_clone\train\images", r"F:\0_work\数据集\datasets_clone\val\images"]
yolo_label_dir_path = [r"F:\0_work\数据集\datasets_clone\train\new_labels", r"F:\0_work\数据集\datasets_clone\val\new_labels"]
class MakeLabel():
def __init__(self):
# 设置模型路径,并加载模型
self.people_model_path = "./weights/people_best.pt"
self.tree_model_path = "./weights/tree_best.pt"
self.people_model = YOLO(self.people_model_path, task='detect')
self.tree_model = YOLO(self.tree_model_path, task='detect')
def check_label_dirs(self, destination_folder):
if not os.path.exists(destination_folder):
os.makedirs(destination_folder)
def get_label_data(self, labels, boxes, id):
for box in boxes:
x = box[0]
y = box[1]
w = box[2]
h = box[3]
labels.append([id, x, y, w, h])
def getlabel(self, img_path):
img = cv2.imread(img_path)
new_label_file_name = img_path.split('\\')[-1].split('.')[0]
# 得到模型检测的结果
people_results = self.people_model(img, verbose=False, imgsz=640, conf=0.1, iou=0.35, show=False)
tree_results = self.tree_model(img, verbose=False, conf=0.7, iou=0.35, imgsz=320)
# 生成模型检测的结果,将数据格式表示为经过比率转换的YOLO标注格式,即[class_id, x*, y*, w*, h*]
labels = []
people_boxes = people_results[0].boxes.xywhn.cpu().numpy().astype(float)
tree_boxes = tree_results[0].boxes.xywhn.cpu().numpy().astype(float)
people_id = 0
tree_id = 1
# 添加 class_id
self.get_label_data(labels, people_boxes, people_id)
self.get_label_data(labels, tree_boxes, tree_id)
# 将检测的结果存入txt
with open(os.path.join(yolo_label_dir_path, new_label_file_name + '.txt'), 'w') as out_file:
print(type(labels))
for line in labels:
out_file.write(" ".join(str(data) for data in line) + '\n')
out_file.close()
if __name__ == '__main__':
makelabel = MakeLabel()
source_path_dir_numbers = len(source_dir_path)
for index in range(source_path_dir_numbers):
makelabel.check_label_dirs(yolo_label_dir_path[index])
for file in os.listdir(source_dir_path[index]):
file_path = os.path.join(source_dir_path, file)
makelabel.getlabel(file_path)