原文 https://segmentfault.com/a/1190000038690313

因為想用yolov5演算法訓練自己資料集識別數字“0-9”，一開始用labeling標註了圖片，生成了大量的xml檔案。因為圖片中0，1比較多，而其他數字偏少，標註到後面，就忽略了大量的0，1。後面發現，漏標註會導致訓練時把目標識別成背景，嚴重影響演算法識別的準確性。然後，我也不想重新去標註圖片了，就想著寫個Python程式根據xml檔案，按照標註框，把目標都裁剪出來。

1、裁剪圖片

首先是根據xml檔案把對應標註圖片，按標註框，裁剪出來。我在的基礎上實現了裁剪圖片按類別儲存到對應資料夾裡面，並在該類別下按順序編號

匯入模組import cv2import xml.etree.ElementTree as ETimport osfrom pathlib import Pathimport numpy as npimport random

原圖片、標籤檔案、裁剪圖片路徑

img_path = r'D:yolov5-3.1cutc_1'

xml_path = r'D:yolov5-3.1cutxml'

obj_img_path = r'D:yolov5-3.1cutc_3'

宣告一個空字典用於儲存裁剪圖片的類別及其數量

Numpic = {}

把原圖片裁剪後，按類別新建資料夾儲存，並在該類別下按順序編號for file in os.listdir(img_path):

if file[-4:] in ['.png', '.jpg']: # 判斷檔案是否為圖片格式

img_filename = os.path.join(img_path, file) # 將圖片路徑與圖片名進行拼接

img_cv = cv2.imread(img_filename) # 讀取圖片

img_name = (os.path.splitext(img_file)[0]) # 分割出圖片名，如“000.png” 圖片名為“000”

xml_name = xml_path + '' + '%s.xml' % name # 利用標籤路徑、圖片名、xml字尾拼接出完整的標籤路徑名

if os.path.exists(xml_name): # 判斷與圖片同名的標籤是否存在，因為圖片不一定每張都打標

root = ET.parse(xml_name).getroot() # 利用ET讀取xml檔案

for obj in root.iter('object'): # 遍歷所有目標框

name = obj.find('name').text # 獲取目標框名稱，即label名

xmlbox = obj.find('bndbox') # 找到框目標

x0 = xmlbox.find('xmin').text # 將框目標的四個頂點座標取出

y0 = xmlbox.find('ymin').text

x1 = xmlbox.find('xmax').text

y1 = xmlbox.find('ymax').text

obj_img = cv[int(y0):int(y1), int(x0):int(x1)] # cv2裁剪出目標框中的圖片

Numpic.setdefault(name, 0) # 判斷字典中有無當前name對應的類別，無則新建

Numpic[name] += 1 # 當前類別對應數量 + 1

my_file = Path(obj_img_path + '' + name) # 判斷當前name對應的類別有無資料夾

if 1 - my_file.is_dir(): # 無則新建

os.mkdir(obj_img_path + '' + str(name))

cv2.imwrite(obj_img_path + '' + name + '' + '%04d' % (Numpic[name]) + '.jpg',

obj_img) # 儲存裁剪圖片，圖片命名4位，不足補0

2、圖片擴容

只是把標註框裁剪出來，跟單網gendan5.com還會有一個問題就是，每個類別的數量不一致，0，1的圖片多，其他數字少，作為訓練集可能不太好。我想，要是每個類別的圖片數量都一致就好了。於是我繼續把裁剪圖片進行擴容，這裡只是透過給圖片增加噪點來擴容。

新建一個圖片加噪點的函式

def random_noise(image,noise_num):

img_noiseimg = cv2.imread(image) # 讀取圖片

ows, cols, chn = noise.shape

for i in range(noise_num):

x = np.random.randint(0, rows)#隨機生成指定範圍的整數

y = np.random.randint(0, cols)

img_noise[x, y, :] = 0 # 0代表黑色，255代表白色

return noise

圖片擴容

max_Numpic = max(Numpic.values()) # 提取裁剪圖片中，類別下數量最大值for name in Numpic:# 遍歷每一個類別

for i in range (Numpic[name] + 1, max_Numpic + 1):# 把其餘類別的圖片數量擴充到，與數量值最大的類別相等（我的資料集裡面“0”這個類別數量是最多的）

Noisenum = random.randint(1, 20)# 生成隨機的噪點數

Num = random.randint(1, Numpic[name])# 隨機選擇該類別下已存在的一個圖片

Noicepic = random_noise(obj_img_path + '' + name + '' + '%04d' % Num + '.jpg', Noisenum)# 給圖片加噪點

cv2.imwrite(obj_img_path + '' + name + '' + '%04d' % (i) + '.jpg', Noicepic)# 儲存圖片