python数据扩充:gamma+重采样+旋转+镜像

训练的时候因为数据出现问题，把做数据扩充的重新整理了一下，应该一步步检查的。

1、gamma变换调节图像对比度

"""
=================================
Gamma and log contrast adjustment
=================================
This example adjusts image contrast by performing a Gamma and a Logarithmic
correction on the input image.

"""

import numpy as np
import cv2
import math
from matplotlib import pyplot as plt

from skimage import exposure
import shutil
import os

def gamma_corrected(savefolder, imgfolder, img, para):

	#pare > 1 Increase brightness
	#pare < 1 Decrease brightness
	imgsavefolder = savefolder + '/JPEG/'
	xmlsavefolder = savefolder + '/XML/'

	if not os.path.exists(imgsavefolder):
		os.makedirs(imgsavefolder)

	if not os.path.exists(xmlsavefolder):
		os.makedirs(xmlsavefolder)

	xmlfolder = imgfolder.replace('JPEG', 'XML')
	xmlfile = xmlfolder + img.replace('jpg', 'xml')

	origimgpath = os.path.join(imgfolder,img)

	im = cv2.imread(os.path.join(imgfolder,img))

	for p in para:
		try:
			gamma_img = exposure.adjust_gamma(im, p)

			gamma_imgsavepath = imgsavefolder + 'g' + str(p) + '_' + img
			orig_imgsavepath = imgsavefolder + img

			gamma_xmlsavepath = xmlsavefolder + 'g' + str(p) + '_' + img.replace('jpg', 'xml')
			orig_xmlsavepath = xmlsavefolder + img.replace('jpg', 'xml')
			
			#save gamma image
			cv2.imwrite(gamma_imgsavepath, gamma_img)
			#sava original image
			shutil.copy(origimgpath,orig_imgsavepath)
			#save gamma xml
			shutil.copy (xmlfile, gamma_xmlsavepath)
			#save original xml
			shutil.copy(xmlfile, orig_xmlsavepath)
		
		except Exception as e:
			orig_imgsavepath = imgsavefolder + img
			orig_xmlsavepath = xmlsavefolder + img.replace('jpg', 'xml')

			#sava original image
			shutil.copy(origimgpath,orig_imgsavepath)
			#save original xml
			shutil.copy(xmlfile, orig_xmlsavepath)
			raise
#log	
# def logarithmic_corrected():
# 	logarithmic_corrected = exposure.adjust_log(img, 0.9)
# 	pass

##Gamma
para = [0.6, 1.5]

img_org_folder = '/Desktop/power_traindata/JPEG/'
save_folder = '/PowerDetectTrainDataExp/Gamma/'

if not os.path.exists(save_folder):
	os.makedirs(save_folder)
	os.makedirs(save_folder + 'JPEG')
	os.makedirs(save_folder + 'XML')

for img in os.listdir(img_org_folder):
	gamma_corrected(save_folder, img_org_folder, img, para)



## Logarithmic

2、上下采样：

#coding:utf-8

import os
import cv2
import numpy as np

try:
	import xml.etree.cElementTree as ET
except ImportEroor:
	import xml.etree.ElementTree as ET
import copy,random
# ==================down sample=========
def img_downsample():
	JPEGImages = os.path.join(folder, 'JPEG')
	saveJPEGImages = os.path.join(save_folder, 'JPEG')
	saveJPEGImages.replace('\\', '/')
	
	if not os.path.exists(saveJPEGImages):
		os.makedirs(saveJPEGImages)
	
	Annotations = os.path.join(folder, 'XML')
	saveAnnotations = os.path.join(save_folder, 'XML')
	saveAnnotations.replace('\\', '/')
	
	if not os.path.exists(saveAnnotations):
		os.makedirs(saveAnnotations)
	
	scales = None
	
	for imgfile in os.listdir(JPEGImages):
		#images
		img = cv2.imread(os.path.join(JPEGImages,imgfile))
		scales = [random.uniform(4,20)*0.1 for _ in range(2)]
		
		for i,scale in enumerate(scales):
			img_downsampled = cv2.resize(img,(int(scale*img.shape[1]),int(scale*img.shape[0])))
			new_name = 's' + str(i) + '_' + imgfile
			# print(new_name)
			cv2.imwrite(os.path.join(saveJPEGImages,new_name),img_downsampled)

			xmlfile = imgfile.replace('jpg','xml')
			Anno_Path = os.path.join(Annotations,xmlfile)

			tree = ET.parse(Anno_Path.replace('\\', '/'))
			root = tree.getroot()
			for obj in root.findall('object'):
				bndbox = obj.find('bndbox')

				xmin = bndbox.find('xmin').text
				ymin = bndbox.find('ymin').text
				xmax = bndbox.find('xmax').text
				ymax = bndbox.find('ymax').text

				bndbox.find('xmin').text = str(int(int(xmin)*scale))
				bndbox.find('ymin').text = str(int(int(ymin)*scale))
				bndbox.find('xmax').text = str(int(int(xmax)*scale))
				bndbox.find('ymax').text = str(int(int(ymax)*scale))

				assert(((int(xmax) - int(xmin))*scale)>0)
				assert(((int(ymax) - int(ymin))*scale)>0)
		
			filename = tree.find('filename')
			filename.text = new_name
			size = tree.find('size')
			height = size.find('height')
			width = size.find('width')
			height.text = str(int(float(height.text)*scale))
			width.text = str(int(float(width.text)*scale))
			tree.write(os.path.join(saveAnnotations,new_name[:-4]+'.xml'))

