SmallObjectAugmentation(小目标数据增强)论文要点及代码实现

效果图

COCO关于目标大小的定义

小目标的两个特点

1.数据集中包含小物体的图片相对较少，这可能会使任何检测模型更倾向于关注中等和大型物体。
2.小物体所占面积较小，意味着小物体位置的缺乏多样性。

解决方案

我们通过对包含小物体的图片进行过采样来解决第一个问题。我们通过在每个包含小物体的图像中多次复制粘贴小物体来解决第二个问题。在复制粘贴每个物体时，我们确保粘贴的物体不与任何现有物体重叠。这增加了小物体的位置多样性，同时确保这些物体出现在正确的上下文中。
copy-paste增加锚框数量。

操作细节

在将对象粘贴到新位置之前，我们对其进行随机变换。我们通过更改物体大小±20%来缩放物体，并将其旋转±15◦。我们只考虑未遮挡的对象，因为将分离的分割掩码分别插入中间未曾见过的部分通常会导致不太逼真的图像。我们确保新粘贴的对象不与任何现有对象重叠，并且距离图像边界至少有5个像素。

我们测试了三种设置。在第一种设置中，我们用复制和粘贴小物体后的图像替换每个带有小物体的原始图像。在第二种设置中，我们复制这些增强图像以模拟过采样。在最后一个设置中，我们保留原始图像和增强后的图像，这相当于将带有小物体的图像过采样两倍，并将复制的副本增加更多的小物体进行增强。

复制和粘贴小物体有不同的方法。我们考虑了三种单独的策略。首先，我们选择一张图像中的一个小物体，并在随机位置复制和粘贴多次。第二，我们选择多个小物体，并将每个小物体粘贴到任意位置。最后，我们将每个小物体在图像中复制和粘贴到多个随机位置。在所有情况下，我们使用上面的第三种增强设置；也就是说，我们保留原始图像和其增强副本。

在粘贴小物体的副本时，有两件事需要考虑。首先，我们必须决定粘贴的对象是否会与任何其他对象重叠。虽然我们选择不引入任何重叠，但我们通过实验证明了它是否是一个好策略。其次，进行平滑处理以使得粘贴对象的边缘更加平滑是一个设计选择。我们尝试使用不同的高斯滤波器大小进行平滑处理，以帮助提高性能。

结论

训练模型对“粘贴”小物体进行了有效的过拟合，但不一定对原始小物体进行了过拟合。我们认为这是由于粘贴造成的伪影，例如不完美的对象掩码和与背景的亮度差异，相对容易被神经网络检测到。最好的结果是通过将过采样和概率为p = 0.5的增强（原始+增强）与原始和增强小物体的比率为2：1相结合来达到的。这种设置产生了比单独进行过采样更好的结果，确认了粘贴小物体所提出的策略的有效性。

代码实现

import aug as am
import Helpers as hp
from util import *
import os
from os.path import join
from tqdm import tqdm
import random

base_dir = os.getcwd()

save_base_dir = join(base_dir, 'save')

check_dir(save_base_dir)
#train.txt存放需要粘贴的图像的位置
imgs_dir = [f.strip() for f in open(join(base_dir, 'train.txt')).readlines()]
labels_dir = hp.replace_labels(imgs_dir)
#small.txt存放粘贴图像的位置
small_imgs_dir = [f.strip() for f in open(join(base_dir, 'small.txt')).readlines()]
#打乱存放被粘贴的图像的顺序
random.shuffle(small_imgs_dir)

for image_dir, label_dir in tqdm(zip(imgs_dir, labels_dir)):
    small_img = []
    #搞八个小目标图像出来
    for x in range(8):
        if small_imgs_dir == []:
            #exit()
            small_imgs_dir = [f.strip() for f in open(join(base_dir,'small.txt')).readlines()]
            random.shuffle(small_imgs_dir)
        small_img.append(small_imgs_dir.pop())
    am.copysmallobjects2(image_dir, label_dir, save_base_dir,small_img)

接下来主要是copysmallobjects2这个函数的调用

