def test_can_save_and_load_with_landmarks(self): source_dataset = Dataset.from_iterable([ DatasetItem( id='name0_0001', subset='test', image=np.ones((2, 5, 3)), annotations=[ Label(0, attributes={'positive_pairs': ['name0/name0_0002']}), Points([0, 4, 3, 3, 2, 2, 1, 0, 3, 0], label=0), ]), DatasetItem(id='name0_0002', subset='test', image=np.ones((2, 5, 3)), annotations=[ Label(0), Points([0, 5, 3, 5, 2, 2, 1, 0, 3, 0], label=0), ]), ], categories=['name0']) with TestDir() as test_dir: LfwConverter.convert(source_dataset, test_dir, save_images=True) parsed_dataset = Dataset.import_from(test_dir, 'lfw') compare_datasets(self, source_dataset, parsed_dataset)
def test_can_import(self): expected_dataset = Dataset.from_iterable( [ DatasetItem(id='n000001/0001_01', subset='train', image=np.ones((10, 15, 3)), annotations=[ Bbox(2, 2, 1, 2, label=0), Points([ 2.787, 2.898, 2.965, 2.79, 2.8, 2.456, 2.81, 2.32, 2.89, 2.3 ], label=0), ]), DatasetItem(id='n000002/0001_01', subset='train', image=np.ones((10, 15, 3)), annotations=[ Bbox(2, 4, 2, 2, label=1), Points([ 2.3, 4.9, 2.9, 4.93, 2.62, 4.745, 2.54, 4.45, 2.76, 4.43 ], label=1) ]), DatasetItem(id='n000002/0002_01', subset='train', image=np.ones((10, 15, 3)), annotations=[ Bbox(1, 3, 1, 1, label=1), Points([ 1.2, 3.8, 1.8, 3.82, 1.51, 3.634, 1.43, 3.34, 1.65, 3.32 ], label=1) ]), DatasetItem( id='n000003/0003_01', subset='test', image=np.ones((10, 15, 3)), annotations=[ Bbox(1, 1, 1, 1, label=2), Points( [0.2, 2.8, 0.8, 2.9, 0.5, 2.6, 0.4, 2.3, 0.6, 2.3], label=2) ]) ], categories={ AnnotationType.label: LabelCategories.from_iterable([('n000001', 'Karl'), ('n000002', 'Jay'), ('n000003', 'Pol')]), }) dataset = Dataset.import_from(DUMMY_DATASET_DIR, 'vgg_face2') compare_datasets(self, expected_dataset, dataset)
def test_can_save_and_load_image_with_arbitrary_extension(self): dataset = Dataset.from_iterable([ DatasetItem('no_label/q/1', image=Image(path='q/1.JPEG', data=np.zeros( (4, 3, 3)))), DatasetItem('a/b/c/2', image=Image(path='a/b/c/2.bmp', data=np.zeros((3, 4, 3))), annotations=[ Bbox(0, 2, 4, 2, label=0), Points([ 4.23, 4.32, 5.34, 4.45, 3.54, 3.56, 4.52, 3.51, 4.78, 3.34 ], label=0), ]), ], categories=['a']) with TestDir() as test_dir: VggFace2Converter.convert(dataset, test_dir, save_images=True) parsed_dataset = Dataset.import_from(test_dir, 'vgg_face2') compare_datasets(self, dataset, parsed_dataset, require_images=True)
def test_can_import_specific_subset(self): expected_dataset = Dataset.from_iterable( [ DatasetItem( id='n000003/0003_01', subset='test', image=np.ones((10, 15, 3)), annotations=[ Bbox(1, 1, 1, 1, label=2), Points( [0.2, 2.8, 0.8, 2.9, 0.5, 2.6, 0.4, 2.3, 0.6, 2.3], label=2) ]) ], categories={ AnnotationType.label: LabelCategories.from_iterable([('n000001', 'Karl'), ('n000002', 'Jay'), ('n000003', 'Pol')]), }) specific_subset = osp.join(DUMMY_DATASET_DIR, 'bb_landmark', 'loose_bb_test.csv') dataset = Dataset.import_from(specific_subset, 'vgg_face2') compare_datasets(self, expected_dataset, dataset)
def test_shapes_to_boxes(self): source_dataset = Dataset.from_iterable([ DatasetItem(id=1, image=np.zeros((5, 5, 3)), annotations=[ Mask(np.array([ [0, 0, 1, 1, 1], [0, 0, 0, 0, 1], [1, 0, 0, 0, 1], [1, 0, 0, 0, 0], [1, 1, 1, 0, 0]], ), id=1), Polygon([1, 1, 4, 1, 4, 4, 1, 4], id=2), PolyLine([1, 1, 2, 1, 2, 2, 1, 2], id=3), Points([2, 2, 4, 2, 4, 4, 2, 4], id=4), ] ), ]) target_dataset = Dataset.from_iterable([ DatasetItem(id=1, image=np.zeros((5, 5, 3)), annotations=[ Bbox(0, 0, 4, 4, id=1), Bbox(1, 1, 3, 3, id=2), Bbox(1, 1, 1, 1, id=3), Bbox(2, 2, 2, 2, id=4), ] ), ]) actual = transforms.ShapesToBoxes(source_dataset) compare_datasets(self, target_dataset, actual)
def test_can_load_image(self): expected_dataset = Dataset.from_iterable([ DatasetItem(id='img0', subset='train', image=np.ones((8, 8, 3)), annotations=[ Bbox(0, 2, 4, 2, label=0, z_order=1, attributes={ 'occluded': True, 'a1': True, 'a2': 'v3', 'a3': '0003', 'a4': 2.4, }), PolyLine([1, 2, 3, 4, 5, 6, 7, 8], attributes={'occluded': False}), ], attributes={'frame': 0}), DatasetItem(id='img1', subset='train', image=np.ones((10, 10, 3)), annotations=[ Polygon([1, 2, 3, 4, 6, 5], z_order=1, attributes={'occluded': False}), Points([1, 2, 3, 4, 5, 6], label=1, z_order=2, attributes={'occluded': False}), ], attributes={'frame': 1}), ], categories={ AnnotationType.label: LabelCategories.from_iterable([ ['label1', '', {'a1', 'a2', 'a3', 'a4'}], ['label2'], ]) }) parsed_dataset = Dataset.import_from(DUMMY_IMAGE_DATASET_DIR, 'cvat') compare_datasets(self, expected_dataset, parsed_dataset)
def test_can_save_dataset_with_no_labels(self): source_dataset = Dataset.from_iterable([ DatasetItem(id='no_label/1', image=np.ones((8, 8, 3)), annotations=[ Bbox(0, 2, 4, 2), Points([ 4.23, 4.32, 5.34, 4.45, 3.54, 3.56, 4.52, 3.51, 4.78, 3.34 ]), ]), DatasetItem(id='no_label/2', image=np.ones((8, 8, 3)), annotations=[ Bbox(2, 2, 4, 2), ]), ], categories=[]) with TestDir() as test_dir: VggFace2Converter.convert(source_dataset, test_dir, save_images=False) parsed_dataset = Dataset.import_from(test_dir, 'vgg_face2') compare_datasets(self, source_dataset, parsed_dataset)
def test_can_import_dataset_wo_numpy_files(self): expected_dataset = Dataset.from_iterable([ DatasetItem(id='000000001', image=np.ones((5, 5, 3)), annotations=[ Points([620.0, 394.0, 616.0, 269.0, 573.0, 185.0, 647.0, 188.0, 661.0, 221.0, 656.0, 231.0, 610.0, 187.0, 647.0, 176.0, 637.02, 189.818, 695.98, 108.182, 606.0, 217.0, 553.0, 161.0, 601.0, 167.0, 692.0, 185.0, 693.0, 240.0, 688.0, 313.0], [1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1], attributes={'center': [594.0, 257.0], 'scale': 3.021}, label=0, group=1) ] ), DatasetItem(id='000000002', image=np.ones((5, 5, 3)), annotations=[ Points([650.0, 424.0, 646.0, 309.0, 603.0, 215.0, 677.0, 218.0, 691.0, 251.0, 686.0, 261.0, 640.0, 217.0, 677.0, 216.0, 667.02, 219.818, 725.98, 138.182, 636.0, 247.0, 583.0, 191.0, 631.0, 197.0, 722.0, 215.0, 723.0, 270.0, 718.0, 343.0], [1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1], attributes={'center': [624.0, 287.0], 'scale': 3.7}, label=0, group=1) ] ), DatasetItem(id='000000003', image=np.ones((5, 5, 3)), annotations=[ Points([590.0, 364.0, 586.0, 239.0, 533.0, 155.0, 617.0, 158.0, 631.0, 191.0, 626.0, 201.0, 580.0, 157.0, 617.0, 146.0, 607.02, 159.818, 645.98, 68.182, 576.0, 187.0, 532.0, 131.0, 571.0, 137.0, 662.0, 155.0, 663.0, 210.0, 658.0, 283.0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1], attributes={'center': [564.0, 227.0], 'scale': 3.2}, label=0, group=1) ] ) ], categories={ AnnotationType.label: LabelCategories.from_iterable(['human']), AnnotationType.points: PointsCategories.from_iterable( [(0, MPII_POINTS_LABELS, MPII_POINTS_JOINTS)]) }) dataset = Dataset.import_from(DUMMY_DATASET_DIR_WO_NUMPY_FILES, 'mpii_json') compare_datasets(self, expected_dataset, dataset, require_images=True)
def test_can_resize(self): small_dataset = Dataset.from_iterable([ DatasetItem(id=i, image=np.ones((4, 4)) * i, annotations=[ Label(1), Bbox(1, 1, 2, 2, label=2), Polygon([1, 1, 1, 2, 2, 2, 2, 1], label=1), PolyLine([1, 1, 1, 2, 2, 2, 2, 1], label=2), Points([1, 1, 1, 2, 2, 2, 2, 1], label=2), Mask(np.array([ [0, 0, 1, 1], [1, 0, 0, 1], [0, 1, 1, 0], [1, 1, 0, 0], ])) ]) for i in range(3) ], categories=['a', 'b', 'c']) big_dataset = Dataset.from_iterable([ DatasetItem(id=i, image=np.ones((8, 8)) * i, annotations=[ Label(1), Bbox(2, 2, 4, 4, label=2), Polygon([2, 2, 2, 4, 4, 4, 4, 2], label=1), PolyLine([2, 2, 2, 4, 4, 4, 4, 2], label=2), Points([2, 2, 2, 4, 4, 4, 4, 2], label=2), Mask(np.array([ [0, 0, 0, 0, 1, 1, 1, 1], [0, 0, 0, 0, 1, 1, 1, 1], [1, 1, 0, 0, 0, 0, 1, 1], [1, 1, 0, 0, 0, 0, 1, 1], [0, 0, 1, 1, 1, 1, 0, 0], [0, 0, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 0, 0, 0, 0], [1, 1, 1, 1, 0, 0, 0, 0], ])) ]) for i in range(3) ], categories=['a', 'b', 'c']) with self.subTest('upscale'): actual = transforms.ResizeTransform(small_dataset, width=8, height=8) compare_datasets(self, big_dataset, actual) with self.subTest('downscale'): actual = transforms.ResizeTransform(big_dataset, width=4, height=4) compare_datasets(self, small_dataset, actual)
def test_can_match_points(self): item1 = DatasetItem(id=1, annotations=[ Points([1, 2, 2, 0, 1, 1], label=0), Points([3, 5, 5, 7, 5, 3], label=0), ]) item2 = DatasetItem(id=2, annotations=[ Points([1.5, 2, 2, 0.5, 1, 1.5], label=0), Points([5, 7, 7, 7, 7, 5], label=0), ]) result = DistanceComparator().match_points(item1, item2) matches, mismatches, a_greater, b_greater = result self.assertEqual(1, len(a_greater)) self.assertEqual(1, len(b_greater)) self.assertEqual(1, len(matches)) self.assertEqual(0, len(mismatches))
def test_can_import_without_people_file(self): expected_dataset = Dataset.from_iterable([ DatasetItem(id='name0_0001', subset='test', image=np.ones((2, 5, 3)), annotations=[ Label(0, attributes={ 'negative_pairs': ['name1/name1_0001', 'name1/name1_0002'] }), Points([0, 4, 3, 3, 2, 2, 1, 0, 3, 0], label=0), ]), DatasetItem(id='name1_0001', subset='test', image=np.ones((2, 5, 3)), annotations=[ Label(1, attributes={ 'positive_pairs': ['name1/name1_0002'], }), Points([1, 6, 4, 6, 3, 3, 2, 1, 4, 1], label=1), ]), DatasetItem(id='name1_0002', subset='test', image=np.ones((2, 5, 3)), annotations=[ Label(1), Points([0, 5, 3, 5, 2, 2, 1, 0, 3, 0], label=1), ]), ], categories=['name0', 'name1']) with TestDir() as test_dir: dataset_path = osp.join(test_dir, 'dataset') shutil.copytree(DUMMY_DATASET_DIR, dataset_path) os.remove( osp.join(dataset_path, 'test', 'annotations', 'people.txt')) dataset = Dataset.import_from(DUMMY_DATASET_DIR, 'lfw') compare_datasets(self, expected_dataset, dataset)
def test_can_save_and_load_with_meta_file(self): source_dataset = Dataset.from_iterable( [ DatasetItem(id='class_0/1', subset='train', image=np.ones((8, 8, 3)), annotations=[ Bbox(0, 2, 4, 2, label=0), Points([ 3.2, 3.12, 4.11, 3.2, 2.11, 2.5, 3.5, 2.11, 3.8, 2.13 ], label=0), ]), DatasetItem(id='class_1/2', subset='train', image=np.ones((10, 10, 3)), annotations=[ Points([ 4.23, 4.32, 5.34, 4.45, 3.54, 3.56, 4.52, 3.51, 4.78, 3.34 ], label=1), ]) ], categories={ AnnotationType.label: LabelCategories.from_iterable([('class_%s' % i) for i in range(5)]), }) with TestDir() as test_dir: VggFace2Converter.convert(source_dataset, test_dir, save_images=True, save_dataset_meta=True) parsed_dataset = Dataset.import_from(test_dir, 'vgg_face2') self.assertTrue(osp.isfile(osp.join(test_dir, 'dataset_meta.json'))) compare_datasets(self, source_dataset, parsed_dataset)
def test_can_import(self): expected_dataset = Dataset.from_iterable([ DatasetItem(id='name0_0001', subset='test', image=np.ones((2, 5, 3)), annotations=[ Label(0, attributes={ 'negative_pairs': ['name1/name1_0001', 'name1/name1_0002'] }), Points([0, 4, 3, 3, 2, 2, 1, 0, 3, 0], label=0), ]), DatasetItem(id='name1_0001', subset='test', image=np.ones((2, 5, 3)), annotations=[ Label(1, attributes={ 'positive_pairs': ['name1/name1_0002'], }), Points([1, 6, 4, 6, 3, 3, 2, 1, 4, 1], label=1), ]), DatasetItem(id='name1_0002', subset='test', image=np.ones((2, 5, 3)), annotations=[ Label(1), Points([0, 5, 3, 5, 2, 2, 1, 0, 3, 0], label=1), ]), ], categories=['name0', 'name1']) dataset = Dataset.import_from(DUMMY_DATASET_DIR, 'lfw') compare_datasets(self, expected_dataset, dataset)
def test_can_save_dataset_with_wrong_number_of_points(self): source_dataset = Dataset.from_iterable([ DatasetItem(id='no_label/1', image=np.ones((8, 8, 3)), annotations=[ Points([4.23, 4.32, 5.34, 3.51, 4.78, 3.34]), ]), ], categories=[]) target_dataset = Dataset.from_iterable([ DatasetItem( id='no_label/1', image=np.ones((8, 8, 3)), annotations=[]), ], categories=[]) with TestDir() as test_dir: VggFace2Converter.convert(source_dataset, test_dir, save_images=True) parsed_dataset = Dataset.import_from(test_dir, 'vgg_face2') compare_datasets(self, target_dataset, parsed_dataset)
def test_can_save_dataset_with_cyrillic_and_spaces_in_filename(self): source_dataset = Dataset.from_iterable([ DatasetItem(id='a/кириллица с пробелом', image=np.ones((8, 8, 3)), annotations=[ Points([ 4.23, 4.32, 5.34, 4.45, 3.54, 3.56, 4.52, 3.51, 4.78, 3.34 ], label=0), ]), ], categories=['a']) with TestDir() as test_dir: VggFace2Converter.convert(source_dataset, test_dir, save_images=True) parsed_dataset = Dataset.import_from(test_dir, 'vgg_face2') compare_datasets(self, source_dataset, parsed_dataset, require_images=True)
def _load_items(self, root_dir): items = {} image_dir = osp.join(root_dir, CelebaPath.IMAGES_DIR) if osp.isdir(image_dir): images = { osp.splitext(osp.relpath(p, image_dir))[0].replace('\\', '/'): p for p in find_images(image_dir, recursive=True) } else: images = {} label_categories = self._categories[AnnotationType.label] labels_path = osp.join(root_dir, CelebaPath.LABELS_FILE) if not osp.isfile(labels_path): raise DatasetImportError("File '%s': was not found" % labels_path) with open(labels_path, encoding='utf-8') as f: for line in f: item_id, item_ann = self.split_annotation(line) label_ids = [int(id) for id in item_ann] anno = [] for label in label_ids: while len(label_categories) <= label: label_categories.add('class-%d' % len(label_categories)) anno.append(Label(label)) items[item_id] = DatasetItem(id=item_id, image=images.get(item_id), annotations=anno) landmark_path = osp.join(root_dir, CelebaPath.LANDMARKS_FILE) if osp.isfile(landmark_path): with open(landmark_path, encoding='utf-8') as f: landmarks_number = int(f.readline().strip()) point_cat = PointsCategories() for i, point_name in enumerate(f.readline().strip().split()): point_cat.add(i, [point_name]) self._categories[AnnotationType.points] = point_cat counter = 0 for counter, line in enumerate(f): item_id, item_ann = self.split_annotation(line) landmarks = [float(id) for id in item_ann] if len(landmarks) != len(point_cat): raise DatasetImportError("File '%s', line %s: " "points do not match the header of this file" % \ (landmark_path, line)) if item_id not in items: raise DatasetImportError("File '%s', line %s: " "for this item are not label in %s " % \ (landmark_path, line, CelebaPath.LABELS_FILE)) anno = items[item_id].annotations label = anno[0].label anno.append(Points(landmarks, label=label)) if landmarks_number - 1 != counter: raise DatasetImportError( "File '%s': the number of " "landmarks does not match the specified number " "at the beginning of the file " % landmark_path) bbox_path = osp.join(root_dir, CelebaPath.BBOXES_FILE) if osp.isfile(bbox_path): with open(bbox_path, encoding='utf-8') as f: bboxes_number = int(f.readline().strip()) if f.readline().strip() != CelebaPath.BBOXES_HEADER: raise DatasetImportError("File '%s': the header " "does not match the expected format '%s'" % \ (bbox_path, CelebaPath.BBOXES_HEADER)) counter = 0 for counter, line in enumerate(f): item_id, item_ann = self.split_annotation(line) bbox = [float(id) for id in item_ann] if item_id not in items: raise DatasetImportError("File '%s', line %s: " "for this item are not label in %s " % \ (bbox_path, line, CelebaPath.LABELS_FILE)) anno = items[item_id].annotations label = anno[0].label anno.append( Bbox(bbox[0], bbox[1], bbox[2], bbox[3], label=label)) if bboxes_number - 1 != counter: raise DatasetImportError( "File '%s': the number of bounding " "boxes does not match the specified number " "at the beginning of the file " % bbox_path) attr_path = osp.join(root_dir, CelebaPath.ATTRS_FILE) if osp.isfile(attr_path): with open(attr_path, encoding='utf-8') as f: attr_number = int(f.readline().strip()) attr_names = f.readline().split() counter = 0 for counter, line in enumerate(f): item_id, item_ann = self.split_annotation(line) if len(attr_names) != len(item_ann): raise DatasetImportError( "File '%s', line %s: " "the number of attributes " "in the line does not match the number at the " "beginning of the file " % (attr_path, line)) attrs = { name: 0 < int(ann) for name, ann in zip(attr_names, item_ann) } if item_id not in items: items[item_id] = DatasetItem(id=item_id, image=images.get(item_id)) items[item_id].attributes = attrs if attr_number - 1 != counter: raise DatasetImportError( "File %s: the number of items " "with attributes does not match the specified number " "at the beginning of the file " % attr_path) subset_path = osp.join(root_dir, CelebaPath.SUBSETS_FILE) if osp.isfile(subset_path): with open(subset_path, encoding='utf-8') as f: for line in f: item_id, item_ann = self.split_annotation(line) subset_id = item_ann[0] subset = CelebaPath.SUBSETS[subset_id] if item_id not in items: items[item_id] = DatasetItem(id=item_id, image=images.get(item_id)) items[item_id].subset = subset if 'default' in self._subsets: self._subsets.pop() self._subsets.append(subset) return items
def _parse_shape_ann(cls, ann, categories): ann_id = ann.get('id', 0) ann_type = ann['type'] attributes = ann.get('attributes') or {} if 'occluded' in categories[AnnotationType.label].attributes: attributes['occluded'] = ann.get('occluded', False) if 'outside' in ann: attributes['outside'] = ann['outside'] if 'keyframe' in ann: attributes['keyframe'] = ann['keyframe'] if 'track_id' in ann: attributes['track_id'] = ann['track_id'] group = ann.get('group') label = ann.get('label') label_id = categories[AnnotationType.label].find(label)[0] z_order = ann.get('z_order', 0) points = ann.get('points', []) if ann_type == 'polyline': return PolyLine(points, label=label_id, z_order=z_order, id=ann_id, attributes=attributes, group=group) elif ann_type == 'polygon': return Polygon(points, label=label_id, z_order=z_order, id=ann_id, attributes=attributes, group=group) elif ann_type == 'points': return Points(points, label=label_id, z_order=z_order, id=ann_id, attributes=attributes, group=group) elif ann_type == 'box': x, y = points[0], points[1] w, h = points[2] - x, points[3] - y return Bbox(x, y, w, h, label=label_id, z_order=z_order, id=ann_id, attributes=attributes, group=group) else: raise NotImplementedError("Unknown annotation type '%s'" % ann_type)
def _load_annotations(self, ann, image_info=None, parsed_annotations=None): if parsed_annotations is None: parsed_annotations = [] ann_id = ann['id'] attributes = ann.get('attributes', {}) if 'score' in ann: attributes['score'] = ann['score'] group = ann_id # make sure all tasks' annotations are merged if self._task is CocoTask.instances or \ self._task is CocoTask.person_keypoints or \ self._task is CocoTask.stuff: label_id = self._get_label_id(ann) attributes['is_crowd'] = bool(ann['iscrowd']) if self._task is CocoTask.person_keypoints: keypoints = ann['keypoints'] points = [] visibility = [] for x, y, v in take_by(keypoints, 3): points.append(x) points.append(y) visibility.append(v) parsed_annotations.append( Points(points, visibility, label=label_id, id=ann_id, attributes=attributes, group=group)) segmentation = ann['segmentation'] if segmentation and segmentation != [[]]: rle = None if isinstance(segmentation, list): if not self._merge_instance_polygons: # polygon - a single object can consist of multiple parts for polygon_points in segmentation: parsed_annotations.append( Polygon(points=polygon_points, label=label_id, id=ann_id, attributes=attributes, group=group)) else: # merge all parts into a single mask RLE rle = self._lazy_merged_mask(segmentation, image_info['height'], image_info['width']) elif isinstance(segmentation['counts'], list): # uncompressed RLE img_h = image_info['height'] img_w = image_info['width'] mask_h, mask_w = segmentation['size'] if img_h == mask_h and img_w == mask_w: rle = self._lazy_merged_mask([segmentation], mask_h, mask_w) else: log.warning( "item #%s: mask #%s " "does not match image size: %s vs. %s. " "Skipping this annotation.", image_info['id'], ann_id, (mask_h, mask_w), (img_h, img_w)) else: # compressed RLE rle = segmentation if rle: parsed_annotations.append( RleMask(rle=rle, label=label_id, id=ann_id, attributes=attributes, group=group)) else: x, y, w, h = ann['bbox'] parsed_annotations.append( Bbox(x, y, w, h, label=label_id, id=ann_id, attributes=attributes, group=group)) elif self._task is CocoTask.labels: label_id = self._get_label_id(ann) parsed_annotations.append( Label(label=label_id, id=ann_id, attributes=attributes, group=group)) elif self._task is CocoTask.captions: caption = ann['caption'] parsed_annotations.append( Caption(caption, id=ann_id, attributes=attributes, group=group)) else: raise NotImplementedError() return parsed_annotations
def _load_annotations(item): parsed = item['annotations'] loaded = [] for ann in parsed: ann_id = ann.get('id') ann_type = AnnotationType[ann['type']] attributes = ann.get('attributes') group = ann.get('group') label_id = ann.get('label_id') z_order = ann.get('z_order') points = ann.get('points') if ann_type == AnnotationType.label: loaded.append( Label(label=label_id, id=ann_id, attributes=attributes, group=group)) elif ann_type == AnnotationType.mask: rle = ann['rle'] rle['counts'] = rle['counts'].encode('ascii') loaded.append( RleMask(rle=rle, label=label_id, id=ann_id, attributes=attributes, group=group, z_order=z_order)) elif ann_type == AnnotationType.polyline: loaded.append( PolyLine(points, label=label_id, id=ann_id, attributes=attributes, group=group, z_order=z_order)) elif ann_type == AnnotationType.polygon: loaded.append( Polygon(points, label=label_id, id=ann_id, attributes=attributes, group=group, z_order=z_order)) elif ann_type == AnnotationType.bbox: x, y, w, h = ann['bbox'] loaded.append( Bbox(x, y, w, h, label=label_id, id=ann_id, attributes=attributes, group=group, z_order=z_order)) elif ann_type == AnnotationType.points: loaded.append( Points(points, label=label_id, id=ann_id, attributes=attributes, group=group, z_order=z_order)) elif ann_type == AnnotationType.caption: caption = ann.get('caption') loaded.append( Caption(caption, id=ann_id, attributes=attributes, group=group)) elif ann_type == AnnotationType.cuboid_3d: loaded.append( Cuboid3d(ann.get('position'), ann.get('rotation'), ann.get('scale'), label=label_id, id=ann_id, attributes=attributes, group=group)) else: raise NotImplementedError() return loaded
def test_can_match_shapes(self): source0 = Dataset.from_iterable([ DatasetItem(1, annotations=[ # unique Bbox(1, 2, 3, 4, label=1), # common Mask(label=2, z_order=2, image=np.array([ [0, 0, 0, 0], [0, 0, 0, 0], [1, 1, 1, 0], [1, 1, 1, 0], ])), Polygon([1, 0, 3, 2, 1, 2]), # an instance with keypoints Bbox(4, 5, 2, 4, label=2, z_order=1, group=1), Points([5, 6], label=0, group=1), Points([6, 8], label=1, group=1), PolyLine([1, 1, 2, 1, 3, 1]), ]), ], categories=['a', 'b', 'c']) source1 = Dataset.from_iterable([ DatasetItem(1, annotations=[ # common Mask(label=2, image=np.array([ [0, 0, 0, 0], [0, 1, 1, 1], [0, 1, 1, 1], [0, 1, 1, 1], ])), Polygon([0, 2, 2, 0, 2, 1]), # an instance with keypoints Bbox(4, 4, 2, 5, label=2, z_order=1, group=2), Points([5.5, 6.5], label=0, group=2), Points([6, 8], label=1, group=2), PolyLine([1, 1.5, 2, 1.5]), ]), ], categories=['a', 'b', 'c']) source2 = Dataset.from_iterable([ DatasetItem(1, annotations=[ # common Mask(label=2, z_order=3, image=np.array([ [0, 0, 1, 1], [0, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 0], ])), Polygon([3, 1, 2, 2, 0, 1]), # an instance with keypoints, one is missing Bbox(3, 6, 2, 3, label=2, z_order=4, group=3), Points([4.5, 5.5], label=0, group=3), PolyLine([1, 1.25, 3, 1, 4, 2]), ]), ], categories=['a', 'b', 'c']) expected = Dataset.from_iterable([ DatasetItem(1, annotations=[ # unique Bbox(1, 2, 3, 4, label=1), # common # nearest to mean bbox Mask(label=2, z_order=3, image=np.array([ [0, 0, 0, 0], [0, 1, 1, 1], [0, 1, 1, 1], [0, 1, 1, 1], ])), Polygon([1, 0, 3, 2, 1, 2]), # an instance with keypoints Bbox(4, 5, 2, 4, label=2, z_order=4, group=1), Points([5, 6], label=0, group=1), Points([6, 8], label=1, group=1), PolyLine([1, 1.25, 3, 1, 4, 2]), ]), ], categories=['a', 'b', 'c']) merger = IntersectMerge(conf={'quorum': 1, 'pairwise_dist': 0.1}) merged = merger([source0, source1, source2]) compare_datasets(self, expected, merged, ignored_attrs={'score'}) self.assertEqual( [ NoMatchingAnnError(item_id=('1', DEFAULT_SUBSET_NAME), sources={2}, ann=source0.get('1').annotations[5]), NoMatchingAnnError(item_id=('1', DEFAULT_SUBSET_NAME), sources={1, 2}, ann=source0.get('1').annotations[0]), ], sorted((e for e in merger.errors if isinstance(e, NoMatchingAnnError)), key=lambda e: len(e.sources)) )
def test_annotation_comparison(self): a = Dataset.from_iterable( [ DatasetItem( id=1, annotations=[ Caption('hello'), # unmatched Caption('world', group=5), Label(2, attributes={ 'x': 1, 'y': '2', }), Bbox(1, 2, 3, 4, label=4, z_order=1, attributes={ 'score': 1.0, }), Bbox(5, 6, 7, 8, group=5), Points([1, 2, 2, 0, 1, 1], label=0, z_order=4), Mask(label=3, z_order=2, image=np.ones((2, 3))), ]), ], categories=['a', 'b', 'c', 'd']) b = Dataset.from_iterable( [ DatasetItem( id=1, annotations=[ Caption('world', group=5), Label(2, attributes={ 'x': 1, 'y': '2', }), Bbox(1, 2, 3, 4, label=4, z_order=1, attributes={ 'score': 1.0, }), Bbox(5, 6, 7, 8, group=5), Bbox(5, 6, 7, 8, group=5), # unmatched Points([1, 2, 2, 0, 1, 1], label=0, z_order=4), Mask(label=3, z_order=2, image=np.ones((2, 3))), ]), ], categories=['a', 'b', 'c', 'd']) comp = ExactComparator() matched, unmatched, _, _, errors = comp.compare_datasets(a, b) self.assertEqual(6, len(matched), matched) self.assertEqual(2, len(unmatched), unmatched) self.assertEqual(0, len(errors), errors)
def test_can_load_video(self): expected_dataset = Dataset.from_iterable([ DatasetItem(id='frame_000010', subset='annotations', image=255 * np.ones((20, 25, 3)), annotations=[ Bbox(3, 4, 7, 1, label=2, id=0, attributes={ 'occluded': True, 'outside': False, 'keyframe': True, 'track_id': 0 }), Points([21.95, 8.00, 2.55, 15.09, 2.23, 3.16], label=0, id=1, attributes={ 'occluded': False, 'outside': False, 'keyframe': True, 'track_id': 1, 'hgl': 'hgkf', }), ], attributes={'frame': 10}), DatasetItem(id='frame_000013', subset='annotations', image=255 * np.ones((20, 25, 3)), annotations=[ Bbox(7, 6, 7, 2, label=2, id=0, attributes={ 'occluded': False, 'outside': True, 'keyframe': True, 'track_id': 0 }), Points([21.95, 8.00, 9.55, 15.09, 5.23, 1.16], label=0, id=1, attributes={ 'occluded': False, 'outside': True, 'keyframe': True, 'track_id': 1, 'hgl': 'jk', }), PolyLine([7.85, 13.88, 3.50, 6.67, 15.90, 2.00, 13.31, 7.21], label=2, id=2, attributes={ 'occluded': False, 'outside': False, 'keyframe': True, 'track_id': 2, }), ], attributes={'frame': 13}), DatasetItem(id='frame_000016', subset='annotations', image=Image(path='frame_0000016.png', size=(20, 25)), annotations=[ Bbox(8, 7, 6, 10, label=2, id=0, attributes={ 'occluded': False, 'outside': True, 'keyframe': True, 'track_id': 0 }), PolyLine([7.85, 13.88, 3.50, 6.67, 15.90, 2.00, 13.31, 7.21], label=2, id=2, attributes={ 'occluded': False, 'outside': True, 'keyframe': True, 'track_id': 2, }), ], attributes={'frame': 16}), ], categories={ AnnotationType.label: LabelCategories.from_iterable([ ['klhg', '', {'hgl'}], ['z U k'], ['II'] ]), }) parsed_dataset = Dataset.import_from(DUMMY_VIDEO_DATASET_DIR, 'cvat') compare_datasets(self, expected_dataset, parsed_dataset)
def test_stats(self): dataset = Dataset.from_iterable([ DatasetItem(id=1, image=np.ones((5, 5, 3)), annotations=[ Caption('hello'), Caption('world'), Label(2, attributes={ 'x': 1, 'y': '2', }), Bbox(1, 2, 2, 2, label=2, attributes={ 'score': 0.5, }), Bbox(5, 6, 2, 2, attributes={ 'x': 1, 'y': '3', 'occluded': True, }), Points([1, 2, 2, 0, 1, 1], label=0), Mask(label=3, image=np.array([ [0, 0, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 1, 1, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], ])), ]), DatasetItem(id=2, image=np.ones((2, 4, 3)), annotations=[ Label(2, attributes={ 'x': 2, 'y': '2', }), Bbox(1, 2, 2, 2, label=3, attributes={ 'score': 0.5, }), Bbox(5, 6, 2, 2, attributes={ 'x': 2, 'y': '3', 'occluded': False, }), ]), DatasetItem(id=3), DatasetItem(id='2.2', image=np.ones((2, 4, 3))), ], categories=['label_%s' % i for i in range(4)]) expected = { 'images count': 4, 'annotations count': 10, 'unannotated images count': 2, 'unannotated images': ['3', '2.2'], 'annotations by type': { 'label': { 'count': 2, }, 'polygon': { 'count': 0, }, 'polyline': { 'count': 0, }, 'bbox': { 'count': 4, }, 'mask': { 'count': 1, }, 'points': { 'count': 1, }, 'caption': { 'count': 2, }, 'cuboid_3d': {'count': 0}, }, 'annotations': { 'labels': { 'count': 6, 'distribution': { 'label_0': [1, 1/6], 'label_1': [0, 0.0], 'label_2': [3, 3/6], 'label_3': [2, 2/6], }, 'attributes': { 'x': { 'count': 2, # annotations with no label are skipped 'values count': 2, 'values present': ['1', '2'], 'distribution': { '1': [1, 1/2], '2': [1, 1/2], }, }, 'y': { 'count': 2, # annotations with no label are skipped 'values count': 1, 'values present': ['2'], 'distribution': { '2': [2, 2/2], }, }, # must not include "special" attributes like "occluded" } }, 'segments': { 'avg. area': (4 * 2 + 9 * 1) / 3, 'area distribution': [ {'min': 4.0, 'max': 4.5, 'count': 2, 'percent': 2/3}, {'min': 4.5, 'max': 5.0, 'count': 0, 'percent': 0.0}, {'min': 5.0, 'max': 5.5, 'count': 0, 'percent': 0.0}, {'min': 5.5, 'max': 6.0, 'count': 0, 'percent': 0.0}, {'min': 6.0, 'max': 6.5, 'count': 0, 'percent': 0.0}, {'min': 6.5, 'max': 7.0, 'count': 0, 'percent': 0.0}, {'min': 7.0, 'max': 7.5, 'count': 0, 'percent': 0.0}, {'min': 7.5, 'max': 8.0, 'count': 0, 'percent': 0.0}, {'min': 8.0, 'max': 8.5, 'count': 0, 'percent': 0.0}, {'min': 8.5, 'max': 9.0, 'count': 1, 'percent': 1/3}, ], 'pixel distribution': { 'label_0': [0, 0.0], 'label_1': [0, 0.0], 'label_2': [4, 4/17], 'label_3': [13, 13/17], }, } }, } actual = compute_ann_statistics(dataset) self.assertEqual(expected, actual)
def test_can_import(self): expected_dataset = Dataset.from_iterable( [ DatasetItem( id='000001', subset='train', image=np.ones((3, 4, 3)), annotations=[ Label(12), Points([69, 109, 106, 113, 77, 142, 73, 152, 108, 154], label=12) ], attributes={ '5_o_Clock_Shadow': False, 'Arched_Eyebrows': True, 'Attractive': True, 'Bags_Under_Eyes': False, 'Bald': False, 'Bangs': False, 'Big_Lips': False, 'Big_Nose': False }), DatasetItem( id='000002', subset='train', image=np.ones((3, 4, 3)), annotations=[ Label(5), Points([69, 110, 107, 112, 81, 135, 70, 151, 108, 153], label=5) ]), DatasetItem( id='000003', subset='val', image=np.ones((3, 4, 3)), annotations=[ Label(2), Points([76, 112, 104, 106, 108, 128, 74, 156, 98, 158], label=2) ], attributes={ '5_o_Clock_Shadow': False, 'Arched_Eyebrows': False, 'Attractive': False, 'Bags_Under_Eyes': True, 'Bald': False, 'Bangs': False, 'Big_Lips': False, 'Big_Nose': True }), DatasetItem( id='000004', subset='test', image=np.ones((3, 4, 3)), annotations=[ Label(10), Points( [72, 113, 108, 108, 101, 138, 71, 155, 101, 151], label=10) ]), DatasetItem( id='000005', subset='test', image=np.ones((3, 4, 3)), annotations=[ Label(7), Points([66, 114, 112, 112, 86, 119, 71, 147, 104, 150], label=7) ]) ], categories={ AnnotationType.label: LabelCategories.from_iterable(f'class-{i}' for i in range(13)), AnnotationType.points: PointsCategories.from_iterable([(0, ['lefteye_x']), (1, ['lefteye_y']), (2, ['righteye_x']), (3, ['righteye_y']), (4, ['nose_x']), (5, ['nose_y']), (6, ['leftmouth_x']), (7, ['leftmouth_y']), (8, ['rightmouth_x']), (9, ['rightmouth_y'])]) }) dataset = Dataset.import_from(DUMMY_ALIGN_DATASET_DIR, 'align_celeba') compare_datasets(self, expected_dataset, dataset, require_images=True)
def _load_items(self, subset): def _get_label(path): label_name = path.split('/')[0] label = None if label_name != VggFace2Path.IMAGES_DIR_NO_LABEL: label = \ self._categories[AnnotationType.label].find(label_name)[0] return label items = {} image_dir = osp.join(self._dataset_dir, subset) if osp.isdir(image_dir): images = { osp.splitext(osp.relpath(p, image_dir))[0].replace('\\', '/'): p for p in find_images(image_dir, recursive=True) } else: images = {} landmarks_path = osp.join( self._dataset_dir, VggFace2Path.ANNOTATION_DIR, VggFace2Path.LANDMARKS_FILE + subset + '.csv') if osp.isfile(landmarks_path): with open(landmarks_path, encoding='utf-8') as content: landmarks_table = list(csv.DictReader(content)) for row in landmarks_table: item_id = row['NAME_ID'] label = None if '/' in item_id: label = _get_label(item_id) if item_id not in items: items[item_id] = DatasetItem(id=item_id, subset=subset, image=images.get( row['NAME_ID'])) annotations = items[item_id].annotations if [a for a in annotations if a.type == AnnotationType.points]: raise Exception("Item %s: an image can have only one " "set of landmarks" % item_id) if len([p for p in row if row[p] == '']) == 0 and len(row) == 11: annotations.append( Points([float(row[p]) for p in row if p != 'NAME_ID'], label=label)) elif label is not None: annotations.append(Label(label=label)) bboxes_path = osp.join(self._dataset_dir, VggFace2Path.ANNOTATION_DIR, VggFace2Path.BBOXES_FILE + subset + '.csv') if osp.isfile(bboxes_path): with open(bboxes_path, encoding='utf-8') as content: bboxes_table = list(csv.DictReader(content)) for row in bboxes_table: item_id = row['NAME_ID'] label = None if '/' in item_id: label = _get_label(item_id) if item_id not in items: items[item_id] = DatasetItem(id=item_id, subset=subset, image=images.get( row['NAME_ID'])) annotations = items[item_id].annotations if [a for a in annotations if a.type == AnnotationType.bbox]: raise Exception("Item %s: an image can have only one " "bbox" % item_id) if len([p for p in row if row[p] == '']) == 0 and len(row) == 5: annotations.append( Bbox(float(row['X']), float(row['Y']), float(row['W']), float(row['H']), label=label)) return items
def test_can_import_dataset_witn_numpy_files(self): expected_dataset = Dataset.from_iterable([ DatasetItem(id='000000001', image=np.ones((5, 5, 3)), annotations=[ Points([620.0, 394.0, 616.0, 269.0, 573.0, 185.0, 647.0, 188.0, 661.0, 221.0, 656.0, 231.0, 610.0, 187.0, 647.0, 176.0, 637.02, 189.818, 695.98, 108.182, 606.0, 217.0, 553.0, 161.0, 601.0, 167.0, 692.0, 185.0, 693.0, 240.0, 688.0, 313.0], [1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1], attributes={'center': [594.0, 257.0], 'scale': 3.021}, label=0, group=1), Bbox(615, 218.65, 288.4, 286.95, label=0, group=1) ] ), DatasetItem(id='000000002', image=np.ones((5, 5, 3)), annotations=[ Points([650.0, 424.0, 646.0, 309.0, 603.0, 215.0, 677.0, 218.0, 691.0, 251.0, 686.0, 261.0, 640.0, 217.0, 677.0, 216.0, 667.02, 219.818, 725.98, 138.182, 636.0, 247.0, 583.0, 191.0, 631.0, 197.0, 722.0, 215.0, 723.0, 270.0, 718.0, 343.0], [1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1], attributes={'center': [624.0, 287.0], 'scale': 3.7}, label=0, group=1), Bbox(101.1, 33.3, 113.9, 81.4, label=0, group=1) ] ), DatasetItem(id='000000003', image=np.ones((5, 5, 3)), annotations=[ Points([590.0, 364.0, 586.0, 239.0, 533.0, 155.0, 617.0, 158.0, 631.0, 191.0, 626.0, 201.0, 580.0, 157.0, 617.0, 146.0, 607.02, 159.818, 645.98, 68.182, 576.0, 187.0, 532.0, 131.0, 571.0, 137.0, 662.0, 155.0, 663.0, 210.0, 658.0, 283.0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1], attributes={'center': [564.0, 227.0], 'scale': 3.2}, label=0, group=1), Bbox(313.3, 512.43, 220.7, 121.57, label=0, group=1), Points([490.0, 264.0, 486.0, 139.0, 433.0, 55.0, 517.0, 58.0, 531.0, 91.0, 526.0, 101.0, 480.0, 57.0, 517.0, 46.0, 507.02, 59.818, 545.98, 8.182, 476.0, 87.0, 432.0, 31.0, 471.0, 37.0, 562.0, 55.0, 563.0, 110.0, 558.0, 183.0], [1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1], attributes={'center': [464.0, 127.0], 'scale': 2.65}, label=0, group=2), Points([690.0, 464.0, 686.0, 339.0, 633.0, 255.0, 717.0, 258.0, 731.0, 291.0, 726.0, 301.0, 680.0, 257.0, 717.0, 246.0, 707.02, 259.818, 745.98, 168.182, 676.0, 287.0, 632.0, 231.0, 671.0, 237.0, 762.0, 255.0, 763.0, 310.0, 758.0, 383.0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1], attributes={'center': [664.0, 327.0], 'scale': 3.9}, label=0, group=3) ] ) ], categories={ AnnotationType.label: LabelCategories.from_iterable(['human']), AnnotationType.points: PointsCategories.from_iterable( [(0, MPII_POINTS_LABELS, MPII_POINTS_JOINTS)]) }) dataset = Dataset.import_from(DUMMY_DATASET_DIR_WITH_NUMPY_FILES, 'mpii_json') compare_datasets(self, expected_dataset, dataset, require_images=True)
def test_can_save_and_load(self): source_dataset = Dataset.from_iterable( [ DatasetItem(id='label_0/1', subset='train', image=np.ones((8, 8, 3)), annotations=[ Bbox(0, 2, 4, 2, label=0), Points([ 3.2, 3.12, 4.11, 3.2, 2.11, 2.5, 3.5, 2.11, 3.8, 2.13 ], label=0), ]), DatasetItem(id='label_1/2', subset='train', image=np.ones((10, 10, 3)), annotations=[ Points([ 4.23, 4.32, 5.34, 4.45, 3.54, 3.56, 4.52, 3.51, 4.78, 3.34 ], label=1), ]), DatasetItem(id='label_2/3', subset='train', image=np.ones((8, 8, 3)), annotations=[Label(2)]), DatasetItem(id='label_3/4', subset='train', image=np.ones((10, 10, 3)), annotations=[ Bbox(0, 2, 4, 2, label=3), Points([ 3.2, 3.12, 4.11, 3.2, 2.11, 2.5, 3.5, 2.11, 3.8, 2.13 ], label=3), ]), DatasetItem(id='no_label/a/5', subset='train', image=np.ones((8, 8, 3)), annotations=[ Bbox(2, 2, 2, 2), ]), DatasetItem( id='no_label/label_0', subset='train', image=np.ones((8, 8, 3)), ), ], categories={ AnnotationType.label: LabelCategories.from_iterable([('label_%s' % i, 'class_%s' % i) for i in range(5)]), }) with TestDir() as test_dir: VggFace2Converter.convert(source_dataset, test_dir, save_images=True) parsed_dataset = Dataset.import_from(test_dir, 'vgg_face2') compare_datasets(self, source_dataset, parsed_dataset)
def _load_items(self, path): items = {} root_dir = osp.dirname(path) data = spio.loadmat(path, struct_as_record=False, squeeze_me=True).get('RELEASE', {}) data = getattr(data, 'annolist', []) for item in data: image = '' annotations = [] group_num = 1 image = getattr(item, 'image', '') if isinstance(image, spio.matlab.mio5_params.mat_struct): image = getattr(image, 'name', '') anno_values = getattr(item, 'annorect', []) if isinstance(anno_values, spio.matlab.mio5_params.mat_struct): anno_values = [anno_values] for val in anno_values: x1 = None x2 = None y1 = None y2 = None keypoints = {} is_visible = {} attributes = {} scale = getattr(val, 'scale', 0.0) if isinstance(scale, float): attributes['scale'] = scale objpos = getattr(val, 'objpos', None) if isinstance(objpos, spio.matlab.mio5_params.mat_struct): attributes['center'] = [ getattr(objpos, 'x', 0), getattr(objpos, 'y', 0) ] annopoints = getattr(val, 'annopoints', None) if isinstance(annopoints, spio.matlab.mio5_params.mat_struct) and \ not isinstance(getattr(annopoints, 'point'), spio.matlab.mio5_params.mat_struct): for point in getattr(annopoints, 'point'): point_id = getattr(point, 'id') keypoints[point_id] = [ getattr(point, 'x', 0), getattr(point, 'y', 0) ] is_visible[point_id] = getattr(point, 'is_visible', 1) if not isinstance(is_visible[point_id], int): is_visible[point_id] = 1 x1 = getattr(val, 'x1', None) if not isinstance(x1, (int, float)): x1 = None x2 = getattr(val, 'x2', None) if not isinstance(x2, (int, float)): x2 = None y1 = getattr(val, 'y1', None) if not isinstance(y1, (int, float)): y1 = None y2 = getattr(val, 'y2', None) if not isinstance(y2, (int, float)): y2 = None if keypoints: points = [0] * (2 * len(keypoints)) vis = [0] * len(keypoints) keypoints = sorted(keypoints.items(), key=lambda x: x[0]) for i, (key, point) in enumerate(keypoints): points[2 * i] = point[0] points[2 * i + 1] = point[1] vis[i] = is_visible.get(key, 1) annotations.append( Points(points, vis, label=0, group=group_num, attributes=attributes)) if x1 is not None and x2 is not None \ and y1 is not None and y2 is not None: annotations.append( Bbox(x1, y1, x2 - x1, y2 - y1, label=0, group=group_num)) group_num += 1 item_id = osp.splitext(image)[0] items[item_id] = DatasetItem( id=item_id, subset=self._subset, image=Image(path=osp.join(root_dir, image)), annotations=annotations) return items
def test_can_save_and_load(self): src_label_cat = LabelCategories(attributes={'occluded', 'common'}) for i in range(10): src_label_cat.add(str(i)) src_label_cat.items[2].attributes.update(['a1', 'a2', 'empty']) source_dataset = Dataset.from_iterable([ DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)), annotations=[ Polygon([0, 0, 4, 0, 4, 4], label=1, group=4, attributes={ 'occluded': True, 'common': 't' }), Points([1, 1, 3, 2, 2, 3], label=2, attributes={ 'a1': 'x', 'a2': 42, 'empty': '', 'unknown': 'bar' }), Label(1), Label(2, attributes={ 'a1': 'y', 'a2': 44 }), ] ), DatasetItem(id=1, subset='s1', annotations=[ PolyLine([0, 0, 4, 0, 4, 4], label=3, id=4, group=4), Bbox(5, 0, 1, 9, label=3, id=4, group=4), ] ), DatasetItem(id=2, subset='s2', image=np.ones((5, 10, 3)), annotations=[ Polygon([0, 0, 4, 0, 4, 4], z_order=1, label=3, group=4, attributes={ 'occluded': False }), PolyLine([5, 0, 9, 0, 5, 5]), # will be skipped as no label ] ), DatasetItem(id=3, subset='s3', image=Image( path='3.jpg', size=(2, 4))), ], categories={ AnnotationType.label: src_label_cat }) target_label_cat = LabelCategories( attributes={'occluded'}) # unable to represent a common attribute for i in range(10): target_label_cat.add(str(i), attributes={'common'}) target_label_cat.items[2].attributes.update(['a1', 'a2', 'empty', 'common']) target_dataset = Dataset.from_iterable([ DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)), annotations=[ Polygon([0, 0, 4, 0, 4, 4], label=1, group=4, attributes={ 'occluded': True, 'common': 't' }), Points([1, 1, 3, 2, 2, 3], label=2, attributes={ 'occluded': False, 'empty': '', 'a1': 'x', 'a2': '42' }), Label(1), Label(2, attributes={ 'a1': 'y', 'a2': '44' }), ], attributes={'frame': 0} ), DatasetItem(id=1, subset='s1', annotations=[ PolyLine([0, 0, 4, 0, 4, 4], label=3, group=4, attributes={ 'occluded': False }), Bbox(5, 0, 1, 9, label=3, group=4, attributes={ 'occluded': False }), ], attributes={'frame': 1} ), DatasetItem(id=2, subset='s2', image=np.ones((5, 10, 3)), annotations=[ Polygon([0, 0, 4, 0, 4, 4], z_order=1, label=3, group=4, attributes={ 'occluded': False }), ], attributes={'frame': 0} ), DatasetItem(id=3, subset='s3', image=Image( path='3.jpg', size=(2, 4)), attributes={'frame': 0}), ], categories={ AnnotationType.label: target_label_cat }) with TestDir() as test_dir: self._test_save_and_load(source_dataset, partial(CvatConverter.convert, save_images=True), test_dir, target_dataset=target_dataset)
def test_can_compare_projects(self): # just a smoke test label_categories1 = LabelCategories.from_iterable(['x', 'a', 'b', 'y']) mask_categories1 = MaskCategories.generate(len(label_categories1)) point_categories1 = PointsCategories() for index, _ in enumerate(label_categories1.items): point_categories1.add(index, ['cat1', 'cat2'], joints=[[0, 1]]) dataset1 = Dataset.from_iterable([ DatasetItem(id=100, subset='train', image=np.ones((10, 6, 3)), annotations=[ Caption('hello', id=1), Caption('world', id=2, group=5), Label(2, id=3, attributes={ 'x': 1, 'y': '2', }), Bbox(1, 2, 3, 4, label=0, id=4, z_order=1, attributes={ 'score': 1.0, }), Bbox(5, 6, 7, 8, id=5, group=5), Points([1, 2, 2, 0, 1, 1], label=0, id=5, z_order=4), Mask(label=3, id=5, z_order=2, image=np.ones((2, 3))), ]), DatasetItem(id=21, subset='train', annotations=[ Caption('test'), Label(2), Bbox(1, 2, 3, 4, label=2, id=42, group=42) ]), DatasetItem(id=2, subset='val', annotations=[ PolyLine([1, 2, 3, 4, 5, 6, 7, 8], id=11, z_order=1), Polygon([1, 2, 3, 4, 5, 6, 7, 8], id=12, z_order=4), ]), DatasetItem(id=42, subset='test', attributes={'a1': 5, 'a2': '42'}), DatasetItem(id=42), DatasetItem(id=43, image=Image(path='1/b/c.qq', size=(2, 4))), ], categories={ AnnotationType.label: label_categories1, AnnotationType.mask: mask_categories1, AnnotationType.points: point_categories1, }) label_categories2 = LabelCategories.from_iterable(['a', 'b', 'x', 'y']) mask_categories2 = MaskCategories.generate(len(label_categories2)) point_categories2 = PointsCategories() for index, _ in enumerate(label_categories2.items): point_categories2.add(index, ['cat1', 'cat2'], joints=[[0, 1]]) dataset2 = Dataset.from_iterable([ DatasetItem(id=100, subset='train', image=np.ones((10, 6, 3)), annotations=[ Caption('hello', id=1), Caption('world', id=2, group=5), Label(2, id=3, attributes={ 'x': 1, 'y': '2', }), Bbox(1, 2, 3, 4, label=1, id=4, z_order=1, attributes={ 'score': 1.0, }), Bbox(5, 6, 7, 8, id=5, group=5), Points([1, 2, 2, 0, 1, 1], label=0, id=5, z_order=4), Mask(label=3, id=5, z_order=2, image=np.ones((2, 3))), ]), DatasetItem(id=21, subset='train', annotations=[ Caption('test'), Label(2), Bbox(1, 2, 3, 4, label=3, id=42, group=42) ]), DatasetItem(id=2, subset='val', annotations=[ PolyLine([1, 2, 3, 4, 5, 6, 7, 8], id=11, z_order=1), Polygon([1, 2, 3, 4, 5, 6, 7, 8], id=12, z_order=4), ]), DatasetItem(id=42, subset='test', attributes={'a1': 5, 'a2': '42'}), DatasetItem(id=42), DatasetItem(id=43, image=Image(path='1/b/c.qq', size=(2, 4))), ], categories={ AnnotationType.label: label_categories2, AnnotationType.mask: mask_categories2, AnnotationType.points: point_categories2, }) with TestDir() as test_dir: with DiffVisualizer(save_dir=test_dir, comparator=DistanceComparator(iou_threshold=0.8), ) as visualizer: visualizer.save(dataset1, dataset2) self.assertNotEqual(0, os.listdir(osp.join(test_dir)))