if __name__=='__main__': 
	global folder, save_folder
	folder = '/PowerDetectTrainDataExp/Gamma/'
	save_folder = '/PowerDetectTrainDataExp/Sampled/'
	img_downsample()

3、旋转：0，90，180, 270

# coding:utf-8

import cv2
import os
import numpy as np
import math
import copy
try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET

def rotate_about_center(src, angle, scale=1.):
    w = src.shape[1]
    h = src.shape[0]

    rangle = np.deg2rad(angle) #angle in radians
    nw = (abs(np.sin(rangle)*h)+abs(np.cos(rangle)*w))*scale
    nh = (abs(np.cos(rangle)*h)+abs(np.sin(rangle)*w))*scale

    rot_mat = cv2.getRotationMatrix2D((nw*0.5,nh*0.5), angle, scale)# rotate with center
    rot_move = np.dot(rot_mat,np.array([(nw-w)*0.5,(nh-h)*0.5,0]))

    rot_mat[0,2] += rot_move[0]
    rot_mat[1,2] += rot_move[1]
    return cv2.warpAffine(src,rot_mat,(int(math.ceil(nw)),int(math.ceil(nh))),flags = cv2.INTER_LANCZOS4)


def enhancement_using_rotation(ImagePath,AnnotationsPath,ImageSavePath,AnnotationsSavePath):

    for imgfile in os.listdir(ImagePath):
        if not os.path.isfile(os.path.join(AnnotationsPath,imgfile[:-4]+'.xml')):
            continue
        print(imgfile)
        img = cv2.imread(os.path.join(ImagePath,imgfile))

        #rotation
        for angle in rotation_angle:

            new_name =  'r'+ str(angle) + '_' + imgfile
            rotate_img = rotate_about_center(img,angle)
            cv2.imwrite(os.path.join(ImageSavePath,new_name),rotate_img)

            center_x = img.shape[1]/2
            center_y = img.shape[0]/2
            new_center_x = rotate_img.shape[1]/2
            new_center_y = rotate_img.shape[0]/2


            if annotation_file is 'txt':            
                ff = open(os.path.join(AnnotationsPath,'labels',imgfile[:-4]+'.txt'),'r')
                f = open(os.path.join(AnnotationsSavePath,'labels',new_name[:-4] +'.txt'),'a')
                for line in ff.readlines():
                    [filename, cls, x1,y1,x2,y2,t] = line.split(' ')
                    
                    final_x1, final_y1, final_x2, final_y2 = \
                                        _rotated_location(x1,y1,x2,y2,center_x,center_y,new_center_x,new_center_y)
                    assert(final_y2-final_y1>0) & (final_x2-final_x1>0)

                    if ifshow == 1:
                        cv2.rectangle(rotate_img,(final_x1,final_y1),(final_x2,final_y2),(0,0,255),2)
                        cv2.putText(rotate_img, cls, (int(final_x1),int(final_y1)),0,1.2,(0,0,255),2)
                    nline = new_name + ' ' + cls + ' ' + str(final_x1) + ' ' + str(final_y1) + ' ' + str(final_x2) + ' ' + str(final_y2) + '\n'
                    f.write(nline)
                f.close()
                ff.close()
            
            elif annotation_file is 'xml':
                Annotations = os.path.join(AnnotationsPath,imgfile[:-4]+'.xml')
                saveAnnotations = os.path.join(AnnotationsSavePath,new_name[:-4]+'.xml')
                tree = ET.parse(Annotations)

                filename = tree.find('filename')
                filename.text = new_name

                size = tree.find('size')
                height = size.find('height')
                height.text = str(rotate_img.shape[0])
                width = size.find('width')
                width.text = str(rotate_img.shape[1])

                root = tree.getroot()
                for obj in tree.findall('object'):
                    class_node = obj.find('name')
                    bndbox = obj.find('bndbox')
                    x1 = bndbox.find('xmin').text
                    y1 = bndbox.find('ymin').text
                    x2 = bndbox.find('xmax').text
                    y2 = bndbox.find('ymax').text
                    if not(int(y2)-int(y1)>0) & (int(x2)-int(x1)>0):
                        print(Annotations)
                        root.remove(obj)
                        continue
                    assert(int(y2)-int(y1)>0) & (int(x2)-int(x1)>0)

                    final_x1, final_y1, final_x2, final_y2 = \
                                        _rotated_location(x1,y1,x2,y2,center_x,center_y,new_center_x,new_center_y,angle)
                    if final_y1<0:
                        print(x1,y1,x2,y2)
                        print(final_x1,final_y1)
                    assert(final_y2-final_y1>0)
                    assert(final_x2-final_x1>0)
                    assert(final_x1>=0)
                    assert(final_y1>=0)
                    assert(final_x2<=int(width.text))
                    assert(final_y2<=int(height.text))

                    if ifshow == 1:
                        cv2.rectangle(rotate_img,(final_x1,final_y1),(final_x2,final_y2),(0,0,255),2)
                        cv2.putText(rotate_img, cls, (int(final_x1),int(final_y1)),0,1.2,(0,0,255),2)
                    bndbox.find('xmin').text = str(final_x1)
                    bndbox.find('ymin').text = str(final_y1)
                    bndbox.find('xmax').text = str(final_x2)
                    bndbox.find('ymax').text = str(final_y2)
              
                tree.write(saveAnnotations)

            if ifshow == 1:
                cv2.imshow(new_name,rotate_img)
                key = cv2.waitKey(500)
                if key ==27:
                        quit()
                cv2.destroyWindow(new_name)


