cvat-cli 2.67.0__tar.gz → 2.69.0__tar.gz

{cvat_cli-2.67.0 → cvat_cli-2.69.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cvat-cli
-Version: 2.67.0
+Version: 2.69.0
 Summary: Command-line client for CVAT
 Author-email: "CVAT.ai Corporation" <support@cvat.ai>
 License-Expression: MIT
@@ -9,7 +9,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
-Requires-Dist: cvat-sdk==2.67.0
+Requires-Dist: cvat-sdk==2.69.0
 Requires-Dist: attrs>=24.2.0
 Requires-Dist: Pillow>=10.3.0
 

{cvat_cli-2.67.0 → cvat_cli-2.69.0}/pyproject.toml

@@ -16,7 +16,7 @@ classifiers = [
 ]
 requires-python = ">=3.10"
 dependencies = [
-    "cvat-sdk==2.67.0",
+    "cvat-sdk==2.69.0",
 
     "attrs>=24.2.0",
     "Pillow>=10.3.0",

{cvat_cli-2.67.0 → cvat_cli-2.69.0}/src/cvat_cli/_internal/agent.py

@@ -14,35 +14,35 @@ import shutil
 import tempfile
 import threading
 import time
-from collections import OrderedDict
-from collections.abc import Callable, Generator, Iterator, Sequence
+from collections.abc import Generator, Iterator
 from datetime import datetime, timedelta, timezone
 from http import HTTPStatus
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, TypeAlias
+from typing import TYPE_CHECKING
 
 import attrs
 import cvat_sdk.auto_annotation as cvataa
 import cvat_sdk.datasets as cvatds
-import PIL.Image
 import urllib3.exceptions
-from cvat_sdk import Client, models
-from cvat_sdk.auto_annotation.driver import (
-    _AnnotationMapper,
-    _DetectionFunctionContextImpl,
-    _SpecNameMapping,
-)
+from cvat_sdk import Client
 from cvat_sdk.datasets.caching import make_cache_manager
 from cvat_sdk.exceptions import ApiException
 
+from .agent_driver import (
+    AgentFunctionDriver,
+    BadArError,
+    IncompatibleFunctionError,
+    set_worker_current_function,
+    worker_current_function,
+)
+from .agent_driver_detection import AgentDetectionFunctionDriver
+from .agent_driver_tracking import AgentTrackingFunctionDriver, TrackingStateIdGenerator
 from .common import CriticalError, FunctionLoader
 
 if TYPE_CHECKING:
     from _typeshed import SupportsReadline
 
 FUNCTION_PROVIDER_NATIVE = "native"
-FUNCTION_KIND_DETECTOR = "detector"
-FUNCTION_KIND_TRACKER = "tracker"
 REQUEST_CATEGORY_BATCH = "batch"
 REQUEST_CATEGORY_INTERACTIVE = "interactive"
 
@@ -55,8 +55,6 @@ _DEFAULT_RETRY_DELAY = timedelta(seconds=5)
 
 _UPDATE_INTERVAL = timedelta(seconds=30)
 
-_MAX_AGE_OF_TRACKING_STATE = timedelta(hours=8)
-
 
 class _ExponentialBackoff:
     def __init__(self, max_delay: timedelta, current_delay: timedelta) -> None:
@@ -72,7 +70,7 @@ class _ExponentialBackoff:
         return delay
 
 
-class _RecoverableExecutor:
+class RecoverableExecutor:
     # A wrapper around ProcessPoolExecutor that recreates the underlying
     # executor when a worker crashes.
     def __init__(self, initializer, initargs):
@@ -107,146 +105,23 @@ class _RecoverableExecutor:
             raise
 
 
-_TrackingStateIdGenerator: TypeAlias = Callable[[], str]
-
-
 def _default_tracking_state_id_generator() -> str:
     # This is defined as a separate function so that tests can monkeypatch it
     # in order to get deterministic state IDs.
     return secrets.token_urlsafe(32)
 
 
-_current_function: cvataa.AutoAnnotationFunction
-_tracking_states: _TrackingStateContainer
-_tracking_state_id_generator: _TrackingStateIdGenerator
-
-
-@attrs.define
-class _ExtendedTrackingState:
-    inner_state: Any  # the state produced by the AA function
-    original_shape_type: str
-    original_task_id: int
-    original_image_dims: tuple[int, int]
-    last_accessed_at: datetime = attrs.field(factory=lambda: datetime.now(tz=timezone.utc))
-
-
-class _TrackingStateContainer:
-    def __init__(self):
-        self._id_to_ext_state: OrderedDict[str, _ExtendedTrackingState] = OrderedDict()
-
-    def store(self, state: Any, shape_type: str, task_id: int, image_dims: tuple[int, int]) -> str:
-        state_id = _tracking_state_id_generator()
-        self._id_to_ext_state[state_id] = _ExtendedTrackingState(
-            inner_state=state,
-            original_shape_type=shape_type,
-            original_task_id=task_id,
-            original_image_dims=image_dims,
-        )
-        return state_id
-
-    def retrieve(self, state_id: str, task_id: int, image_dims: tuple[int, int]) -> Any:
-        ext_state = self._id_to_ext_state.get(state_id)
-
-        if not ext_state:
-            raise _BadArError(f"Tracking state {state_id!r} not found - possibly expired")
-
-        if ext_state.original_task_id != task_id:
-            # This is a defense-in-depth measure. State IDs are supposed to be unguessable,
-            # but even if an attacker manages to obtain one, they will not be able to use it
-            # to get any information about a task they don't have access to.
-            raise _BadArError(f"Tracking state {state_id!r} is not for task #{task_id}")
-
-        if image_dims != ext_state.original_image_dims:
-            raise _BadArError(f"Image sizes of the start frame and the current frame are different")
-
-        ext_state.last_accessed_at = datetime.now(tz=timezone.utc)
-        self._id_to_ext_state.move_to_end(state_id)
-
-        return ext_state.inner_state, ext_state.original_shape_type
-
-    def prune(self) -> None:
-        cutoff = datetime.now(tz=timezone.utc) - _MAX_AGE_OF_TRACKING_STATE
-
-        while (
-            self._id_to_ext_state
-            and next(iter(self._id_to_ext_state.values())).last_accessed_at < cutoff
-        ):
-            self._id_to_ext_state.popitem(last=False)
-
-
-def _worker_init(function_loader: FunctionLoader, state_id_generator):
-    global _current_function
-    _current_function = function_loader.load()
-
-    if isinstance(_current_function.spec, cvataa.TrackingFunctionSpec):
-        global _tracking_states
-        _tracking_states = _TrackingStateContainer()
+def _worker_init(
+    function_loader: FunctionLoader, state_id_generator: TrackingStateIdGenerator
+) -> None:
+    current_function = function_loader.load()
+    set_worker_current_function(current_function)
 
-        global _tracking_state_id_generator
-        _tracking_state_id_generator = state_id_generator
+    get_function_driver_class(current_function.spec).init_worker(state_id_generator)
 
 
 def _worker_job_get_function_spec():
-    return _current_function.spec
-
-
-def _worker_job_detect(
-    context: _DetectionFunctionContextImpl, image: PIL.Image.Image
-) -> list[cvataa.DetectionAnnotation]:
-    return _current_function.detect(context, image)
-
-
-def _worker_job_init_tracking(
-    task_id: int,
-    image: PIL.Image.Image,
-    shapes: list[cvataa.TrackableShape],
-) -> list[str]:
-    _tracking_states.prune()
-
-    if hasattr(_current_function, "preprocess_image"):
-        pp_image = _current_function.preprocess_image(_TrackingFunctionContextImpl(), image)
-    else:
-        pp_image = image
-
-    return [
-        _tracking_states.store(
-            state=_current_function.init_tracking_state(
-                _TrackingFunctionShapeContextImpl(original_shape_type=shape.type), pp_image, shape
-            ),
-            shape_type=shape.type,
-            task_id=task_id,
-            image_dims=image.size,
-        )
-        for shape in shapes
-    ]
-
-
-def _worker_job_track(
-    task_id: int, image: PIL.Image.Image, states: list[str]
-) -> list[cvataa.TrackableShape | None]:
-    _tracking_states.prune()
-
-    pp_image = _current_function.preprocess_image(_TrackingFunctionContextImpl(), image)
-
-    def track(state_id):
-        inner_state, original_shape_type = _tracking_states.retrieve(
-            state_id=state_id, task_id=task_id, image_dims=image.size
-        )
-
-        output_shape = _current_function.track(
-            _TrackingFunctionShapeContextImpl(original_shape_type=original_shape_type),
-            pp_image,
-            inner_state,
-        )
-
-        if output_shape and output_shape.type != original_shape_type:
-            raise cvataa.BadFunctionError(
-                f"function output shape of type {output_shape.type!r}, "
-                f"but original shape was of type {original_shape_type!r}"
-            )
-        return output_shape
-
-    return list(map(track, states))
+    return worker_current_function().spec
 
 
 @attrs.frozen
@@ -263,64 +138,41 @@ class _NewReconnectionDelay:
 class _TaskCacheLimiter:
     """
     This class deletes least-recently used tasks from the dataset cache,
-    so that at any time the cache contains at most _MAX_TASKS_WITH_CHUNKS
-    tasks with downloaded chunks, and at most _MAX_TASKS_WITHOUT_CHUNKS without.
+    so that at any time the cache contains at most _MAX_CACHED_TASKS tasks.
 
     This helps manage disk usage, since agents may run indefinitely, and
     we don't want the dataset cache to keep growing.
     """
 
-    _MAX_TASKS_WITH_CHUNKS = 1
-    _MAX_TASKS_WITHOUT_CHUNKS = 10
+    _MAX_CACHED_TASKS = 10
 
     def __init__(self, client: Client) -> None:
         self._client = client
         self._cache_manager = make_cache_manager(client, cvatds.UpdatePolicy.IF_MISSING_OR_STALE)
 
-        self._cached_with_chunks_task_ids = []
-        self._cached_without_chunks_task_ids = []
+        self._cached_task_ids = []
 
         self._task_ids_in_use = set()
 
     @contextlib.contextmanager
-    def using_cache_for_task(
-        self, task_id: int, *, with_chunks: bool
-    ) -> Generator[None, None, None]:
+    def using_cache_for_task(self, task_id: int) -> Generator[None, None, None]:
         if task_id in self._task_ids_in_use:
-            # If with_chunks is True, we would have to ensure that task_id is returned to the
-            # "with chunks" list after it leaves _task_ids_in_use, regardless of the value of
-            # with_chunks in the call that initially put it in. That would be tricky to implement,
-            # and we don't have a use case for it, so just ban it.
-            assert not with_chunks
             yield
             return
 
-        if task_id in self._cached_with_chunks_task_ids:
-            # If the task already had cached chunks, we have to return it back to
-            # _cached_with_chunks_task_ids in the end.
-            with_chunks = True
-
-            self._cached_with_chunks_task_ids.remove(task_id)
-        elif task_id in self._cached_without_chunks_task_ids:
-            self._cached_without_chunks_task_ids.remove(task_id)
+        if task_id in self._cached_task_ids:
+            self._cached_task_ids.remove(task_id)
 
         self._task_ids_in_use.add(task_id)
 
-        if with_chunks:
-            cached_task_ids = self._cached_with_chunks_task_ids
-            max_cached_tasks = self._MAX_TASKS_WITH_CHUNKS
-        else:
-            cached_task_ids = self._cached_without_chunks_task_ids
-            max_cached_tasks = self._MAX_TASKS_WITHOUT_CHUNKS
-
-        if len(cached_task_ids) + len(self._task_ids_in_use) > max_cached_tasks:
-            self._delete_task_cache(cached_task_ids.pop(0))
+        if len(self._cached_task_ids) + len(self._task_ids_in_use) > self._MAX_CACHED_TASKS:
+            self._delete_task_cache(self._cached_task_ids.pop(0))
 
         try:
             yield
         finally:
             self._task_ids_in_use.remove(task_id)
-            cached_task_ids.append(task_id)
+            self._cached_task_ids.append(task_id)
 
     def _delete_task_cache(self, task_id: int) -> None:
         self._client.logger.info("Deleting task %d from the cache to make room...", task_id)
@@ -371,33 +223,25 @@ def _parse_event_stream(
                 yield _NewReconnectionDelay(timedelta(milliseconds=int(field_value)))
 
 
-class _BadArError(Exception):
-    pass
-
+def get_function_driver_class(function_spec: object) -> type[AgentFunctionDriver]:
+    if isinstance(function_spec, cvataa.DetectionFunctionSpec):
+        return AgentDetectionFunctionDriver
 
-class _IncompatibleFunctionError(Exception):
-    # This should only be thrown from inside _validate_X_function_compatibility methods.
-    pass
+    if isinstance(function_spec, cvataa.TrackingFunctionSpec):
+        return AgentTrackingFunctionDriver
 
-
-class _TrackingFunctionContextImpl(cvataa.TrackingFunctionContext):
-    pass
-
-
-@attrs.frozen(kw_only=True)
-class _TrackingFunctionShapeContextImpl(cvataa.TrackingFunctionShapeContext):
-    original_shape_type: str
+    raise CriticalError(f"Unsupported function spec type: {type(function_spec).__name__}")
 
 
 class _Agent:
-    def __init__(self, client: Client, executor: _RecoverableExecutor, function_id: int):
+    def __init__(self, client: Client, executor: RecoverableExecutor, function_id: int):
         self._rng = random.Random()  # nosec
 
         self._client = client
-        self._executor = executor
         self._function_id = function_id
-        self._function_spec = self._executor.result(
-            self._executor.submit(_worker_job_get_function_spec)
+        function_spec = executor.result(executor.submit(_worker_job_get_function_spec))
+        self._function_driver = get_function_driver_class(function_spec)(
+            client, executor, function_spec
         )
 
         _, response = self._client.api_client.call_api(
@@ -445,91 +289,18 @@ class _Agent:
             )
 
         try:
-            if isinstance(self._function_spec, cvataa.DetectionFunctionSpec):
-                self._validate_detection_function_compatibility(remote_function)
-                self._calculate_result_for_ar = self._calculate_result_for_detection_ar
-            elif isinstance(self._function_spec, cvataa.TrackingFunctionSpec):
-                self._validate_tracking_function_compatibility(remote_function)
-                self._calculate_result_for_ar = self._calculate_result_for_tracking_ar
-            else:
-                raise CriticalError(
-                    f"Unsupported function spec type: {type(self._function_spec).__name__}"
+            if remote_function["kind"] != self._function_driver.FUNCTION_KIND:
+                raise IncompatibleFunctionError(
+                    f"kind is {remote_function['kind']!r} "
+                    f"(expected {self._function_driver.FUNCTION_KIND!r})."
                 )
-        except _IncompatibleFunctionError as ex:
+
+            self._function_driver.validate_function_compatibility(remote_function)
+        except IncompatibleFunctionError as ex:
             raise CriticalError(
                 f"Function #{function_id} is incompatible with function object: {ex}"
             ) from ex
 
-    def _validate_detection_function_compatibility(self, remote_function: dict) -> None:
-        self._validate_remote_function_kind(remote_function, FUNCTION_KIND_DETECTOR)
-
-        labels_by_name = {label.name: label for label in self._function_spec.labels}
-
-        for remote_label in remote_function["labels_v2"]:
-            label_desc = f"label {remote_label['name']!r}"
-            label = labels_by_name.get(remote_label["name"])
-
-            self._validate_sublabel_compatibility(remote_label, label, label_desc)
-
-            sublabels_by_name = {sl.name: sl for sl in getattr(label, "sublabels", [])}
-
-            for remote_sl in remote_label.get("sublabels", []):
-                sl_desc = f"sublabel {remote_sl['name']!r} of {label_desc}"
-                sl = sublabels_by_name.get(remote_sl["name"])
-
-                self._validate_sublabel_compatibility(remote_sl, sl, sl_desc)
-
-    def _validate_sublabel_compatibility(
-        self, remote_sl: dict, sl: models.Sublabel | None, sl_desc: str
-    ):
-        if not sl:
-            raise _IncompatibleFunctionError(f"{sl_desc} is not supported.")
-
-        if remote_sl["type"] not in {"any", "unknown"} and remote_sl["type"] != sl.type:
-            raise _IncompatibleFunctionError(
-                f"{sl_desc} has type {remote_sl['type']!r}, "
-                f"but the function object declares type {sl.type!r}."
-            )
-
-        attrs_by_name = {attr.name: attr for attr in getattr(sl, "attributes", [])}
-
-        for remote_attr in remote_sl["attributes"]:
-            attr_desc = f"attribute {remote_attr['name']!r} of {sl_desc}"
-            attr = attrs_by_name.get(remote_attr["name"])
-
-            if not attr:
-                raise _IncompatibleFunctionError(f"{attr_desc} is not supported.")
-
-            if remote_attr["input_type"] != attr.input_type.value:
-                raise _IncompatibleFunctionError(
-                    f"{attr_desc} has input type {remote_attr['input_type']!r},"
-                    f" but the function object declares input type {attr.input_type.value!r}."
-                )
-
-            if remote_attr["values"] != attr.values:
-                raise _IncompatibleFunctionError(
-                    f"{attr_desc} has values {remote_attr['values']!r},"
-                    f" but the function object declares values {attr.values!r}."
-                )
-
-    def _validate_tracking_function_compatibility(self, remote_function: dict) -> None:
-        self._validate_remote_function_kind(remote_function, FUNCTION_KIND_TRACKER)
-
-        remote_supported_shape_types = frozenset(remote_function["supported_shape_types"])
-        unsupported = remote_supported_shape_types - self._function_spec.supported_shape_types
-
-        if unsupported:
-            raise _IncompatibleFunctionError(
-                "the function object does not support the following shape types: "
-                + ", ".join(map(repr, unsupported))
-            )
-
-    def _validate_remote_function_kind(self, remote_function: dict, expected_kind: str) -> None:
-        if remote_function["kind"] != expected_kind:
-            raise _IncompatibleFunctionError(
-                f"kind is {remote_function['kind']!r} (expected {expected_kind!r})."
-            )
-
     def _wait_between_polls(self):
         # offset the interval randomly to avoid synchronization between workers
         timeout_multiplier = self._rng.uniform(1 - _JITTER_AMOUNT, 1 + _JITTER_AMOUNT)
@@ -701,9 +472,23 @@ class _Agent:
         )
         self._client.logger.debug("AR %r parameters: %r", ar_id, ar_params)
 
-        try:
-            result = self._calculate_result_for_ar(ar_id, ar_params)
+        last_update_timestamp = datetime.now(tz=timezone.utc)
+
+        def check_in(*, current_progress: float) -> None:
+            nonlocal last_update_timestamp
+            current_timestamp = datetime.now(tz=timezone.utc)
 
+            if current_timestamp >= last_update_timestamp + _UPDATE_INTERVAL:
+                self._update_ar(ar_id, current_progress)
+                last_update_timestamp = current_timestamp
+
+            # Interactive requests are time sensitive, so if there are any,
+            # we have to put the current AR on hold and process them ASAP.
+            self._process_available_ars(REQUEST_CATEGORY_INTERACTIVE)
+
+        try:
+            with self._task_cache_limiter.using_cache_for_task(ar_params["task"]):
+                result = self._function_driver.calculate_result_for_ar(ar_params, check_in)
             self._complete_ar(ar_id, result)
         except Exception as ex:
             self._client.logger.error("Failed to process AR %r", ar_id, exc_info=True)
@@ -729,7 +514,7 @@ class _Agent:
                 error_message = "Failed to make an HTTP request"
             elif isinstance(ex, cvataa.BadFunctionError):
                 error_message = "Underlying function returned incorrect result: " + str(ex)
-            elif isinstance(ex, _BadArError):
+            elif isinstance(ex, BadArError):
                 error_message = "Invalid annotation request: " + str(ex)
             elif isinstance(ex, concurrent.futures.BrokenExecutor):
                 error_message = "Worker process crashed"
@@ -807,164 +592,6 @@ class _Agent:
         response_data = json.loads(response.data)
         return response_data["ar_assignment"]
 
-    def _calculate_result_for_detection_ar(self, ar_id: str, ar_params) -> dict[str, Any]:
-        if ar_params["type"] == "annotate_task":
-            with self._task_cache_limiter.using_cache_for_task(ar_params["task"], with_chunks=True):
-                return self._calculate_result_for_annotate_task_ar(ar_id, ar_params)
-        elif ar_params["type"] == "annotate_frame":
-            with self._task_cache_limiter.using_cache_for_task(
-                ar_params["task"], with_chunks=False
-            ):
-                return self._calculate_result_for_annotate_frame_ar(ar_id, ar_params)
-        else:
-            raise _BadArError(f"unsupported type: {ar_params['type']!r}")
-
-    def _create_annotation_mapper_for_detection_ar(
-        self, ar_params: dict, ds_labels: Sequence[models.ILabel]
-    ) -> _AnnotationMapper:
-        spec_nm = _SpecNameMapping.from_api(
-            {
-                k: models.LabelMappingEntryRequest._from_openapi_data(**v)
-                for k, v in ar_params["mapping"].items()
-            }
-        )
-
-        return _AnnotationMapper(
-            self._client.logger,
-            self._function_spec.labels,
-            ds_labels,
-            allow_unmatched_labels=False,
-            spec_nm=spec_nm,
-            conv_mask_to_poly=ar_params["conv_mask_to_poly"],
-        )
-
-    def _create_detection_function_context(
-        self, ar_params: dict, frame_name: str
-    ) -> cvataa.DetectionFunctionContext:
-        return _DetectionFunctionContextImpl(
-            frame_name=frame_name,
-            conf_threshold=ar_params["threshold"],
-            conv_mask_to_poly=ar_params["conv_mask_to_poly"],
-        )
-
-    def _calculate_result_for_annotate_task_ar(self, ar_id: str, ar_params) -> dict[str, Any]:
-        ds = cvatds.TaskDataset(self._client, ar_params["task"], load_annotations=False)
-
-        # Fetching the dataset might take a while, so do a progress update to let the server
-        # know we're still alive.
-        self._update_ar(ar_id, 0)
-        last_update_timestamp = datetime.now(tz=timezone.utc)
-
-        mapper = self._create_annotation_mapper_for_detection_ar(ar_params, ds.labels)
-
-        all_annotations = models.PatchedLabeledDataRequest(tags=[], shapes=[])
-
-        for sample_index, sample in enumerate(ds.samples):
-            context = self._create_detection_function_context(ar_params, sample.frame_name)
-            annotations = self._executor.result(
-                self._executor.submit(_worker_job_detect, context, sample.media.load_image())
-            )
-
-            tags, shapes = mapper.validate_and_remap(annotations, sample.frame_index)
-            all_annotations.tags.extend(tags)
-            all_annotations.shapes.extend(shapes)
-
-            current_timestamp = datetime.now(tz=timezone.utc)
-
-            if current_timestamp >= last_update_timestamp + _UPDATE_INTERVAL:
-                self._update_ar(ar_id, (sample_index + 1) / len(ds.samples))
-                last_update_timestamp = current_timestamp
-
-            # Interactive requests are time sensitive, so if there are any,
-            # we have to put the current AR on hold and process them ASAP.
-            self._process_available_ars(REQUEST_CATEGORY_INTERACTIVE)
-
-        return {"annotations": all_annotations}
-
-    def _calculate_result_for_annotate_frame_ar(self, ar_id: str, ar_params) -> dict[str, Any]:
-        sample, ds_labels = self._get_sample_from_ar_params(ar_params)
-
-        mapper = self._create_annotation_mapper_for_detection_ar(ar_params, ds_labels)
-
-        context = self._create_detection_function_context(ar_params, sample.frame_name)
-
-        annotations = self._executor.result(
-            self._executor.submit(_worker_job_detect, context, sample.media.load_image())
-        )
-
-        tags, shapes = mapper.validate_and_remap(annotations, sample.frame_index)
-        return {"annotations": models.PatchedLabeledDataRequest(tags=tags, shapes=shapes)}
-
-    def _calculate_result_for_tracking_ar(self, ar_id: str, ar_params) -> dict[str, Any]:
-        if ar_params["type"] == "init_tracking":
-            with self._task_cache_limiter.using_cache_for_task(
-                ar_params["task"], with_chunks=False
-            ):
-                return self._calculate_result_for_init_tracking_ar(ar_id, ar_params)
-        elif ar_params["type"] == "track":
-            with self._task_cache_limiter.using_cache_for_task(
-                ar_params["task"], with_chunks=False
-            ):
-                return self._calculate_result_for_track_ar(ar_id, ar_params)
-        else:
-            raise _BadArError(f"unsupported type: {ar_params['type']!r}")
-
-    def _calculate_result_for_init_tracking_ar(self, ar_id: str, ar_params) -> dict[str, Any]:
-        sample, _ = self._get_sample_from_ar_params(ar_params)
-
-        def convert_shape(shape: dict) -> cvataa.TrackableShape:
-            if shape["type"] not in self._function_spec.supported_shape_types:
-                raise _BadArError(f"Unsupported shape type {shape['type']!r}")
-            return cvataa.TrackableShape(type=shape["type"], points=shape["points"])
-
-        shapes = list(map(convert_shape, ar_params["shapes"]))
-
-        states = self._executor.result(
-            self._executor.submit(
-                _worker_job_init_tracking,
-                ar_params["task"],
-                sample.media.load_image(),
-                shapes,
-            )
-        )
-
-        return {"states": states}
-
-    def _calculate_result_for_track_ar(self, ar_id: str, ar_params) -> dict[str, Any]:
-        sample, _ = self._get_sample_from_ar_params(ar_params)
-
-        states = ar_params["states"]
-        shapes = self._executor.result(
-            self._executor.submit(
-                _worker_job_track, ar_params["task"], sample.media.load_image(), states
-            )
-        )
-
-        return {
-            "states": states,
-            "shapes": [attrs.asdict(shape) if shape else None for shape in shapes],
-        }
-
-    def _get_sample_from_ar_params(self, ar_params):
-        ds = cvatds.TaskDataset(
-            self._client,
-            ar_params["task"],
-            load_annotations=False,
-            media_download_policy=cvatds.MediaDownloadPolicy.FETCH_FRAMES_ON_DEMAND,
-        )
-
-        frame_index = ar_params["frame"]
-
-        # Since ds.samples excludes deleted frames, we can't just do sample = ds.samples[frame_index].
-        # Once we drop Python 3.9, we can change this to use bisect instead of the linear search.
-        for sample in ds.samples:
-            if sample.frame_index == frame_index:
-                break
-        else:
-            raise _BadArError(f"Frame with index {frame_index} does not exist in the task")
-
-        return sample, ds.labels
-
     def _update_ar(self, ar_id: str, progress: float) -> None:
         self._client.logger.info("Updating AR %r progress to %.2f%%...", ar_id, progress * 100)
 
@@ -1015,7 +642,7 @@ def run_agent(
     client: Client, function_loader: FunctionLoader, function_id: int, *, burst: bool
 ) -> None:
     with (
-        _RecoverableExecutor(
+        RecoverableExecutor(
             initializer=_worker_init,
             initargs=[function_loader, _default_tracking_state_id_generator],
         ) as executor,

cvat_cli-2.69.0/src/cvat_cli/_internal/agent_driver.py

@@ -0,0 +1,95 @@
+# Copyright (C) CVAT.ai Corporation
+#
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, ClassVar, Protocol
+
+import cvat_sdk.auto_annotation as cvataa
+import cvat_sdk.datasets as cvatds
+from cvat_sdk import Client
+from typing_extensions import Self
+
+if TYPE_CHECKING:
+    from .agent import RecoverableExecutor
+    from .agent_driver_tracking import TrackingStateIdGenerator
+
+
+class BadArError(Exception):
+    pass
+
+
+class IncompatibleFunctionError(Exception):
+    # This should only be thrown from inside validate_function_compatibility methods.
+    pass
+
+
+_current_function: cvataa.AutoAnnotationFunction
+
+
+def worker_current_function() -> cvataa.AutoAnnotationFunction:
+    return _current_function
+
+
+def set_worker_current_function(func: cvataa.AutoAnnotationFunction) -> None:
+    global _current_function
+    _current_function = func
+
+
+class CheckInCallback(Protocol):
+    def __call__(self, *, current_progress: float) -> None:
+        """
+        Temporarily suspends the processing of the current AR to report progress back to the server
+        and optionally process any pending interactive ARs. Should only be called during batch AR
+        processing (interactive ARs should be processed quickly enough to not require this).
+        """
+
+
+class AgentFunctionDriver:
+    FUNCTION_KIND: ClassVar[str]
+
+    def __init__(self, client: Client, executor: RecoverableExecutor, function_spec: object):
+        self._client = client
+        self._executor = executor
+        self._function_spec = function_spec
+
+    @classmethod
+    def init_worker(cls, state_id_generator: TrackingStateIdGenerator) -> None:
+        pass
+
+    def validate_function_compatibility(self, remote_function: dict) -> None:
+        raise NotImplementedError
+
+    _CALCULATE_RESULT_PER_AR_TYPE: ClassVar[
+        dict[str, Callable[[Self, dict[str, Any], CheckInCallback], dict[str, Any]]]
+    ]
+
+    def calculate_result_for_ar(
+        self, ar_params: dict[str, Any], check_in: CheckInCallback
+    ) -> dict[str, Any]:
+        if calc := self._CALCULATE_RESULT_PER_AR_TYPE.get(ar_params["type"]):
+            return calc(self, ar_params, check_in)
+
+        raise BadArError(f"unsupported type: {ar_params['type']!r}")
+
+    def _get_sample_from_ar_params(self, ar_params):
+        ds = cvatds.TaskDataset(
+            self._client,
+            ar_params["task"],
+            load_annotations=False,
+            media_download_policy=cvatds.MediaDownloadPolicy.FETCH_FRAMES_ON_DEMAND,
+        )
+
+        frame_index = ar_params["frame"]
+
+        # Since ds.samples excludes deleted frames, we can't just do sample = ds.samples[frame_index].
+        # Once we drop Python 3.9, we can change this to use bisect instead of the linear search.
+        for sample in ds.samples:
+            if sample.frame_index == frame_index:
+                break
+        else:
+            raise BadArError(f"Frame with index {frame_index} does not exist in the task")
+
+        return sample, ds.labels

cvat_cli-2.69.0/src/cvat_cli/_internal/agent_driver_detection.py

@@ -0,0 +1,167 @@
+# Copyright (C) CVAT.ai Corporation
+#
+# SPDX-License-Identifier: MIT
+
+from collections.abc import Sequence
+from typing import Any, cast
+
+import cvat_sdk.auto_annotation as cvataa
+import cvat_sdk.datasets as cvatds
+import PIL.Image
+from cvat_sdk import models
+from cvat_sdk.auto_annotation.driver import (
+    _AnnotationMapper,
+    _DetectionFunctionContextImpl,
+    _SpecNameMapping,
+)
+
+from .agent_driver import (
+    AgentFunctionDriver,
+    CheckInCallback,
+    IncompatibleFunctionError,
+    worker_current_function,
+)
+
+
+def _worker_job_detect(
+    context: _DetectionFunctionContextImpl, image: PIL.Image.Image
+) -> Sequence[cvataa.DetectionAnnotation]:
+    current_function = cast(cvataa.DetectionFunction, worker_current_function())
+    return current_function.detect(context, image)
+
+
+class AgentDetectionFunctionDriver(AgentFunctionDriver):
+    FUNCTION_KIND = "detector"
+    _function_spec: cvataa.DetectionFunctionSpec
+
+    def _validate_sublabel_compatibility(
+        self, remote_sl: dict, sl: models.Sublabel | None, sl_desc: str
+    ):
+        if not sl:
+            raise IncompatibleFunctionError(f"{sl_desc} is not supported.")
+
+        if remote_sl["type"] not in {"any", "unknown"} and remote_sl["type"] != sl.type:
+            raise IncompatibleFunctionError(
+                f"{sl_desc} has type {remote_sl['type']!r}, "
+                f"but the function object declares type {sl.type!r}."
+            )
+
+        attrs_by_name = {attr.name: attr for attr in getattr(sl, "attributes", [])}
+
+        for remote_attr in remote_sl["attributes"]:
+            attr_desc = f"attribute {remote_attr['name']!r} of {sl_desc}"
+            attr = attrs_by_name.get(remote_attr["name"])
+
+            if not attr:
+                raise IncompatibleFunctionError(f"{attr_desc} is not supported.")
+
+            if remote_attr["input_type"] != attr.input_type.value:
+                raise IncompatibleFunctionError(
+                    f"{attr_desc} has input type {remote_attr['input_type']!r},"
+                    f" but the function object declares input type {attr.input_type.value!r}."
+                )
+
+            if remote_attr["values"] != attr.values:
+                raise IncompatibleFunctionError(
+                    f"{attr_desc} has values {remote_attr['values']!r},"
+                    f" but the function object declares values {attr.values!r}."
+                )
+
+    def validate_function_compatibility(self, remote_function: dict) -> None:
+        labels_by_name = {label.name: label for label in self._function_spec.labels}
+
+        for remote_label in remote_function["labels_v2"]:
+            label_desc = f"label {remote_label['name']!r}"
+            label = labels_by_name.get(remote_label["name"])
+
+            self._validate_sublabel_compatibility(remote_label, label, label_desc)
+
+            sublabels_by_name = {sl.name: sl for sl in getattr(label, "sublabels", [])}
+
+            for remote_sl in remote_label.get("sublabels", []):
+                sl_desc = f"sublabel {remote_sl['name']!r} of {label_desc}"
+                sl = sublabels_by_name.get(remote_sl["name"])
+
+                self._validate_sublabel_compatibility(remote_sl, sl, sl_desc)
+
+    def _create_annotation_mapper_for_detection_ar(
+        self, ar_params: dict, ds_labels: Sequence[models.ILabel]
+    ) -> _AnnotationMapper:
+        spec_nm = _SpecNameMapping.from_api(
+            {
+                k: models.LabelMappingEntryRequest._from_openapi_data(**v)
+                for k, v in ar_params["mapping"].items()
+            }
+        )
+
+        return _AnnotationMapper(
+            self._client.logger,
+            self._function_spec.labels,
+            ds_labels,
+            allow_unmatched_labels=False,
+            spec_nm=spec_nm,
+            conv_mask_to_poly=ar_params["conv_mask_to_poly"],
+        )
+
+    def _create_detection_function_context(
+        self, ar_params: dict, frame_name: str
+    ) -> cvataa.DetectionFunctionContext:
+        return _DetectionFunctionContextImpl(
+            frame_name=frame_name,
+            conf_threshold=ar_params["threshold"],
+            conv_mask_to_poly=ar_params["conv_mask_to_poly"],
+        )
+
+    def _calculate_result_for_annotate_task_ar(
+        self, ar_params: dict[str, Any], check_in: CheckInCallback
+    ) -> dict[str, Any]:
+        ds = cvatds.TaskDataset(
+            self._client,
+            ar_params["task"],
+            load_annotations=False,
+            media_download_policy=cvatds.MediaDownloadPolicy.FETCH_CHUNKS_ON_DEMAND,
+        )
+
+        # Fetching the dataset might take a while, so check in to let the server
+        # know we're still alive.
+        check_in(current_progress=0)
+
+        mapper = self._create_annotation_mapper_for_detection_ar(ar_params, ds.labels)
+
+        all_annotations = models.PatchedLabeledDataRequest(tags=[], shapes=[])
+
+        with ds.iter_samples(temporary_chunks=True) as samples:
+            for sample_index, sample in enumerate(samples):
+                context = self._create_detection_function_context(ar_params, sample.frame_name)
+                annotations = self._executor.result(
+                    self._executor.submit(_worker_job_detect, context, sample.media.load_image())
+                )
+
+                tags, shapes = mapper.validate_and_remap(annotations, sample.frame_index)
+                all_annotations.tags.extend(tags)
+                all_annotations.shapes.extend(shapes)
+
+                check_in(current_progress=(sample_index + 1) / len(ds.samples))
+
+        return {"annotations": all_annotations}
+
+    def _calculate_result_for_annotate_frame_ar(
+        self, ar_params: dict[str, Any], check_in: object
+    ) -> dict[str, Any]:
+        sample, ds_labels = self._get_sample_from_ar_params(ar_params)
+
+        mapper = self._create_annotation_mapper_for_detection_ar(ar_params, ds_labels)
+
+        context = self._create_detection_function_context(ar_params, sample.frame_name)
+
+        annotations = self._executor.result(
+            self._executor.submit(_worker_job_detect, context, sample.media.load_image())
+        )
+
+        tags, shapes = mapper.validate_and_remap(annotations, sample.frame_index)
+        return {"annotations": models.PatchedLabeledDataRequest(tags=tags, shapes=shapes)}
+
+    _CALCULATE_RESULT_PER_AR_TYPE = {
+        "annotate_task": _calculate_result_for_annotate_task_ar,
+        "annotate_frame": _calculate_result_for_annotate_frame_ar,
+    }

cvat_cli-2.69.0/src/cvat_cli/_internal/agent_driver_tracking.py

@@ -0,0 +1,214 @@
+# Copyright (C) CVAT.ai Corporation
+#
+# SPDX-License-Identifier: MIT
+
+from collections import OrderedDict
+from datetime import datetime, timedelta, timezone
+from typing import Any, Callable, TypeAlias, cast
+
+import attrs
+import cvat_sdk.auto_annotation as cvataa
+import PIL.Image
+
+from .agent_driver import (
+    AgentFunctionDriver,
+    BadArError,
+    IncompatibleFunctionError,
+    worker_current_function,
+)
+
+_MAX_AGE_OF_TRACKING_STATE = timedelta(hours=8)
+
+
+@attrs.define
+class _ExtendedTrackingState:
+    inner_state: Any  # the state produced by the AA function
+    original_shape_type: str
+    original_task_id: int
+    original_image_dims: tuple[int, int]
+    last_accessed_at: datetime = attrs.field(factory=lambda: datetime.now(tz=timezone.utc))
+
+
+class _TrackingStateContainer:
+    def __init__(self):
+        self._id_to_ext_state: OrderedDict[str, _ExtendedTrackingState] = OrderedDict()
+
+    def store(self, state: Any, shape_type: str, task_id: int, image_dims: tuple[int, int]) -> str:
+        state_id = _tracking_state_id_generator()
+        self._id_to_ext_state[state_id] = _ExtendedTrackingState(
+            inner_state=state,
+            original_shape_type=shape_type,
+            original_task_id=task_id,
+            original_image_dims=image_dims,
+        )
+        return state_id
+
+    def retrieve(self, state_id: str, task_id: int, image_dims: tuple[int, int]) -> Any:
+        ext_state = self._id_to_ext_state.get(state_id)
+
+        if not ext_state:
+            raise BadArError(f"Tracking state {state_id!r} not found - possibly expired")
+
+        if ext_state.original_task_id != task_id:
+            # This is a defense-in-depth measure. State IDs are supposed to be unguessable,
+            # but even if an attacker manages to obtain one, they will not be able to use it
+            # to get any information about a task they don't have access to.
+            raise BadArError(f"Tracking state {state_id!r} is not for task #{task_id}")
+
+        if image_dims != ext_state.original_image_dims:
+            raise BadArError("Image sizes of the start frame and the current frame are different")
+
+        ext_state.last_accessed_at = datetime.now(tz=timezone.utc)
+        self._id_to_ext_state.move_to_end(state_id)
+
+        return ext_state.inner_state, ext_state.original_shape_type
+
+    def prune(self) -> None:
+        cutoff = datetime.now(tz=timezone.utc) - _MAX_AGE_OF_TRACKING_STATE
+
+        while (
+            self._id_to_ext_state
+            and next(iter(self._id_to_ext_state.values())).last_accessed_at < cutoff
+        ):
+            self._id_to_ext_state.popitem(last=False)
+
+
+TrackingStateIdGenerator: TypeAlias = Callable[[], str]
+
+
+_tracking_states: _TrackingStateContainer
+_tracking_state_id_generator: TrackingStateIdGenerator
+
+
+def _worker_job_init_tracking(
+    task_id: int,
+    image: PIL.Image.Image,
+    shapes: list[cvataa.TrackableShape],
+) -> list[str]:
+    _tracking_states.prune()
+
+    current_function = cast(cvataa.TrackingFunction, worker_current_function())
+
+    if hasattr(current_function, "preprocess_image"):
+        pp_image = current_function.preprocess_image(_TrackingFunctionContextImpl(), image)
+    else:
+        pp_image = image
+
+    return [
+        _tracking_states.store(
+            state=current_function.init_tracking_state(
+                _TrackingFunctionShapeContextImpl(original_shape_type=shape.type), pp_image, shape
+            ),
+            shape_type=shape.type,
+            task_id=task_id,
+            image_dims=image.size,
+        )
+        for shape in shapes
+    ]
+
+
+def _worker_job_track(
+    task_id: int, image: PIL.Image.Image, states: list[str]
+) -> list[cvataa.TrackableShape | None]:
+    _tracking_states.prune()
+
+    current_function = cast(cvataa.TrackingFunction, worker_current_function())
+
+    pp_image = current_function.preprocess_image(_TrackingFunctionContextImpl(), image)
+
+    def track(state_id):
+        inner_state, original_shape_type = _tracking_states.retrieve(
+            state_id=state_id, task_id=task_id, image_dims=image.size
+        )
+
+        output_shape = current_function.track(
+            _TrackingFunctionShapeContextImpl(original_shape_type=original_shape_type),
+            pp_image,
+            inner_state,
+        )
+
+        if output_shape and output_shape.type != original_shape_type:
+            raise cvataa.BadFunctionError(
+                f"function output shape of type {output_shape.type!r}, "
+                f"but original shape was of type {original_shape_type!r}"
+            )
+        return output_shape
+
+    return list(map(track, states))
+
+
+class _TrackingFunctionContextImpl(cvataa.TrackingFunctionContext):
+    pass
+
+
+@attrs.frozen(kw_only=True)
+class _TrackingFunctionShapeContextImpl(cvataa.TrackingFunctionShapeContext):
+    original_shape_type: str
+
+
+class AgentTrackingFunctionDriver(AgentFunctionDriver):
+    FUNCTION_KIND = "tracker"
+    _function_spec: cvataa.TrackingFunctionSpec
+
+    @classmethod
+    def init_worker(cls, state_id_generator: TrackingStateIdGenerator) -> None:
+        global _tracking_states
+        _tracking_states = _TrackingStateContainer()
+
+        global _tracking_state_id_generator
+        _tracking_state_id_generator = state_id_generator
+
+    def validate_function_compatibility(self, remote_function: dict) -> None:
+        remote_supported_shape_types = frozenset(remote_function["supported_shape_types"])
+        unsupported = remote_supported_shape_types - self._function_spec.supported_shape_types
+
+        if unsupported:
+            raise IncompatibleFunctionError(
+                "the function object does not support the following shape types: "
+                + ", ".join(map(repr, unsupported))
+            )
+
+    def _calculate_result_for_init_tracking_ar(
+        self, ar_params: dict[str, Any], check_in: object
+    ) -> dict[str, Any]:
+        sample, _ = self._get_sample_from_ar_params(ar_params)
+
+        def convert_shape(shape: dict) -> cvataa.TrackableShape:
+            if shape["type"] not in self._function_spec.supported_shape_types:
+                raise BadArError(f"Unsupported shape type {shape['type']!r}")
+            return cvataa.TrackableShape(type=shape["type"], points=shape["points"])
+
+        shapes = list(map(convert_shape, ar_params["shapes"]))
+
+        states = self._executor.result(
+            self._executor.submit(
+                _worker_job_init_tracking,
+                ar_params["task"],
+                sample.media.load_image(),
+                shapes,
+            )
+        )
+
+        return {"states": states}
+
+    def _calculate_result_for_track_ar(
+        self, ar_params: dict[str, Any], check_in: object
+    ) -> dict[str, Any]:
+        sample, _ = self._get_sample_from_ar_params(ar_params)
+
+        states = ar_params["states"]
+        shapes = self._executor.result(
+            self._executor.submit(
+                _worker_job_track, ar_params["task"], sample.media.load_image(), states
+            )
+        )
+
+        return {
+            "states": states,
+            "shapes": [attrs.asdict(shape) if shape else None for shape in shapes],
+        }
+
+    _CALCULATE_RESULT_PER_AR_TYPE = {
+        "init_tracking": _calculate_result_for_init_tracking_ar,
+        "track": _calculate_result_for_track_ar,
+    }

{cvat_cli-2.67.0 → cvat_cli-2.69.0}/src/cvat_cli/_internal/commands_functions.py

@@ -11,12 +11,9 @@ from typing import Any
 import cvat_sdk.auto_annotation as cvataa
 from cvat_sdk import Client, models
 
-from .agent import (
-    FUNCTION_KIND_DETECTOR,
-    FUNCTION_KIND_TRACKER,
-    FUNCTION_PROVIDER_NATIVE,
-    run_agent,
-)
+from .agent import FUNCTION_PROVIDER_NATIVE, run_agent
+from .agent_driver_detection import AgentDetectionFunctionDriver
+from .agent_driver_tracking import AgentTrackingFunctionDriver
 from .command_base import CommandGroup
 from .common import FunctionLoader, configure_function_implementation_arguments
 
@@ -85,7 +82,7 @@ class FunctionCreateNative:
         spec = function.spec
 
         if isinstance(spec, cvataa.DetectionFunctionSpec):
-            remote_function["kind"] = FUNCTION_KIND_DETECTOR
+            remote_function["kind"] = AgentDetectionFunctionDriver.FUNCTION_KIND
             remote_function["labels_v2"] = []
 
             for label_spec in spec.labels:
@@ -96,7 +93,7 @@ class FunctionCreateNative:
                         self._dump_sublabel_spec(sublabel) for sublabel in sublabels
                     ]
         elif isinstance(spec, cvataa.TrackingFunctionSpec):
-            remote_function["kind"] = FUNCTION_KIND_TRACKER
+            remote_function["kind"] = AgentTrackingFunctionDriver.FUNCTION_KIND
             remote_function["supported_shape_types"] = sorted(spec.supported_shape_types)
         else:
             raise cvataa.BadFunctionError(f"Unsupported function spec type: {type(spec).__name__}")

cvat_cli-2.69.0/src/cvat_cli/version.py

@@ -0,0 +1 @@
+VERSION = "2.69.0"

{cvat_cli-2.67.0 → cvat_cli-2.69.0}/src/cvat_cli.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cvat-cli
-Version: 2.67.0
+Version: 2.69.0
 Summary: Command-line client for CVAT
 Author-email: "CVAT.ai Corporation" <support@cvat.ai>
 License-Expression: MIT
@@ -9,7 +9,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
-Requires-Dist: cvat-sdk==2.67.0
+Requires-Dist: cvat-sdk==2.69.0
 Requires-Dist: attrs>=24.2.0
 Requires-Dist: Pillow>=10.3.0
 

{cvat_cli-2.67.0 → cvat_cli-2.69.0}/src/cvat_cli.egg-info/SOURCES.txt

@@ -11,6 +11,9 @@ src/cvat_cli.egg-info/requires.txt
 src/cvat_cli.egg-info/top_level.txt
 src/cvat_cli/_internal/__init__.py
 src/cvat_cli/_internal/agent.py
+src/cvat_cli/_internal/agent_driver.py
+src/cvat_cli/_internal/agent_driver_detection.py
+src/cvat_cli/_internal/agent_driver_tracking.py
 src/cvat_cli/_internal/command_base.py
 src/cvat_cli/_internal/commands_all.py
 src/cvat_cli/_internal/commands_functions.py

{cvat_cli-2.67.0 → cvat_cli-2.69.0}/src/cvat_cli.egg-info/requires.txt

@@ -1,3 +1,3 @@
-cvat-sdk==2.67.0
+cvat-sdk==2.69.0
 attrs>=24.2.0
 Pillow>=10.3.0

cvat_cli-2.67.0/src/cvat_cli/version.py

@@ -1 +0,0 @@
-VERSION = "2.67.0"