devbox
代码探险家
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

【Python数据增强】图像数据集扩充

作者：代码探险家 | 2024-07-26 10:04:03

赞

踩

【Python数据增强】图像数据集扩充

 前言：该脚本用于图像数据增强，特别是目标检测任务中的图像和标签数据增强。通过应用一系列数据增强技术（如旋转、平移、裁剪、加噪声、改变亮度、cutout、翻转等），生成多样化的图像数据集，以提高目标检测模型的鲁棒性和准确性。

效果：img存的原始图像168张图片，img2扩充的数量为5040张图片

目录

1.环境准备

2.显示图片函数

3.数据增强类

3.1类初始化

3.2数据增强方法

3.3 数据增强主方法

4.XML解析工具类 

4.1 解析XML

4.2 保存图片 

4.3 保存XML 

5. 主函数

完整程序

1.环境准备

这段代码导入了脚本所需的库，用于图像处理（cv2、numpy）、随机操作（random）、文件操作（os）、XML解析（etree）等。

  1. # -*- coding=utf-8 -*-
  2. import time
  3. import random
  4. import copy
  5. import cv2
  6. import os
  7. import math
  8. import numpy as np
  9. from skimage.util import random_noise
  10. from lxml import etree, objectify
  11. import xml.etree.ElementTree as ET
  12. import argparse

2.显示图片函数

该函数用于显示图片，并在图片上绘制边界框（bounding box）。

  1. def show_pic(img, bboxes=None):
  2.     '''
  3.     输入:
  4.         img: 图像array
  5.         bboxes: 图像的所有bounding box list, 格式为[[x_min, y_min, x_max, y_max]....]
  6.     '''
  7.     for i in range(len(bboxes)):
  8.         bbox = bboxes[i]
  9.         x_min = bbox[0]
  10.         y_min = bbox[1]
  11.         x_max = bbox[2]
  12.         y_max = bbox[3]
  13.         cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 3)
  14.     cv2.namedWindow('pic', 0)
  15.     cv2.moveWindow('pic', 0, 0)
  16.     cv2.resizeWindow('pic', 1200, 800)
  17.     cv2.imshow('pic', img)
  18.     cv2.waitKey(0)
  19.     cv2.destroyAllWindows()

3.数据增强类

3.1类初始化

该类初始化函数设置了数据增强的各种参数和是否启用某种增强方式的标志。

  1. class DataAugmentForObjectDetection():
  2.     def __init__(self, rotation_rate=0.5, max_rotation_angle=5,
  3.                  crop_rate=0.5, shift_rate=0.5, change_light_rate=0.5,
  4.                  add_noise_rate=0.5, flip_rate=0.5,
  5.                  cutout_rate=0.5, cut_out_length=50, cut_out_holes=1, cut_out_threshold=0.5,
  6.                  is_addNoise=True, is_changeLight=True, is_cutout=True, is_rotate_img_bbox=True,
  7.                  is_crop_img_bboxes=True, is_shift_pic_bboxes=True, is_filp_pic_bboxes=True):
  8.  
  9.         self.rotation_rate = rotation_rate
  10.         self.max_rotation_angle = max_rotation_angle
  11.         self.crop_rate = crop_rate
  12.         self.shift_rate = shift_rate
  13.         self.change_light_rate = change_light_rate
  14.         self.add_noise_rate = add_noise_rate
  15.         self.flip_rate = flip_rate
  16.         self.cutout_rate = cutout_rate
  17.  
  18.         self.cut_out_length = cut_out_length
  19.         self.cut_out_holes = cut_out_holes
  20.         self.cut_out_threshold = cut_out_threshold
  21.  
  22.         self.is_addNoise = is_addNoise
  23.         self.is_changeLight = is_changeLight
  24.         self.is_cutout = is_cutout
  25.         self.is_rotate_img_bbox = is_rotate_img_bbox
  26.         self.is_crop_img_bboxes = is_crop_img_bboxes
  27.         self.is_shift_pic_bboxes = is_shift_pic_bboxes
  28.         self.is_filp_pic_bboxes = is_filp_pic_bboxes

3.2数据增强方法

加噪声。为图像添加高斯噪声。

  1. def _addNoise(self, img):
  2.     return random_noise(img, mode='gaussian', clip=True) * 255

改变亮度。随机改变图像亮度。

  1. def _changeLight(self, img):
  2.     alpha = random.uniform(0.35, 1)
  3.     blank = np.zeros(img.shape, img.dtype)
  4.     return cv2.addWeighted(img, alpha, blank, 1 - alpha, 0)

