在样本量不足的情况下,我们通常会采用mixup或者填鸭式的方法来进行数据增强。其中mixup是将正负样本融合成新的一组样本,使得样本量翻倍。填鸭式是将原本样本里的目标抠出来,随机复制粘贴到其他地方。(个人理解)
import os
import json
import numpy as np
import pandas as pd
import cv2
import glob
import random
from PIL import Image
import time
from sklearn import metrics as mr
random.seed(2019)
aug_name=['target']
defect_img_root='C:\\Users\\Administrator\\Desktop\\new_front'
save_dir='...'
if not os.path.exists(save_dir):
os.mkdir(save_dir)
anno_file='...'
anno_result= pd.read_json(open(anno_file,"r"))
name_list=anno_result["name"].unique()
ring_width=10# default is 5
result=[]
last_result_length=0
img_name_count=0
for root,paths,files in os.walk(defect_img_root):
for path in paths:
img_name=path+'.jpg'
img_anno = anno_result[anno_result["name"] == img_name]
bboxs = img_anno["bbox"].tolist()
# print(bboxs)
img_name_count+=1
defect_names = img_anno["defect_name"].tolist()
# defect_names = [defect_name2label[x] for x in defect_names]
print(defect_names)
assert img_anno["name"].unique()[0] == img_name
# testimg=cv2.imread(root+path+'/'+img_name)
testimg=Image.open(root+path+'/'+img_name)
template_img_name='template_'+path.split('_')[0]+'.jpg'
# temp_img=cv2.imread(root+path+'/'+template_img_name)
temp_img=Image.open(root+path+'/'+template_img_name)
save_temp_name='template_'+path.split('_')[0]+str(img_name_count)+'.jpg'
for idx in range(len(bboxs)):
pts=bboxs[idx]
d_name=defect_names[idx]
xmin=pts[0]
ymin=pts[1]
xmax=pts[2]
ymax=pts[3]
defect_h=abs(xmax-xmin)
defect_w=abs(ymax-ymin)
w_h=round(defect_w/defect_h,2)
h_w=round(defect_h/defect_w,2)
if w_h > 5 or h_w >5 or (ymax-ymin)*(xmax-xmin)>300000:#TODO: 这里的两个5感觉有问题,300000的应该丢弃
left_top_x=random.randint(1,3)
left_top_y=random.randint(1,testimg.size[1])
else:
left_top_x=random.randint(1,testimg.size[0])
left_top_y=random.randint(1,testimg.size[1])
# print(left_top_x,left_top_y)
mask=np.zeros_like(temp_img)
if d_name in aug_name:
scale=random.randint(3,5)
mask[int(scale*(left_top_y-ring_width)):int(scale*(left_top_y+defect_w+ring_width)),int(scale*(left_top_x-ring_width)):int(scale*(left_top_x+defect_h+ring_width))]=255
mask[int(left_top_y):int(scale*(left_top_y+defect_w)),int(left_top_x):int(scale*(left_top_x+defect_h))]=0
patch=testimg.crop((xmin,ymin,xmax,ymax))
#====相似度计算==============================================================================================#
patch1=patch.copy()
patch2=temp_img.crop((left_top_x,left_top_y,int(left_top_x+patch1.size[0]),int(left_top_y+patch1.size[1])))
patch2.resize((patch1.size[0],patch1.size[1]))
patch1=np.resize(patch1,-1)
patch2=np.resize(patch2,-1)
mutual_infor=mr.mutual_info_score(patch1,patch2)
print(mutual_infor)
if mutual_infor>0.8:
print(patch.size)
patch=patch.resize((patch.size[0]*scale,patch.size[1]*scale))
print(patch.size)
temp_img.paste(patch,(left_top_x,left_top_y))
temp_img = cv2.cvtColor(np.asarray(temp_img),cv2.COLOR_RGB2BGR)
temp_img = cv2.inpaint(temp_img,mask[:,:,0],3,cv2.INPAINT_TELEA)
temp_img = Image.fromarray(cv2.cvtColor(temp_img,cv2.COLOR_BGR2RGB))
result.append({'name': save_temp_name, 'defect_name': d_name, 'bbox': [left_top_x,left_top_y,left_top_x+defect_h*scale,left_top_y+defect_w*scale]})
else:
continue
else:
mask[int(left_top_y-ring_width):int(left_top_y+defect_w+ring_width),int(left_top_x-ring_width):int(left_top_x+defect_h+ring_width)]=255
mask[int(left_top_y):int(left_top_y+defect_w),int(left_top_x):int(left_top_x+defect_h)]=0
patch=testimg.crop((xmin,ymin,xmax,ymax))
patch1=patch.copy()
patch2=temp_img.crop((left_top_x,left_top_y,int(left_top_x+patch1.size[0]),int(left_top_y+patch1.size[1])))
if mutual_infor>0.8:
temp_img.paste(patch,(left_top_x,left_top_y))
temp_img = cv2.cvtColor(np.asarray(temp_img),cv2.COLOR_RGB2BGR)
temp_img = cv2.inpaint(temp_img,mask[:,:,0],3,cv2.INPAINT_TELEA)
temp_img = Image.fromarray(cv2.cvtColor(temp_img,cv2.COLOR_BGR2RGB))
result.append({'name': save_temp_name, 'defect_name': d_name, 'bbox': [left_top_x,left_top_y,left_top_x+defect_h,left_top_y+defect_w]})
else:
continue
temp_img.save(save_dir+save_temp_name)
json_name='./Duck_inject_normal.json'
with open(json_name,'w') as fp:
json.dump(result, fp, indent = 4, separators=(',', ': '))