编程语言
首页 > 编程语言> > pytorch源码解析系列-yolov4最核心技巧代码详解(3)- 数据处理以及图像增强

pytorch源码解析系列-yolov4最核心技巧代码详解(3)- 数据处理以及图像增强

作者:互联网

输入是啥?

现在 我们知道了yolo的模型,知道了模型获取的推理输出,在进入训练代码之前,我们需要了解到yolo数据的处理方式,以及用了何种图像增强方法。(不然连输入是啥都不知道,你怎么看的懂训练过程呢)

数据载入

其实数据载入这块没有必要单独拿出来说的,但是学会如何写collate function是很重要的
我们先看下数据源

## 如果你多进程报错 num_workers要改成1 大部分CPU训练机器都会出这个问题
## batchsize = batch // subdivisions 所以你就知道超参数这块实际上最终的batchsize是由这两个数据决定的
train_loader = DataLoader(train_dataset, batch_size=config.batch // config.subdivisions, shuffle=True,
                              num_workers=8, pin_memory=True, drop_last=True, collate_fn=collate)

只需要在collate_fn指定自己写的方法就行了,自己写可以把很多数据预处理都搞定在这个方法里

def collate(batch):
    images = []
    bboxes = []
    for img, box in batch:
        images.append([img])
        bboxes.append([box])
    #拿torch和numpy写都行,反正都要最后转成torch的
    #这边就是数据的标准预处理了,把格式处理成[B,C,H,W] 再除以255
    images = np.concatenate(images, axis=0)
    images = images.transpose(0, 3, 1, 2)
    images = torch.from_numpy(images).div(255.0)
    bboxes = np.concatenate(bboxes, axis=0)
    bboxes = torch.from_numpy(bboxes)
    return images, bboxes

Mosaic图像增强

小目标的AP一般比中目标和大目标低很多,而且小目标的分布并不均匀,所以作者用Mosaic处理这种问题
CutMix数据增强:两张图片进行拼接
然后给他变形一下↓
Mosaic数据增强:4张图片,随机缩放、随机裁剪、随机排布的方式进行拼接。
好处就是,一次把4个图片拧一块,batch直接4倍了,岂不美哉?
看看代码如何实现的:(注:代码实现较老,没有用torchvision)

#yolo的Dataset写法不会的建议先学下基础 这里的代码实际上没必要,
#dataset的label文件是指定一个txt的 你如果跑过yolo4的训练就知道要自己制作train.txt
class Yolo_dataset(Dataset):
    def __init__(self, lable_path, cfg, train=True):
        super(Yolo_dataset, self).__init__()
    #这边参数的意思是 mixup=4是用 mosaic+cutmix, 2是只用mixup,3是mosaic 。
        if cfg.mixup == 2:
            print("cutmix=1 - isn't supported for Detector")
            raise
        elif cfg.mixup == 2 and cfg.letter_box:
            print("Combination: letter_box=1 & mosaic=1 - isn't supported, use only 1 of these parameters")
            raise

        self.cfg = cfg
        self.train = train
		# label载入 每张图片的里面分类都载进字典
        truth = {}
        f = open(lable_path, 'r', encoding='utf-8')
        for line in f.readlines():
            data = line.split(" ")
            truth[data[0]] = []
            for i in data[1:]:
                truth[data[0]].append([int(float(j)) for j in i.split(',')])

        self.truth = truth
        self.imgs = list(self.truth.keys())

