【利用Yolov8实现多个检测模型融合】

1 背景

1.1 案例背景

现在有Yolov8两个模型，一个模型用于识别人，一个模型用于识别树。现要求将识别人和识别树的两个模型合并成一个模型，仅利用一个模型就能同时识别人和树。

1.2 解决思路

将同时存在人和树的数据集重新标注，使label文件里存在具有人和树的类别，然后再次训练出新的模型。

1.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)