如何用Python讀取xml檔案後，裁剪標註圖片和擴容資料

因為想用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.1\cut\c_1'

xml_path = r'D:\yolov5-3.1\cut\xml'

obj_img_path = r'D:\yolov5-3.1\cut\c_3'

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

Numpic = {}

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

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

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

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

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

xml_name = xml_path + '\\' + '%s.xml' % img_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 = 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、圖片擴容

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

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

def random_noise(image,noise_num):

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

ows, cols, chn = img_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 img_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)# 儲存圖片