cutout。随机在图像中遮挡某些部分（cutout），避免遮挡太多目标。

  1. def _cutout(self, img, bboxes, length=100, n_holes=1, threshold=0.5):
  2.     def cal_iou(boxA, boxB):
  3.         xA = max(boxA[0], boxB[0])
  4.         yA = max(boxA[1], boxB[1])
  5.         xB = min(boxA[2], boxB[2])
  6.         yB = min(boxA[3], boxB[3])
  7.         if xB <= xA or yB <= yA:
  8.             return 0.0
  9.         interArea = (xB - xA + 1) * (yB - yA + 1)
  10.         boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
  11.         boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
  12.         iou = interArea / float(boxBArea)
  13.         return iou
  14.  
  15.     if img.ndim == 3:
  16.         h, w, c = img.shape
  17.     else:
  18.         _, h, w, c = img.shape
  19.     mask = np.ones((h, w, c), np.float32)
  20.     for n in range(n_holes):
  21.         chongdie = True
  22.         while chongdie:
  23.             y = np.random.randint(h)
  24.             x = np.random.randint(w)
  25.             y1 = np.clip(y - length // 2, 0, h)
  26.             y2 = np.clip(y + length // 2, 0, h)
  27.             x1 = np.clip(x - length // 2, 0, w)
  28.             x2 = np.clip(x + length // 2, 0, w)
  29.             chongdie = False
  30.             for box in bboxes:
  31.                 if cal_iou([x1, y1, x2, y2], box) > threshold:
  32.                     chongdie = True
  33.                     break
  34.         mask[y1: y2, x1: x2, :] = 0.
  35.     img = img * mask
  36.     return img

旋转。旋转图像和对应的边界框。

  1. def _rotate_img_bbox(self, img, bboxes, angle=5, scale=1.):
  2.     w, h = img.shape[1], img.shape[0]
  3.     rangle = np.deg2rad(angle)
  4.     nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scale
  5.     nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale
  6.     rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)
  7.     rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
  8.     rot_mat[0, 2] += rot_move[0]
  9.     rot_mat[1, 2] += rot_move[1]
  10.     rot_img = cv2.warpAffine(img, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4)
  11.  
  12.     rot_bboxes = []
  13.     for bbox in bboxes:
  14.         points = np.array([[bbox[0], bbox[1]], [bbox[2], bbox[1]], [bbox[2], bbox[3]], [bbox[0], bbox[3]]])
  15.         new_points = cv2.transform(points[None, :, :], rot_mat)[0]
  16.         rx, ry, rw, rh = cv2.boundingRect(new_points)
  17.         corrected_bbox = [max(0, rx), max(0, ry), min(nw, rx + rw), min(nh, ry + rh)]
  18.         corrected_bbox = [int(val) for val in corrected_bbox]
  19.         rot_bboxes.append(corrected_bbox)
  20.     return rot_img, rot_bboxes

裁剪。随机裁剪图像，同时裁剪对应的边界框。

  1. def _crop_img_bboxes(self, img, bboxes):
  2.     w = img.shape[1]
  3.     h = img.shape[0]
  4.     x_min = w
  5.     x_max = 0
  6.     y_min = h
  7.     y_max = 0
  8.     for bbox in bboxes:
  9.         x_min = min(x_min, bbox[0])
  10.         y_min = min(y_min, bbox[1])
  11.         x_max = max(x_max, bbox[2])
  12.         y_max = max(y_max, bbox[3])
  13.  
  14.     d_to_left = x_min
  15.     d_to_right = w - x_max
  16.     d_to_top = y_min
  17.     d_to_bottom = h - y_max
  18.  
  19.     crop_x_min = int(x_min - random.uniform(0, d_to_left))
  20.     crop_y_min = int(y_min - random.uniform(0, d_to_top))
  21.     crop_x_max = int(x_max + random.uniform(0, d_to_right))
  22.     crop_y_max = int(y_max + random.uniform(0, d_to_bottom))
  23.  
  24.     crop_x_min = max(0, crop_x_min)
  25.     crop_y_min = max(0, crop_y_min)
  26.     crop_x_max = min(w, crop_x_max)
  27.     crop_y_max = min(h, crop_y_max)
  28.  
  29.     crop_img = img[crop_y_min:crop_y_max, crop_x_min:crop_x_max]
  30.  
  31.     crop_bboxes = list()
  32.     for bbox in bboxes:
  33.         crop_bboxes.append([bbox[0] - crop_x_min, bbox[1] - crop_y_min, bbox[2] - crop_x_min, bbox[3] - crop_y_min])
  34.  
  35.     return crop_img, crop_bboxes

平移。随机平移图像和对应的边界框。

  1. def _shift_pic_bboxes(self, img, bboxes):
  2.     h, w = img.shape[:2]
  3.     x = random.uniform(-w * 0.2, w * 0.2)
  4.     y = random.uniform(-h * 0.2, h * 0.2)
  5.     M = np.float32([[1, 0, x], [0, 1, y]])
  6.     shift_img = cv2.warpAffine(img, M, (w, h))
  7.  
  8.     shift_bboxes = []
  9.     for bbox in bboxes:
  10.         new_bbox = [bbox[0] + x, bbox[1] + y, bbox[2] + x, bbox[3] + y]
  11.         corrected_bbox = [max(0, new_bbox[0]), max(0, new_bbox[1]), min(w, new_bbox[2]), min(h, new_bbox[3])]
  12.         corrected_bbox = [int(val) for val in corrected_bbox]
  13.         shift_bboxes.append(corrected_bbox)
  14.     return shift_img, shift_bboxes

 翻转。随机翻转图像和对应的边界框。

  1. def _filp_pic_bboxes(self, img, bboxes):
  2.     flipCode = random.choice([-1, 0, 1])
  3.     flip_img = cv2.flip(img, flipCode)
  4.     h, w, _ = img.shape
  5.     flip_bboxes = []
  6.  
  7.     for bbox in bboxes:
  8.         x_min, y_min, x_max, y_max = bbox
  9.         if flipCode == 0:
  10.             new_bbox = [x_min, h - y_max, x_max, h - y_min]
  11.         elif flipCode == 1:
  12.             new_bbox = [w - x_max, y_min, w - x_min, y_max]
  13.         else:
  14.             new_bbox = [w - x_max, h - y_max, w - x_min, h - y_min]
  15.         flip_bboxes.append(new_bbox)
  16.  
  17.     return flip_img, flip_bboxes

3.3 数据增强主方法

综合应用各种数据增强方法，对输入图像和边界框进行增强。

  1. def dataAugment(self, img, bboxes):
  2.     change_num = 0
  3.     while change_num < 1:
  4.         if self.is_rotate_img_bbox:
  5.             if random.random() > self.rotation_rate:
  6.                 change_num += 1
  7.                 angle = random.uniform(-self.max_rotation_angle, self.max_rotation_angle)
  8.                 scale = random.uniform(0.7, 0.8)
  9.                 img, bboxes = self._rotate_img_bbox(img, bboxes, angle, scale)
  10.  
  11.         if self.is_shift_pic_bboxes:
  12.             if random.random() < self.shift_rate:
  13.                 change_num += 1
  14.                 img, bboxes = self._shift_pic_bboxes(img, bboxes)
  15.  
  16.         if self.is_changeLight:
  17.             if random.random() > self.change_light_rate:
  18.                 change_num += 1
  19.                 img = self._changeLight(img)
  20.  
  21.         if self.is_addNoise:
  22.             if random.random() < self.add_noise_rate:
  23.                 change_num += 1
  24.                 img = self._addNoise(img)
  25.  
  26.         if self.is_cutout:
  27.             if random.random() < self.cutout_rate:
  28.                 change_num += 1
  29.                 img = self._cutout(img, bboxes, length=self.cut_out_length, n_holes=self.cut_out_holes,
  30.                                    threshold=self.cut_out_threshold)
  31.  
  32.         if self.is_filp_pic_bboxes:
  33.             if random.random() < self.flip_rate:
  34.                 change_num += 1
  35.                 img, bboxes = self._filp_pic_bboxes(img, bboxes)
  36.  
  37.     return img, bboxes

4.XML解析工具类 

4.1 解析XML

从XML文件中提取边界框信息。

  1. class ToolHelper():
  2.     def parse_xml(self, path):
  3.         tree = ET.parse(path)
  4.         root = tree.getroot()
  5.         objs = root.findall('object')
  6.         coords = list()
  7.         for ix, obj in enumerate(objs):
  8.             name = obj.find('name').text
  9.             box = obj.find('bndbox')
  10.             x_min = int(box[0].text)
  11.             y_min = int(box[1].text)
  12.             x_max = int(box[2].text)
  13.             y_max = int(box[3].text)
  14.             coords.append([x_min, y_min, x_max, y_max, name])
  15.         return coords

4.2 保存图片 

保存增强后的图片。

  1. def save_img(self, file_name, save_folder, img):
  2.     cv2.imwrite(os.path.join(save_folder, file_name), img)

4.3 保存XML 

保存增强后的XML文件。

  1. def save_xml(self, file_name, save_folder, img_info, height, width, channel, bboxs_info):
  2.     folder_name, img_name = img_info
  3.  
  4.     E = objectify.ElementMaker(annotate=False)
  5.  
  6.     anno_tree = E.annotation(
  7.         E.folder(folder_name),
  8.         E.filename(img_name),
  9.         E.path(os.path.join(folder_name, img_name)),
  10.         E.source(
  11.             E.database('Unknown'),
  12.         ),
  13.         E.size(
  14.             E.width(width),
  15.             E.height(height),
  16.             E.depth(channel)
  17.         ),
  18.         E.segmented(0),
  19.     )
  20.  
  21.     labels, bboxs = bboxs_info
  22.     for label, box in zip(labels, bboxs):
  23.         anno_tree.append(
  24.             E.object(
  25.                 E.name(label),
  26.                 E.pose('Unspecified'),
  27.                 E.truncated('0'),
  28.                 E.difficult('0'),
  29.                 E.bndbox(
  30.                     E.xmin(box[0]),
  31.                     E.ymin(box[1]),
  32.                     E.xmax(box[2]),
  33.                     E.ymax(box[3])
  34.                 )
  35.             ))
  36.  
  37.     etree.ElementTree(anno_tree).write(os.path.join(save_folder, file_name), pretty_print=True)

5. 主函数

首先新建几个文件夹，修改主函数里相应的文件路径，即可。

  • img 用于存放自己手里已有的数据集图片
  • img2 用于存放增强后的数据集图片
  • xml 用于存放自己手里已有的数据集图片对应的标签（这里必须是VOC格式）
  • xml2 用于存放增强后的数据集图片对应的标签
  • txt 用于存放将xml2中的voc格式的标签转换成txt格式（yolov识别txt格式的标签）

修改每个图片的增强次数即可决定增强图片的数量。 

主函数：

  • 解析命令行参数，获取图片和XML文件路径。
  • 创建保存路径文件夹（如果不存在）。
  • 遍历源图片路径，读取图片和对应的XML文件。
  • 应用数据增强，保存增强后的图片和XML文件。
  1. if __name__ == '__main__':
  2.     need_aug_num = 30  # 每张图片需要增强的次数
  3.  
  4.     is_endwidth_dot = True  # 文件是否以.jpg或者png结尾
  5.  
  6.     dataAug = DataAugmentForObjectDetection()  # 数据增强工具类
  7.  
  8.     toolhelper = ToolHelper()  # 工具
  9.  
  10.     # 获取相关参数
  11.     parser = argparse.ArgumentParser()
  12.     parser.add_argument('--source_img_path', type=str, default='D:/lenovo/Archie/shujukuochongv1.0/img')
  13.     parser.add_argument('--source_xml_path', type=str, default='D:/lenovo/Archie/shujukuochongv1.0/xml')
  14.     parser.add_argument('--save_img_path', type=str, default='D:/lenovo/Archie/shujukuochongv1.0/img2')
  15.     parser.add_argument('--save_xml_path', type=str, default='D:/lenovo/Archie/shujukuochongv1.0/xml2')
  16.     args = parser.parse_args()
  17.     source_img_path = args.source_img_path  # 图片原始位置
  18.     source_xml_path = args.source_xml_path  # xml的原始位置
  19.  
  20.     save_img_path = args.save_img_path  # 图片增强结果保存文件
  21.     save_xml_path = args.save_xml_path  # xml增强结果保存文件
  22.  
  23.     if not os.path.exists(save_img_path):
  24.         os.mkdir(save_img_path)
  25.  
  26.     if not os.path.exists(save_xml_path):
  27.         os.mkdir(save_xml_path)
  28.  
  29.     for parent, _, files in os.walk(source_img_path):
  30.         files.sort()
  31.         for file in files:
  32.             cnt = 0
  33.             pic_path = os.path.join(parent, file)
  34.             xml_path = os.path.join(source_xml_path, file[:-4] + '.xml')
  35.             values = toolhelper.parse_xml(xml_path)
  36.             coords = [v[:4] for v in values]
  37.             labels = [v[-1] for v in values]
  38.  
  39.             if is_endwidth_dot:
  40.                 dot_index = file.rfind('.')
  41.                 _file_prefix = file[:dot_index]
  42.                 _file_suffix = file[dot_index:]
  43.             img = cv2.imread(pic_path)
  44.  
  45.             while cnt < need_aug_num:
  46.                 auged_img, auged_bboxes = dataAug.dataAugment(img, coords)
  47.                 auged_bboxes_int = np.array(auged_bboxes).astype(np.int32)
  48.                 height, width, channel = auged_img.shape
  49.                 img_name = '{}_{}{}'.format(_file_prefix, cnt + 1, _file_suffix)
  50.                 tool

完整程序

该脚本用于对图像数据进行各种数据增强操作，并保存增强后的图像和标签数据。通过这些增强操作，可以生成大量多样化的训练数据，提升目标检测模型的鲁棒性和准确性。

  1. # -*- coding=utf-8 -*-
  2.  
  3. import time
  4. import random
  5. import copy
  6. import cv2
  7. import os
  8. import math
  9. import numpy as np
  10. from skimage.util import random_noise
  11. from lxml import etree, objectify
  12. import xml.etree.ElementTree as ET
  13. import argparse
  14.  
  15.  
  16. # 显示图片
  17. def show_pic(img, bboxes=None):
  18.     '''
  19.     输入:
  20.         img:图像array
  21.         bboxes:图像的所有boudning box list, 格式为[[x_min, y_min, x_max, y_max]....]
  22.         names:每个box对应的名称
  23.     '''
  24.     for i in range(len(bboxes)):
  25.         bbox = bboxes[i]
  26.         x_min = bbox[0]
  27.         y_min = bbox[1]
  28.         x_max = bbox[2]
  29.         y_max = bbox[3]
  30.         cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)), (0, 255, 0), 3)
  31.     cv2.namedWindow('pic', 0)  # 1表示原图
  32.     cv2.moveWindow('pic', 0, 0)
  33.     cv2.resizeWindow('pic', 1200, 800)  # 可视化的图片大小
  34.     cv2.imshow('pic', img)
  35.     cv2.waitKey(0)
  36.     cv2.destroyAllWindows()
  37.  
  38.  
  39. # 图像均为cv2读取
  40. class DataAugmentForObjectDetection():
  41.     def __init__(self, rotation_rate=0.5, max_rotation_angle=5,
  42.                  crop_rate=0.5, shift_rate=0.5, change_light_rate=0.5,
  43.                  add_noise_rate=0.5, flip_rate=0.5,
  44.                  cutout_rate=0.5, cut_out_length=50, cut_out_holes=1, cut_out_threshold=0.5,
  45.                  is_addNoise=True, is_changeLight=True, is_cutout=True, is_rotate_img_bbox=True,
  46.                  is_crop_img_bboxes=True, is_shift_pic_bboxes=True, is_filp_pic_bboxes=True):
  47.  
  48.         # 配置各个操作的属性
  49.         self.rotation_rate = rotation_rate
  50.         self.max_rotation_angle = max_rotation_angle
  51.         self.crop_rate = crop_rate
  52.         self.shift_rate = shift_rate
  53.         self.change_light_rate = change_light_rate
  54.         self.add_noise_rate = add_noise_rate
  55.         self.flip_rate = flip_rate
  56.         self.cutout_rate = cutout_rate
  57.  
  58.         self.cut_out_length = cut_out_length
  59.         self.cut_out_holes = cut_out_holes
  60.         self.cut_out_threshold = cut_out_threshold
  61.  
  62.         # 是否使用某种增强方式
  63.         self.is_addNoise = is_addNoise
  64.         self.is_changeLight = is_changeLight
  65.         self.is_cutout = is_cutout
  66.         self.is_rotate_img_bbox = is_rotate_img_bbox
  67.         self.is_crop_img_bboxes = is_crop_img_bboxes
  68.         self.is_shift_pic_bboxes = is_shift_pic_bboxes
  69.         self.is_filp_pic_bboxes = is_filp_pic_bboxes
  70.  
  71.     # ----1.加噪声---- #
  72.     def _addNoise(self, img):
  73.         '''
  74.         输入:
  75.             img:图像array
  76.         输出:
  77.             加噪声后的图像array,由于输出的像素是在[0,1]之间,所以得乘以255
  78.         '''
  79.         # return cv2.GaussianBlur(img, (11, 11), 0)
  80.         return random_noise(img, mode='gaussian', clip=True) * 255
  81.  
  82.     # ---2.调整亮度--- #
  83.     def _changeLight(self, img):
  84.         alpha = random.uniform(0.35, 1)
  85.         blank = np.zeros(img.shape, img.dtype)
  86.         return cv2.addWeighted(img, alpha, blank, 1 - alpha, 0)
  87.  
  88.     # ---3.cutout--- #
  89.     def _cutout(self, img, bboxes, length=100, n_holes=1, threshold=0.5):
  90.         '''
  91.         原版本：https://github.com/uoguelph-mlrg/Cutout/blob/master/util/cutout.py
  92.         Randomly mask out one or more patches from an image.
  93.         Args:
  94.             img : a 3D numpy array,(h,w,c)
  95.             bboxes : 框的坐标
  96.             n_holes (int): Number of patches to cut out of each image.
  97.             length (int): The length (in pixels) of each square patch.
  98.         '''
  99.  
  100.         def cal_iou(boxA, boxB):
  101.             '''
  102.             boxA, boxB为两个框，返回iou
  103.             boxB为bouding box
  104.             '''
  105.             # determine the (x, y)-coordinates of the intersection rectangle
  106.             xA = max(boxA[0], boxB[0])
  107.             yA = max(boxA[1], boxB[1])
  108.             xB = min(boxA[2], boxB[2])
  109.             yB = min(boxA[3], boxB[3])
  110.  
  111.             if xB <= xA or yB <= yA:
  112.                 return 0.0
  113.  
  114.             # compute the area of intersection rectangle
  115.             interArea = (xB - xA + 1) * (yB - yA + 1)
  116.  
  117.             # compute the area of both the prediction and ground-truth
  118.             # rectangles
  119.             boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
  120.             boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
  121.             iou = interArea / float(boxBArea)
  122.             return iou
  123.  
  124.         # 得到h和w
  125.         if img.ndim == 3:
  126.             h, w, c = img.shape
  127.         else:
  128.             _, h, w, c = img.shape
  129.         mask = np.ones((h, w, c), np.float32)
  130.         for n in range(n_holes):
  131.             chongdie = True  # 看切割的区域是否与box重叠太多
  132.             while chongdie:
  133.                 y = np.random.randint(h)
  134.                 x = np.random.randint(w)
  135.  
  136.                 y1 = np.clip(y - length // 2, 0,
  137.                              h)  # numpy.clip(a, a_min, a_max, out=None), clip这个函数将将数组中的元素限制在a_min, a_max之间，大于a_max的就使得它等于 a_max，小于a_min,的就使得它等于a_min
  138.                 y2 = np.clip(y + length // 2, 0, h)
  139.                 x1 = np.clip(x - length // 2, 0, w)
  140.                 x2 = np.clip(x + length // 2, 0, w)
  141.  
  142.                 chongdie = False
  143.                 for box in bboxes:
  144.                     if cal_iou([x1, y1, x2, y2], box) > threshold:
  145.                         chongdie = True
  146.                         break
  147.             mask[y1: y2, x1: x2, :] = 0.
  148.         img = img * mask
  149.         return img
  150.  
  151.     # ---4.旋转--- #
  152.     def _rotate_img_bbox(self, img, bboxes, angle=5, scale=1.):
  153.         w, h = img.shape[1], img.shape[0]
  154.         rangle = np.deg2rad(angle)  # angle in radians
  155.         nw = (abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)) * scale
  156.         nh = (abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)) * scale
  157.         rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)
  158.         rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
  159.         rot_mat[0, 2] += rot_move[0]
  160.         rot_mat[1, 2] += rot_move[1]
  161.         rot_img = cv2.warpAffine(img, rot_mat, (int(math.ceil(nw)), int(math.ceil(nh))), flags=cv2.INTER_LANCZOS4)
  162.  
  163.         rot_bboxes = []
  164.         for bbox in bboxes:
  165.             points = np.array([[bbox[0], bbox[1]], [bbox[2], bbox[1]], [bbox[2], bbox[3]], [bbox[0], bbox[3]]])
  166.             new_points = cv2.transform(points[None, :, :], rot_mat)[0]
  167.             rx, ry, rw, rh = cv2.boundingRect(new_points)
  168.             corrected_bbox = [max(0, rx), max(0, ry), min(nw, rx + rw), min(nh, ry + rh)]
  169.             corrected_bbox = [int(val) for val in corrected_bbox]  # Convert to int and correct order if necessary
  170.             rot_bboxes.append(corrected_bbox)
  171.         return rot_img, rot_bboxes
  172.  
  173.     # ---5.裁剪--- #
  174.     def _crop_img_bboxes(self, img, bboxes):
  175.         '''
  176.         裁剪后的图片要包含所有的框
  177.         输入:
  178.             img:图像array
  179.             bboxes:该图像包含的所有boundingboxs,一个list,每个元素为[x_min, y_min, x_max, y_max],要确保是数值
  180.         输出:
  181.             crop_img:裁剪后的图像array
  182.             crop_bboxes:裁剪后的bounding box的坐标list
  183.         '''
  184.         # 裁剪图像
  185.         w = img.shape[1]
  186.         h = img.shape[0]
  187.         x_min = w  # 裁剪后的包含所有目标框的最小的框
  188.         x_max = 0
  189.         y_min = h
  190.         y_max = 0
  191.         for bbox in bboxes:
  192.             x_min = min(x_min, bbox[0])
  193.             y_min = min(y_min, bbox[1])
  194.             x_max = max(x_max, bbox[2])
  195.             y_max = max(y_max, bbox[3])
  196.  
  197.         d_to_left = x_min  # 包含所有目标框的最小框到左边的距离
  198.         d_to_right = w - x_max  # 包含所有目标框的最小框到右边的距离
  199.         d_to_top = y_min  # 包含所有目标框的最小框到顶端的距离
  200.         d_to_bottom = h - y_max  # 包含所有目标框的最小框到底部的距离
  201.  
  202.         # 随机扩展这个最小框
  203.         crop_x_min = int(x_min - random.uniform(0, d_to_left))
  204.         crop_y_min = int(y_min - random.uniform(0, d_to_top))
  205.         crop_x_max = int(x_max + random.uniform(0, d_to_right))
  206.         crop_y_max = int(y_max + random.uniform(0, d_to_bottom))
  207.  
  208.         # 随机扩展这个最小框 , 防止别裁的太小
  209.         # crop_x_min = int(x_min - random.uniform(d_to_left//2, d_to_left))
  210.         # crop_y_min = int(y_min - random.uniform(d_to_top//2, d_to_top))
  211.         # crop_x_max = int(x_max + random.uniform(d_to_right//2, d_to_right))
  212.         # crop_y_max = int(y_max + random.uniform(d_to_bottom//2, d_to_bottom))
  213.  
  214.         # 确保不要越界
  215.         crop_x_min = max(0, crop_x_min)
  216.         crop_y_min = max(0, crop_y_min)
  217.         crop_x_max = min(w, crop_x_max)
  218.         crop_y_max = min(h, crop_y_max)
  219.  
  220.         crop_img = img[crop_y_min:crop_y_max, crop_x_min:crop_x_max]
  221.  
  222.         # 裁剪boundingbox
  223.         # 裁剪后的boundingbox坐标计算
  224.         crop_bboxes = list()
  225.         for bbox in bboxes:
  226.             crop_bboxes.append([bbox[0] - crop_x_min, bbox[1] - crop_y_min, bbox[2] - crop_x_min, bbox[3] - crop_y_min])
  227.  
  228.         return crop_img, crop_bboxes
  229.  
  230.     # ---6.平移--- #
  231.     def _shift_pic_bboxes(self, img, bboxes):
  232.         h, w = img.shape[:2]
  233.         x = random.uniform(-w * 0.2, w * 0.2)
  234.         y = random.uniform(-h * 0.2, h * 0.2)
  235.         M = np.float32([[1, 0, x], [0, 1, y]])
  236.         shift_img = cv2.warpAffine(img, M, (w, h))
  237.  
  238.         shift_bboxes = []
  239.         for bbox in bboxes:
  240.             new_bbox = [bbox[0] + x, bbox[1] + y, bbox[2] + x, bbox[3] + y]
  241.             corrected_bbox = [max(0, new_bbox[0]), max(0, new_bbox[1]), min(w, new_bbox[2]), min(h, new_bbox[3])]
  242.             corrected_bbox = [int(val) for val in corrected_bbox]  # Convert to int and correct order if necessary
  243.             shift_bboxes.append(corrected_bbox)
  244.         return shift_img, shift_bboxes
  245.  
  246.     # ---7.镜像--- #
  247.     def _filp_pic_bboxes(self, img, bboxes):
  248.         # Randomly decide the flip method
  249.         flipCode = random.choice([-1, 0, 1])  # -1: both; 0: vertical; 1: horizontal
  250.         flip_img = cv2.flip(img, flipCode)  # Apply the flip
  251.         h, w, _ = img.shape
  252.         flip_bboxes = []
  253.  
  254.         for bbox in bboxes:
  255.             x_min, y_min, x_max, y_max = bbox
  256.             if flipCode == 0:  # Vertical flip
  257.                 new_bbox = [x_min, h - y_max, x_max, h - y_min]
  258.             elif flipCode == 1:  # Horizontal flip
  259.                 new_bbox = [w - x_max, y_min, w - x_min, y_max]
  260.             else:  # Both flips
  261.                 new_bbox = [w - x_max, h - y_max, w - x_min, h - y_min]
  262.             flip_bboxes.append(new_bbox)
  263.  
  264.         return flip_img, flip_bboxes
  265.  
  266.     # 图像增强方法
  267.     def dataAugment(self, img, bboxes):
  268.         '''
  269.         图像增强
  270.         输入:
  271.             img:图像array
  272.             bboxes:该图像的所有框坐标
  273.         输出:
  274.             img:增强后的图像
  275.             bboxes:增强后图片对应的box
  276.         '''
  277.         change_num = 0  # 改变的次数
  278.         # print('------')
  279.         while change_num < 1:  # 默认至少有一种数据增强生效
  280.  
  281.             if self.is_rotate_img_bbox:
  282.                 if random.random() > self.rotation_rate:  # 旋转
  283.                     change_num += 1
  284.                     angle = random.uniform(-self.max_rotation_angle, self.max_rotation_angle)
  285.                     scale = random.uniform(0.7, 0.8)
  286.                     img, bboxes = self._rotate_img_bbox(img, bboxes, angle, scale)
  287.  
  288.             if self.is_shift_pic_bboxes:
  289.                 if random.random() < self.shift_rate:  # 平移
  290.                     change_num += 1
  291.                     img, bboxes = self._shift_pic_bboxes(img, bboxes)
  292.  
  293.             if self.is_changeLight:
  294.                 if random.random() > self.change_light_rate:  # 改变亮度
  295.                     change_num += 1
  296.                     img = self._changeLight(img)
  297.  
  298.             if self.is_addNoise:
  299.                 if random.random() < self.add_noise_rate:  # 加噪声
  300.                     change_num += 1
  301.                     img = self._addNoise(img)
  302.             if self.is_cutout:
  303.                 if random.random() < self.cutout_rate:  # cutout
  304.                     change_num += 1
  305.                     img = self._cutout(img, bboxes, length=self.cut_out_length, n_holes=self.cut_out_holes,
  306.                                        threshold=self.cut_out_threshold)
  307.             if self.is_filp_pic_bboxes:
  308.                 if random.random() < self.flip_rate:  # 翻转
  309.                     change_num += 1
  310.                     img, bboxes = self._filp_pic_bboxes(img, bboxes)
  311.  
  312.         return img, bboxes
  313.  
  314.  
  315. # xml解析工具
  316. class ToolHelper():
  317.     # 从xml文件中提取bounding box信息, 格式为[[x_min, y_min, x_max, y_max, name]]
  318.     def parse_xml(self, path):
  319.         '''
  320.         输入：
  321.             xml_path: xml的文件路径
  322.         输出：
  323.             从xml文件中提取bounding box信息, 格式为[[x_min, y_min, x_max, y_max, name]]
  324.         '''
  325.         tree = ET.parse(path)
  326.         root = tree.getroot()
  327.         objs = root.findall('object')
  328.         coords = list()
  329.         for ix, obj in enumerate(objs):
  330.             name = obj.find('name').text
  331.             box = obj.find('bndbox')
  332.             x_min = int(box[0].text)
  333.             y_min = int(box[1].text)
  334.             x_max = int(box[2].text)
  335.             y_max = int(box[3].text)
  336.             coords.append([x_min, y_min, x_max, y_max, name])
  337.         return coords
  338.  
  339.     # 保存图片结果
  340.     def save_img(self, file_name, save_folder, img):
  341.         cv2.imwrite(os.path.join(save_folder, file_name), img)
  342.  
  343.     # 保持xml结果
  344.     def save_xml(self, file_name, save_folder, img_info, height, width, channel, bboxs_info):
  345.         '''
  346.         :param file_name:文件名
  347.         :param save_folder:#保存的xml文件的结果
  348.         :param height:图片的信息
  349.         :param width:图片的宽度
  350.         :param channel:通道
  351.         :return:
  352.         '''
  353.         folder_name, img_name = img_info  # 得到图片的信息
  354.  
  355.         E = objectify.ElementMaker(annotate=False)
  356.  
  357.         anno_tree = E.annotation(
  358.             E.folder(folder_name),
  359.             E.filename(img_name),
  360.             E.path(os.path.join(folder_name, img_name)),
  361.             E.source(
  362.                 E.database('Unknown'),
  363.             ),
  364.             E.size(
  365.                 E.width(width),
  366.                 E.height(height),
  367.                 E.depth(channel)
  368.             ),
  369.             E.segmented(0),
  370.         )
  371.  
  372.         labels, bboxs = bboxs_info  # 得到边框和标签信息
  373.         for label, box in zip(labels, bboxs):
  374.             anno_tree.append(
  375.                 E.object(
  376.                     E.name(label),
  377.                     E.pose('Unspecified'),
  378.                     E.truncated('0'),
  379.                     E.difficult('0'),
  380.                     E.bndbox(
  381.                         E.xmin(box[0]),
  382.                         E.ymin(box[1]),
  383.                         E.xmax(box[2]),
  384.                         E.ymax(box[3])
  385.                     )
  386.                 ))
  387.  
  388.         etree.ElementTree(anno_tree).write(os.path.join(save_folder, file_name), pretty_print=True)
  389.  
  390.  
  391. if __name__ == '__main__':
  392.  
  393.     need_aug_num = 30  # 每张图片需要增强的次数
  394.  
  395.     is_endwidth_dot = True  # 文件是否以.jpg或者png结尾
  396.  
  397.     dataAug = DataAugmentForObjectDetection()  # 数据增强工具类
  398.  
  399.     toolhelper = ToolHelper()  # 工具
  400.  
  401.     # 获取相关参数
  402.     parser = argparse.ArgumentParser()
  403.     parser.add_argument('--source_img_path', type=str, default='D:/lenovo/Archie/shujukuochongv1.0/img')
  404.     parser.add_argument('--source_xml_path', type=str, default='D:/lenovo/Archie/shujukuochongv1.0/xml')
  405.     parser.add_argument('--save_img_path', type=str, default='D:/lenovo/Archie/shujukuochongv1.0/img2')
  406.     parser.add_argument('--save_xml_path', type=str, default='D:/lenovo/Archie/shujukuochongv1.0/xml2')
  407.     args = parser.parse_args()
  408.     source_img_path = args.source_img_path  # 图片原始位置
  409.     source_xml_path = args.source_xml_path  # xml的原始位置
  410.  
  411.     save_img_path = args.save_img_path  # 图片增强结果保存文件
  412.     save_xml_path = args.save_xml_path  # xml增强结果保存文件
  413.  
  414.     # 如果保存文件夹不存在就创建
  415.     if not os.path.exists(save_img_path):
  416.         os.mkdir(save_img_path)
  417.  
  418.     if not os.path.exists(save_xml_path):
  419.         os.mkdir(save_xml_path)
  420.  
  421.     for parent, _, files in os.walk(source_img_path):
  422.         files.sort()
  423.         for file in files:
  424.             cnt = 0
  425.             pic_path = os.path.join(parent, file)
  426.             xml_path = os.path.join(source_xml_path, file[:-4] + '.xml')
  427.             values = toolhelper.parse_xml(xml_path)  # 解析得到box信息，格式为[[x_min,y_min,x_max,y_max,name]]
  428.             coords = [v[:4] for v in values]  # 得到框
  429.             labels = [v[-1] for v in values]  # 对象的标签
  430.  
  431.             # 如果图片是有后缀的
  432.             if is_endwidth_dot:
  433.                 # 找到文件的最后名字
  434.                 dot_index = file.rfind('.')
  435.                 _file_prefix = file[:dot_index]  # 文件名的前缀
  436.                 _file_suffix = file[dot_index:]  # 文件名的后缀
  437.             img = cv2.imread(pic_path)
  438.  
  439.             # show_pic(img, coords)  # 显示原图
  440.             while cnt < need_aug_num:  # 继续增强
  441.                 auged_img, auged_bboxes = dataAug.dataAugment(img, coords)
  442.                 auged_bboxes_int = np.array(auged_bboxes).astype(np.int32)
  443.                 height, width, channel = auged_img.shape  # 得到图片的属性
  444.                 img_name = '{}_{}{}'.format(_file_prefix, cnt + 1, _file_suffix)  # 图片保存的信息
  445.                 toolhelper.save_img(img_name, save_img_path,
  446.                                     auged_img)  # 保存增强图片
  447.  
  448.                 toolhelper.save_xml('{}_{}.xml'.format(_file_prefix, cnt + 1),
  449.                                     save_xml_path, (save_img_path, img_name), height, width, channel,
  450.                                     (labels, auged_bboxes_int))  # 保存xml文件
  451.                 # show_pic(auged_img, auged_bboxes)  # 强化后的图
  452.                 print(img_name)
  453.                 cnt += 1  # 继续增强下一张
  454.
声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/代码探险家/article/detail/884714
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号