本文的文字及圖片來(lái)源于網(wǎng)絡(luò),僅供學(xué)習(xí)撞鹉、交流使用,不具有任何商業(yè)用途,,版權(quán)歸原作者所有,如有問題請(qǐng)及時(shí)聯(lián)系我們以作處理
作者:Gendan
原文https://segmentfault.com/a/1190000038690313
想要獲取更多Python學(xué)習(xí)資料嵌巷,了解更多關(guān)于Python的知識(shí),可以加Q群630390733踴躍發(fā)言室抽,大家一起來(lái)學(xué)習(xí)討論吧搪哪!
因?yàn)橄胗脃olov5算法訓(xùn)練自己數(shù)據(jù)集識(shí)別數(shù)字“0-9”,一開始用labeling標(biāo)注了圖片坪圾,生成了大量的xml文件晓折。因?yàn)閳D片中0,1比較多兽泄,而其他數(shù)字偏少漓概,標(biāo)注到后面,就忽略了大量的0病梢,1胃珍。后面發(fā)現(xiàn),漏標(biāo)注會(huì)導(dǎo)致訓(xùn)練時(shí)把目標(biāo)識(shí)別成背景蜓陌,嚴(yán)重影響算法識(shí)別的準(zhǔn)確性觅彰。然后,我也不想重新去標(biāo)注圖片了钮热,就想著寫個(gè)Python程序根據(jù)xml文件缔莲,按照標(biāo)注框,把目標(biāo)都裁剪出來(lái)霉旗。
1痴奏、裁剪圖片
首先是根據(jù)xml文件把對(duì)應(yīng)標(biāo)注圖片蛀骇,按標(biāo)注框,裁剪出來(lái)读拆。我在的基礎(chǔ)上實(shí)現(xiàn)了裁剪圖片按類別保存到對(duì)應(yīng)文件夾里面擅憔,并在該類別下按順序編號(hào)
導(dǎo)入模塊import cv2import xml.etree.ElementTree as ETimport osfrom pathlib import Pathimport numpy as npimport random
原圖片、標(biāo)簽文件檐晕、裁剪圖片路徑
img_path = r'D:yolov5-3.1cutc_1'
xml_path = r'D:yolov5-3.1cutxml'
obj_img_path = r'D:yolov5-3.1cutc_3'
聲明一個(gè)空字典用于儲(chǔ)存裁剪圖片的類別及其數(shù)量
Numpic = {}
把原圖片裁剪后暑诸,按類別新建文件夾保存,并在該類別下按順序編號(hào)for img_file in os.listdir(img_path):
if img_file[-4:] in ['.png', '.jpg']: ?# 判斷文件是否為圖片格式
img_filename = os.path.join(img_path, img_file) ?# 將圖片路徑與圖片名進(jìn)行拼接
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 ?# 利用標(biāo)簽路徑个榕、圖片名、xml后綴拼接出完整的標(biāo)簽路徑名
if os.path.exists(xml_name): ?# 判斷與圖片同名的標(biāo)簽是否存在芥喇,因?yàn)閳D片不一定每張都打標(biāo)
root = ET.parse(xml_name).getroot() ?# 利用ET讀取xml文件
for obj in root.iter('object'): ?# 遍歷所有目標(biāo)框
name = obj.find('name').text ?# 獲取目標(biāo)框名稱西采,即label名
xmlbox = obj.find('bndbox') ?# 找到框目標(biāo)
x0 = xmlbox.find('xmin').text ?# 將框目標(biāo)的四個(gè)頂點(diǎn)坐標(biāo)取出
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裁剪出目標(biāo)框中的圖片
Numpic.setdefault(name, 0) ?# 判斷字典中有無(wú)當(dāng)前name對(duì)應(yīng)的類別,無(wú)則新建
Numpic[name] += 1 ?# 當(dāng)前類別對(duì)應(yīng)數(shù)量 + 1
my_file = Path(obj_img_path + '' + name) ?# 判斷當(dāng)前name對(duì)應(yīng)的類別有無(wú)文件夾
if 1 - my_file.is_dir(): ?# 無(wú)則新建
os.mkdir(obj_img_path + '' + str(name))
cv2.imwrite(obj_img_path + '' + name + '' + '%04d' % (Numpic[name]) + '.jpg',
obj_img) ?# 保存裁剪圖片继控,圖片命名4位械馆,不足補(bǔ)0
2、圖片擴(kuò)容
只是把標(biāo)注框裁剪出來(lái)武通,跟單網(wǎng)gendan5.com還會(huì)有一個(gè)問題就是霹崎,每個(gè)類別的數(shù)量不一致,0冶忱,1的圖片多尾菇,其他數(shù)字少,作為訓(xùn)練集可能不太好囚枪。我想错沽,要是每個(gè)類別的圖片數(shù)量都一致就好了。于是我繼續(xù)把裁剪圖片進(jìn)行擴(kuò)容眶拉,這里只是通過給圖片增加噪點(diǎn)來(lái)擴(kuò)容。
新建一個(gè)圖片加噪點(diǎn)的函數(shù)
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)#隨機(jī)生成指定范圍的整數(shù)
y = np.random.randint(0, cols)
img_noise[x, y, :] = 0 # 0代表黑色憔儿,255代表白色
return img_noise
圖片擴(kuò)容
max_Numpic = max(Numpic.values()) # 提取裁剪圖片中忆植,類別下數(shù)量最大值for name in Numpic:# 遍歷每一個(gè)類別
for i in range (Numpic[name] + 1, max_Numpic + 1):# 把其余類別的圖片數(shù)量擴(kuò)充到,與數(shù)量值最大的類別相等(我的數(shù)據(jù)集里面“0”這個(gè)類別數(shù)量是最多的)
Noisenum = random.randint(1, 20)# 生成隨機(jī)的噪點(diǎn)數(shù)
Num = random.randint(1, Numpic[name])# 隨機(jī)選擇該類別下已存在的一個(gè)圖片
Noicepic = random_noise(obj_img_path + '' + name + '' + '%04d' % Num + '.jpg', Noisenum)# 給圖片加噪點(diǎn)
cv2.imwrite(obj_img_path + '' + name + '' + '%04d' % (i) + '.jpg', Noicepic)# 保存圖片