def test_inplace_save_writes_only_updated_data(self): with TestDir() as path: dataset = Dataset.from_iterable([ DatasetItem(id='frame1', annotations=[ Cuboid3d(position=[3.5, 9.8, 0.3], label=0, attributes={'track_id': 1}) ], point_cloud=self.pcd1, related_images=[self.image1], attributes={'frame': 0}) ], categories=['car', 'bus']) dataset.export(path, 'kitti_raw', save_images=True) dataset.put( DatasetItem('frame2', annotations=[ Cuboid3d(position=[1, 2, 0], label=1, attributes={'track_id': 1}) ], point_cloud=self.pcd2, related_images=[self.image2], attributes={'frame': 1})) dataset.remove('frame1') dataset.save(save_images=True) self.assertEqual({'frame2.png'}, set(os.listdir(osp.join(path, 'image_00', 'data')))) self.assertEqual( {'frame2.pcd'}, set(os.listdir(osp.join(path, 'velodyne_points', 'data'))))
def test_reindex_allows_single_annotations(self): source_dataset = Dataset.from_iterable([ DatasetItem(id='abc', annotations=[ Cuboid3d(position=[0.4, -1, 2.24], label=0), ]) ], categories=['dog']) expected_dataset = Dataset.from_iterable([ DatasetItem(id='abc', annotations=[ Cuboid3d(position=[0.4, -1, 2.24], label=0, attributes={ 'track_id': 1, 'occluded': False }), ], attributes={'frame': 0}) ], categories=['dog']) with TestDir() as test_dir: self._test_save_and_load(source_dataset, partial(KittiRawConverter.convert, reindex=True), test_dir, target_dataset=expected_dataset)
def test_can_save_and_load_with_pointcloud(self): source_dataset = Dataset.from_iterable([ DatasetItem(id=1, subset='test', point_cloud='1.pcd', related_images=[ Image(data=np.ones((5, 5, 3)), path='1/a.jpg'), Image(data=np.ones((5, 4, 3)), path='1/b.jpg'), Image(size=(5, 3), path='1/c.jpg'), '1/d.jpg', ], annotations=[ Cuboid3d([2, 2, 2], [1, 1, 1], [3, 3, 1], id=1, group=1, label=0, attributes={'x': True}) ]), ], categories=['label']) with TestDir() as test_dir: target_dataset = Dataset.from_iterable([ DatasetItem( id=1, subset='test', point_cloud=osp.join(test_dir, 'point_clouds', 'test', '1.pcd'), related_images=[ Image(data=np.ones((5, 5, 3)), path=osp.join(test_dir, 'related_images', 'test', '1', 'image_0.jpg')), Image(data=np.ones((5, 4, 3)), path=osp.join(test_dir, 'related_images', 'test', '1', 'image_1.jpg')), Image(size=(5, 3), path=osp.join(test_dir, 'related_images', 'test', '1', 'image_2.jpg')), osp.join(test_dir, 'related_images', 'test', '1', 'image_3.jpg'), ], annotations=[ Cuboid3d([2, 2, 2], [1, 1, 1], [3, 3, 1], id=1, group=1, label=0, attributes={'x': True}) ]), ], categories=['label']) self._test_save_and_load(source_dataset, partial(DatumaroConverter.convert, save_images=True), test_dir, target_dataset, compare=None, dimension=Dimensions.dim_3d)
def test_can_save_and_load_multiple_related_images(self): source_dataset = Dataset.from_iterable([ DatasetItem(id='a/d', annotations=[ Cuboid3d(position=[1, 2, 3], label=0, attributes={'track_id': 1}) ], point_cloud=self.pcd1, related_images=[self.image1, self.image2, self.image3], attributes={'frame': 3}), ], categories=['cat']) with TestDir() as test_dir: target_label_cat = LabelCategories(attributes={'occluded'}) target_label_cat.add('cat') target_dataset = Dataset.from_iterable([ DatasetItem( id='a/d', annotations=[ Cuboid3d(position=[1, 2, 3], label=0, attributes={ 'track_id': 1, 'occluded': False }) ], point_cloud=osp.join(test_dir, 'velodyne_points', 'data', 'a', 'd.pcd'), related_images=[ osp.join(test_dir, 'image_00', 'data', 'a', 'd.png'), osp.join(test_dir, 'image_01', 'data', 'a', 'd.png'), osp.join(test_dir, 'image_02', 'data', 'a', 'd.png'), ], attributes={'frame': 3}), ], categories={ AnnotationType.label: target_label_cat }) self._test_save_and_load(source_dataset, partial(KittiRawConverter.convert, save_images=True), test_dir, target_dataset=target_dataset, require_point_cloud=True) self.assertTrue( osp.isfile(osp.join(test_dir, 'image_00', 'data', 'a', 'd.png'))) self.assertTrue( osp.isfile(osp.join(test_dir, 'image_01', 'data', 'a', 'd.png'))) self.assertTrue( osp.isfile(osp.join(test_dir, 'image_02', 'data', 'a', 'd.png')))
def test_can_keep_undeclared_attributes(self): src_label_cat = LabelCategories(attributes={'occluded'}) src_label_cat.add('label1', attributes={'a'}) source_dataset = Dataset.from_iterable( [ DatasetItem(id='frame_000000', annotations=[ Cuboid3d(id=206, label=0, position=[320.86, 979.18, 1.04], attributes={ 'track_id': 1, 'occluded': False, 'a': 5, 'undeclared': 'y' }), ], attributes={'frame': 0}), ], categories={AnnotationType.label: src_label_cat}) with TestDir() as test_dir: self._test_save_and_load( source_dataset, partial(SuperviselyPointCloudConverter.convert, save_images=True, allow_undeclared_attrs=True), test_dir, ignored_attrs=['description'])
def _parse_track(cls, track_id, track, categories): common_attrs = { k: cls._parse_attr(v) for k, v in track['attributes'].items() } scale = [track['scale'][k] for k in ['w', 'h', 'l']] label = categories[AnnotationType.label].find(track['label'])[0] kf_occluded = False for shape in track['shapes']: occluded = shape['occluded'] in { OcclusionStates.FULLY, OcclusionStates.PARTLY} if shape['occluded_kf']: kf_occluded = occluded elif shape['occluded'] == OcclusionStates.OCCLUSION_UNSET: occluded = kf_occluded if shape['truncated'] in {TruncationStates.OUT_IMAGE, TruncationStates.BEHIND_IMAGE}: # skip these frames continue local_attrs = { k: cls._parse_attr(v) for k, v in shape['attributes'].items() } local_attrs['occluded'] = occluded local_attrs['track_id'] = track_id attrs = dict(common_attrs) attrs.update(local_attrs) position = [shape['points'][k] for k in ['tx', 'ty', 'tz']] rotation = [shape['points'][k] for k in ['rx', 'ry', 'rz']] yield Cuboid3d(position, rotation, scale, label=label, attributes=attrs)
def test_inplace_save_writes_only_updated_data(self): with TestDir() as path: dataset = Dataset.from_iterable([ DatasetItem(id='frame1', annotations=[ Cuboid3d(id=215, position=[320.59, 979.48, 1.03], label=0) ], point_cloud=self.pcd1, related_images=[self.image1], attributes={'frame': 0}) ], categories=['car', 'bus']) dataset.export(path, 'sly_pointcloud', save_images=True) dataset.put( DatasetItem(id='frame2', annotations=[ Cuboid3d(id=216, position=[0.59, 14.41, -0.61], label=1) ], point_cloud=self.pcd2, related_images=[self.image2], attributes={'frame': 1})) dataset.remove('frame1') dataset.save(save_images=True) self.assertEqual({'frame2.pcd.json'}, set(os.listdir(osp.join(path, 'ds0', 'ann')))) self.assertEqual({'frame2.pcd'}, set( os.listdir(osp.join(path, 'ds0', 'pointcloud')))) self.assertTrue( osp.isfile( osp.join(path, 'ds0', 'related_images', 'frame2_pcd', 'img1.png'))) self.assertFalse( osp.isfile( osp.join(path, 'ds0', 'related_images', 'frame1_pcd', 'img2.png')))
def test_can_save_and_load_attributes(self): source_dataset = Dataset.from_iterable([ DatasetItem(id='0000000000', annotations=[ Cuboid3d(position=[13.54, -9.41, 0.24], label=0, attributes={ 'track_id': 1, 'occluded': True, 'a': 'w', 'b': 5 }) ], attributes={'frame': 0}) ], categories=['cat']) target_label_cat = LabelCategories(attributes={'occluded'}) target_label_cat.add('cat', attributes=['a', 'b']) target_dataset = Dataset.from_iterable( [ DatasetItem(id='0000000000', annotations=[ Cuboid3d(position=[13.54, -9.41, 0.24], label=0, attributes={ 'track_id': 1, 'occluded': True, 'a': 'w', 'b': 5 }) ], attributes={'frame': 0}) ], categories={AnnotationType.label: target_label_cat}) with TestDir() as test_dir: self._test_save_and_load(source_dataset, partial(KittiRawConverter.convert, allow_attrs=True), test_dir, target_dataset=target_dataset)
def test_requires_track_id(self): source_dataset = Dataset.from_iterable([ DatasetItem(id='abc', annotations=[ Cuboid3d(position=[0.4, -1, 2.24], label=0), ]) ], categories=['dog']) with TestDir() as test_dir: with self.assertRaisesRegex(Exception, 'track_id'): KittiRawConverter.convert(source_dataset, test_dir)
def test_can_save_and_load(self): src_label_cat = LabelCategories(attributes={'occluded'}) src_label_cat.add('car', attributes=['x']) src_label_cat.add('bus') source_dataset = Dataset.from_iterable( [ DatasetItem(id='frame_1', annotations=[ Cuboid3d(id=206, label=0, position=[320.86, 979.18, 1.04], attributes={ 'occluded': False, 'track_id': 1, 'x': 1 }), Cuboid3d(id=207, label=1, position=[318.19, 974.65, 1.29], attributes={ 'occluded': True, 'track_id': 2 }), ], point_cloud=self.pcd1, attributes={ 'frame': 0, 'description': 'zzz' }), DatasetItem(id='frm2', annotations=[ Cuboid3d(id=208, label=1, position=[23.04, 8.75, -0.78], attributes={ 'occluded': False, 'track_id': 2 }) ], point_cloud=self.pcd2, related_images=[self.image2], attributes={'frame': 1}), ], categories={AnnotationType.label: src_label_cat}) with TestDir() as test_dir: target_label_cat = LabelCategories(attributes={'occluded'}) target_label_cat.add('car', attributes=['x']) target_label_cat.add('bus') target_dataset = Dataset.from_iterable([ DatasetItem(id='frame_1', annotations=[ Cuboid3d(id=206, label=0, position=[320.86, 979.18, 1.04], attributes={ 'occluded': False, 'track_id': 1, 'x': 1 }), Cuboid3d(id=207, label=1, position=[318.19, 974.65, 1.29], attributes={ 'occluded': True, 'track_id': 2 }), ], point_cloud=osp.join(test_dir, 'ds0', 'pointcloud', 'frame_1.pcd'), attributes={ 'frame': 0, 'description': 'zzz' }), DatasetItem(id='frm2', annotations=[ Cuboid3d(id=208, label=1, position=[23.04, 8.75, -0.78], attributes={ 'occluded': False, 'track_id': 2 }), ], point_cloud=osp.join(test_dir, 'ds0', 'pointcloud', 'frm2.pcd'), related_images=[ osp.join(test_dir, 'ds0', 'related_images', 'frm2_pcd', 'img1.png') ], attributes={ 'frame': 1, 'description': '' }) ], categories={ AnnotationType.label: target_label_cat }) self._test_save_and_load( source_dataset, partial(SuperviselyPointCloudConverter.convert, save_images=True), test_dir, target_dataset=target_dataset, require_point_cloud=True)
def _parse(cls, rootpath): mapping = parse_json_file( osp.join(rootpath, PointCloudPath.KEY_ID_FILE)) meta = parse_json_file(osp.join(rootpath, PointCloudPath.META_FILE)) label_cat = LabelCategories() for label in meta.get('classes', []): label_cat.add(label['title']) tags = {} for tag in meta.get('tags', []): # See reference at: # https://github.com/supervisely/supervisely/blob/047e52ebe407cfee61464c1bd0beb9c906892253/supervisely_lib/annotation/tag_meta.py#L139 tags[tag['name']] = tag applicable_to = tag.get('applicable_type', 'all') if applicable_to == 'imagesOnly': continue # an image attribute elif applicable_to not in {'all', 'objectsOnly'}: raise Exception("Unexpected tag 'applicable_type' value '%s'" % \ applicable_to) applicable_classes = tag.get('classes', []) if not applicable_classes: label_cat.attributes.add(tag['name']) else: for label_name in applicable_classes: _, label = label_cat.find(label_name) if label is None: raise Exception("Unknown class for tag '%s'" % \ label_name) label.attributes.add(tag['name']) categories = {AnnotationType.label: label_cat} def _get_label_attrs(label_id): attrs = set(label_cat.attributes) attrs.update(label_cat[label_id].attributes) return attrs def _parse_tag(tag): if tag['value'] == 'true': value = True elif tag['value'] == 'false': value = False else: value = tag['value'] return value ann_dir = osp.join(rootpath, PointCloudPath.BASE_DIR, PointCloudPath.ANNNOTATION_DIR) items = {} for ann_file in iglob(osp.join(ann_dir, '**', '*.json'), recursive=True): ann_data = parse_json_file(ann_file) objects = {} for obj in ann_data['objects']: obj['id'] = mapping['objects'][obj['key']] objects[obj['key']] = obj frame_attributes = {'description': ann_data.get('description', '')} for tag in ann_data['tags']: frame_attributes[tag['name']] = _parse_tag(tag) frame = mapping['videos'][ann_data['key']] frame_desc = items.setdefault( frame, { 'name': osp.splitext(osp.relpath(ann_file, ann_dir))[0], 'annotations': [], 'attributes': frame_attributes, }) for figure in ann_data['figures']: geometry = { dst_field: [ float(figure['geometry'][src_field][axis]) for axis in ['x', 'y', 'z'] ] for src_field, dst_field in { 'position': 'position', 'rotation': 'rotation', 'dimensions': 'scale' }.items() } ann_id = mapping['figures'][figure['key']] obj = objects[figure['objectKey']] label = categories[AnnotationType.label].find( obj['classTitle'])[0] attributes = {} attributes['track_id'] = obj['id'] for tag in obj.get('tags', []): attributes[tag['name']] = _parse_tag(tag) for attr in _get_label_attrs(label): if attr in attributes: continue if tags[attr]['value_type'] == 'any_string': value = '' elif tags[attr]['value_type'] == 'oneof_string': value = (tags[attr]['values'] or [''])[0] elif tags[attr]['value_type'] == 'any_number': value = 0 else: value = None attributes[attr] = value shape = Cuboid3d(**geometry, label=label, id=ann_id, attributes=attributes) frame_desc['annotations'].append(shape) return items, categories
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_convert_to_kitti_raw(self): with TestDir() as test_dir: export_dir = osp.join(test_dir, 'export_dir') expected_label_cat = LabelCategories(attributes={'occluded'}) expected_label_cat.add('bus') expected_label_cat.add('car') expected_dataset = Dataset.from_iterable([ DatasetItem(id='0000000000', annotations=[ Cuboid3d(position=[1, 2, 3], scale=[7.95, -3.62, -1.03], label=1, attributes={'occluded': False, 'track_id': 1}), Cuboid3d(position=[1, 1, 0], scale=[8.34, 23.01, -0.76], label=0, attributes={'occluded': False, 'track_id': 2}) ], point_cloud=osp.join(export_dir, 'ds0', 'pointcloud', '0000000000.pcd'), related_images=[osp.join(export_dir, 'ds0', 'related_images', '0000000000_pcd', '0000000000.png') ], attributes={'frame': 0, 'description': ''} ), DatasetItem(id='0000000001', annotations=[ Cuboid3d(position=[0, 1, 0], scale=[8.34, 23.01, -0.76], rotation=[1, 1, 3], label=0, attributes={'occluded': True, 'track_id': 2}) ], point_cloud=osp.join(export_dir, 'ds0', 'pointcloud', '0000000001.pcd'), related_images=[osp.join(export_dir, 'ds0', 'related_images', '0000000001_pcd', '0000000001.png') ], attributes={'frame': 1, 'description': ''} ), DatasetItem(id='0000000002', annotations=[ Cuboid3d(position=[1, 2, 3], scale=[-9.41, 13.54, 0.24], label=1, attributes={'occluded': False, 'track_id': 3}) ], point_cloud=osp.join(export_dir, 'ds0', 'pointcloud', '0000000002.pcd'), related_images=[osp.join(export_dir, 'ds0', 'related_images', '0000000002_pcd', '0000000002.png') ], attributes={'frame': 2, 'description': ''} ), ], categories={AnnotationType.label: expected_label_cat}) run(self, 'convert', '-if', 'kitti_raw', '-i', DUMMY_DATASET_DIR, '-f', 'sly_pointcloud', '-o', export_dir, '--', '--save-images') parsed_dataset = Dataset.import_from(export_dir, format='sly_pointcloud') compare_datasets_3d(self, expected_dataset, parsed_dataset, require_point_cloud=True)
def test_can_save_and_load_with_meta_file(self): source_dataset = Dataset.from_iterable([ DatasetItem(id='0000000000', annotations=[ Cuboid3d(position=[13.54, -9.41, 0.24], label=0, attributes={ 'occluded': False, 'track_id': 1 }) ], point_cloud=self.pcd1, related_images=[self.image1], attributes={'frame': 0}), DatasetItem( id='0000000001', annotations=[ Cuboid3d(position=[1.4, 2.1, 1.4], label=1, attributes={'track_id': 2}) ], ) ], categories=['cat', 'dog']) with TestDir() as test_dir: target_label_cat = LabelCategories(attributes={'occluded'}) target_label_cat.add('cat') target_label_cat.add('dog') target_dataset = Dataset.from_iterable([ DatasetItem(id='0000000000', annotations=[ Cuboid3d(position=[13.54, -9.41, 0.24], label=0, attributes={ 'occluded': False, 'track_id': 1 }) ], point_cloud=osp.join(test_dir, 'velodyne_points', 'data', '0000000000.pcd'), related_images=[ osp.join(test_dir, 'image_00', 'data', '0000000000.png') ], attributes={'frame': 0}), DatasetItem(id='0000000001', annotations=[ Cuboid3d(position=[1.4, 2.1, 1.4], label=1, attributes={ 'occluded': False, 'track_id': 2 }) ], attributes={'frame': 1}) ], categories={ AnnotationType.label: target_label_cat }) self._test_save_and_load(source_dataset, partial(KittiRawConverter.convert, save_images=True, save_dataset_meta=True), test_dir, target_dataset=target_dataset, require_point_cloud=True) self.assertTrue(osp.isfile(osp.join(test_dir, 'dataset_meta.json')))
def test_can_load(self): pcd1 = osp.join(DUMMY_DATASET_DIR, 'ds0', 'pointcloud', 'frame1.pcd') pcd2 = osp.join(DUMMY_DATASET_DIR, 'ds0', 'pointcloud', 'frame2.pcd') image1 = osp.join(DUMMY_DATASET_DIR, 'ds0', 'related_images', 'frame1_pcd', 'img2.png') image2 = osp.join(DUMMY_DATASET_DIR, 'ds0', 'related_images', 'frame2_pcd', 'img1.png') label_cat = LabelCategories(attributes={'tag1', 'tag3'}) label_cat.add('car') label_cat.add('bus') expected_dataset = Dataset.from_iterable( [ DatasetItem(id='frame1', annotations=[ Cuboid3d(id=755220128, label=0, position=[0.47, 0.23, 0.79], scale=[0.01, 0.01, 0.01], attributes={ 'track_id': 231825, 'tag1': 'fd', 'tag3': '4s' }), Cuboid3d(id=755337225, label=0, position=[0.36, 0.64, 0.93], scale=[0.01, 0.01, 0.01], attributes={ 'track_id': 231831, 'tag1': 'v12', 'tag3': '' }), ], point_cloud=pcd1, related_images=[image1], attributes={ 'frame': 0, 'description': '', 'tag1': '25dsd', 'tag2': 65 }), DatasetItem(id='frame2', annotations=[ Cuboid3d(id=216, label=1, position=[0.59, 14.41, -0.61], attributes={ 'track_id': 36, 'tag1': '', 'tag3': '' }) ], point_cloud=pcd2, related_images=[image2], attributes={ 'frame': 1, 'description': '' }), ], categories={AnnotationType.label: label_cat}) parsed_dataset = Dataset.import_from(DUMMY_DATASET_DIR, 'sly_pointcloud') compare_datasets_3d(self, expected_dataset, parsed_dataset, require_point_cloud=True)
def test_can_convert_to_kitti_raw(self): with TestDir() as test_dir: export_dir = osp.join(test_dir, 'export_dir') expected_label_cat = LabelCategories(attributes={'occluded'}) expected_label_cat.add('bus', attributes={'tag1', 'tag3'}) expected_label_cat.add('car', attributes={'tag1', 'tag3'}) expected_dataset = Dataset.from_iterable([ DatasetItem(id='frame1', annotations=[ Cuboid3d(label=1, position=[0.47, 0.23, 0.79], scale=[0.01, 0.01, 0.01], attributes={ 'track_id': 2, 'tag1': 'fd', 'tag3': '4s', 'occluded': False }), Cuboid3d(label=1, position=[0.36, 0.64, 0.93], scale=[0.01, 0.01, 0.01], attributes={ 'track_id': 3, 'tag1': 'v12', 'tag3': '', 'occluded': False }), ], point_cloud=osp.join(export_dir, 'velodyne_points', 'data', 'frame1.pcd'), related_images=[ osp.join(export_dir, 'image_00', 'data', 'frame1.png') ], attributes={'frame': 0}), DatasetItem(id='frame2', annotations=[ Cuboid3d(label=0, position=[0.59, 14.41, -0.61], attributes={ 'track_id': 1, 'tag1': '', 'tag3': '', 'occluded': False }) ], point_cloud=osp.join(export_dir, 'velodyne_points', 'data', 'frame2.pcd'), related_images=[ osp.join(export_dir, 'image_00', 'data', 'frame2.png') ], attributes={'frame': 1}), ], categories={ AnnotationType.label: expected_label_cat }) run(self, 'convert', '-if', 'sly_pointcloud', '-i', DUMMY_DATASET_DIR, '-f', 'kitti_raw', '-o', export_dir, '--', '--save-images', '--allow-attrs') parsed_dataset = Dataset.import_from(export_dir, format='kitti_raw') compare_datasets_3d(self, expected_dataset, parsed_dataset, require_point_cloud=True)
def test_can_load(self): pcd1 = osp.join(DUMMY_DATASET_DIR, 'velodyne_points', 'data', '0000000000.pcd') pcd2 = osp.join(DUMMY_DATASET_DIR, 'velodyne_points', 'data', '0000000001.pcd') pcd3 = osp.join(DUMMY_DATASET_DIR, 'velodyne_points', 'data', '0000000002.pcd') image1 = osp.join(DUMMY_DATASET_DIR, 'IMAGE_00', 'data', '0000000000.png') image2 = osp.join(DUMMY_DATASET_DIR, 'IMAGE_00', 'data', '0000000001.png') image3 = osp.join(DUMMY_DATASET_DIR, 'IMAGE_00', 'data', '0000000002.png') expected_label_cat = LabelCategories(attributes={'occluded'}) expected_label_cat.add('bus') expected_label_cat.add('car') expected_dataset = Dataset.from_iterable( [ DatasetItem(id='0000000000', annotations=[ Cuboid3d(position=[1, 2, 3], scale=[7.95, -3.62, -1.03], label=1, attributes={ 'occluded': False, 'track_id': 1 }), Cuboid3d(position=[1, 1, 0], scale=[8.34, 23.01, -0.76], label=0, attributes={ 'occluded': False, 'track_id': 2 }) ], point_cloud=pcd1, related_images=[image1], attributes={'frame': 0}), DatasetItem(id='0000000001', annotations=[ Cuboid3d(position=[0, 1, 0], scale=[8.34, 23.01, -0.76], rotation=[1, 1, 3], label=0, attributes={ 'occluded': True, 'track_id': 2 }) ], point_cloud=pcd2, related_images=[image2], attributes={'frame': 1}), DatasetItem(id='0000000002', annotations=[ Cuboid3d(position=[1, 2, 3], scale=[-9.41, 13.54, 0.24], label=1, attributes={ 'occluded': False, 'track_id': 3 }) ], point_cloud=pcd3, related_images=[image3], attributes={'frame': 2}) ], categories={AnnotationType.label: expected_label_cat}) parsed_dataset = Dataset.import_from(DUMMY_DATASET_DIR, 'kitti_raw') compare_datasets_3d(self, expected_dataset, parsed_dataset, require_point_cloud=True)
def test_dataset(self): label_categories = LabelCategories(attributes={'a', 'b', 'score'}) for i in range(5): label_categories.add('cat' + str(i), attributes={'x', 'y'}) mask_categories = MaskCategories( generate_colormap(len(label_categories.items))) points_categories = PointsCategories() for index, _ in enumerate(label_categories.items): points_categories.add(index, ['cat1', 'cat2'], joints=[[0, 1]]) return 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=4, id=4, z_order=1, attributes={ 'score': 1.0, }), Bbox(5, 6, 7, 8, id=5, group=5, attributes={ 'a': 1.5, 'b': 'text', }), Points([1, 2, 2, 0, 1, 1], label=0, id=5, z_order=4, attributes={ 'x': 1, 'y': '2', }), Mask(label=3, id=5, z_order=2, image=np.ones((2, 3)), attributes={ 'x': 1, 'y': '2', }), ]), DatasetItem(id=21, subset='train', annotations=[ Caption('test'), Label(2), Bbox(1, 2, 3, 4, label=5, 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=1, subset='test', annotations=[ Cuboid3d([1.0, 2.0, 3.0], [2.0, 2.0, 4.0], [1.0, 3.0, 4.0], id=6, label=0, attributes={'occluded': True}, group=6) ]), 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_categories, AnnotationType.mask: mask_categories, AnnotationType.points: points_categories, })