基于 imageAug 做图像增广

# -*- coding: utf-8 -*-
"""
Created on Fri May 31 13:41:00 2019
对 500* 500 的TTK 图片进行增广，因为 增广之前的信号灯都比较小
@author: panghaiyan
"""

import cv2
import os
import imgaug as ia
from imgaug import augmenters as iaa


# 增广方式
seq = iaa.SomeOf(2, [
    #iaa.Affine(rotate=45),
    #iaa.Affine(rotate=(-10, 10)),
    #iaa.GaussianBlur(sigma=(0, 0.05)), # blur images with a sigma of 0 to 3.0,
    iaa.Affine(scale=(0.15, 0.3)),
], random_order=True)

# 输入数据
NewTxtPah = "E:\\CPlusDev\\TTKYOLO\\NewTag3\\"
NewJpgsPah = "E:\\CPlusDev\\TTKYOLO\\NewLight3\\"

# 输出数据
AugJpgPath = "E:\\CPlusDev\\TTKYOLO\\TTK500Aug3\\"
AugTxtPath = "E:\\CPlusDev\\TTKYOLO\\TTK500AugTxt3\\"

files = os.listdir(NewTxtPah)

iCtrl = 0
counter = 0

NEWSIZE = 500

iCtrl = 0

# 每张图片增广次数
m = 2

# 所有信号灯都是18这个 ID
lightId = "18"

for file in files:
    if not os.path.isdir(file):
        
        # 获取图像
        srcImg = cv2.imread(NewJpgsPah + file.split(".")[0] + ".jpg")
        # 应该都是 500*500
        height, width, channels = srcImg.shape
        
        # 获取一张图片的所有标注信息
        srcTxt = open(NewTxtPah+file)
        iter_f = iter(srcTxt)
        
        bbs = []
        # 遍历每个框
        for line in iter_f:
            row = line.split(" ")
            print(row)
            classIdx = int(row[0])

            xcenter = int(float(row[1]) * width)
            ycenter = int(float(row[2]) * height)
            lightW  = int(float(row[3]) * width)
            lightH  = int(float(row[4]) * height)
            
            # 过大的信号灯就不要了
            if(lightW > 300 or lightH > 300):
                continue
                
            # 获取某个灯的位置
            leftX = max(0, int((xcenter-lightW/2)))
            leftY = max(0, int((ycenter-lightH/2)))
            
            rightX = min(width-1, int((xcenter+lightW/2)))
            leftBottomY = min(height-1, int((ycenter+lightH/2)))     
            
            bb = ia.BoundingBox(x1=leftX, y1=leftY, x2=rightX, y2=leftBottomY);
            bbs.append(bb)
        
        #总共增广 m 次
        for j in range(m):
            bbsi = ia.BoundingBoxesOnImage(bbs,shape=srcImg.shape)
            seq_det=seq.to_deterministic()#保持坐标和图像同步
            
            images = [srcImg]
            images_aug = seq_det.augment_images(images) 
            
            baseName = file.split(".")[0] + "_small_" + str(j)
            outputName = baseName + ".jpg"
            # 保存图片
            cv2.imwrite(AugJpgPath + outputName, images_aug[0])
        
            bbs_aug = seq_det.augment_bounding_boxes(bbsi)
            # 保存对应的标注框信息
            newTxt = open(AugTxtPath + baseName+".txt", "w")
            for i in range(len(bbs_aug.bounding_boxes)):
                after = bbs_aug.bounding_boxes[i]
                
                xcenter = (after.x1 + after.x2) / 2 / width
                ycenter = (after.y1 + after.y2) / 2 / height
                lightW  = (after.x2 - after.x1) / width
                lightH = (after.y2 - after.y1 ) / height
                
                #TODO: 增广过后，太小的检测框 就不添加 标注信息了
                info = lightId + " " + str(xcenter) + " " + str(ycenter) + " " + str(lightW) + " " + str(lightH) + "\n"
                newTxt.write(info)
            
            newTxt.close()
            
#    iCtrl = iCtrl+1
#    if iCtrl > 3:
#        break