545 lines
21 KiB
Python
545 lines
21 KiB
Python
|
# -*- coding=utf-8 -*-
|
|||
|
|
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 显示图片
|
|||
|
|
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) # 1表示原图
|
|||
|
|
cv2.moveWindow('pic', 0, 0)
|
|||
|
|
cv2.resizeWindow('pic', 640, 640) # 可视化的图片大小
|
|||
|
|
cv2.imshow('pic', img)
|
|||
|
|
cv2.waitKey(0)
|
|||
|
|
cv2.destroyAllWindows()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 图像均为cv2读取
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
# ----1.加噪声---- #
|
|||
|
|
def _addNoise(self, img):
|
|||
|
|
'''
|
|||
|
|
输入:
|
|||
|
|
img:图像array
|
|||
|
|
输出:
|
|||
|
|
加噪声后的图像array,由于输出的像素是在[0,1]之间,所以得乘以255
|
|||
|
|
'''
|
|||
|
|
# return cv2.GaussianBlur(img, (11, 11), 0)
|
|||
|
|
return random_noise(img, mode='gaussian', seed=int(time.time()), clip=True) * 255
|
|||
|
|
|
|||
|
|
# ---2.调整亮度--- #
|
|||
|
|
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)
|
|||
|
|
|
|||
|
|
# ---3.cutout--- #
|
|||
|
|
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
|
|||
|
|
'''
|
|||
|
|
# determine the (x, y)-coordinates of the intersection rectangle
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
# compute the area of intersection rectangle
|
|||
|
|
interArea = (xB - xA + 1) * (yB - yA + 1)
|
|||
|
|
|
|||
|
|
# compute the area of both the prediction and ground-truth
|
|||
|
|
# rectangles
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
# 得到h和w
|
|||
|
|
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 # 看切割的区域是否与box重叠太多
|
|||
|
|
while chongdie:
|
|||
|
|
y = np.random.randint(h)
|
|||
|
|
x = np.random.randint(w)
|
|||
|
|
|
|||
|
|
y1 = np.clip(y - length // 2, 0,
|
|||
|
|
h) # numpy.clip(a, a_min, a_max, out=None), clip这个函数将将数组中的元素限制在a_min, a_max之间，大于a_max的就使得它等于 a_max，小于a_min,的就使得它等于a_min
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
# ---4.旋转--- #
|
|||
|
|
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) # angle in radians
|
|||
|
|
# now calculate new image width and height
|
|||
|
|
nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scale
|
|||
|
|
nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale
|
|||
|
|
# ask OpenCV for the rotation matrix
|
|||
|
|
rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)
|
|||
|
|
# calculate the move from the old center to the new center combined
|
|||
|
|
# with the rotation
|
|||
|
|
rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
|
|||
|
|
# the move only affects the translation, so update the translation
|
|||
|
|
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)
|
|||
|
|
|
|||
|
|
# 矫正bbox坐标
|
|||
|
|
# rot_mat是最终的旋转矩阵
|
|||
|
|
# 获取原始bbox的四个中点，然后将这四个点转换到旋转后的坐标系下
|
|||
|
|
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]))
|
|||
|
|
# 合并np.array
|
|||
|
|
concat = np.vstack((point1, point2, point3, point4))
|
|||
|
|
# 改变array类型
|
|||
|
|
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
|
|||
|
|
# 加入list中
|
|||
|
|
rot_bboxes.append([rx_min, ry_min, rx_max, ry_max])
|
|||
|
|
|
|||
|
|
return rot_img, rot_bboxes
|
|||
|
|
|
|||
|
|
# ---5.裁剪--- #
|
|||
|
|
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 = int(x_min - random.uniform(d_to_left//2, d_to_left))
|
|||
|
|
# crop_y_min = int(y_min - random.uniform(d_to_top//2, d_to_top))
|
|||
|
|
# crop_x_max = int(x_max + random.uniform(d_to_right//2, d_to_right))
|
|||
|
|
# crop_y_max = int(y_max + random.uniform(d_to_bottom//2, 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]
|
|||
|
|
|
|||
|
|
# 裁剪boundingbox
|
|||
|
|
# 裁剪后的boundingbox坐标计算
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
# ---6.平移--- #
|
|||
|
|
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]]) # x为向左或右移动的像素值,正为向右负为向左; y为向上或者向下移动的像素值,正为向下负为向上
|
|||
|
|
shift_img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
|
|||
|
|
|
|||
|
|
# 平移boundingbox
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
# ---7.镜像--- #
|
|||
|
|
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: # 0.33的概率水平翻转，0.33的概率垂直翻转,0.33是对角反转
|
|||
|
|
flip_img = cv2.flip(flip_img, 0) # _flip_x
|
|||
|
|
inver = 0
|
|||
|
|
elif 0.33 < sed < 0.66:
|
|||
|
|
flip_img = cv2.flip(flip_img, 1) # _flip_y
|
|||
|
|
inver = 1
|
|||
|
|
else:
|
|||
|
|
flip_img = cv2.flip(flip_img, -1) # flip_x_y
|
|||
|
|
inver = -1
|
|||
|
|
|
|||
|
|
# 调整boundingbox
|
|||
|
|
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:
|
|||
|
|
# 0：垂直翻转
|
|||
|
|
flip_bboxes.append([x_min, h - y_max, x_max, h - y_min])
|
|||
|
|
elif inver == 1:
|
|||
|
|
# 1：水平翻转
|
|||
|
|
flip_bboxes.append([w - x_max, y_min, w - x_min, y_max])
|
|||
|
|
elif inver == -1:
|
|||
|
|
# -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 # 改变的次数
|
|||
|
|
# print('------')
|
|||
|
|
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: # cutout
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
|
|||
|
|
# xml解析工具
|
|||
|
|
class ToolHelper():
|
|||
|
|
# 从xml文件中提取bounding box信息, 格式为[[x_min, y_min, x_max, y_max, name]]
|
|||
|
|
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)
|
|||
|
|
|
|||
|
|
# 保持xml结果
|
|||
|
|
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 # 文件是否以.jpg或者png结尾
|
|||
|
|
|
|||
|
|
dataAug = DataAugmentForObjectDetection() # 数据增强工具类
|
|||
|
|
|
|||
|
|
toolhelper = ToolHelper() # 工具
|
|||
|
|
|
|||
|
|
# 获取相关参数
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument('--source_img_path', type=str, default=r'D:\1BNS_projects\yolo_safehat_num_v8\datasets\2m')
|
|||
|
|
parser.add_argument('--source_xml_path', type=str, default=r'D:\1BNS_projects\yolo_safehat_num_v8\datasets\labels')
|
|||
|
|
parser.add_argument('--save_img_path', type=str, default=r'D:\1BNS_projects\yolo_safehat_num_v8\datasets\2m1')
|
|||
|
|
parser.add_argument('--save_xml_path', type=str, default=r'D:\1BNS_projects\yolo_safehat_num_v8\datasets\labels1')
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
source_img_path = args.source_img_path # 图片原始位置
|
|||
|
|
source_xml_path = args.source_xml_path # xml的原始位置
|
|||
|
|
|
|||
|
|
save_img_path = args.save_img_path # 图片增强结果保存文件
|
|||
|
|
save_xml_path = args.save_xml_path # xml增强结果保存文件
|
|||
|
|
|
|||
|
|
# 如果保存文件夹不存在就创建
|
|||
|
|
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 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) # 解析得到box信息，格式为[[x_min,y_min,x_max,y_max,name]]
|
|||
|
|
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)
|
|||
|
|
|
|||
|
|
# show_pic(img, coords) # 显示原图
|
|||
|
|
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)) # 保存xml文件
|
|||
|
|
# show_pic(auged_img, auged_bboxes) # 强化后的图
|
|||
|
|
print(img_name)
|
|||
|
|
cnt += 1 # 继续增强下一张
|