重点在下面mosiac的实现:

	# 每次从dataset拿数据的时候,都会随机抽4张图片做mosiac
    def __getitem__(self, index):
        if not self.train:
            return self._get_val_item(index)
        img_path = self.imgs[index]
        #从label里拿到类别bboxes
        bboxes = np.array(self.truth.get(img_path), dtype=np.float)
        img_path = os.path.join(self.cfg.dataset_dir, img_path)
        use_mixup = self.cfg.mixup
        # 一半概率不用增强
        if random.randint(0, 1):
            use_mixup = 0
		# mixup==3是使用mosaic 
        if use_mixup == 3:
            min_offset = 0.2 #指定剪切率之后 随机从宽高剪切一个区域
            cut_x = random.randint(int(self.cfg.w * min_offset), int(self.cfg.w * (1 - min_offset)))
            cut_y = random.randint(int(self.cfg.h * min_offset), int(self.cfg.h * (1 - min_offset)))

        r1, r2, r3, r4, r_scale = 0, 0, 0, 0, 0
        dhue, dsat, dexp, flip, blur = 0, 0, 0, 0, 0
        gaussian_noise = 0

        out_img = np.zeros([self.cfg.h, self.cfg.w, 3])
        out_bboxes = []
		#这边就知道 use_mixup里数字的意思了,0就是不用 只循环1次,3是mosiac,循环4次
        for i in range(use_mixup + 1):
            if i != 0:
            	#不是第一张图片 就从图片库里随便抽一张做拼接
                img_path = random.choice(list(self.truth.keys()))
                bboxes = np.array(self.truth.get(img_path), dtype=np.float)
                img_path = os.path.join(self.cfg.dataset_dir, img_path)
            img = cv2.imread(img_path) 
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) #bgr2rgb处理 你不会这个都不知道吧?
            if img is None:
                continue
            oh, ow, oc = img.shape
            # 默认jitter=0.2 hue = 0.1 sat = 1.5 exp = 1.5
            # 这边三个增强方法写在最后,有兴趣看看,主要是随机获取HSV中三个随机增强参数
            dh, dw, dc = np.array(np.array([oh, ow, oc]) * self.cfg.jitter, dtype=np.int)
            dhue = rand_uniform_strong(-self.cfg.hue, self.cfg.hue)
            dsat = rand_scale(self.cfg.saturation)
            dexp = rand_scale(self.cfg.exposure)
			# 下面很长 有兴趣的可以看看 主要是图像增强的其他一系列旋转、滤波、拼接操作
            pleft = random.randint(-dw, dw)
            pright = random.randint(-dw, dw)
            ptop = random.randint(-dh, dh)
            pbot = random.randint(-dh, dh)
			#旋转
            flip = random.randint(0, 1) if self.cfg.flip else 0
			#高斯滤波
            if (self.cfg.blur):
                tmp_blur = random.randint(0, 2)  # 0 - disable, 1 - blur background, 2 - blur the whole image
                if tmp_blur == 0:
                    blur = 0
                elif tmp_blur == 1:
                    blur = 1
                else:
                    blur = self.cfg.blur
			#高斯
            if self.cfg.gaussian and random.randint(0, 1):
                gaussian_noise = self.cfg.gaussian
            else:
                gaussian_noise = 0
			#补灰(类似pad)
            if self.cfg.letter_box:
                img_ar = ow / oh
                net_ar = self.cfg.w / self.cfg.h
                result_ar = img_ar / net_ar
                if result_ar > 1:  # 补sheight
                    oh_tmp = ow / net_ar
                    delta_h = (oh_tmp - oh) / 2
                    ptop = ptop - delta_h
                    pbot = pbot - delta_h
                else:  # 补swidth
                    ow_tmp = oh * net_ar
                    delta_w = (ow_tmp - ow) / 2
                    pleft = pleft - delta_w
                    pright = pright - delta_w

            swidth = ow - pleft - pright
            sheight = oh - ptop - pbot
			#这个代码贴在最后 有兴趣看一下 主要是拿剪切后的label
            truth, min_w_h = fill_truth_detection(bboxes, self.cfg.boxes, self.cfg.classes, flip, pleft, ptop, swidth,
                                                  sheight, self.cfg.w, self.cfg.h)
            if (min_w_h / 8) < blur and blur > 1:  # 如果太小就不用滤波了
                blur = min_w_h / 8
			# 主要图片增强处理 核心内容是图片转成HSV格式,用参数对其增强,有兴趣看最后的代码
            ai = image_data_augmentation(img, self.cfg.w, self.cfg.h, pleft, ptop, swidth, sheight, flip,
                                         dhue, dsat, dexp, gaussian_noise, blur, truth)

            if use_mixup == 0:
                out_img = ai
                out_bboxes = truth
            ## 1是cutmix 拼接两个图像
            if use_mixup == 1:
                if i == 0:
                    old_img = ai.copy()
                    old_truth = truth.copy()
                elif i == 1:
                    out_img = cv2.addWeighted(ai, 0.5, old_img, 0.5)
                    out_bboxes = np.concatenate([old_truth, truth], axis=0)
             ## 3 是mosaic:旋转、滤波
            elif use_mixup == 3:
                if flip:
                    tmp = pleft
                    pleft = pright
                    pright = tmp

                left_shift = int(min(cut_x, max(0, (-int(pleft) * self.cfg.w / swidth))))
                top_shift = int(min(cut_y, max(0, (-int(ptop) * self.cfg.h / sheight))))

                right_shift = int(min((self.cfg.w - cut_x), max(0, (-int(pright) * self.cfg.w / swidth))))
                bot_shift = int(min(self.cfg.h - cut_y, max(0, (-int(pbot) * self.cfg.h / sheight))))
				#blend_truth_mosaic这个代码也贴在最后 有兴趣看看 主要是拿剪切后的图片和label
                out_img, out_bbox = blend_truth_mosaic(out_img, ai, truth.copy(), self.cfg.w, self.cfg.h, cut_x,
                                                       cut_y, i, left_shift, right_shift, top_shift, bot_shift)
                out_bboxes.append(out_bbox)
        #后面就是拼接图片和label了
        if use_mixup == 3:
            out_bboxes = np.concatenate(out_bboxes, axis=0)
        out_bboxes1 = np.zeros([self.cfg.boxes, 5])
        out_bboxes1[:min(out_bboxes.shape[0], self.cfg.boxes)] = out_bboxes[:min(out_bboxes.shape[0], self.cfg.boxes)]
        return out_img, out_bboxes1

