如何使用PyTorch Dataloader从Mosaic增强中获取类标签?

问题描述

注意:我想不出更好的 Title 名称,欢迎您 编辑或建议。

更新

直接Colab Link。只需获取给定的虚拟数据集并将其加载到colab。


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

在我的训练机制中,由于扩充机制会随机选择样本的子部分,因此我有点无法正确检索每个类别的类别标签。但是,以下是我们到目前为止通过相关边界框实现的马赛克增强的结果。

enter image description here

数据集

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

enter image description here

它总共有 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,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 个样本以创建此类扩充,但图像 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,labels = self.load_image_and_Boxes(index)

            if i == 0:    # top left
                x1a,y1a,x2a,y2a =  0,xc,yc
                x1b,y1b,x2b,y2b = s - xc,s - yc,s # from bottom right
            elif i == 1:  # top right
                x1a,y2a = xc,y2b = 0,s - xc,s # from bottom left
            elif i == 2:  # bottom left
                x1a,y2a = 0,yc,s
                x1b,s-yc   # from top right
            elif i == 3:  # bottom right
                x1a,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,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:],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

就是这样。希望我能使查询更清楚。您的建议将不胜感激。


通过此查询,我还更新了几天前提出的另一个非常相关查询,但没有得到足够的答复。我也更新了该查询,并使其更加清晰。如果您有兴趣,请链接Stratified K-Fold For Multi-Class Object Detection?

解决方法

已解决-)

问题已解决。最初,我以非常困难的方式考虑它,但是,我只需要同时解析bounding boxclass label信息。 J 撇开,我失去了100个赏金> _

无论如何,下面是我们现在实现的输出。如果您有兴趣尝试使用自己的数据集进行尝试,这里是colab notebook入门指南。编码愉快-)

enter image description here

相关问答

Selenium Web驱动程序和Java。元素在(x,y)点处不可单击。其...
Python-如何使用点“。” 访问字典成员?
Java 字符串是不可变的。到底是什么意思?
Java中的“ final”关键字如何工作?(我仍然可以修改对象。...
“loop:”在Java代码中。这是什么,为什么要编译?
java.lang.ClassNotFoundException:sun.jdbc.odbc.JdbcOdbc...