yolo系列的模型在目標(biāo)檢測(cè)領(lǐng)域里面受眾非常廣，也十分流行，但是在使用yolo進(jìn)行目標(biāo)檢測(cè)訓(xùn)練的時(shí)候，往往要將VOC格式的數(shù)據(jù)集轉(zhuǎn)化為yolo專(zhuān)屬的數(shù)據(jù)集，而yolo的訓(xùn)練數(shù)據(jù)集制作方法呢，最常見(jiàn)的也是有兩種，下面我們只講述一種最常用的方法，也是我最常使用的。

1. voc轉(zhuǎn)yolo格式

我最常使用的目標(biāo)檢測(cè)數(shù)據(jù)集為VOC格式，而它的格式一般如下所示：

- dataset|- annotations|  |- image1.xml|  |- image2.xml|  |- ...||- images|  |- image1.jpg|  |- image2.jpg|  |- ...

• dataset 是數(shù)據(jù)集的根目錄。
• annotations 目錄包含每個(gè)圖像對(duì)應(yīng)的 XML 注釋文件。
• images 目錄包含每個(gè)圖像文件。

而我們要轉(zhuǎn)換的yolo格式如下所示:

- dataset|- images|  |- image1.jpg|  |- image2.jpg|  |- ...||- labels|  |- image1.txt|  |- image2.txt|  |- ...

• dataset 是數(shù)據(jù)集的根目錄。
• images 目錄包含每個(gè)圖像文件，通常是以 .jpg 或 .png 等格式保存的圖像文件。
• labels 目錄包含每個(gè)圖像對(duì)應(yīng)的標(biāo)簽文件，通常是以 .txt 格式保存的文本文件。

而 labels 里面的內(nèi)容填寫(xiě)格式為下圖所示：

通常，每行的格式為：class x_center y_center width height，其中class代表的是圖片中目標(biāo)所對(duì)應(yīng)的類(lèi)別，x_center, y_center是邊界框的中心點(diǎn)坐標(biāo)相對(duì)于圖像寬度和高度的歸一化值，width 和 height 是邊界框的寬度和高度相對(duì)于圖像寬度和高度的歸一化值。
舉例如下：

轉(zhuǎn)換代碼：

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import joindef convert(size, box):x_center = (box[0] + box[1]) / 2.0y_center = (box[2] + box[3]) / 2.0x = x_center / size[0]y = y_center / size[1]w = (box[1] - box[0]) / size[0]h = (box[3] - box[2]) / size[1]return (x, y, w, h)def convert_annotation(xml_files_path, save_txt_files_path, classes):xml_files = os.listdir(xml_files_path)for xml_name in xml_files:xml_file = os.path.join(xml_files_path, xml_name)out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')out_txt_f = open(out_txt_path, 'w')tree = ET.parse(xml_file)root = tree.getroot()size = root.find('size')w = int(size.find('width').text)h = int(size.find('height').text)for obj in root.iter('object'):#difficult = obj.find('difficult').textcls = obj.find('name').text#if cls not in classes or int(difficult) == 1:#continuecls_id = classes.index(cls)xmlbox = obj.find('bndbox')b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),float(xmlbox.find('ymax').text))# b=(xmin, xmax, ymin, ymax)# print(w, h, b)bb = convert((w, h), b)out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')if __name__ == "__main__":# 把forklift_pallet的voc的xml標(biāo)簽文件轉(zhuǎn)化為yolo的txt標(biāo)簽文件# 1、需要轉(zhuǎn)化的類(lèi)別，這里我直接用數(shù)字代表類(lèi)別，由于我是八類(lèi)，所以從0到7classes = ['0', '1', '2', '3', '4', '5', '6', '7']# 2、voc格式的xml標(biāo)簽文件路徑xml_files1 = 'annotations'# 3、轉(zhuǎn)化為yolo格式的txt標(biāo)簽文件存儲(chǔ)路徑save_txt_files1 = 'labels'convert_annotation(xml_files1, save_txt_files1, classes)

