目标检测中数据增强的方法——填鸭式:针对小样本

在样本量不足的情况下,我们通常会采用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=(',', ': '))


你可能感兴趣的:(计算机视觉)