cvat/tests/python/rest_api/_test_base.py

import io
import math
from contextlib import closing
from functools import partial
from typing import Generator, Optional, Sequence

import numpy as np
import pytest
from cvat_sdk.api_client import models
from PIL import Image
from pytest_cases import fixture, fixture_ref, parametrize

import shared.utils.s3 as s3
from rest_api.utils import calc_end_frame, create_task, unique
from shared.tasks.enums import SourceDataType
from shared.tasks.interface import ITaskSpec
from shared.tasks.types import ImagesTaskSpec, VideoTaskSpec
from shared.tasks.utils import parse_frame_step
from shared.utils.config import make_api_client
from shared.utils.helpers import generate_image_files, generate_video_file


class TestTasksBase:
    _USERNAME = "admin1"

    def _image_task_fxt_base(
        self,
        request: pytest.FixtureRequest,
        *,
        frame_count: Optional[int] = 10,
        image_files: Optional[Sequence[io.BytesIO]] = None,
        start_frame: Optional[int] = None,
        stop_frame: Optional[int] = None,
        step: Optional[int] = None,
        segment_size: Optional[int] = None,
        server_files: Optional[Sequence[str]] = None,
        cloud_storage_id: Optional[int] = None,
        job_replication: Optional[int] = None,
        **data_kwargs,
    ) -> Generator[tuple[ImagesTaskSpec, int], None, None]:
        task_params = {
            "name": f"{request.node.name}[{request.fixturename}]",
            "labels": [{"name": "a"}],
            **({"segment_size": segment_size} if segment_size else {}),
            **({"consensus_replicas": job_replication} if job_replication else {}),
        }

        if server_files is not None:
            assert (
                image_files is not None
            ), "'server_files' must be used together with 'image_files'"
        else:
            assert bool(image_files) ^ bool(
                frame_count
            ), "Expected only one of 'image_files' and 'frame_count'"
            if not image_files:
                image_files = generate_image_files(frame_count)

        images_data = [f.getvalue() for f in image_files]

        resulting_task_size = len(
            range(start_frame or 0, (stop_frame or len(images_data) - 1) + 1, step or 1)
        )

        data_params = {
            "image_quality": 70,
            "sorting_method": "natural",
            "chunk_size": max(1, (segment_size or resulting_task_size) // 2),
            **(
                {
                    "server_files": server_files,
                    "cloud_storage_id": cloud_storage_id,
                }
                if server_files
                else {"client_files": image_files}
            ),
        }
        data_params.update(data_kwargs)

        if start_frame is not None:
            data_params["start_frame"] = start_frame

        if stop_frame is not None:
            data_params["stop_frame"] = stop_frame

        if step is not None:
            data_params["frame_filter"] = f"step={step}"

        def get_frame(i: int) -> bytes:
            return images_data[i]

        task_id, _ = create_task(self._USERNAME, spec=task_params, data=data_params)
        yield ImagesTaskSpec(
            models.TaskWriteRequest._from_openapi_data(**task_params),
            models.DataRequest._from_openapi_data(**data_params),
            get_frame=get_frame,
            size=resulting_task_size,
        ), task_id

    @pytest.fixture(scope="class")
    def fxt_uploaded_images_task(
        self, request: pytest.FixtureRequest
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._image_task_fxt_base(request=request)

    @pytest.fixture(scope="class")
    def fxt_uploaded_images_task_with_segments(
        self, request: pytest.FixtureRequest
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._image_task_fxt_base(request=request, segment_size=4)

    @fixture(scope="class")
    @parametrize("step", [2, 5])
    @parametrize("stop_frame", [15, 26])
    @parametrize("start_frame", [3, 7])
    def fxt_uploaded_images_task_with_segments_start_stop_step(
        self, request: pytest.FixtureRequest, start_frame: int, stop_frame: Optional[int], step: int
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._image_task_fxt_base(
            request=request,
            frame_count=30,
            segment_size=4,
            start_frame=start_frame,
            stop_frame=stop_frame,
            step=step,
        )

    @pytest.fixture(scope="class")
    def fxt_uploaded_images_task_with_segments_and_consensus(
        self, request: pytest.FixtureRequest
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._image_task_fxt_base(request=request, segment_size=4, job_replication=2)

    def _image_task_with_honeypots_and_segments_base(
        self,
        request: pytest.FixtureRequest,
        *,
        start_frame: Optional[int] = None,
        step: Optional[int] = None,
        random_seed: int = 42,
        image_files: Optional[Sequence[io.BytesIO]] = None,
        server_files: Optional[Sequence[str]] = None,
        cloud_storage_id: Optional[int] = None,
        **kwargs,
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        validation_params = models.DataRequestValidationParams._from_openapi_data(
            mode="gt_pool",
            frame_selection_method="random_uniform",
            random_seed=random_seed,
            frame_count=5,
            frames_per_job_count=2,
        )

        used_frames_count = 15
        total_frame_count = (start_frame or 0) + used_frames_count * (step or 1)
        base_segment_size = 4
        regular_frame_count = used_frames_count - validation_params.frame_count
        final_segment_size = base_segment_size + validation_params.frames_per_job_count
        final_task_size = (
            regular_frame_count
            + validation_params.frames_per_job_count
            * math.ceil(regular_frame_count / base_segment_size)
            + validation_params.frame_count
        )

        if image_files:
            if len(image_files) != total_frame_count:
                raise ValueError(
                    f"If provided, image_files must contain {total_frame_count} images"
                )
        else:
            image_files = generate_image_files(total_frame_count)

        with closing(
            self._image_task_fxt_base(
                request=request,
                frame_count=None,
                image_files=image_files,
                segment_size=base_segment_size,
                sorting_method="random",
                start_frame=start_frame,
                step=step,
                validation_params=validation_params,
                server_files=server_files,
                cloud_storage_id=cloud_storage_id,
                **kwargs,
            )
        ) as task_gen:
            for task_spec, task_id in task_gen:
                # Get the actual frame order after the task is created
                with make_api_client(self._USERNAME) as api_client:
                    (task_meta, _) = api_client.tasks_api.retrieve_data_meta(task_id)
                    frame_map = [
                        next(i for i, f in enumerate(image_files) if f.name == frame_info.name)
                        for frame_info in task_meta.frames
                    ]

                _get_frame = task_spec._get_frame
                task_spec._get_frame = lambda i: _get_frame(frame_map[i])

                task_spec.size = final_task_size
                task_spec._params.segment_size = final_segment_size

                # These parameters are not applicable to the resulting task,
                # they are only effective during task creation
                if start_frame or step:
                    task_spec._data_params.start_frame = 0
                    task_spec._data_params.stop_frame = task_spec.size
                    task_spec._data_params.frame_filter = ""

                yield task_spec, task_id

    @fixture(scope="class")
    def fxt_uploaded_images_task_with_honeypots_and_segments(
        self, request: pytest.FixtureRequest
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._image_task_with_honeypots_and_segments_base(request)

    @fixture(scope="class")
    @parametrize("start_frame, step", [(2, 3)])
    def fxt_uploaded_images_task_with_honeypots_and_segments_start_step(
        self, request: pytest.FixtureRequest, start_frame: Optional[int], step: Optional[int]
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._image_task_with_honeypots_and_segments_base(
            request, start_frame=start_frame, step=step
        )

    def _images_task_with_honeypots_and_changed_real_frames_base(
        self, request: pytest.FixtureRequest, **kwargs
    ):
        with closing(
            self._image_task_with_honeypots_and_segments_base(
                request, start_frame=2, step=3, **kwargs
            )
        ) as gen_iter:
            task_spec, task_id = next(gen_iter)

            with make_api_client(self._USERNAME) as api_client:
                validation_layout, _ = api_client.tasks_api.retrieve_validation_layout(task_id)
                validation_frames = validation_layout.validation_frames

                new_honeypot_real_frames = [
                    validation_frames[(validation_frames.index(f) + 1) % len(validation_frames)]
                    for f in validation_layout.honeypot_real_frames
                ]
                api_client.tasks_api.partial_update_validation_layout(
                    task_id,
                    patched_task_validation_layout_write_request=(
                        models.PatchedTaskValidationLayoutWriteRequest(
                            frame_selection_method="manual",
                            honeypot_real_frames=new_honeypot_real_frames,
                        )
                    ),
                )

                # Get the new frame order
                frame_map = dict(zip(validation_layout.honeypot_frames, new_honeypot_real_frames))

                _get_frame = task_spec._get_frame
                task_spec._get_frame = lambda i: _get_frame(frame_map.get(i, i))

            yield task_spec, task_id

    @fixture(scope="class")
    @parametrize("random_seed", [1, 2, 5])
    def fxt_uploaded_images_task_with_honeypots_and_changed_real_frames(
        self, request: pytest.FixtureRequest, random_seed: int
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._images_task_with_honeypots_and_changed_real_frames_base(
            request, random_seed=random_seed
        )

    @fixture(scope="class")
    @parametrize(
        "cloud_storage_id",
        [pytest.param(2, marks=[pytest.mark.with_external_services, pytest.mark.timeout(60)])],
    )
    def fxt_cloud_images_task_with_honeypots_and_changed_real_frames(
        self, request: pytest.FixtureRequest, cloud_storages, cloud_storage_id: int
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        cloud_storage = cloud_storages[cloud_storage_id]
        s3_client = s3.make_client(bucket=cloud_storage["resource"])

        image_files = generate_image_files(47)

        for image in image_files:
            image.name = f"test/{image.name}"
            image.seek(0)

            s3_client.create_file(data=image, filename=image.name)
            request.addfinalizer(partial(s3_client.remove_file, filename=image.name))

        server_files = [f.name for f in image_files]

        for image in image_files:
            image.seek(0)

        yield from self._images_task_with_honeypots_and_changed_real_frames_base(
            request,
            image_files=image_files,
            server_files=server_files,
            cloud_storage_id=cloud_storage_id,
            # FIXME: random sorting with frame filter and cloud images (and, optionally, honeypots)
            # doesn't work with static cache
            # https://github.com/cvat-ai/cvat/issues/9021
            use_cache=True,
        )

    def _uploaded_images_task_with_gt_and_segments_base(
        self,
        request: pytest.FixtureRequest,
        *,
        start_frame: Optional[int] = None,
        step: Optional[int] = None,
        frame_selection_method: str = "random_uniform",
        job_replication: Optional[int] = None,
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        used_frames_count = 16
        total_frame_count = (start_frame or 0) + used_frames_count * (step or 1)
        segment_size = 5
        image_files = generate_image_files(total_frame_count)

        validation_params_kwargs = {"frame_selection_method": frame_selection_method}

        if "random" in frame_selection_method:
            validation_params_kwargs["random_seed"] = 42

        if frame_selection_method == "random_uniform":
            validation_frames_count = 10
            validation_params_kwargs["frame_count"] = validation_frames_count
        elif frame_selection_method == "random_per_job":
            frames_per_job_count = 3
            validation_params_kwargs["frames_per_job_count"] = frames_per_job_count
            validation_frames_count = used_frames_count // segment_size + min(
                used_frames_count % segment_size, frames_per_job_count
            )
        elif frame_selection_method == "manual":
            validation_frames_count = 10

            valid_frame_ids = range(
                (start_frame or 0), (start_frame or 0) + used_frames_count * (step or 1), step or 1
            )
            rng = np.random.Generator(np.random.MT19937(seed=42))
            validation_params_kwargs["frames"] = rng.choice(
                [f.name for i, f in enumerate(image_files) if i in valid_frame_ids],
                validation_frames_count,
                replace=False,
            ).tolist()
        else:
            raise NotImplementedError

        validation_params = models.DataRequestValidationParams._from_openapi_data(
            mode="gt",
            **validation_params_kwargs,
        )

        yield from self._image_task_fxt_base(
            request=request,
            frame_count=None,
            image_files=image_files,
            segment_size=segment_size,
            sorting_method="natural",
            start_frame=start_frame,
            step=step,
            validation_params=validation_params,
            job_replication=job_replication,
        )

    @pytest.fixture(scope="class")
    def fxt_uploaded_images_task_with_gt_and_segments_and_consensus(
        self, request: pytest.FixtureRequest
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._uploaded_images_task_with_gt_and_segments_base(
            request=request, job_replication=2
        )

    @fixture(scope="class")
    @parametrize("start_frame, step", [(2, 3)])
    @parametrize("frame_selection_method", ["random_uniform", "random_per_job", "manual"])
    def fxt_uploaded_images_task_with_gt_and_segments_start_step(
        self,
        request: pytest.FixtureRequest,
        start_frame: Optional[int],
        step: Optional[int],
        frame_selection_method: str,
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._uploaded_images_task_with_gt_and_segments_base(
            request,
            start_frame=start_frame,
            step=step,
            frame_selection_method=frame_selection_method,
        )

    def _uploaded_video_task_fxt_base(
        self,
        request: pytest.FixtureRequest,
        *,
        frame_count: int = 10,
        segment_size: Optional[int] = None,
        start_frame: Optional[int] = None,
        stop_frame: Optional[int] = None,
        step: Optional[int] = None,
    ) -> Generator[tuple[VideoTaskSpec, int], None, None]:
        task_params = {
            "name": f"{request.node.name}[{request.fixturename}]",
            "labels": [{"name": "a"}],
        }
        if segment_size:
            task_params["segment_size"] = segment_size

        resulting_task_size = len(
            range(start_frame or 0, (stop_frame or frame_count - 1) + 1, step or 1)
        )

        video_file = generate_video_file(frame_count)
        video_data = video_file.getvalue()
        data_params = {
            "image_quality": 70,
            "client_files": [video_file],
            "chunk_size": max(1, (segment_size or resulting_task_size) // 2),
        }

        if start_frame is not None:
            data_params["start_frame"] = start_frame

        if stop_frame is not None:
            data_params["stop_frame"] = stop_frame

        if step is not None:
            data_params["frame_filter"] = f"step={step}"

        def get_video_file() -> io.BytesIO:
            return io.BytesIO(video_data)

        task_id, _ = create_task(self._USERNAME, spec=task_params, data=data_params)
        yield VideoTaskSpec(
            models.TaskWriteRequest._from_openapi_data(**task_params),
            models.DataRequest._from_openapi_data(**data_params),
            get_video_file=get_video_file,
            size=resulting_task_size,
        ), task_id

    @pytest.fixture(scope="class")
    def fxt_uploaded_video_task(
        self,
        request: pytest.FixtureRequest,
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._uploaded_video_task_fxt_base(request=request)

    @pytest.fixture(scope="class")
    def fxt_uploaded_video_task_with_segments(
        self, request: pytest.FixtureRequest
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._uploaded_video_task_fxt_base(request=request, segment_size=4)

    @fixture(scope="class")
    @parametrize("step", [2, 5])
    @parametrize("stop_frame", [15, 26])
    @parametrize("start_frame", [3, 7])
    def fxt_uploaded_video_task_with_segments_start_stop_step(
        self, request: pytest.FixtureRequest, start_frame: int, stop_frame: Optional[int], step: int
    ) -> Generator[tuple[ITaskSpec, int], None, None]:
        yield from self._uploaded_video_task_fxt_base(
            request=request,
            frame_count=30,
            segment_size=4,
            start_frame=start_frame,
            stop_frame=stop_frame,
            step=step,
        )

    def _compute_annotation_segment_params(self, task_spec: ITaskSpec) -> list[tuple[int, int]]:
        segment_params = []
        frame_step = task_spec.frame_step
        segment_size = getattr(task_spec, "segment_size", 0) or task_spec.size * frame_step
        start_frame = getattr(task_spec, "start_frame", 0)
        stop_frame = getattr(task_spec, "stop_frame", None) or (
            start_frame + (task_spec.size - 1) * frame_step
        )
        end_frame = calc_end_frame(start_frame, stop_frame, frame_step)

        validation_params = getattr(task_spec, "validation_params", None)
        if validation_params and validation_params.mode.value == "gt_pool":
            end_frame = min(
                end_frame, (task_spec.size - validation_params.frame_count) * frame_step
            )
            segment_size = min(segment_size, end_frame - 1)

        overlap = min(
            (
                getattr(task_spec, "overlap", None) or 0
                if task_spec.source_data_type == SourceDataType.images
                else 5
            ),
            segment_size // 2,
        )
        segment_start = start_frame
        while segment_start < end_frame:
            if start_frame < segment_start:
                segment_start -= overlap * frame_step

            segment_end = segment_start + frame_step * segment_size

            segment_params.append((segment_start, min(segment_end, end_frame) - frame_step))
            segment_start = segment_end

        return segment_params

    @staticmethod
    def _compare_images(
        expected: Image.Image, actual: Image.Image, *, must_be_identical: bool = True
    ):
        expected_pixels = np.array(expected)
        chunk_frame_pixels = np.array(actual)
        assert expected_pixels.shape == chunk_frame_pixels.shape

        if not must_be_identical:
            # video chunks can have slightly changed colors, due to codec specifics
            # compressed images can also be distorted
            assert np.allclose(chunk_frame_pixels, expected_pixels, atol=2)
        else:
            assert np.array_equal(chunk_frame_pixels, expected_pixels)

    def _get_job_abs_frame_set(self, job_meta: models.DataMetaRead) -> Sequence[int]:
        if job_meta.included_frames:
            return job_meta.included_frames
        else:
            return range(
                job_meta.start_frame,
                job_meta.stop_frame + 1,
                parse_frame_step(job_meta.frame_filter),
            )

    _tasks_with_honeypots_cases = [
        fixture_ref("fxt_uploaded_images_task_with_honeypots_and_segments"),
        fixture_ref("fxt_uploaded_images_task_with_honeypots_and_segments_start_step"),
        fixture_ref("fxt_uploaded_images_task_with_honeypots_and_changed_real_frames"),
        fixture_ref("fxt_cloud_images_task_with_honeypots_and_changed_real_frames"),
    ]

    _tasks_with_simple_gt_job_cases = [
        fixture_ref("fxt_uploaded_images_task_with_gt_and_segments_start_step"),
        fixture_ref("fxt_uploaded_images_task_with_gt_and_segments_and_consensus"),
    ]

    _tasks_with_consensus_cases = [
        fixture_ref("fxt_uploaded_images_task_with_segments_and_consensus"),
        fixture_ref("fxt_uploaded_images_task_with_gt_and_segments_and_consensus"),
    ]

    # Keep in mind that these fixtures are generated eagerly
    # (before each depending test or group of tests),
    # e.g. a failing task creation in one the fixtures will fail all the depending tests cases.
    _all_task_cases = unique(
        [
            fixture_ref("fxt_uploaded_images_task"),
            fixture_ref("fxt_uploaded_images_task_with_segments"),
            fixture_ref("fxt_uploaded_images_task_with_segments_start_stop_step"),
            fixture_ref("fxt_uploaded_video_task"),
            fixture_ref("fxt_uploaded_video_task_with_segments"),
            fixture_ref("fxt_uploaded_video_task_with_segments_start_stop_step"),
        ]
        + _tasks_with_honeypots_cases
        + _tasks_with_simple_gt_job_cases
        + _tasks_with_consensus_cases,
        key=lambda fxt_ref: fxt_ref.fixture,
    )