如何从对象检测数据加载器中的马赛克增强中获取类标签?

Posted

技术标签:

【中文标题】如何从对象检测数据加载器中的马赛克增强中获取类标签?【英文标题】:How can I get Class label from Mosaic augmentation in Object Detection Dataloader? 【发布时间】:2021-01-27 19:44:16 【问题描述】:

我正在尝试针对多类问题训练对象检测模型。在我的训练中,我使用Mosaic augmentation、Paper 来完成这项任务。

在我的训练机制中,我有点坚持正确检索每个类别的类标签,因为增强机制会随机选择样本的子部分。然而,下面是我们迄今为止使用相关边界框实现的马赛克增强的结果。

数据集

我创建了一个虚拟数据集。 df.head()

共有4类df.object.value_counts():

human    23
car      13
cat       5
dog       3

数据加载器和马赛克增强

数据加载器定义如下。但是,马赛克增强应该在内部定义,但现在,我将创建一个单独的代码 sn-p 以便更好地演示:


IMG_SIZE = 2000

class DatasetRetriever(Dataset):

    def __init__(self, main_df, image_ids, transforms=None, test=False):
        super().__init__()

        self.image_ids = image_ids
        self.main_df = main_df
        self.transforms = transforms
        self.size_limit = 1
        self.test = test

    def __getitem__(self, index: int):
        image_id = self.image_ids[index] 
        image, boxes, labels = self.load_mosaic_image_and_boxes(index)
        
        # labels = torch.tensor(labels, dtype=torch.int64) # for multi-class 
        labels = torch.ones((boxes.shape[0],), dtype=torch.int64) # for single-class 
         
        target = 
        target['boxes'] = boxes
        target['cls'] = labels
        target['image_id'] = torch.tensor([index])

        if self.transforms:
            for i in range(10):
                sample = self.transforms(**
                    'image' : image,
                    'bboxes': target['boxes'],
                    'labels': target['cls'] 
                )
                
                assert len(sample['bboxes']) == target['cls'].shape[0], 'not equal!'
                if len(sample['bboxes']) > 0:
                    # image
                    image = sample['image']
                    
                    # box
                    target['boxes'] = torch.tensor(sample['bboxes'])
                    target['boxes'][:,[0,1,2,3]] = target['boxes'][:,[1,0,3,2]]
                    
                    # label
                    target['cls'] = torch.stack(sample['labels'])
                    break
                    
        return image, target

    def __len__(self) -> int:
        return self.image_ids.shape[0]

基本变换

def get_transforms():
    return A.Compose(
        [
            A.Resize(height=IMG_SIZE, width=IMG_SIZE, p=1.0),
            ToTensorV2(p=1.0),
        ], 
        p=1.0, 
        bbox_params=A.BboxParams(
            format='pascal_voc',
            min_area=0, 
            min_visibility=0,
            label_fields=['labels']
        )
    )

马赛克增强

注意,它应该在数据加载器内部定义。主要问题是,在这个增强中,在迭代所有 4 个样本以创建这种增强时,imagebounding_box 被重新缩放如下:

mosaic_image[y1a:y2a, x1a:x2a] = image[y1b:y2b, x1b:x2b]

offset_x = x1a - x1b
offset_y = y1a - y1b
boxes[:, 0] += offset_x
boxes[:, 1] += offset_y
boxes[:, 2] += offset_x
boxes[:, 3] += offset_y

这样,我如何为那些被选中的bounding_box选择相关的类标签?请看下面的完整代码:

def load_mosaic_image_and_boxes(self, index, s=3000, 
                                    minfrac=0.25, maxfrac=0.75):
        self.mosaic_size = s
        xc, yc = np.random.randint(s * minfrac, s * maxfrac, (2,))

        # random other 3 sample 
        indices = [index] + random.sample(range(len(self.image_ids)), 3) 

        mosaic_image = np.zeros((s, s, 3), dtype=np.float32)
        final_boxes  = [] # box for the sub-region
        final_labels = [] # relevant class labels
        
        for i, index in enumerate(indices):
            image, boxes, labels = self.load_image_and_boxes(index)

            if i == 0:    # top left
                x1a, y1a, x2a, y2a =  0,  0, xc, yc
                x1b, y1b, x2b, y2b = s - xc, s - yc, s, s # from bottom right
            elif i == 1:  # top right
                x1a, y1a, x2a, y2a = xc, 0, s , yc
                x1b, y1b, x2b, y2b = 0, s - yc, s - xc, s # from bottom left
            elif i == 2:  # bottom left
                x1a, y1a, x2a, y2a = 0, yc, xc, s
                x1b, y1b, x2b, y2b = s - xc, 0, s, s-yc   # from top right
            elif i == 3:  # bottom right
                x1a, y1a, x2a, y2a = xc, yc,  s, s
                x1b, y1b, x2b, y2b = 0, 0, s-xc, s-yc    # from top left

            # calculate and apply box offsets due to replacement            
            offset_x = x1a - x1b
            offset_y = y1a - y1b
            boxes[:, 0] += offset_x
            boxes[:, 1] += offset_y
            boxes[:, 2] += offset_x
            boxes[:, 3] += offset_y
            
            # cut image, save boxes
            mosaic_image[y1a:y2a, x1a:x2a] = image[y1b:y2b, x1b:x2b]
            final_boxes.append(boxes)

            '''
            ATTENTION: 
            Need some mechanism to get relevant class labels
            '''
            final_labels.append(labels)

        # collect boxes
        final_boxes  = np.vstack(final_boxes)
        final_labels = np.hstack(final_labels)

        # clip boxes to the image area
        final_boxes[:, 0:] = np.clip(final_boxes[:, 0:], 0, s).astype(np.int32)
        w = (final_boxes[:,2] - final_boxes[:,0])
        h = (final_boxes[:,3] - final_boxes[:,1])
        
        # discard boxes where w or h <10
        final_boxes = final_boxes[(w>=self.size_limit) & (h>=self.size_limit)]

        return mosaic_image, final_boxes, final_labels

【问题讨论】:

【参考方案1】:

我同时解析了bounding boxclass label信息。

以下是我们实现的输出。要尝试使用您自己的数据集,请先。

【讨论】:

以上是关于如何从对象检测数据加载器中的马赛克增强中获取类标签?的主要内容,如果未能解决你的问题,请参考以下文章

如何使用Albumentations 对目标检测任务做增强

如何从 pytorch 数据加载器中获取批迭代的总数?

如何在模块化的 java 11 应用程序中动态加载 Libreoffice jar,而不从自定义类加载器中获取 ClassCastException

对象的创建

如何从 jQuery 移动加载器中删除 <h1> 标签?

如何从 csv 文件中提取图像、标签并使用 Torch 创建训练集?