def copysmallobjects2(image_dir, label_dir, save_base_dir, small_img_dir):
    image = cv2.imread(image_dir)
    labels = read_label_txt(label_dir)
    if len(labels) == 0:
        return
    rescale_labels = rescale_yolo_labels(labels, image.shape)  # 转换坐标表示
    all_boxes = []
    for _, rescale_label in enumerate(rescale_labels):
        all_boxes.append(rescale_label)

    for small_img_dirs in small_img_dir:
        image_bbox = cv2.imread(small_img_dirs)
        #roi = image_bbox
        roi = suo_fang(image_bbox,area_max=3000,area_min=1500)

        new_bboxes = random_add_patches2(roi.shape, rescale_labels, image.shape, paste_number=1, iou_thresh=0)
        count = 0
        for new_bbox in new_bboxes:
            count += 1

            cl, bbox_left, bbox_top, bbox_right, bbox_bottom = new_bbox[0], new_bbox[1], new_bbox[2], new_bbox[3], \
                                                               new_bbox[4]
            #roi = GaussianBlurImg(roi)  # 高斯模糊
            height, width, channels = roi.shape
            center = (int(width / 2),int(height / 2))
            #ran_point = (int((bbox_top+bbox_bottom)/2),int((bbox_left+bbox_right)/2))
            mask = 255 * np.ones(roi.shape, roi.dtype)

            try:
                if count > 1:
                    roi = flip_bbox(roi)
                #image[bbox_top:bbox_bottom, bbox_left:bbox_right] = roi
                #image[bbox_top:bbox_bottom, bbox_left:bbox_right] = cv2.addWeighted(image[bbox_top:bbox_bottom, bbox_left:bbox_right],
                #                                                                    0.5,roi,0.5,0) #图片融合

                # 泊松融合
                #image = cv2.seamlessClone(roi, image, mask, ran_point, cv2.NORMAL_CLONE)
                #print(str(bbox_bottom-bbox_top) + "|" + str(bbox_right-bbox_left))
                #print(roi.shape)
                #print(mask.shape)
                image[bbox_top:bbox_bottom, bbox_left:bbox_right] = cv2.seamlessClone(roi, image[bbox_top:bbox_bottom, bbox_left:bbox_right],
                                                                                      mask, center, cv2.NORMAL_CLONE)
                all_boxes.append(new_bbox)
                rescale_labels.append(new_bbox)
            except ValueError:
                print("---")
                continue
    dir_name = find_str(image_dir)
    save_dir = join(save_base_dir, dir_name)
    check_dir(save_dir)
    yolo_txt_dir = join(save_dir, basename(image_dir.replace('.jpg', '_augment.txt')))
    cv2.imwrite(join(save_dir, basename(image_dir).replace('.jpg', '_augment.jpg')), image)
    convert_all_boxes(image.shape, all_boxes, yolo_txt_dir)

具体解析如下：

def copysmallobjects2(image_dir, label_dir, save_base_dir, small_img_dir):
# 读取原始图像和标签文件
image = cv2.imread(image_dir)
labels = read_label_txt(label_dir)
# 如果标签为空，直接返回
if len(labels) == 0:
return
# 调用rescale_yolo_labels函数，将坐标值从Yolo格式转换为图像尺寸格式
rescale_labels = rescale_yolo_labels(labels, image.shape)
# 定义一个空列表
all_boxes = []
# 遍历每个物体框及其对应的标签
for _, rescale_label in enumerate(rescale_labels):
# 将物体框信息添加到all_boxes列表中
all_boxes.append(rescale_label)

# 遍历小物体路径列表
for small_img_dirs in small_img_dir:
    # 读取小物体图像
    image_bbox = cv2.imread(small_img_dirs)
    # 使用suo_fang函数进行图像缩放
    roi = suo_fang(image_bbox, area_max=3000, area_min=1500) # 将小物体图像缩放至合适的大小

    # 使用random_add_patches2函数，在原始图像上随机添加小物体
    new_bboxes = random_add_patches2(roi.shape, rescale_labels, image.shape, paste_number=1, iou_thresh=0)
    count = 0
    # 对每个新添加的小物体进行处理
    for new_bbox in new_bboxes:
        count += 1
        cl, bbox_left, bbox_top, bbox_right, bbox_bottom = new_bbox[0], new_bbox[1], new_bbox[2], new_bbox[3], \
                                                           new_bbox[4]
        #roi = GaussianBlurImg(roi)  # 高斯模糊
        height, width, channels = roi.shape
        center = (int(width / 2),int(height / 2))
        #ran_point = (int((bbox_top+bbox_bottom)/2),int((bbox_left+bbox_right)/2))
        mask = 255 * np.ones(roi.shape, roi.dtype)

        try:
            # 如果是第二次插入图片，则翻转图片
            if count > 1:
                roi = flip_bbox(roi)
            # 将小物体复制到原始图像上
            #image[bbox_top:bbox_bottom, bbox_left:bbox_right] = roi
            #image[bbox_top:bbox_bottom, bbox_left:bbox_right] = cv2.addWeighted(image[bbox_top:bbox_bottom, bbox_left:bbox_right],
            #                                                                    0.5,roi,0.5,0) # 图片融合

            # 使用泊松融合方法，将小物体复制到原始图像上
            #image = cv2.seamlessClone(roi, image, mask, ran_point, cv2.NORMAL_CLONE)

            # 泊松融合方法可能抛出异常，需要进行捕获处理
            image[bbox_top:bbox_bottom, bbox_left:bbox_right] = cv2.seamlessClone(roi, image[bbox_top:bbox_bottom, bbox_left:bbox_right],
                                                                                  mask, center, cv2.NORMAL_CLONE)
            # 将新的物体框信息添加到all_boxes列表和rescale_labels列表中
            all_boxes.append(new_bbox)
            rescale_labels.append(new_bbox)
        except ValueError:
            print("---")
            continue