上面的内容 用一句话概括:
50%的概率使用图像增强,增强方式是对4张图片采用随机缩放,裁剪,旋转,滤波,最后拼在一起的方式搞成一张图片
在这里插入图片描述

选看

下面的代码可以看,可以不看,主要是之前代码里面图像增强的具体实现方式,如果你真的有手撕代码的需要,可以用来参考

从简单的来看,HSE三个随机参数:
没啥东西 就是随机增强

def rand_uniform_strong(min, max):
    if min > max:
        swap = min
        min = max
        max = swap
    return random.random() * (max - min) + min
def rand_scale(s):
    scale = rand_uniform_strong(1, s)
    if random.randint(0, 1) % 2:
        return scale
    return 1. / scale
def rand_precalc_random(min, max, random_part):
    if max < min:
        swap = min
        min = max
        max = swap
    return (random_part * (max - min)) + min

然后是处理图片裁剪和缩放后的label:

#truth, min_w_h = fill_truth_detection(bboxes, self.cfg.boxes, self.cfg.classes, flip, pleft, ptop, swidth,sheight, self.cfg.w, self.cfg.h)
#这边的dx,dy,sx,sy对应之前的pleft,ptop,swidth,sheight 了解了这个再往后看
def fill_truth_detection(bboxes, num_boxes, classes, flip, dx, dy, sx, sy, net_w, net_h):
    if bboxes.shape[0] == 0:
        return bboxes, 10000
    np.random.shuffle(bboxes)
    bboxes[:, 0] -= dx
    bboxes[:, 2] -= dx
    bboxes[:, 1] -= dy
    bboxes[:, 3] -= dy
	#上面减去随机位置后,用下面的参数保证长宽落在[0,sx]中间
    bboxes[:, 0] = np.clip(bboxes[:, 0], 0, sx)
    bboxes[:, 2] = np.clip(bboxes[:, 2], 0, sx)
	#同理
    bboxes[:, 1] = np.clip(bboxes[:, 1], 0, sy)
    bboxes[:, 3] = np.clip(bboxes[:, 3], 0, sy)
	#过滤掉out_box,也就是说过滤掉长宽都同时落在边缘(如果图片裁剪区域里没有目标)
    out_box = list(np.where(((bboxes[:, 1] == sy) & (bboxes[:, 3] == sy)) |
                            ((bboxes[:, 0] == sx) & (bboxes[:, 2] == sx)) |
                            ((bboxes[:, 1] == 0) & (bboxes[:, 3] == 0)) |
                            ((bboxes[:, 0] == 0) & (bboxes[:, 2] == 0)))[0])
    list_box = list(range(bboxes.shape[0]))
    for i in out_box:
        list_box.remove(i)
    bboxes = bboxes[list_box]

    if bboxes.shape[0] == 0:
        return bboxes, 10000
	#我感觉这一步多余的 类别判断为啥在这里
    bboxes = bboxes[np.where((bboxes[:, 4] < classes) & (bboxes[:, 4] >= 0))[0]]

    if bboxes.shape[0] > num_boxes:
        bboxes = bboxes[:num_boxes]
	#最短宽高
    min_w_h = np.array([bboxes[:, 2] - bboxes[:, 0], bboxes[:, 3] - bboxes[:, 1]]).min()
	#这个net参数是自己cfg设的,控制缩放比率
    bboxes[:, 0] *= (net_w / sx)
    bboxes[:, 2] *= (net_w / sx)
    bboxes[:, 1] *= (net_h / sy)
    bboxes[:, 3] *= (net_h / sy)
	#是否翻转
    if flip:
        temp = net_w - bboxes[:, 0]
        bboxes[:, 0] = net_w - bboxes[:, 2]
        bboxes[:, 2] = temp
	#最后是返回处理后的图片label和最短宽高
    return bboxes, min_w_h

