From 094dc27e2238cb3f3d32706c0aad96f8b738d32f Mon Sep 17 00:00:00 2001 From: Dmitrii Lavrukhin Date: Mon, 10 Feb 2025 16:43:16 +0400 Subject: [PATCH] streaming export for task or job data, and for yolo format --- cvat/apps/dataset_manager/bindings.py | 39 ++++++++++++++--------- cvat/apps/dataset_manager/formats/yolo.py | 7 ++-- 2 files changed, 29 insertions(+), 17 deletions(-) diff --git a/cvat/apps/dataset_manager/bindings.py b/cvat/apps/dataset_manager/bindings.py index 6015715b9c2..60388725a8b 100644 --- a/cvat/apps/dataset_manager/bindings.py +++ b/cvat/apps/dataset_manager/bindings.py @@ -1660,11 +1660,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): @@ -1678,15 +1675,21 @@ 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 } - 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 @@ -1694,7 +1697,7 @@ def __init__( 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 ) @@ -1705,7 +1708,7 @@ def __init__( 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, @@ -1713,8 +1716,8 @@ def __init__( 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"] @@ -1731,9 +1734,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() @@ -1747,6 +1751,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, diff --git a/cvat/apps/dataset_manager/formats/yolo.py b/cvat/apps/dataset_manager/formats/yolo.py index 4bb696c4d62..fae68629052 100644 --- a/cvat/apps/dataset_manager/formats/yolo.py +++ b/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)