Streaming export for YOLO and COCO formats (#9084)

Author: Eldies

changelog.d/20250304_215108_dmitrii.lavrukhin_stream_export.md

@@ -0,0 +1,4 @@
+### Changed
+
+- Optimized memory usage on export with YOLO and COCO formats for tasks
+  (<https://github.com/cvat-ai/cvat/pull/9084>)

cvat/apps/dataset_manager/bindings.py

@@ -1661,11 +1661,8 @@ def __init__(
         self._user = self._load_user_info(instance_meta) if dimension == DimensionType.DIM_3D else {}
         self._dimension = dimension
         self._format_type = format_type
-
-        is_video = instance_meta['mode'] == 'interpolation'
-        ext = ''
-        if is_video:
-            ext = TaskFrameProvider.VIDEO_FRAME_EXT
+        self._instance_data = instance_data
+        self._include_images = include_images
 
         if dimension == DimensionType.DIM_3D or include_images:
             if isinstance(instance_data, TaskData):
@@ -1679,46 +1676,64 @@ def __init__(
                 {0: MediaSource(db_task)}
             )
 
-        dm_items: list[dm.DatasetItem] = []
-        for frame_data in instance_data.group_by_frame(include_empty=True):
+    def __iter__(self):
+        instance_meta = self._instance_data.meta[self._instance_data.META_FIELD]
+        is_video = instance_meta['mode'] == 'interpolation'
+        ext = ''
+        if is_video:
+            ext = TaskFrameProvider.VIDEO_FRAME_EXT
+
+        for frame_data in self._instance_data.group_by_frame(include_empty=True):
             dm_media_args = {
                 'path': frame_data.name + ext,
                 'ext': ext or frame_data.name.rsplit(osp.extsep, maxsplit=1)[1],
             }
-            if dimension == DimensionType.DIM_3D:
+            if self._dimension == DimensionType.DIM_3D:
                 dm_media: dm.PointCloud = self._media_provider.get_media_for_frame(
                     0, frame_data.id, **dm_media_args
                 )
 
-                if not include_images:
+                if not self._include_images:
                     dm_media_args["extra_images"] = [
                         dm.Image.from_file(path=osp.basename(image.path))
                         for image in dm_media.extra_images
                     ]
                     dm_media = dm.PointCloud.from_file(**dm_media_args)
             else:
                 dm_media_args['size'] = (frame_data.height, frame_data.width)
-                if include_images:
+                if self._include_images:
                     dm_media: dm.Image = self._media_provider.get_media_for_frame(
                         0, frame_data.idx, **dm_media_args
                     )
                 else:
                     dm_media = dm.Image.from_file(**dm_media_args)
 
+            # do not keep parsed lazy list data after this iteration
+            frame_data = frame_data._replace(
+                labeled_shapes=[
+                    (
+                        shape._replace(points=shape.points.lazy_copy())
+                        if isinstance(shape.points, LazyList) and not shape.points.is_parsed
+                        else shape
+                    )
+                    for shape in frame_data.labeled_shapes
+                ]
+            )
+
             dm_anno = self._read_cvat_anno(frame_data, instance_meta['labels'])
 
             dm_attributes = {'frame': frame_data.frame}
 
-            if dimension == DimensionType.DIM_2D:
+            if self._dimension == DimensionType.DIM_2D:
                 dm_item = dm.DatasetItem(
                     id=osp.splitext(frame_data.name)[0],
                     subset=frame_data.subset,
                     annotations=dm_anno,
                     media=dm_media,
                     attributes=dm_attributes,
                 )
-            elif dimension == DimensionType.DIM_3D:
-                if format_type == "sly_pointcloud":
+            elif self._dimension == DimensionType.DIM_3D:
+                if self._format_type == "sly_pointcloud":
                     dm_attributes["name"] = self._user["name"]
                     dm_attributes["createdAt"] = self._user["createdAt"]
                     dm_attributes["updatedAt"] = self._user["updatedAt"]
@@ -1735,9 +1750,10 @@ def __init__(
                     attributes=dm_attributes,
                 )
 
-            dm_items.append(dm_item)
+            yield dm_item
 
-        self._items = dm_items
+    def __len__(self):
+        return len(self._instance_data)
 
     def _read_cvat_anno(self, cvat_frame_anno: CommonData.Frame, labels: list):
         categories = self.categories()
@@ -1751,6 +1767,11 @@ def map_label(name, parent=''): return label_cat.find(name, parent)[0]
         return self.convert_annotations(cvat_frame_anno,
             label_attrs, map_label, self._format_type, self._dimension)
 
+    @property
+    def is_stream(self) -> bool:
+        return True
+
+
 class CVATProjectDataExtractor(dm.DatasetBase, CVATDataExtractorMixin):
     def __init__(
         self,

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
         )

cvat/apps/dataset_manager/formats/yolo.py

@@ -7,6 +7,7 @@
 from typing import Callable, Optional
 
 from datumaro.components.annotation import AnnotationType
+from datumaro.components.dataset import StreamDataset
 from datumaro.components.dataset_base import DatasetItem
 from datumaro.components.project import Dataset
 from pyunpack import Archive
@@ -36,7 +37,8 @@ def _export_common(
     **kwargs,
 ):
     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, format_name, save_media=save_images, **kwargs)
 
     make_zip_archive(temp_dir, dst_file)
@@ -109,7 +111,8 @@ def _export_yolo_ultralytics_oriented_boxes(*args, **kwargs):
 @exporter(name="Ultralytics YOLO Segmentation", ext="ZIP", version="1.0")
 def _export_yolo_ultralytics_segmentation(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 = dataset.transform("masks_to_polygons")
         dataset.export(temp_dir, "yolo_ultralytics_segmentation", save_media=save_images)
 

cvat/apps/engine/lazy_list.py

@@ -268,3 +268,19 @@ def __setstate__(self, state):
         self._parsed = state["parsed"]
         if self._parsed:
             self.extend(state["parsed_elements"])
+
+    def lazy_copy(self) -> list[T]:
+        """
+        Makes a copy without parsing elements.
+        Only works if elements have not been parsed yet.
+        """
+        assert not self._parsed
+        return LazyList(
+            string=self._string,
+            separator=self._separator,
+            converter=self._converter,
+        )
+
+    @property
+    def is_parsed(self):
+        return self._parsed