这类代码写起来比看要容易,因为当你有明确的目标后,你再去写就有思路,光看是很难看出思路的

上面获取偏转缩放后的label,经过mosaic还要处理一次,核心思想就是4张图片分别用何种处理获得mosiac后的label区间

def blend_truth_mosaic(out_img, img, bboxes, w, h, cut_x, cut_y, i_mixup,
                       left_shift, right_shift, top_shift, bot_shift):
    left_shift = min(left_shift, w - cut_x)
    top_shift = min(top_shift, h - cut_y)
    right_shift = min(right_shift, cut_x)
    bot_shift = min(bot_shift, cut_y)
	#i_mixup 对应外面4张图片的4个循环
    if i_mixup == 0:
    #filter_truth这个代码不写了,主要就是fill_truth_detection里OUTBOX做的事情,过滤掉处理后没有目标的图片
        bboxes = filter_truth(bboxes, left_shift, top_shift, cut_x, cut_y, 0, 0)
        out_img[:cut_y, :cut_x] = img[top_shift:top_shift + cut_y, left_shift:left_shift + cut_x]
    if i_mixup == 1:
        bboxes = filter_truth(bboxes, cut_x - right_shift, top_shift, w - cut_x, cut_y, cut_x, 0)
        out_img[:cut_y, cut_x:] = img[top_shift:top_shift + cut_y, cut_x - right_shift:w - right_shift]
    if i_mixup == 2:
        bboxes = filter_truth(bboxes, left_shift, cut_y - bot_shift, cut_x, h - cut_y, 0, cut_y)
        out_img[cut_y:, :cut_x] = img[cut_y - bot_shift:h - bot_shift, left_shift:left_shift + cut_x]
    if i_mixup == 3:
        bboxes = filter_truth(bboxes, cut_x - right_shift, cut_y - bot_shift, w - cut_x, h - cut_y, cut_x, cut_y)
        out_img[cut_y:, cut_x:] = img[cut_y - bot_shift:h - bot_shift, cut_x - right_shift:w - right_shift]

    return out_img, bboxes

最后的代码是图片的处理,主要是转成HSV处理增强 再转回来


