cvat-ai--cvat
wehub-resource-sync
164 行
4.6 KiB
Python
164 行
4.6 KiB
Python
# Copyright (C) CVAT.ai Corporation
|
|
#
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from types import TracebackType
|
|
from unittest.mock import patch
|
|
|
|
from datumaro.components.dataset import StreamDataset
|
|
from datumaro.util.os_util import rmfile, rmtree
|
|
|
|
from cvat.apps.dataset_manager.bindings import CvatDataExtractor
|
|
from cvat.apps.dataset_manager.task import JobAnnotation
|
|
from cvat.apps.dataset_manager.util import current_function_name
|
|
|
|
|
|
class FileRemover:
|
|
def __init__(self, path: str, is_dir: bool = False):
|
|
self.path = path
|
|
self.is_dir = is_dir
|
|
|
|
def __enter__(self) -> str:
|
|
return self.path
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_value: BaseException | None,
|
|
traceback: TracebackType | None,
|
|
) -> None:
|
|
if self.is_dir:
|
|
try:
|
|
rmtree(self.path)
|
|
except unittest.SkipTest:
|
|
# Suppress skip test errors from git.util.rmtree
|
|
if not exc_type:
|
|
raise
|
|
else:
|
|
rmfile(self.path)
|
|
|
|
|
|
class TestDir(FileRemover):
|
|
"""
|
|
Creates a temporary directory for a test. Uses the name of
|
|
the test function to name the directory.
|
|
|
|
Usage:
|
|
|
|
.. code-block::
|
|
|
|
with TestDir() as test_dir:
|
|
...
|
|
"""
|
|
|
|
def __init__(self, path: str | None = None, frame_id: int = 2):
|
|
if not path:
|
|
prefix = f"temp_{current_function_name(frame_id)}-"
|
|
else:
|
|
prefix = None
|
|
self._prefix = prefix
|
|
|
|
super().__init__(path, is_dir=True)
|
|
|
|
def __enter__(self) -> str:
|
|
"""
|
|
Creates a test directory.
|
|
|
|
Returns: path to the directory
|
|
"""
|
|
|
|
path = self.path
|
|
|
|
if path is None:
|
|
path = tempfile.mkdtemp(dir=os.getcwd(), prefix=self._prefix)
|
|
self.path = path
|
|
else:
|
|
os.makedirs(path, exist_ok=False)
|
|
|
|
return path
|
|
|
|
|
|
def ensure_streaming_importers(cls):
|
|
original_import_from = StreamDataset.import_from
|
|
|
|
def mock_import_from(*args, **kwargs):
|
|
dataset = original_import_from(*args, **kwargs)
|
|
extractor = dataset._data._source
|
|
|
|
annotation_initiation_counter = 0
|
|
|
|
class ExtractorWrapper:
|
|
def __getattr__(self, attr):
|
|
return getattr(extractor, attr)
|
|
|
|
def __len__(self):
|
|
return len(extractor)
|
|
|
|
def __iter__(self):
|
|
nonlocal annotation_initiation_counter
|
|
for item in extractor:
|
|
yield item
|
|
if item.annotations_are_initialized:
|
|
annotation_initiation_counter += 1
|
|
# annotations should be initialized once per item and no more
|
|
assert annotation_initiation_counter <= len(extractor)
|
|
|
|
dataset._data._source = ExtractorWrapper()
|
|
|
|
return dataset
|
|
|
|
return patch.object(StreamDataset, "import_from", mock_import_from)(cls)
|
|
|
|
|
|
def ensure_extractors_efficiency(cls):
|
|
class MockExtractor(CvatDataExtractor):
|
|
def __init__(self, *args, **kwargs):
|
|
self.ann_init_counter = 0
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def _read_cvat_anno(self, *args, **kwargs):
|
|
self.ann_init_counter += 1
|
|
# annotations should be initialized once per item and no more
|
|
assert self.ann_init_counter <= len(self)
|
|
return super()._read_cvat_anno(*args, **kwargs)
|
|
|
|
class MockJobAnnotation(JobAnnotation):
|
|
def _init_shapes_from_db(self, *, streaming: bool = False):
|
|
super()._init_shapes_from_db(streaming=streaming)
|
|
|
|
assert isinstance(self.ir_data.shapes, list) != streaming
|
|
|
|
if streaming:
|
|
# should only generate shapes once
|
|
shapes = self.ir_data.shapes
|
|
|
|
already_iterated = False
|
|
|
|
class SinglePass:
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
nonlocal already_iterated
|
|
assert not already_iterated
|
|
try:
|
|
return next(shapes)
|
|
except StopIteration:
|
|
already_iterated = True
|
|
raise
|
|
|
|
self.ir_data.shapes = SinglePass()
|
|
|
|
cls = patch(
|
|
"cvat.apps.dataset_manager.bindings.CvatDataExtractor",
|
|
MockExtractor,
|
|
)(cls)
|
|
cls = patch(
|
|
"cvat.apps.dataset_manager.task.JobAnnotation",
|
|
MockJobAnnotation,
|
|
)(cls)
|
|
return cls
|