VOC数据集到RRPN所需数据格式转换脚本

- 日理万妓 2022-02-13 09:47 279阅读 0赞

# 1. 前言 #

在之前的[文章][Link 1]中已经介绍了过了RRPN的原理，在给出的代码里面也写了相关数据转换的脚本，其实只需要理解了其内部所需数据的类型就可以按照自己的意愿将现有的数据集转换过去，而且该算法可以做到多类别检测，只是需要很小的修改就可以实现。

# 2. 转换脚本 #

def get_PascalVOC_2007():
        DATASET_DIR = 'xxx/data/VOCdevkit2007'  # VOC数据集路径
    
        img_dir = os.path.join(DATASET_DIR, 'VOC2007', 'JPEGImages')
        gt_dir = os.path.join(DATASET_DIR, 'VOC2007', 'Annotations')
        img_set = os.path.join(DATASET_DIR, 'VOC2007', 'ImageSets', 'Main', 'trainval.txt')
    
    
        num_classes = 5  # 分类的数目
        box_classes = ('__background__', 'class1', 'class2', 'class3', 'class4')
        box_classes_to_ind = dict(zip(box_classes, np.arange(num_classes)))
    
        im_infos = []  # 标注信息
    
        with open(img_set, 'r') as f:
            gt_file_list = f.readlines()
        print("buddy: all labeled image size: {}".format(len(gt_file_list)))
    
        # gt_file_list = gt_file_list[0:100]
    
        for index, gt_file in enumerate(gt_file_list):
            if index % 10000 == 0:
                print('buddy: load process-{}/{}'.format(index, len(gt_file_list)))
            gt_file = gt_file.strip('\n')
            gt_fobj = open(os.path.join(gt_dir, gt_file + '.xml'))
            tree = ET.parse(gt_fobj)
            objs = tree.findall('object')  # 找到所有的标注框节点
    
            img_name = os.path.join(img_dir, gt_file + '.jpg')
            img = cv2.imread(img_name)
    
            len_of_bboxes = len(objs)
            gt_boxes = np.zeros((len_of_bboxes, 5), dtype=np.int16)
            gt_classes = np.zeros((len_of_bboxes), dtype=np.int32)
            overlaps = np.zeros((len_of_bboxes, num_classes), dtype=np.float32)  # text or non-text
            seg_areas = np.zeros((len_of_bboxes), dtype=np.float32)
    
            # Load object bounding boxes into a data frame.
            for idx, obj in enumerate(objs):
                bbox = obj.find('bndbox')
                # Make pixel indexes 0-based
                x1 = float(bbox.find('xmin').text) - 1
                y1 = float(bbox.find('ymin').text) - 1
                x2 = float(bbox.find('xmax').text) - 1
                y2 = float(bbox.find('ymax').text) - 1
                cls = box_classes_to_ind[obj.find('name').text.lower().strip()]
                width = x2 - x1
                height = y2 - y1
                x_ctr = x1 + width/2
                y_ctr = y1 + height/2
                gt_boxes[idx, :] = [x_ctr, y_ctr, height, width, 0]
                gt_classes[idx] = cls
                overlaps[idx, cls] = 1.0
                seg_areas[idx] = (x2 - x1 + 1) * (y2 - y1 + 1)
    
            max_overlaps = overlaps.max(axis=1)
            # gt class that had the max overlap
            max_classes = overlaps.argmax(axis=1)
    
            im_info = { 
                'gt_classes': gt_classes,
                'max_classes': max_classes,
                'image': os.path.join(img_dir, gt_file + '.jpg'),
                'boxes': gt_boxes,
                'flipped': False,
                'gt_overlaps': overlaps,
                'seg_areas': seg_areas,
                'height': img.shape[0],
                'width': img.shape[1],
                'max_overlaps': max_overlaps,
                'rotated': True  # 使用旋转增广数据
            }
            im_infos.append(im_info)
        return im_infos

[Link 1]: https://blog.csdn.net/m_buddy/article/details/89416473