一、代码
import time
import random
import copy
import cv2
import os
import math
import numpy as np
from skimage.util import random_noise
from lxml import etree, objectify
import xml.etree.ElementTree as ET
import argparse
from tqdm import tqdm
def show_pic(img, bboxes=None):
'''
输入:
img:图像array
bboxes:图像的所有boudning box list, 格式为[[x_min, y_min, x_max, y_max]....]
names:每个box对应的名称
'''
for i in range(len(bboxes)):
bbox = bboxes[i]
x_min = bbox[0]
y_min = bbox[1]
x_max = bbox[2]
y_max = bbox[3]
cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 3)
cv2.namedWindow('pic', 0)
cv2.moveWindow('pic', 0, 0)
cv2.resizeWindow('pic', 1200, 800)
cv2.imshow('pic', img)
cv2.waitKey(0)
cv2.destroyAllWindows()
class DataAugmentForObjectDetection():
def __init__(self, rotation_rate=0.5, max_rotation_angle=5,
crop_rate=0.5, shift_rate=0.5, change_light_rate=0.5,
add_noise_rate=0.5, flip_rate=0.5,
cutout_rate=0.5, cut_out_length=50, cut_out_holes=1, cut_out_threshold=0.5,
is_addNoise=True, is_changeLight=True, is_cutout=True, is_rotate_img_bbox=True,
is_crop_img_bboxes=True, is_shift_pic_bboxes=True, is_filp_pic_bboxes=True):
self.rotation_rate = rotation_rate
self.max_rotation_angle = max_rotation_angle
self.crop_rate = crop_rate
self.shift_rate = shift_rate
self.change_light_rate = change_light_rate
self.add_noise_rate = add_noise_rate
self.flip_rate = flip_rate
self.cutout_rate = cutout_rate
self.cut_out_length = cut_out_length
self.cut_out_holes = cut_out_holes
self.cut_out_threshold = cut_out_threshold
self.is_addNoise = is_addNoise
self.is_changeLight = is_changeLight
self.is_cutout = is_cutout
self.is_rotate_img_bbox = is_rotate_img_bbox
self.is_crop_img_bboxes = is_crop_img_bboxes
self.is_shift_pic_bboxes = is_shift_pic_bboxes
self.is_filp_pic_bboxes = is_filp_pic_bboxes
def _addNoise(self, img):
'''
输入:
img:图像array
输出:
加噪声后的图像array,由于输出的像素是在[0,1]之间,所以得乘以255
'''
return random_noise(img, mode='gaussian', seed=int(time.time()), clip=True) * 255
def _changeLight(self, img):
alpha = random.uniform(0.35, 1)
blank = np.zeros(img.shape, img.dtype)
return cv2.addWeighted(img, alpha, blank, 1 - alpha, 0)
def _cutout(self, img, bboxes, length=100, n_holes=1, threshold=0.5):
'''
原版本:https://github.com/uoguelph-mlrg/Cutout/blob/master/util/cutout.py
Randomly mask out one or more patches from an image.
Args:
img : a 3D numpy array,(h,w,c)
bboxes : 框的坐标
n_holes (int): Number of patches to cut out of each image.
length (int): The length (in pixels) of each square patch.
'''
def cal_iou(boxA, boxB):
'''
boxA, boxB为两个框,返回iou
boxB为bouding box
'''
xA = max(boxA[0], boxB[0])
yA = max(boxA[1], boxB[1])
xB = min(boxA[2], boxB[2])
yB = min(boxA[3], boxB[3])
if xB <= xA or yB <= yA:
return 0.0
interArea = (xB - xA + 1) * (yB - yA + 1)
boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
iou = interArea / float(boxBArea)
return iou
if img.ndim == 3:
h, w, c = img.shape
else:
_, h, w, c = img.shape
mask = np.ones((h, w, c), np.float32)
for n in range(n_holes):
chongdie = True
while chongdie:
y = np.random.randint(h)
x = np.random.randint(w)
y1 = np.clip(y - length // 2, 0,
h)
y2 = np.clip(y + length // 2, 0, h)
x1 = np.clip(x - length // 2, 0, w)
x2 = np.clip(x + length // 2, 0, w)
chongdie = False
for box in bboxes:
if cal_iou([x1, y1, x2, y2], box) > threshold:
chongdie = True
break
mask[y1: y2, x1: x2, :] = 0.
img = img * mask
return img
def _rotate_img_bbox(self, img, bboxes, angle=5, scale=1.):
'''
参考:https://blog.csdn.net/u014540717/article/details/53301195crop_rate
输入:
img:图像array,(h,w,c)
bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
angle:旋转角度
scale:默认1
输出:
rot_img:旋转后的图像array
rot_bboxes:旋转后的boundingbox坐标list
'''
w = img.shape[1]
h = img.shape[0]
rangle = np.deg2rad(angle)
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)
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]
rot_img = cv2.warpAffine(img, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4)
rot_bboxes = list()
for bbox in bboxes:
xmin = bbox[0]
ymin = bbox[1]
xmax = bbox[2]
ymax = bbox[3]
point1 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymin, 1]))
point2 = np.dot(rot_mat, np.array([xmax, (ymin + ymax) / 2, 1]))
point3 = np.dot(rot_mat, np.array([(xmin + xmax) / 2, ymax, 1]))
point4 = np.dot(rot_mat, np.array([xmin, (ymin + ymax) / 2, 1]))
concat = np.vstack((point1, point2, point3, point4))
concat = concat.astype(np.int32)
rx, ry, rw, rh = cv2.boundingRect(concat)
rx_min = rx
ry_min = ry
rx_max = rx + rw
ry_max = ry + rh
rot_bboxes.append([rx_min, ry_min, rx_max, ry_max])
return rot_img, rot_bboxes
def _crop_img_bboxes(self, img, bboxes):
'''
裁剪后的图片要包含所有的框
输入:
img:图像array
bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
输出:
crop_img:裁剪后的图像array
crop_bboxes:裁剪后的bounding box的坐标list
'''
w = img.shape[1]
h = img.shape[0]
x_min = w
x_max = 0
y_min = h
y_max = 0
for bbox in bboxes:
x_min = min(x_min, bbox[0])
y_min = min(y_min, bbox[1])
x_max = max(x_max, bbox[2])
y_max = max(y_max, bbox[3])
d_to_left = x_min
d_to_right = w - x_max
d_to_top = y_min
d_to_bottom = h - y_max
crop_x_min = int(x_min - random.uniform(0, d_to_left))
crop_y_min = int(y_min - random.uniform(0, d_to_top))
crop_x_max = int(x_max + random.uniform(0, d_to_right))
crop_y_max = int(y_max + random.uniform(0, d_to_bottom))
crop_x_min = max(0, crop_x_min)
crop_y_min = max(0, crop_y_min)
crop_x_max = min(w, crop_x_max)
crop_y_max = min(h, crop_y_max)
crop_img = img[crop_y_min:crop_y_max, crop_x_min:crop_x_max]
crop_bboxes = list()
for bbox in bboxes:
crop_bboxes.append([bbox[0] - crop_x_min, bbox[1] - crop_y_min, bbox[2] - crop_x_min, bbox[3] - crop_y_min])
return crop_img, crop_bboxes
def _shift_pic_bboxes(self, img, bboxes):
'''
平移后的图片要包含所有的框
输入:
img:图像array
bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
输出:
shift_img:平移后的图像array
shift_bboxes:平移后的bounding box的坐标list
'''
w = img.shape[1]
h = img.shape[0]
x_min = w
x_max = 0
y_min = h
y_max = 0
for bbox in bboxes:
x_min = min(x_min, bbox[0])
y_min = min(y_min, bbox[1])
x_max = max(x_max, bbox[2])
y_max = max(y_max, bbox[3])
d_to_left = x_min
d_to_right = w - x_max
d_to_top = y_min
d_to_bottom = h - y_max
x = random.uniform(-(d_to_left - 1) / 3, (d_to_right - 1) / 3)
y = random.uniform(-(d_to_top - 1) / 3, (d_to_bottom - 1) / 3)
M = np.float32([[1, 0, x], [0, 1, y]])
shift_img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
shift_bboxes = list()
for bbox in bboxes:
shift_bboxes.append([bbox[0] + x, bbox[1] + y, bbox[2] + x, bbox[3] + y])
return shift_img, shift_bboxes
def _filp_pic_bboxes(self, img, bboxes):
'''
平移后的图片要包含所有的框
输入:
img:图像array
bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
输出:
flip_img:平移后的图像array
flip_bboxes:平移后的bounding box的坐标list
'''
flip_img = copy.deepcopy(img)
h, w, _ = img.shape
sed = random.random()
if 0 < sed < 0.33:
flip_img = cv2.flip(flip_img, 0)
inver = 0
elif 0.33 < sed < 0.66:
flip_img = cv2.flip(flip_img, 1)
inver = 1
else:
flip_img = cv2.flip(flip_img, -1)
inver = -1
flip_bboxes = list()
for box in bboxes:
x_min = box[0]
y_min = box[1]
x_max = box[2]
y_max = box[3]
if inver == 0:
flip_bboxes.append([x_min, h - y_max, x_max, h - y_min])
elif inver == 1:
flip_bboxes.append([w - x_max, y_min, w - x_min, y_max])
elif inver == -1:
flip_bboxes.append([w - x_max, h - y_max, w - x_min, h - y_min])
return flip_img, flip_bboxes
def dataAugment(self, img, bboxes):
'''
图像增强
输入:
img:图像array
bboxes:该图像的所有框坐标
输出:
img:增强后的图像
bboxes:增强后图片对应的box
'''
change_num = 0
while change_num < 1:
if self.is_rotate_img_bbox:
if random.random() > self.rotation_rate:
change_num += 1
angle = random.uniform(-self.max_rotation_angle, self.max_rotation_angle)
scale = random.uniform(0.7, 0.8)
img, bboxes = self._rotate_img_bbox(img, bboxes, angle, scale)
if self.is_shift_pic_bboxes:
if random.random() < self.shift_rate:
change_num += 1
img, bboxes = self._shift_pic_bboxes(img, bboxes)
if self.is_changeLight:
if random.random() > self.change_light_rate:
change_num += 1
img = self._changeLight(img)
if self.is_addNoise:
if random.random() < self.add_noise_rate:
change_num += 1
img = self._addNoise(img)
if self.is_cutout:
if random.random() < self.cutout_rate:
change_num += 1
img = self._cutout(img, bboxes, length=self.cut_out_length, n_holes=self.cut_out_holes,
threshold=self.cut_out_threshold)
if self.is_filp_pic_bboxes:
if random.random() < self.flip_rate:
change_num += 1
img, bboxes = self._filp_pic_bboxes(img, bboxes)
return img, bboxes
class ToolHelper():
def parse_xml(self, path):
'''
输入:
xml_path: xml的文件路径
输出:
从xml文件中提取bounding box信息, 格式为[[x_min, y_min, x_max, y_max, name]]
'''
tree = ET.parse(path)
root = tree.getroot()
objs = root.findall('object')
coords = list()
for ix, obj in enumerate(objs):
name = obj.find('name').text
box = obj.find('bndbox')
x_min = int(box[0].text)
y_min = int(box[1].text)
x_max = int(box[2].text)
y_max = int(box[3].text)
coords.append([x_min, y_min, x_max, y_max, name])
return coords
def save_img(self, file_name, save_folder, img):
cv2.imwrite(os.path.join(save_folder, file_name), img)
def save_xml(self, file_name, save_folder, img_info, height, width, channel, bboxs_info):
'''
:param file_name:文件名
:param save_folder:#保存的xml文件的结果
:param height:图片的信息
:param width:图片的宽度
:param channel:通道
:return:
'''
folder_name, img_name = img_info
E = objectify.ElementMaker(annotate=False)
anno_tree = E.annotation(
E.folder(folder_name),
E.filename(img_name),
E.path(os.path.join(folder_name, img_name)),
E.source(
E.database('Unknown'),
),
E.size(
E.width(width),
E.height(height),
E.depth(channel)
),
E.segmented(0),
)
labels, bboxs = bboxs_info
for label, box in zip(labels, bboxs):
anno_tree.append(
E.object(
E.name(label),
E.pose('Unspecified'),
E.truncated('0'),
E.difficult('0'),
E.bndbox(
E.xmin(box[0]),
E.ymin(box[1]),
E.xmax(box[2]),
E.ymax(box[3])
)
))
etree.ElementTree(anno_tree).write(os.path.join(save_folder, file_name), pretty_print=True)
if __name__ == '__main__':
need_aug_num = 5
is_endwidth_dot = True
dataAug = DataAugmentForObjectDetection()
toolhelper = ToolHelper()
parser = argparse.ArgumentParser()
parser.add_argument('--source_img_path', type=str, default='F:/ship-voc/images')
parser.add_argument('--source_xml_path', type=str, default='F:/ship-voc/annotation')
parser.add_argument('--save_img_path', type=str, default='F:/ship-enhence/images')
parser.add_argument('--save_xml_path', type=str, default='F:/ship-enhence/xmls')
args = parser.parse_args()
source_img_path = args.source_img_path
source_xml_path = args.source_xml_path
save_img_path = args.save_img_path
save_xml_path = args.save_xml_path
if not os.path.exists(save_img_path):
os.mkdir(save_img_path)
if not os.path.exists(save_xml_path):
os.mkdir(save_xml_path)
for parent, _, files in tqdm(os.walk(source_img_path)):
files.sort()
for file in files:
cnt = 0
pic_path = os.path.join(parent, file)
xml_path = os.path.join(source_xml_path, file[:-4] + '.xml')
values = toolhelper.parse_xml(xml_path)
coords = [v[:4] for v in values]
labels = [v[-1] for v in values]
if is_endwidth_dot:
dot_index = file.rfind('.')
_file_prefix = file[:dot_index]
_file_suffix = file[dot_index:]
img = cv2.imread(pic_path)
while cnt < need_aug_num:
auged_img, auged_bboxes = dataAug.dataAugment(img, coords)
auged_bboxes_int = np.array(auged_bboxes).astype(np.int32)
height, width, channel = auged_img.shape
img_name = '{}_{}{}'.format(_file_prefix, cnt + 1, _file_suffix)
toolhelper.save_img(img_name, save_img_path,
auged_img)
toolhelper.save_xml('{}_{}.xml'.format(_file_prefix, cnt + 1),
save_xml_path, (save_img_path, img_name), height, width, channel,
(labels, auged_bboxes_int))
print(img_name)
cnt += 1
二、算法介绍
2.1 主要增强方法及其效果:
- 添加噪声(_addNoise): 向图像中添加随机的高斯噪声。
- 改变亮度(_changeLight): 随机调整图像的亮度。
- Cutout(_cutout): 在图像中随机选择区域进行遮蔽(Cutout)。
- 旋转(_rotate_img_bbox): 以随机角度旋转图像,并相应地调整边界框。
- 裁剪(_crop_img_bboxes): 通过随机裁剪图像的一部分,并确保裁剪后的图像仍然包含所有的目标框。
- 平移(_shift_pic_bboxes): 在水平或垂直方向上随机平移图像,并相应地调整边界框。
- 镜像(_filp_pic_bboxes): 随机选择水平翻转、垂直翻转或对角线翻转图像,并相应地调整边界框。
2.2 用法
DataAugmentForObjectDetection
类中定义了一系列可能的增强方法,并且每次调用dataAugment
方法时,都会根据各自设置的概率率(rate)来决定是否应用上述的某种增强方法。例如,如果rotation_rate
设置为0.5,则有50%的概率会应用旋转增强。如果随机数生成的结果满足条件(即随机数小于对应的概率率),则执行相应的增强方法。如果调用dataAugment
方法时,七个增强方法全部被触发,则在该图像的这次增强中混合了七种增强方法。- 在主函数中,设置
need_aug_num:
确保了每张图像至少会被增强need_aug_num
次,每次增强可能会随机使用上述效果的一种或几种。 - 在
ToolHelper
类的 save_img
和 save_xml
方法分别用于保存增强后的图像和更新后的XML标注文件。