上面代碼中注釋了一部分內(nèi)容，比如difficult這一項(xiàng)，由于我xml文件里面沒(méi)有difficult，所以就注釋掉了，大家按照自己的需求進(jìn)行使用即可。

劃分?jǐn)?shù)據(jù)集

在我們進(jìn)行yolo目標(biāo)檢測(cè)模型訓(xùn)練之前，需要先將數(shù)據(jù)集進(jìn)行合理的劃分，比如說(shuō)劃分為訓(xùn)練集：驗(yàn)證集=8:2，或者訓(xùn)練集：驗(yàn)證集：測(cè)試集=7:2：1。不過(guò)我一般習(xí)慣只劃分訓(xùn)練集和驗(yàn)證集，也就是按8:2的比例進(jìn)行劃分，代碼如下所示：

import os
import shutil
import random# 定義數(shù)據(jù)集文件夾路徑
dataset_path = 'dataset'
images_path = os.path.join(dataset_path, 'images')
labels_path = os.path.join(dataset_path, 'labels')# 定義劃分后的文件夾路徑
new_path = 'mydata'
train_path = os.path.join(new_path, 'train')
val_path = os.path.join(new_path, 'val')# 創(chuàng)建train和val文件夾
os.makedirs(os.path.join(train_path, 'images'), exist_ok=True)
os.makedirs(os.path.join(train_path, 'labels'), exist_ok=True)
os.makedirs(os.path.join(val_path, 'images'), exist_ok=True)
os.makedirs(os.path.join(val_path, 'labels'), exist_ok=True)# 獲取所有圖片文件的文件名
image_files = os.listdir(images_path)
# 隨機(jī)打亂文件順序
random.shuffle(image_files)# 定義驗(yàn)證集所占比例
val_split = 0.1
# 計(jì)算驗(yàn)證集大小
num_val = int(len(image_files) * val_split)# 將數(shù)據(jù)集按照比例劃分到train和val文件夾中
for i, image_file in enumerate(image_files):src_image = os.path.join(images_path, image_file)src_label = os.path.join(labels_path, image_file.replace('.jpg', '.txt'))if i < num_val:dst_image = os.path.join(val_path, 'images', image_file)dst_label = os.path.join(val_path, 'labels', image_file.replace('.jpg', '.txt'))else:dst_image = os.path.join(train_path, 'images', image_file)dst_label = os.path.join(train_path, 'labels', image_file.replace('.jpg', '.txt'))shutil.copy(src_image, dst_image)shutil.copy(src_label, dst_label)

劃分完成以后的文件夾格式為：

- mydata|- train|  |- images|  |- labels||- val|  |- images|  |- labels

images和labels分別是對(duì)應(yīng)的數(shù)據(jù)集圖片和txt標(biāo)簽。

數(shù)據(jù)增強(qiáng)

在我們參加一些目標(biāo)檢測(cè)類(lèi)比賽的時(shí)候，往往會(huì)遇見(jiàn)比賽訓(xùn)練集不足的情況，這將極大程度上影響我們的模型精度，這時(shí)候可能就需要用到一些數(shù)據(jù)增強(qiáng)方法，如翻轉(zhuǎn)、隨機(jī)裁剪等等。當(dāng)然，yolo系列的模型一般都自帶有數(shù)據(jù)增強(qiáng)，但是我們也可以嘗試訓(xùn)練前進(jìn)行增強(qiáng)看看效果。
代碼如下：

