yolov5数据增强代码

YOLOv5是目前最先进的目标检测模型之一，他的数据增强代码优秀的贡献了YOLOv5的精度和鲁棒性。数据增强是深度学习中最核心的手段之一，对数据集进行平移、旋转、缩放、颜色增强等操作，使得数据集变得更加丰富多样，这可以有效提高模型训练的效果。下面介绍YOLOv5的数据增强代码

在YOLOv5的data.py文件中，增加了以下的增强操作，代码如下：

1.随机缩放：可以将图像随机缩放到指定的尺寸，这样可以增加模型对不同尺度物体的适应性。

def random_scale(img, scale_range=(0.1, 2), target_size=None):
if random.random() < 0.5:
if target_size is None:
h, w = img.shape[:2]
target_size = (int(w * random.uniform(scale_range[0], scale_range[1])),
int(h * random.uniform(scale_range[0], scale_range[1])))
img = cv2.resize(img, target_size, interpolation=cv2.INTER_LINEAR)
return img

2.随机翻转：可以将图像随机左右翻转或上下翻转，扩增数据集，减小过拟合的风险。

def random_horizontal_flip(img, prob=0.5):
if random.random() < prob:
return cv2.flip(img, 1)
return img

def random_vertical_flip(img, prob=0.5):
if random.random() < prob:
return cv2.flip(img, 0)
return img

3.随机裁剪：可以将图像随机裁剪到指定大小，这样可以增加模型对不同物体大小的适应性。

def random_crop(img, bboxes, crop_size, max_trial=50, min_scale=0.3):
h, w = img.shape[:2]
crop_h, crop_w = crop_size

for i in range(max_trial):
scale = random.uniform(min_scale, 1)
aspect_ratio = random.uniform(0.5, 2)
crop_h_scale = int(round(crop_h * math.sqrt(aspect_ratio) * scale))
crop_w_scale = int(round(crop_w / math.sqrt(aspect_ratio) * scale))

if crop_h_scale <= h and crop_w_scale <= w:
y = random.randint(0, h - crop_h_scale)
x = random.randint(0, w - crop_w_scale)

crop_img = img[y:y + crop_h_scale, x:x + crop_w_scale]
new_bboxes = bboxes.copy()
new_bboxes[:, [0, 2]] = bboxes[:, [0, 2]] - x
new_bboxes[:, [1, 3]] = bboxes[:, [1, 3]] - y
new_bboxes[:, [0, 2]] = np.clip(new_bboxes[:, [0, 2]], 0, crop_w_scale - 1)
new_bboxes[:, [1, 3]] = np.clip(new_bboxes[:, [1, 3]], 0, crop_h_scale - 1)
new_bboxes = new_bboxes[(new_bboxes[:, 2] - new_bboxes[:, 0] >= 1) &
(new_bboxes[:, 3] - new_bboxes[:, 1] >= 1)]
if new_bboxes.shape[0] != 0:
return crop_img, new_bboxes

return img, bboxes

4.颜色变化：可以对图像进行颜色变化，增加模型对不同光照环境下的适应性。

def random_color_distort(img, brightness_delta=32, contrast_low=0.5, contrast_high=1.5,
saturation_low=0.5, saturation_high=1.5, hue_delta=18):
if random.random() < 0.5:
img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
img = np.array(img, dtype=np.float32)
img[:, :, 0] = (img[:, :, 0] + random.uniform(-hue_delta, hue_delta)) % 180
img[:, :, 1] = np.clip(img[:, :, 1] * random.uniform(saturation_low, saturation_high), 0, 255)
img[:, :, 2] = np.clip(img[:, :, 2] * random.uniform(contrast_low, contrast_high) + random.uniform(-brightness_delta, brightness_delta), 0, 255)
img = np.array(img, dtype=np.uint8)
img = cv2.cvtColor(img, cv2.COLOR_HSV2BGR)
return img

以上是YOLOv5的数据增强代码，可以通过上述增强手段扩大数据集，增加对物体姿态、角度、亮度等复杂情况的适应能力，提高模型性能。