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

Question

我正在尝试针对多类 问题训练对象检测模型。在我的培训中，我使用 Mosaic augmentation , Paper , 对于这个任务。

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

数据集

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

它总共有4个类和df.object.value_counts():

human    23
car      13
cat       5
dog       3

数据加载器和马赛克增强

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

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 样本以创建此类增强时，image 和 bounding_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

Answer 1

我同时解析了bounding box和class label信息。

下面是我们已经实现的输出。要使用您自己的数据集进行尝试，对于初学者。