def image_data_augmentation(mat, w, h, pleft, ptop, swidth, sheight, flip, dhue, dsat, dexp, gaussian_noise, blur,
                            truth):
    try:
        img = mat
        oh, ow, _ = img.shape
        #剪切缩放
        pleft, ptop, swidth, sheight = int(pleft), int(ptop), int(swidth), int(sheight)
        # 剪切目标位置
        src_rect = [pleft, ptop, swidth + pleft, sheight + ptop]  # x1,y1,x2,y2
        img_rect = [0, 0, ow, oh]
        new_src_rect = rect_intersection(src_rect, img_rect)  # 交集 就是获取[minx, miny, maxx, maxy]

        dst_rect = [max(0, -pleft), max(0, -ptop), max(0, -pleft) + new_src_rect[2] - new_src_rect[0],
                    max(0, -ptop) + new_src_rect[3] - new_src_rect[1]]
        # cv2.Mat sized

        if (src_rect[0] == 0 and src_rect[1] == 0 and src_rect[2] == img.shape[0] and src_rect[3] == img.shape[1]):
            sized = cv2.resize(img, (w, h), cv2.INTER_LINEAR)
        else:
            cropped = np.zeros([sheight, swidth, 3])
            cropped[:, :, ] = np.mean(img, axis=(0, 1))

            cropped[dst_rect[1]:dst_rect[3], dst_rect[0]:dst_rect[2]] = \
                img[new_src_rect[1]:new_src_rect[3], new_src_rect[0]:new_src_rect[2]]

            # 线性差值resize
            sized = cv2.resize(cropped, (w, h), cv2.INTER_LINEAR)

        # 翻转
        if flip:
            sized = cv2.flip(sized, 1)  # 0 - x-axis, 1 - y-axis, -1 - both axes (x & y)

        # HSV增强,主要用下面机种方式把图转成HSV
        # cv2.COLOR_BGR2HSV, cv2.COLOR_RGB2HSV, cv2.COLOR_HSV2BGR, cv2.COLOR_HSV2RGB
        if dsat != 1 or dexp != 1 or dhue != 0:
            if img.shape[2] >= 3:
                hsv_src = cv2.cvtColor(sized.astype(np.float32), cv2.COLOR_RGB2HSV) 
                hsv = cv2.split(hsv_src)
                #之前的三个参数增强对hsv增强
                hsv[1] *= dsat
                hsv[2] *= dexp
                hsv[0] += 179 * dhue
                hsv_src = cv2.merge(hsv)
                sized = np.clip(cv2.cvtColor(hsv_src, cv2.COLOR_HSV2RGB), 0, 255)  # 转回RGB,记得合法范围[0,255]
            else:
                sized *= dexp
		#高斯滤波 平滑
        if blur:
            if blur == 1:
                dst = cv2.GaussianBlur(sized, (17, 17), 0)
                # cv2.bilateralFilter(sized, dst, 17, 75, 75)
            else:
                ksize = (blur / 2) * 2 + 1
                dst = cv2.GaussianBlur(sized, (ksize, ksize), 0)

            if blur == 1:
                img_rect = [0, 0, sized.cols, sized.rows]
                for b in truth:
                    left = (b.x - b.w / 2.) * sized.shape[1]
                    width = b.w * sized.shape[1]
                    top = (b.y - b.h / 2.) * sized.shape[0]
                    height = b.h * sized.shape[0]
                    roi(left, top, width, height)
                    roi = roi & img_rect
                    dst[roi[0]:roi[0] + roi[2], roi[1]:roi[1] + roi[3]] = sized[roi[0]:roi[0] + roi[2],
                                                                          roi[1]:roi[1] + roi[3]]

            sized = dst

        if gaussian_noise:
            noise = np.array(sized.shape)
            gaussian_noise = min(gaussian_noise, 127)
            gaussian_noise = max(gaussian_noise, 0)
            cv2.randn(noise, 0, gaussian_noise)  # mean and variance
            sized = sized + noise
    except:
        print("OpenCV can't augment image: " + str(w) + " x " + str(h))
        sized = mat

    return sized

标签:yolov4,img,cfg,self,cut,源码,bboxes,shift,图像增强
来源: https://blog.csdn.net/weixin_48174100/article/details/120482685