def _rotated_location(x1,y1,x2,y2,center_x,center_y,new_center_x,new_center_y,angle):
    x1 = float(x1) - center_x
    y1 = -(float(y1) - center_y)
    x2 = float(x2) - center_x
    y2 = -(float(y2) - center_y)

    rangle = np.deg2rad(angle)
    rotated_x1 = np.cos(rangle)*x1 - np.sin(rangle)*y1
    rotated_y1 = np.cos(rangle)*y1 + np.sin(rangle)*x1
    rotated_x2 = np.cos(rangle)*x2 - np.sin(rangle)*y2
    rotated_y2 = np.cos(rangle)*y2 + np.sin(rangle)*x2

    rotated_x1 = int(rotated_x1 + new_center_x)
    rotated_y1 = int(-rotated_y1 + new_center_y)
    rotated_x2 = int(rotated_x2 + new_center_x)
    rotated_y2 = int(-rotated_y2 + new_center_y)

    final_x1 = int(min(rotated_x1,rotated_x2))
    final_y1 = int(min(rotated_y1,rotated_y2))
    final_x2 = int(max(rotated_x1,rotated_x2))
    final_y2 = int(max(rotated_y1,rotated_y2))
    return final_x1, final_y1, final_x2, final_y2


def enhancement_using_mirror(ImagePath,AnnotationsPath,ImageSavePath,AnnotationsSavePath):

    #mirror
    for imgfile in os.listdir(ImagePath):
        img = cv2.imread(os.path.join(ImagePath,imgfile))
        new_name = 'm_' + imgfile 
        # print(new_name)

        mirror_img  = mirroir_hierachically(img)
        cv2.imwrite(os.path.join(ImageSavePath,new_name),mirror_img)


        w = img.shape[1]

        if annotation_file is 'txt':
            ff = open(os.path.join(txtpath,imgfile[:-4]+'.txt'),'r')
            f = open(os.path.join(txtpath,new_name[:-4]+'.txt'),'a')
            for line in ff.readlines():
                [filename,cls,x1,y1,x2,y2] = line.strip().split(' ')
                x2 = abs(int(x1)-w)
                x1 = abs(int(x2)-w)


                final_x1 = min(x1,x2)
                final_x2 = max(x1,x2)
                assert(final_x2-final_x1>0)

                nline = new_name + ' ' + cls + ' ' + str(final_x1) + ' ' + y1 + ' ' + str(final_x2) +  ' ' + y2 + '\n'
                f.write(nline)

                if ifshow == 1:
                    cv2.rectangle(mirror_img,(final_x1,int(y1)),(final_x2,int(y2)),(0,0,255),2)
                    cv2.putText(mirror_img, cls, (final_x1,int(y1)),0,1.2,(0,0,255),2)           
            
            f.close()
            ff.close()
        
        elif annotation_file is 'xml':
            Annotations = os.path.join(AnnotationsPath, imgfile[:-4]+'.xml')
            saveAnnotations = os.path.join(AnnotationsSavePath,new_name[:-4]+'.xml')
            tree = ET.parse(Annotations)


            filename = tree.find('filename')
            filename.text = new_name


            for obj in tree.findall('object'):
                bndbox = obj.find('bndbox')
                x1 = bndbox.find('xmin').text
                y1 = bndbox.find('ymin').text
                x2 = bndbox.find('xmax').text
                y2 = bndbox.find('ymax').text

                new_x1 = abs(int(x2)-w)
                new_x2 = abs(int(x1)-w)

                final_x1 = int(min(new_x1,new_x2))
                final_x2 = int(max(new_x1,new_x2))
                assert(final_x2-final_x1>0)

                cls = obj.find('name').text 
     
                if ifshow == 1:
                    cv2.rectangle(mirror_img,(final_x1,int(y1)),(final_x2,int(y2)),(0,0,255),2)
                    cv2.putText(mirror_img, cls, (int(final_x1),int(y1)),0,1.2,(0,0,255),2)
                                           
                bndbox.find('xmin').text = str(final_x1)
                bndbox.find('xmax').text = str(final_x2)


            tree.write(saveAnnotations)
            # tree.clear()


        if ifshow == 1:
            cv2.imshow(new_name,mirror_img)
            key = cv2.waitKey(500)
            if key ==27:
                quit()
            cv2.destroyWindow(new_name)