# 根据原始图像的名称，在保存路径下创建相应目录
dir_name = find_str(image_dir)
save_dir = join(save_base_dir, dir_name)
check_dir(save_dir)
# 生成新的标签文件，并将其保存到save_dir下
yolo_txt_dir = join(save_dir, basename(image_dir.replace('.jpg', '_augment.txt')))
convert_all_boxes(image.shape, all_boxes, yolo_txt_dir)
# 将新的图像保存到save_dir下
cv2.imwrite(join(save_dir, basename(image_dir).replace('.jpg', '_augment.jpg')), image)

rescale_abels

其中suo_fang函数

def suo_fang(image, area_max=2000, area_min=1000):
    # 改变图片大小
    height, width, channels = image.shape

    while (height*width) > area_max:
        image = cv2.resize(image, (int(width * 0.9),int(height * 0.9)))
        height, width, channels = image.shape
        height,width = int(height*0.9),int(width*0.9)

    while (height*width) < area_min:
        image = cv2.resize(image, (int(width * 1.1),int(height * 1.1)))
        height, width, channels = image.shape
        height,width = int(height*1.1),int(width*1.1)

    return image

def suo_fang(image, area_max=2000, area_min=1000):
# 获取原始图像的高度、宽度、通道数信息
height, width, channels = image.shape
# 当图像的面积大于area_max时，将图像缩小至面积的0.9倍
while (heightwidth) > area_max:
image = cv2.resize(image, (int(width * 0.9),int(height * 0.9)))
height, width, channels = image.shape
height,width = int(height0.9),int(width0.9)
# 当图像的面积小于area_min时，将图像放大至面积的1.1倍
while (heightwidth) < area_min:
image = cv2.resize(image, (int(width * 1.1),int(height * 1.1)))
height, width, channels = image.shape
height,width = int(height1.1),int(width1.1)
# 返回缩放后的图像
return image

def random_add_patches2(bbox_img, rescale_boxes, shape, paste_number, iou_thresh):
# 复制一份rescale_boxes，命名为temp
temp = []
for rescale_bbox in rescale_boxes:
temp.append(rescale_bbox)
# 获取目标图像的高度、宽度、通道数信息
bbox_h, bbox_w, bbox_c = bbox_img
img_h,img_w,img_c = shape
# 在图像内部选取一定数量的随机点
center_search_space = sampling_new_bbox_center_point2(shape, bbox_img)
success_num = 0
new_bboxes = []
cl = 1
# 随机生成指定数量的小物体，并将其加入到new_bboxes列表中
while success_num < paste_number:
# 随机生成点坐标
new_bbox_x_center, new_bbox_y_center = norm_sampling(center_search_space)
# 如果生成的小物体超出了目标图像的范围，则直接跳过
if new_bbox_x_center-0.5bbox_w < 0 or new_bbox_x_center+0.5bbox_w > img_w:
continue
if new_bbox_y_center-0.5bbox_h < 0 or new_bbox_y_center+0.5bbox_h > img_h:
continue
# 根据生成的随机点，计算出新的物体框坐标信息
new_bbox_x_left, new_bbox_y_left, new_bbox_x_right, new_bbox_y_right = new_bbox_x_center - 0.5 * bbox_w, 

new_bbox_y_center - 0.5 * bbox_h, 

new_bbox_x_center + 0.5 * bbox_w, 

new_bbox_y_center + 0.5 * bbox_h
new_bbox = [cl, int(new_bbox_x_left), int(new_bbox_y_left), int(new_bbox_x_right), int(new_bbox_y_right)]

    # 计算新物体框与已有物体框的交并比，如果IOU值小于阈值，则将新物体框加入new_bboxes列表中
    ious = [bbox_iou(new_bbox, bbox_t) for bbox_t in rescale_boxes]
    ious2 = [bbox_iou(new_bbox,bbox_t1) for bbox_t1 in new_bboxes]
    if ious2 == []:
        ious2.append(0)
    if max(ious) <= iou_thresh and max(ious2) <= iou_thresh:
        success_num += 1
        temp.append(new_bbox)
        new_bboxes.append(new_bbox)
    else:
        continue

return new_bboxes