# -*- coding: utf-8 -*-
"""
Created on Fri May 31 13:41:00 2019
对 500* 500 的TTK 图片进行增广,因为 增广之前的信号灯都比较小
@author: haithink
"""
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
基于 imageAug 做图像增广
最新推荐文章于 2024-08-10 07:00:35 发布