def mirroir_hierachically(src):
    w = src.shape[1]
    h = src.shape[0]
    ll = src.shape[2]
    mirror_img = copy.deepcopy(src)
    for wi in range(w):
        mirror_img[:,w-wi-1] = src[:,wi]
    return mirror_img


def main(dirpath):
    # ##################################### PART 1 ######################################################## 
    savepath = dirpath + 'Rotated/'
    ImageSavePath = savepath + '/JPEG/'
    AnnotationsSavePath = savepath + '/XML/'
    # ImageSetsSavePath = savepath + '/ImageSets/Main/'

    if not os.path.exists(savepath):
        os.makedirs(savepath)  
    if not os.path.exists(ImageSavePath):
        os.makedirs(ImageSavePath)    
    if not os.path.exists(AnnotationsSavePath):
        os.makedirs(AnnotationsSavePath)

  # ##################################### PART 2 ########################################################

    #augument trainning data by rotating
    print('rotating...')
    ImagePath = dirpath + 'Sampled/JPEG/'
    AnnotationsPath = dirpath + 'Sampled/XML/'
    enhancement_using_rotation(ImagePath,AnnotationsPath,ImageSavePath,AnnotationsSavePath)

    #augument trainning data by mirroring, mirroring after rotating
    print('mirroring...') 
    mir_ImagePath = dirpath + 'Rotated/JPEG/'
    mir_AnnotationsPath = dirpath + 'Rotated/XML/'
    mir_savepath = dirpath + 'mirrored/'
    mir_AnnotationsSavePath = mir_savepath + 'XML/'
    mir_imgsavepath = mir_savepath + 'JPEG/'

    if not os.path.exists(mir_savepath):
        os.makedirs(mir_savepath)   
    if not os.path.exists(mir_imgsavepath):
        os.makedirs(mir_imgsavepath) 
    if not os.path.exists(mir_AnnotationsSavePath):
        os.makedirs(mir_AnnotationsSavePath)

    enhancement_using_mirror(mir_ImagePath,mir_AnnotationsPath,mir_imgsavepath,mir_AnnotationsSavePath)

if __name__ == '__main__':
    global annotation_file, rotation_angle, ifshow
    annotation_file = 'xml'#'txt',xml'
    rotation_angle=[0,90,180,270]
    ifshow = 0
    dirpath = 'E:/PowerDetectTrainDataExp/'
    main(dirpath)