YOLOv8目标检测离线数据增强，检查标签是否越界，标签可视化！

# -*- coding=utf-8 -*- ############################################################## # description: # data augmentation for obeject detection # 参考：https://github.com/maozezhong/CV_ToolBox/blob/master/DataAugForObjectDetection ############################################################## # 包括: # 1. 裁剪(需改变bbox) # 2. 平移(需改变bbox) # 3. 改变亮度 # 4. 加噪声 # 5. 旋转角度(需要改变bbox) # 6. 镜像(需要改变bbox) # 7. cutout # 注意: # random.seed(),相同的seed,产生的随机数是一样的!! 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', 1200, 800) # 可视化的图片大小 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 = True self.is_changeLight = True self.is_cutout = True self.is_rotate_img_bbox = True self.is_crop_img_bboxes = True self.is_shift_pic_bboxes = True self.is_filp_pic_bboxes = True # 加噪声 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 # 调整亮度 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) # 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 # 旋转 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