目标检测中图像和检测框bbox同时进行变换
检测和分类不同,检测过程中,图片处理以后,bbox往往也需要同步变换。
[1] 知乎:目标检测-数据增强(Data Augmentation) (推荐,但要注意代码中的个别错误)
[2] CSDN: 目标检测中的数据增强,包括bbox的变换 (写的还可以)
[3] 知乎:一种目标检测任务中图像-标注对增强方法 (介绍了图像数据增强库albumentations)
下面的代码来自上面的参考资料1,github链接:Data-Augmentation/data_augmentation.py at main · REN-HT/Data-Augmentation · GitHub
import torch from PIL import Image from PIL import ImageDraw from torchvision import transforms import numpy as np class DataAugmentation(object): def __init__(self): super(DataAugmentation,self).__init__() def resize(self,img,boxes,size): #--------------------------------------------------------- # 类型为 img=Image.open(path),boxes:Tensor,size:int # 功能为:将图像长和宽缩放到指定值size,并且相应调整boxes #--------------------------------------------------------- w,h=img.size sw=size/w sh=size/h return img.resize((size, size), Image.BILINEAR), boxes*torch.Tensor([sw, sh, sw, sh]) def resize_(self, img, boxes, size): #----------------------------------------------------------- # 类型为 img=Image.open(path),boxes:Tensor,size:int # 功能为:将图像短边缩放到指定值size,保持原有比例不变,并且相应调整boxes #----------------------------------------------------------- w,h=img.size min_size=min(w,h) sw=sh=size/min_size ow=int(sw*w+0.5) oh=int(sh*h+0.5) return img.resize((ow, oh), Image.BILINEAR), boxes*torch.Tensor([sw,sh,sw,sh]) def random_flip_horizon(self, img, boxes): #------------------------------------- # 随机水平翻转 #------------------------------------- if np.random.random() > 0.5: transform=transforms.RandomHorizontalFlip() img=transform(img) w = img.width xmin=w-boxes[:,2] xmax=w-boxes[:,0] boxes[:,0]=xmin boxes[:,2]=xmax return img, boxes def random_flip_vertical(self, img, boxes): #------------------------------------- # 随机垂直翻转 #------------------------------------- if np.random.random()>0.5: transform=transforms.RandomVerticalFlip() img=transform(img) h=img.height ymin=h-boxes[:, 3] ymax=h-boxes[:, 1] boxes[:, 1]=ymin boxes[:, 3]=ymax return img, boxes def center_crop(self, img, boxes): #------------------------------------- # 中心裁剪 #------------------------------------- w,h=img.size size=min(w,h) transform=transforms.CenterCrop(size) img=transform(img) sw,sh=size/w, size/h boxes=boxes*torch.Tensor([sw, sh, sw, sh]) boxes=boxes.clamp(min=0, max=size) return img, boxes #------------------------------------------------------ # 以下img皆为Tensor类型 #------------------------------------------------------ def random_bright(self, img, u=32): #------------------------------------- # 随机亮度变换 #------------------------------------- if np.random.random()>0.5: alpha=np.random.uniform(-u, u)/255 img+=alpha img=img.clamp(min=0.0, max=1.0) return img def random_contrast(self, img, lower=0.5, upper=1.5): #------------------------------------- # 随机增强对比度 #------------------------------------- if np.random.random()>0.5: alpha=np.random.uniform(lower, upper) img*=alpha img=img.clamp(min=0, max=1.0) return img def random_saturation(self, img,lower=0.5, upper=1.5): #----------------------------------------------- # 随机饱和度变换,针对彩色三通道图像,中间通道乘以一个值 #----------------------------------------------- if np.random.random()>0.5: alpha=np.random.uniform(lower, upper) img[1]=img[1]*alpha img[1]=img[1].clamp(min=0,max=1.0) return img def add_gasuss_noise(self, img, mean=0, std=0.1): noise=torch.normal(mean,std,img.shape) img+=noise img=img.clamp(min=0, max=1.0) return img def add_salt_noise(self, img): noise=torch.rand(img.shape) alpha=np.random.random() img[noise[:,:,:]>alpha]=1.0 return img def add_pepper_noise(self, img): noise=torch.rand(img.shape) alpha=np.random.random() img[noise[:,:,:]>alpha]=0 return img def mixup(self, img1, img2, box1, box2,alpha=2): weight=np.random.beta(alpha, alpha) print(weight) miximg=weight*img1+(1-weight)*img2 if weight > 0.5: return miximg, box1 return miximg, box2 def draw_img(self, img, boxes): draw=ImageDraw.Draw(img) for box in boxes: draw.rectangle(list(box), outline='yellow', width=2) img.show() D=DataAugmentation() img1=Image.open('C:/AllProgram/testimage/superResolution/bird_GT.bmp') img2=Image.open('C:/AllProgram/testimage/superResolution/lenna.bmp') box1=torch.Tensor([[92,20,261,190]]) box2=torch.Tensor([[23,18,97,123], [45,65,194,225]]) to_tensor=transforms.ToTensor() to_image=transforms.ToPILImage() img1,box1=D.resize(img1,box1,256) img2,box2=D.resize(img2,box2,256) img1=to_tensor(img1) img2=to_tensor(img2) miximg,box=D.mixup(img1,img2,box1,box2) miximg=to_image(miximg) D.draw_img(miximg,box)