目标检测的数据增强 -- 代码实现
与分类任务不同的是，目标检测任务的数据集有boxes位置的标注，而若进行上下翻转、裁剪等操作，坐标位置也就随之改变，需要在transform时对boxes坐标作出对应的转换。
import math
import random
import torch
from PIL import Image, ImageDraw
1、resize
def resize(img, boxes, size, max_size=1000): # 调整大小缩放
w, h = img.size # 输入img w h (480, 364)
if isinstance(size, int): # 按照短边长度，等比例缩放另一边
size_min = min(w, h) # 输入size=300
size_max = max(w, h) # 首先找到短边，缩放的边是364 h
sw = sh = float(size) / size_min # 计算出短边h缩放的比例，将长边w缩放到同等比例
if sw * size_max > max_size: # 放置缩放过大
sw = sh = float(max_size) / size_max # 缩放至最大程度
ow = int(w * sw + 0.5)
oh = int(h * sh + 0.5) # 向上取整
else: # 直接固定长度缩放
ow, oh = size # 输入size（300,300）
sw = float(ow) / w
sh = float(oh) / h # 计算高h和宽w的缩放比例，boxes的y和x乘以这个比例
# 双线性resize
# x1,y1,x2,y2
return img.resize((ow, oh), Image.BILINEAR), boxes*torch.Tensor([sw, sh, sw, sh])
def draw(img, boxes):
draw = ImageDraw.Draw(img) # 绘图
for box in boxes:
draw.rectangle(list(box), outline='red') # 画一个红线矩形框
img.show() # 显示
def test():
img = Image.open('test.jpg')
boxes = torch.Tensor([[167, 63, 280, 202]])
img, boxes = resize(img, boxes, (300, 300))
draw(img, boxes)
test()
2、center_crop
def center_crop(img, boxes, size): # 中部裁剪
w, h = img.size # 原大小
ow, oh = size # 裁剪后的大小
i = int(round((h - oh) / 2.)) # round 四舍五入
j = int(round((w - ow) / 2.)) # 同样我们只需要将照片上下左右两边需要减掉的高和宽计算出来
img = img.crop((j, i, j+ow, i+oh)) # 利用自带的图像处理，选取图像固定位置 x1,y1,x2,y2
boxes -= torch.Tensor([j, i, j, i]) # 将boxes减去裁剪大小，就可以得到正确的boxes位置信息
boxes[:, 0::2].clamp_(min=0, max=ow-1) # clamp_函数是用来防止超出边界
boxes[:, 1::2].clamp_(min=0, max=oh-1)
return img, boxes
3、random_flip
def random_flip(img, boxes): # 随机水平翻转
if random.random() < 0.5: # 0.5的概率随机翻转
img = img.transpose(Image.FLIP_LEFT_RIGHT) # 左右镜像图像
w = img.width # 宽w
xmin = w - boxes[:, 2] # 计算左上角的位置 w-x2 y的位置都不变
xmax = w - boxes[:, 0] # 计算右下角的位置 w-x1
boxes[:, 0] = xmin # 新的x1 = w - 原来的x2
boxes[:, 2] = xmax # 新的x2 = w - 原来的x1
return img, boxes
4、random_crop
def random_crop(img, boxes):
success = False
for attempt in range(10): # 尝试截取位置，如果尝试10次都失败的话就截取中间部分
area = img.size[0] * img.size[1]
target_area = random.uniform(0.56, 1.0) * area#生成随机大小区域
aspect_ratio = random.uniform(3. / 4, 4. / 3) #再次随机大小
w = int(round(math.sqrt(target_area * aspect_ratio)))#计算高和宽
h = int(round(math.sqrt(target_area / aspect_ratio)))
if random.random() < 0.5:
w, h = h, w
if w <= img.size[0] and h <= img.size[1]: #截取规范，不能超出边界
x = random.randint(0, img.size[0] - w)
y = random.randint(0, img.size[1] - h)
success = True
break
# Fallback 截取中间部分
if not success:#尝试10次都不满足的话就直接中间截取
w = h = min(img.size[0], img.size[1])
x = (img.size[0] - w) // 2
y = (img.size[1] - h) // 2
img = img.crop((x, y, x+w, y+h))
boxes -= torch.Tensor([x,y,x,y]) # 原坐标减去截取的
boxes[:,0::2].clamp_(min=0, max=w-1)
boxes[:,1::2].clamp_(min=0, max=h-1)
return img, boxes
参考：pytorch 目标检测数据增强 详细讲解_视觉盛宴的博客-CSDN博客_pytorch目标检测数据增强