|
9 | 9 | from datumaro.components.annotation import AnnotationType |
10 | 10 | from datumaro.plugins.coco_format.importer import CocoImporter |
11 | 11 |
|
12 | | -from cvat.apps.dataset_manager.bindings import ( |
13 | | - GetCVATDataExtractor, |
14 | | - detect_dataset, |
15 | | - import_dm_annotations, |
16 | | -) |
| 12 | +from cvat.apps.dataset_manager.bindings import GetCVATDataExtractor, detect_dataset, \ |
| 13 | + import_dm_annotations |
17 | 14 | from cvat.apps.dataset_manager.util import make_zip_archive |
18 | 15 |
|
19 | 16 | from .registry import dm_env, exporter, importer |
20 | 17 |
|
21 | | - |
22 | | -@exporter(name="COCO", ext="ZIP", version="1.0") |
| 18 | +@exporter(name='COCO', ext='ZIP', version='1.0') |
23 | 19 | def _export(dst_file, temp_dir, instance_data, save_images=False): |
24 | 20 | with GetCVATDataExtractor(instance_data, include_images=save_images) as extractor: |
25 | 21 | dataset = Dataset.from_extractors(extractor, env=dm_env) |
26 | | - dataset.export( |
27 | | - temp_dir, "coco_instances", save_images=save_images, merge_images=True |
28 | | - ) |
| 22 | + dataset.export(temp_dir, 'coco_instances', save_images=save_images, |
| 23 | + merge_images=True) |
29 | 24 |
|
30 | 25 | make_zip_archive(temp_dir, dst_file) |
31 | 26 |
|
32 | | - |
33 | | -@importer(name="COCO", ext="JSON, ZIP", version="1.0") |
| 27 | +@importer(name='COCO', ext='JSON, ZIP', version='1.0') |
34 | 28 | def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs): |
35 | 29 | if zipfile.is_zipfile(src_file): |
36 | 30 | zipfile.ZipFile(src_file).extractall(temp_dir) |
37 | 31 | # We use coco importer because it gives better error message |
38 | | - detect_dataset(temp_dir, format_name="coco", importer=CocoImporter) |
39 | | - dataset = Dataset.import_from(temp_dir, "coco_instances", env=dm_env) |
| 32 | + detect_dataset(temp_dir, format_name='coco', importer=CocoImporter) |
| 33 | + dataset = Dataset.import_from(temp_dir, 'coco_instances', env=dm_env) |
40 | 34 | if load_data_callback is not None: |
41 | 35 | load_data_callback(dataset, instance_data) |
42 | 36 | import_dm_annotations(dataset, instance_data) |
43 | 37 | else: |
44 | | - dataset = Dataset.import_from(src_file.name, "coco_instances", env=dm_env) |
| 38 | + dataset = Dataset.import_from(src_file.name, |
| 39 | + 'coco_instances', env=dm_env) |
45 | 40 | import_dm_annotations(dataset, instance_data) |
46 | 41 |
|
47 | | - |
48 | | -@exporter(name="COCO Keypoints", ext="ZIP", version="1.0") |
| 42 | +@exporter(name='COCO Keypoints', ext='ZIP', version='1.0') |
49 | 43 | def _export(dst_file, temp_dir, instance_data, save_images=False): |
50 | 44 | with GetCVATDataExtractor(instance_data, include_images=save_images) as extractor: |
51 | 45 | dataset = Dataset.from_extractors(extractor, env=dm_env) |
52 | | - dataset.export( |
53 | | - temp_dir, |
54 | | - "coco_person_keypoints", |
55 | | - save_images=save_images, |
56 | | - merge_images=True, |
57 | | - ) |
| 46 | + dataset.export(temp_dir, 'coco_person_keypoints', save_images=save_images, |
| 47 | + merge_images=True) |
58 | 48 |
|
59 | 49 | make_zip_archive(temp_dir, dst_file) |
60 | 50 |
|
61 | | - |
62 | | -@importer(name="COCO Keypoints", ext="JSON, ZIP", version="1.0") |
| 51 | +@importer(name='COCO Keypoints', ext='JSON, ZIP', version='1.0') |
63 | 52 | def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs): |
64 | 53 | def remove_extra_annotations(dataset): |
65 | 54 | for item in dataset: |
66 | | - annotations = [ |
67 | | - ann for ann in item.annotations if ann.type != AnnotationType.bbox |
68 | | - ] |
| 55 | + annotations = [ann for ann in item.annotations |
| 56 | + if ann.type != AnnotationType.bbox] |
69 | 57 | item.annotations = annotations |
70 | 58 |
|
71 | 59 | if zipfile.is_zipfile(src_file): |
72 | 60 | zipfile.ZipFile(src_file).extractall(temp_dir) |
73 | 61 | # We use coco importer because it gives better error message |
74 | | - detect_dataset(temp_dir, format_name="coco", importer=CocoImporter) |
75 | | - dataset = Dataset.import_from(temp_dir, "coco_person_keypoints", env=dm_env) |
| 62 | + detect_dataset(temp_dir, format_name='coco', importer=CocoImporter) |
| 63 | + dataset = Dataset.import_from(temp_dir, 'coco_person_keypoints', env=dm_env) |
76 | 64 | remove_extra_annotations(dataset) |
77 | 65 | if load_data_callback is not None: |
78 | 66 | load_data_callback(dataset, instance_data) |
79 | 67 | import_dm_annotations(dataset, instance_data) |
80 | 68 | else: |
81 | | - dataset = Dataset.import_from( |
82 | | - src_file.name, "coco_person_keypoints", env=dm_env |
83 | | - ) |
| 69 | + dataset = Dataset.import_from(src_file.name, |
| 70 | + 'coco_person_keypoints', env=dm_env) |
84 | 71 | remove_extra_annotations(dataset) |
85 | 72 | import_dm_annotations(dataset, instance_data) |
0 commit comments