# -*- coding=utf-8 -*-import time
import random
import copy
import cv2
import os
import math
import numpy as np
from skimage.util import random_noise
from lxml import etree, objectify
import xml.etree.ElementTree as ET
import argparse# 顯示圖片
def show_pic(img, bboxes=None):'''輸入:img:圖像arraybboxes:圖像的所有boudning box list, 格式為[[x_min, y_min, x_max, y_max]....]names:每個(gè)box對(duì)應(yīng)的名稱(chēng)'''for i in range(len(bboxes)):bbox = bboxes[i]x_min = bbox[0]y_min = bbox[1]x_max = bbox[2]y_max = bbox[3]cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 3)cv2.namedWindow('pic', 0)  # 1表示原圖cv2.moveWindow('pic', 0, 0)cv2.resizeWindow('pic', 1200, 800)  # 可視化的圖片大小cv2.imshow('pic', img)cv2.waitKey(0)cv2.destroyAllWindows()# 圖像均為cv2讀取
class DataAugmentForObjectDetection():def __init__(self, rotation_rate=0.5, max_rotation_angle=5,crop_rate=0.5, shift_rate=0.5, change_light_rate=0.5,add_noise_rate=0.5, flip_rate=0.5,cutout_rate=0.5, cut_out_length=50, cut_out_holes=1, cut_out_threshold=0.5,is_addNoise=True, is_changeLight=True, is_cutout=True, is_rotate_img_bbox=True,is_crop_img_bboxes=True, is_shift_pic_bboxes=True, is_filp_pic_bboxes=True):# 配置各個(gè)操作的屬性self.rotation_rate = rotation_rateself.max_rotation_angle = max_rotation_angleself.crop_rate = crop_rateself.shift_rate = shift_rateself.change_light_rate = change_light_rateself.add_noise_rate = add_noise_rateself.flip_rate = flip_rateself.cutout_rate = cutout_rateself.cut_out_length = cut_out_lengthself.cut_out_holes = cut_out_holesself.cut_out_threshold = cut_out_threshold# 是否使用某種增強(qiáng)方式self.is_addNoise = is_addNoiseself.is_changeLight = is_changeLightself.is_cutout = is_cutoutself.is_rotate_img_bbox = is_rotate_img_bboxself.is_crop_img_bboxes = is_crop_img_bboxesself.is_shift_pic_bboxes = is_shift_pic_bboxesself.is_filp_pic_bboxes = is_filp_pic_bboxes# ----1.加噪聲---- #def _addNoise(self, img):'''輸入:img:圖像array輸出:加噪聲后的圖像array,由于輸出的像素是在[0,1]之間,所以得乘以255'''# return cv2.GaussianBlur(img, (11, 11), 0)return random_noise(img, mode='gaussian', seed=int(time.time()), clip=True) * 255# ---2.調(diào)整亮度--- #def _changeLight(self, img):alpha = random.uniform(0.35, 1)blank = np.zeros(img.shape, img.dtype)return cv2.addWeighted(img, alpha, blank, 1 - alpha, 0)# ---3.cutout--- #def _cutout(self, img, bboxes, length=100, n_holes=1, threshold=0.5):'''原版本：https://github.com/uoguelph-mlrg/Cutout/blob/master/util/cutout.pyRandomly mask out one or more patches from an image.Args:img : a 3D numpy array,(h,w,c)bboxes : 框的坐標(biāo)n_holes (int): Number of patches to cut out of each image.length (int): The length (in pixels) of each square patch.'''def cal_iou(boxA, boxB):'''boxA, boxB為兩個(gè)框，返回iouboxB為bouding box'''# determine the (x, y)-coordinates of the intersection rectanglexA = max(boxA[0], boxB[0])yA = max(boxA[1], boxB[1])xB = min(boxA[2], boxB[2])yB = min(boxA[3], boxB[3])if xB <= xA or yB <= yA:return 0.0# compute the area of intersection rectangleinterArea = (xB - xA + 1) * (yB - yA + 1)# compute the area of both the prediction and ground-truth# rectanglesboxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)iou = interArea / float(boxBArea)return iou# 得到h和wif img.ndim == 3:h, w, c = img.shapeelse:_, h, w, c = img.shapemask = np.ones((h, w, c), np.float32)for n in range(n_holes):chongdie = True  # 看切割的區(qū)域是否與box重疊太多while chongdie:y = np.random.randint(h)x = np.random.randint(w)y1 = np.clip(y - length // 2, 0,h)  # numpy.clip(a, a_min, a_max, out=None), clip這個(gè)函數(shù)將將數(shù)組中的元素限制在a_min, a_max之間，大于a_max的就使得它等于 a_max，小于a_min,的就使得它等于a_miny2 = np.clip(y + length // 2, 0, h)x1 = np.clip(x - length // 2, 0, w)x2 = np.clip(x + length // 2, 0, w)chongdie = Falsefor box in bboxes:if cal_iou([x1, y1, x2, y2], box) > threshold:chongdie = Truebreakmask[y1: y2, x1: x2, :] = 0.img = img * maskreturn img# ---4.旋轉(zhuǎn)--- #def _rotate_img_bbox(self, img, bboxes, angle=5, scale=1.):'''參考:https://blog.csdn.net/u014540717/article/details/53301195crop_rate輸入:img:圖像array,(h,w,c)bboxes:該圖像包含的所有boundingboxs,一個(gè)list,每個(gè)元素為[x_min, y_min, x_max, y_max],要確保是數(shù)值angle:旋轉(zhuǎn)角度scale:默認(rèn)1輸出:rot_img:旋轉(zhuǎn)后的圖像arrayrot_bboxes:旋轉(zhuǎn)后的boundingbox坐標(biāo)list'''# 旋轉(zhuǎn)圖像w = img.shape[1]h = img.shape[0]# 角度變弧度rangle = np.deg2rad(angle)  # angle in radians# now calculate new image width and heightnw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scalenh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale# ask OpenCV for the rotation matrixrot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)# calculate the move from the old center to the new center combined# with the rotationrot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))# the move only affects the translation, so update the translationrot_mat[0, 2] += rot_move[0]rot_mat[1, 2] += rot_move[1]# 仿射變換rot_img = cv2.warpAffine(img, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4)# 矯正bbox坐標(biāo)# rot_mat是最終的旋轉(zhuǎn)矩陣# 獲取原始bbox的四個(gè)中點(diǎn)，然后將這四個(gè)點(diǎn)轉(zhuǎn)換到旋轉(zhuǎn)后的坐標(biāo)系下rot_bboxes = list()for bbox in bboxes:xmin = bbox[0]ymin = bbox[1]xmax = bbox[2]ymax = bbox[3]point1 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymin, 1]))point2 = np.dot(rot_mat, np.array([xmax, (ymin + ymax) / 2, 1]))point3 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymax, 1]))point4 = np.dot(rot_mat, np.array([xmin, (ymin + ymax) / 2, 1]))# 合并np.arrayconcat = np.vstack((point1, point2, point3, point4))# 改變array類(lèi)型concat = concat.astype(np.int32)# 得到旋轉(zhuǎn)后的坐標(biāo)rx, ry, rw, rh = cv2.boundingRect(concat)rx_min = rxry_min = ryrx_max = rx + rwry_max = ry + rh# 加入list中rot_bboxes.append([rx_min, ry_min, rx_max, ry_max])return rot_img, rot_bboxes# ---5.裁剪--- #def _crop_img_bboxes(self, img, bboxes):'''裁剪后的圖片要包含所有的框輸入:img:圖像arraybboxes:該圖像包含的所有boundingboxs,一個(gè)list,每個(gè)元素為[x_min, y_min, x_max, y_max],要確保是數(shù)值輸出:crop_img:裁剪后的圖像arraycrop_bboxes:裁剪后的bounding box的坐標(biāo)list'''# 裁剪圖像w = img.shape[1]h = img.shape[0]x_min = w  # 裁剪后的包含所有目標(biāo)框的最小的框x_max = 0y_min = hy_max = 0for bbox in bboxes:x_min = min(x_min, bbox[0])y_min = min(y_min, bbox[1])x_max = max(x_max, bbox[2])y_max = max(y_max, bbox[3])d_to_left = x_min  # 包含所有目標(biāo)框的最小框到左邊的距離d_to_right = w - x_max  # 包含所有目標(biāo)框的最小框到右邊的距離d_to_top = y_min  # 包含所有目標(biāo)框的最小框到頂端的距離d_to_bottom = h - y_max  # 包含所有目標(biāo)框的最小框到底部的距離# 隨機(jī)擴(kuò)展這個(gè)最小框crop_x_min = int(x_min - random.uniform(0, d_to_left))crop_y_min = int(y_min - random.uniform(0, d_to_top))crop_x_max = int(x_max + random.uniform(0, d_to_right))crop_y_max = int(y_max + random.uniform(0, d_to_bottom))# 隨機(jī)擴(kuò)展這個(gè)最小框 , 防止別裁的太小# crop_x_min = int(x_min - random.uniform(d_to_left//2, d_to_left))# crop_y_min = int(y_min - random.uniform(d_to_top//2, d_to_top))# crop_x_max = int(x_max + random.uniform(d_to_right//2, d_to_right))# crop_y_max = int(y_max + random.uniform(d_to_bottom//2, d_to_bottom))# 確保不要越界crop_x_min = max(0, crop_x_min)crop_y_min = max(0, crop_y_min)crop_x_max = min(w, crop_x_max)crop_y_max = min(h, crop_y_max)crop_img = img[crop_y_min:crop_y_max, crop_x_min:crop_x_max]# 裁剪boundingbox# 裁剪后的boundingbox坐標(biāo)計(jì)算crop_bboxes = list()for bbox in bboxes:crop_bboxes.append([bbox[0] - crop_x_min, bbox[1] - crop_y_min, bbox[2] - crop_x_min, bbox[3] - crop_y_min])return crop_img, crop_bboxes# ---6.平移--- #def _shift_pic_bboxes(self, img, bboxes):'''平移后的圖片要包含所有的框輸入:img:圖像arraybboxes:該圖像包含的所有boundingboxs,一個(gè)list,每個(gè)元素為[x_min, y_min, x_max, y_max],要確保是數(shù)值輸出:shift_img:平移后的圖像arrayshift_bboxes:平移后的bounding box的坐標(biāo)list'''# 平移圖像w = img.shape[1]h = img.shape[0]x_min = w  # 裁剪后的包含所有目標(biāo)框的最小的框x_max = 0y_min = hy_max = 0for bbox in bboxes:x_min = min(x_min, bbox[0])y_min = min(y_min, bbox[1])x_max = max(x_max, bbox[2])y_max = max(y_max, bbox[3])d_to_left = x_min  # 包含所有目標(biāo)框的最大左移動(dòng)距離d_to_right = w - x_max  # 包含所有目標(biāo)框的最大右移動(dòng)距離d_to_top = y_min  # 包含所有目標(biāo)框的最大上移動(dòng)距離d_to_bottom = h - y_max  # 包含所有目標(biāo)框的最大下移動(dòng)距離x = random.uniform(-(d_to_left - 1) / 3, (d_to_right - 1) / 3)y = random.uniform(-(d_to_top - 1) / 3, (d_to_bottom - 1) / 3)M = np.float32([[1, 0, x], [0, 1, y]])  # x為向左或右移動(dòng)的像素值,正為向右負(fù)為向左; y為向上或者向下移動(dòng)的像素值,正為向下負(fù)為向上shift_img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))#  平移boundingboxshift_bboxes = list()for bbox in bboxes:shift_bboxes.append([bbox[0] + x, bbox[1] + y, bbox[2] + x, bbox[3] + y])return shift_img, shift_bboxes# ---7.鏡像--- #def _filp_pic_bboxes(self, img, bboxes):'''平移后的圖片要包含所有的框輸入:img:圖像arraybboxes:該圖像包含的所有boundingboxs,一個(gè)list,每個(gè)元素為[x_min, y_min, x_max, y_max],要確保是數(shù)值輸出:flip_img:平移后的圖像arrayflip_bboxes:平移后的bounding box的坐標(biāo)list'''# 翻轉(zhuǎn)圖像flip_img = copy.deepcopy(img)h, w, _ = img.shapesed = random.random()if 0 < sed < 0.33:  # 0.33的概率水平翻轉(zhuǎn)，0.33的概率垂直翻轉(zhuǎn),0.33是對(duì)角反轉(zhuǎn)flip_img = cv2.flip(flip_img, 0)  # _flip_xinver = 0elif 0.33 < sed < 0.66:flip_img = cv2.flip(flip_img, 1)  # _flip_yinver = 1else:flip_img = cv2.flip(flip_img, -1)  # flip_x_yinver = -1# 調(diào)整boundingboxflip_bboxes = list()for box in bboxes:x_min = box[0]y_min = box[1]x_max = box[2]y_max = box[3]if inver == 0:# 0：垂直翻轉(zhuǎn)flip_bboxes.append([x_min, h - y_max, x_max, h - y_min])elif inver == 1:# 1：水平翻轉(zhuǎn)flip_bboxes.append([w - x_max, y_min, w - x_min, y_max])elif inver == -1:# -1：水平垂直翻轉(zhuǎn)flip_bboxes.append([w - x_max, h - y_max, w - x_min, h - y_min])return flip_img, flip_bboxes# 圖像增強(qiáng)方法def dataAugment(self, img, bboxes):'''圖像增強(qiáng)輸入:img:圖像arraybboxes:該圖像的所有框坐標(biāo)輸出:img:增強(qiáng)后的圖像bboxes:增強(qiáng)后圖片對(duì)應(yīng)的box'''change_num = 0  # 改變的次數(shù)# print('------')while change_num < 1:  # 默認(rèn)至少有一種數(shù)據(jù)增強(qiáng)生效if self.is_rotate_img_bbox:if random.random() > self.rotation_rate:  # 旋轉(zhuǎn)change_num += 1angle = random.uniform(-self.max_rotation_angle, self.max_rotation_angle)scale = random.uniform(0.7, 0.8)img, bboxes = self._rotate_img_bbox(img, bboxes, angle, scale)if self.is_shift_pic_bboxes:if random.random() < self.shift_rate:  # 平移change_num += 1img, bboxes = self._shift_pic_bboxes(img, bboxes)if self.is_changeLight:if random.random() > self.change_light_rate:  # 改變亮度change_num += 1img = self._changeLight(img)if self.is_addNoise:if random.random() < self.add_noise_rate:  # 加噪聲change_num += 1img = self._addNoise(img)if self.is_cutout:if random.random() < self.cutout_rate:  # cutoutchange_num += 1img = self._cutout(img, bboxes, length=self.cut_out_length, n_holes=self.cut_out_holes,threshold=self.cut_out_threshold)if self.is_filp_pic_bboxes:if random.random() < self.flip_rate:  # 翻轉(zhuǎn)change_num += 1img, bboxes = self._filp_pic_bboxes(img, bboxes)return img, bboxes# xml解析工具
class ToolHelper():# 從xml文件中提取bounding box信息, 格式為[[x_min, y_min, x_max, y_max, name]]def parse_xml(self, path):'''輸入：xml_path: xml的文件路徑輸出：從xml文件中提取bounding box信息, 格式為[[x_min, y_min, x_max, y_max, name]]'''tree = ET.parse(path)root = tree.getroot()objs = root.findall('object')coords = list()for ix, obj in enumerate(objs):name = obj.find('name').textbox = obj.find('bndbox')x_min = int(box[0].text)y_min = int(box[1].text)x_max = int(box[2].text)y_max = int(box[3].text)coords.append([x_min, y_min, x_max, y_max, name])return coords# 保存圖片結(jié)果def save_img(self, file_name, save_folder, img):cv2.imwrite(os.path.join(save_folder, file_name), img)# 保持xml結(jié)果def save_xml(self, file_name, save_folder, img_info, height, width, channel, bboxs_info):''':param file_name:文件名:param save_folder:#保存的xml文件的結(jié)果:param height:圖片的信息:param width:圖片的寬度:param channel:通道:return:'''folder_name, img_name = img_info  # 得到圖片的信息E = objectify.ElementMaker(annotate=False)anno_tree = E.annotation(E.folder(folder_name),E.filename(img_name),E.path(os.path.join(folder_name, img_name)),E.source(E.database('Unknown'),),E.size(E.width(width),E.height(height),E.depth(channel)),E.segmented(0),)labels, bboxs = bboxs_info  # 得到邊框和標(biāo)簽信息for label, box in zip(labels, bboxs):anno_tree.append(E.object(E.name(label),E.pose('Unspecified'),E.truncated('0'),E.difficult('0'),E.bndbox(E.xmin(box[0]),E.ymin(box[1]),E.xmax(box[2]),E.ymax(box[3]))))etree.ElementTree(anno_tree).write(os.path.join(save_folder, file_name), pretty_print=True)if __name__ == '__main__':need_aug_num = 5  # 每張圖片需要增強(qiáng)的次數(shù)is_endwidth_dot = True  # 文件是否以.jpg或者png結(jié)尾dataAug = DataAugmentForObjectDetection()  # 數(shù)據(jù)增強(qiáng)工具類(lèi)toolhelper = ToolHelper()  # 工具# 獲取相關(guān)參數(shù)parser = argparse.ArgumentParser()parser.add_argument('--source_img_path', type=str, default='images')parser.add_argument('--source_xml_path', type=str, default='Annotations')parser.add_argument('--save_img_path', type=str, default='enhance_images')parser.add_argument('--save_xml_path', type=str, default='enhance_Annotations')args = parser.parse_args()source_img_path = args.source_img_path  # 圖片原始位置source_xml_path = args.source_xml_path  # xml的原始位置save_img_path = args.save_img_path  # 圖片增強(qiáng)結(jié)果保存文件save_xml_path = args.save_xml_path  # xml增強(qiáng)結(jié)果保存文件# 如果保存文件夾不存在就創(chuàng)建if not os.path.exists(save_img_path):os.mkdir(save_img_path)if not os.path.exists(save_xml_path):os.mkdir(save_xml_path)for parent, _, files in os.walk(source_img_path):files.sort()for file in files:cnt = 0pic_path = os.path.join(parent, file)xml_path = os.path.join(source_xml_path, file[:-4] + '.xml')values = toolhelper.parse_xml(xml_path)  # 解析得到box信息，格式為[[x_min,y_min,x_max,y_max,name]]coords = [v[:4] for v in values]  # 得到框labels = [v[-1] for v in values]  # 對(duì)象的標(biāo)簽# 如果圖片是有后綴的if is_endwidth_dot:# 找到文件的最后名字dot_index = file.rfind('.')_file_prefix = file[:dot_index]  # 文件名的前綴_file_suffix = file[dot_index:]  # 文件名的后綴img = cv2.imread(pic_path)# show_pic(img, coords)  # 顯示原圖while cnt < need_aug_num:  # 繼續(xù)增強(qiáng)auged_img, auged_bboxes = dataAug.dataAugment(img, coords)auged_bboxes_int = np.array(auged_bboxes).astype(np.int32)height, width, channel = auged_img.shape  # 得到圖片的屬性img_name = '{}_{}{}'.format(_file_prefix, cnt + 1, _file_suffix)  # 圖片保存的信息toolhelper.save_img(img_name, save_img_path,auged_img)  # 保存增強(qiáng)圖片toolhelper.save_xml('{}_{}.xml'.format(_file_prefix, cnt + 1),save_xml_path, (save_img_path, img_name), height, width, channel,(labels, auged_bboxes_int))  # 保存xml文件# show_pic(auged_img, auged_bboxes)  # 強(qiáng)化后的圖print(img_name)cnt += 1  # 繼續(xù)增強(qiáng)下一張

增強(qiáng)后的效果圖如下所示：

詳細(xì)實(shí)戰(zhàn)使用操作請(qǐng)看：基于yolov8的車(chē)牌檢測(cè)訓(xùn)練全流程