streaming coco

cvat-ai · Feb 13, 2025 · 1830690 · 1830690
1 parent 440b6f1
commit 1830690
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/cvat/apps/dataset_manager/formats/coco.py b/cvat/apps/dataset_manager/formats/coco.py
@@ -6,12 +6,13 @@
 import zipfile
 
 from datumaro.components.annotation import AnnotationType
-from datumaro.components.dataset import Dataset
+from datumaro.components.dataset import Dataset, StreamDataset
 from datumaro.plugins.data_formats.coco.importer import CocoImporter
 
 from cvat.apps.dataset_manager.bindings import (
     GetCVATDataExtractor,
     NoMediaInAnnotationFileError,
+    ProjectData,
     detect_dataset,
     import_dm_annotations,
 )
@@ -23,7 +24,8 @@
 @exporter(name="COCO", ext="ZIP", version="1.0")
 def _export(dst_file, temp_dir, instance_data, save_images=False):
     with GetCVATDataExtractor(instance_data, include_images=save_images) as extractor:
-        dataset = Dataset.from_extractors(extractor, env=dm_env)
+        dataset_cls = Dataset if isinstance(instance_data, ProjectData) else StreamDataset
+        dataset = dataset_cls.from_extractors(extractor, env=dm_env)
         dataset.export(temp_dir, "coco_instances", save_media=save_images, merge_images=False)
 
     make_zip_archive(temp_dir, dst_file)
@@ -50,7 +52,8 @@ def _import(src_file, temp_dir, instance_data, load_data_callback=None, **kwargs
 @exporter(name="COCO Keypoints", ext="ZIP", version="1.0")
 def _export(dst_file, temp_dir, instance_data, save_images=False):
     with GetCVATDataExtractor(instance_data, include_images=save_images) as extractor:
-        dataset = Dataset.from_extractors(extractor, env=dm_env)
+        dataset_cls = Dataset if isinstance(instance_data, ProjectData) else StreamDataset
+        dataset = dataset_cls.from_extractors(extractor, env=dm_env)
         dataset.export(
             temp_dir, "coco_person_keypoints", save_media=save_images, merge_images=False
         )