From b32a9ebd3b79953954bb9c5b741a656939e1cc7b Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 25 Jul 2024 19:39:14 +0300 Subject: [PATCH 001/115] Update frame provider and media cache --- cvat/apps/engine/cache.py | 435 +++++++++++++++++---------- cvat/apps/engine/frame_provider.py | 367 +++++++++++++--------- cvat/apps/engine/media_extractors.py | 4 + cvat/apps/engine/pyproject.toml | 12 + 4 files changed, 517 insertions(+), 301 deletions(-) create mode 100644 cvat/apps/engine/pyproject.toml diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 2603c2fd5a1..988e7676121 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -1,68 +1,88 @@ # Copyright (C) 2020-2022 Intel Corporation -# Copyright (C) 2022-2023 CVAT.ai Corporation +# Copyright (C) 2022-2024 CVAT.ai Corporation # # SPDX-License-Identifier: MIT +from __future__ import annotations + import io import os -import zipfile -from datetime import datetime, timezone -from io import BytesIO +import pickle # nosec import shutil import tempfile +import zipfile import zlib - -from typing import Optional, Tuple +from contextlib import contextmanager +from datetime import datetime, timezone +from typing import Any, Callable, Optional, Sequence, Tuple, Type import cv2 import PIL.Image -import pickle # nosec +import PIL.ImageOps from django.conf import settings from django.core.cache import caches from rest_framework.exceptions import NotFound, ValidationError -from cvat.apps.engine.cloud_provider import (Credentials, - db_storage_to_storage_instance, - get_cloud_storage_instance) +from cvat.apps.engine import models +from cvat.apps.engine.cloud_provider import ( + Credentials, + db_storage_to_storage_instance, + get_cloud_storage_instance, +) from cvat.apps.engine.log import ServerLogManager -from cvat.apps.engine.media_extractors import (ImageDatasetManifestReader, - Mpeg4ChunkWriter, - Mpeg4CompressedChunkWriter, - VideoDatasetManifestReader, - ZipChunkWriter, - ZipCompressedChunkWriter) +from cvat.apps.engine.media_extractors import ( + FrameQuality, + IChunkWriter, + ImageDatasetManifestReader, + Mpeg4ChunkWriter, + Mpeg4CompressedChunkWriter, + VideoDatasetManifestReader, + ZipChunkWriter, + ZipCompressedChunkWriter, +) from cvat.apps.engine.mime_types import mimetypes -from cvat.apps.engine.models import (DataChoice, DimensionType, Job, Image, - StorageChoice, CloudStorage) from cvat.apps.engine.utils import md5_hash, preload_images from utils.dataset_manifest import ImageManifestManager slogger = ServerLogManager(__name__) + +DataWithMime = Tuple[io.BytesIO, str] +_CacheItem = Tuple[io.BytesIO, str, int] + + class MediaCache: - def __init__(self, dimension=DimensionType.DIM_2D): - self._dimension = dimension - self._cache = caches['media'] - - def _get_or_set_cache_item(self, key, create_function): - def create_item(): - slogger.glob.info(f'Starting to prepare chunk: key {key}') - item = create_function() - slogger.glob.info(f'Ending to prepare chunk: key {key}') - - if item[0]: - item = (item[0], item[1], zlib.crc32(item[0].getbuffer())) + def __init__(self) -> None: + self._cache = caches["media"] + + # TODO migrate keys (check if they will be removed) + + def get_checksum(self, value: bytes) -> int: + return zlib.crc32(value) + + def _get_or_set_cache_item( + self, key: str, create_callback: Callable[[], DataWithMime] + ) -> DataWithMime: + def create_item() -> _CacheItem: + slogger.glob.info(f"Starting to prepare chunk: key {key}") + item_data = create_callback() + slogger.glob.info(f"Ending to prepare chunk: key {key}") + + if item_data[0]: + item = (item_data[0], item_data[1], self.get_checksum(item_data[0].getbuffer())) self._cache.set(key, item) + else: + item = (item_data[0], item_data[1], None) return item - slogger.glob.info(f'Starting to get chunk from cache: key {key}') + slogger.glob.info(f"Starting to get chunk from cache: key {key}") try: item = self._cache.get(key) except pickle.UnpicklingError: - slogger.glob.error(f'Unable to get item from cache: key {key}', exc_info=True) + slogger.glob.error(f"Unable to get item from cache: key {key}", exc_info=True) item = None - slogger.glob.info(f'Ending to get chunk from cache: key {key}, is_cached {bool(item)}') + slogger.glob.info(f"Ending to get chunk from cache: key {key}, is_cached {bool(item)}") if not item: item = create_item() @@ -70,113 +90,130 @@ def create_item(): # compare checksum item_data = item[0].getbuffer() if isinstance(item[0], io.BytesIO) else item[0] item_checksum = item[2] if len(item) == 3 else None - if item_checksum != zlib.crc32(item_data): - slogger.glob.info(f'Recreating cache item {key} due to checksum mismatch') + if item_checksum != self.get_checksum(item_data): + slogger.glob.info(f"Recreating cache item {key} due to checksum mismatch") item = create_item() return item[0], item[1] - def get_task_chunk_data_with_mime(self, chunk_number, quality, db_data): - item = self._get_or_set_cache_item( - key=f'{db_data.id}_{chunk_number}_{quality}', - create_function=lambda: self._prepare_task_chunk(db_data, quality, chunk_number), - ) + def _get(self, key: str) -> Optional[DataWithMime]: + slogger.glob.info(f"Starting to get chunk from cache: key {key}") + try: + item = self._cache.get(key) + except pickle.UnpicklingError: + slogger.glob.error(f"Unable to get item from cache: key {key}", exc_info=True) + item = None + slogger.glob.info(f"Ending to get chunk from cache: key {key}, is_cached {bool(item)}") return item - def get_selective_job_chunk_data_with_mime(self, chunk_number, quality, job): - item = self._get_or_set_cache_item( - key=f'job_{job.id}_{chunk_number}_{quality}', - create_function=lambda: self.prepare_selective_job_chunk(job, quality, chunk_number), + def get_segment_chunk( + self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality + ) -> DataWithMime: + return self._get_or_set_cache_item( + key=f"segment_{db_segment.id}_{chunk_number}_{quality}", + create_callback=lambda: self.prepare_segment_chunk( + db_segment, chunk_number, quality=quality + ), ) - return item - - def get_local_preview_with_mime(self, frame_number, db_data): - item = self._get_or_set_cache_item( - key=f'data_{db_data.id}_{frame_number}_preview', - create_function=lambda: self._prepare_local_preview(frame_number, db_data), + def get_selective_job_chunk( + self, db_job: models.Job, chunk_number: int, *, quality: FrameQuality + ) -> DataWithMime: + return self._get_or_set_cache_item( + key=f"job_{db_job.id}_{chunk_number}_{quality}", + create_callback=lambda: self.prepare_masked_range_segment_chunk( + db_job.segment, chunk_number, quality=quality + ), ) - return item - - def get_cloud_preview_with_mime( - self, - db_storage: CloudStorage, - ) -> Optional[Tuple[io.BytesIO, str]]: - key = f'cloudstorage_{db_storage.id}_preview' - return self._cache.get(key) - - def get_or_set_cloud_preview_with_mime( - self, - db_storage: CloudStorage, - ) -> Tuple[io.BytesIO, str]: - key = f'cloudstorage_{db_storage.id}_preview' - - item = self._get_or_set_cache_item( - key, create_function=lambda: self._prepare_cloud_preview(db_storage) + def get_local_preview(self, db_data: models.Data, frame_number: int) -> DataWithMime: + return self._get_or_set_cache_item( + key=f"data_{db_data.id}_{frame_number}_preview", + create_callback=lambda: self._prepare_local_preview(frame_number, db_data), ) - return item + def get_cloud_preview(self, db_storage: models.CloudStorage) -> Optional[DataWithMime]: + return self._get(f"cloudstorage_{db_storage.id}_preview") - def get_frame_context_images(self, db_data, frame_number): - item = self._get_or_set_cache_item( - key=f'context_image_{db_data.id}_{frame_number}', - create_function=lambda: self._prepare_context_image(db_data, frame_number) + def get_or_set_cloud_preview(self, db_storage: models.CloudStorage) -> DataWithMime: + return self._get_or_set_cache_item( + f"cloudstorage_{db_storage.id}_preview", + create_callback=lambda: self._prepare_cloud_preview(db_storage), ) - return item + def get_frame_context_images(self, db_data: models.Data, frame_number: int) -> DataWithMime: + return self._get_or_set_cache_item( + key=f"context_image_{db_data.id}_{frame_number}", + create_callback=lambda: self._prepare_context_image(db_data, frame_number), + ) - @staticmethod - def _get_frame_provider_class(): - from cvat.apps.engine.frame_provider import \ - FrameProvider # TODO: remove circular dependency - return FrameProvider + def get_task_preview(self, db_task: models.Task) -> Optional[DataWithMime]: + return self._get(f"task_{db_task.data_id}_preview") - from contextlib import contextmanager + def get_segment_preview(self, db_segment: models.Segment) -> Optional[DataWithMime]: + return self._get(f"segment_{db_segment.id}_preview") - @staticmethod @contextmanager - def _get_images(db_data, chunk_number, dimension): - images = [] + def _read_raw_frames(self, db_task: models.Task, frames: Sequence[int]): + db_data = db_task.data + + media = [] tmp_dir = None - upload_dir = { - StorageChoice.LOCAL: db_data.get_upload_dirname(), - StorageChoice.SHARE: settings.SHARE_ROOT, - StorageChoice.CLOUD_STORAGE: db_data.get_upload_dirname(), + raw_data_dir = { + models.StorageChoice.LOCAL: db_data.get_upload_dirname(), + models.StorageChoice.SHARE: settings.SHARE_ROOT, + models.StorageChoice.CLOUD_STORAGE: db_data.get_upload_dirname(), }[db_data.storage] - try: - if hasattr(db_data, 'video'): - source_path = os.path.join(upload_dir, db_data.video.path) + dimension = db_task.dimension - reader = VideoDatasetManifestReader(manifest_path=db_data.get_manifest_path(), - source_path=source_path, chunk_number=chunk_number, - chunk_size=db_data.chunk_size, start=db_data.start_frame, - stop=db_data.stop_frame, step=db_data.get_frame_step()) + # TODO + try: + if hasattr(db_data, "video"): + source_path = os.path.join(raw_data_dir, db_data.video.path) + + # TODO: refactor to allow non-manifest videos + reader = VideoDatasetManifestReader( + manifest_path=db_data.get_manifest_path(), + source_path=source_path, + chunk_number=chunk_number, + chunk_size=db_data.chunk_size, + start=db_data.start_frame, + stop=db_data.stop_frame, + step=db_data.get_frame_step(), + ) for frame in reader: - images.append((frame, source_path, None)) + media.append((frame, source_path, None)) else: - reader = ImageDatasetManifestReader(manifest_path=db_data.get_manifest_path(), - chunk_number=chunk_number, chunk_size=db_data.chunk_size, - start=db_data.start_frame, stop=db_data.stop_frame, - step=db_data.get_frame_step()) + reader = ImageDatasetManifestReader( + manifest_path=db_data.get_manifest_path(), + chunk_number=chunk_number, + chunk_size=db_data.chunk_size, + start=db_data.start_frame, + stop=db_data.stop_frame, + step=db_data.get_frame_step(), + ) if db_data.storage == StorageChoice.CLOUD_STORAGE: db_cloud_storage = db_data.cloud_storage - assert db_cloud_storage, 'Cloud storage instance was deleted' + assert db_cloud_storage, "Cloud storage instance was deleted" credentials = Credentials() - credentials.convert_from_db({ - 'type': db_cloud_storage.credentials_type, - 'value': db_cloud_storage.credentials, - }) + credentials.convert_from_db( + { + "type": db_cloud_storage.credentials_type, + "value": db_cloud_storage.credentials, + } + ) details = { - 'resource': db_cloud_storage.resource, - 'credentials': credentials, - 'specific_attributes': db_cloud_storage.get_specific_attributes() + "resource": db_cloud_storage.resource, + "credentials": credentials, + "specific_attributes": db_cloud_storage.get_specific_attributes(), } - cloud_storage_instance = get_cloud_storage_instance(cloud_provider=db_cloud_storage.provider_type, **details) + cloud_storage_instance = get_cloud_storage_instance( + cloud_provider=db_cloud_storage.provider_type, **details + ) - tmp_dir = tempfile.mkdtemp(prefix='cvat') + tmp_dir = tempfile.mkdtemp(prefix="cvat") files_to_download = [] checksums = [] for item in reader: @@ -184,51 +221,95 @@ def _get_images(db_data, chunk_number, dimension): fs_filename = os.path.join(tmp_dir, file_name) files_to_download.append(file_name) - checksums.append(item.get('checksum', None)) - images.append((fs_filename, fs_filename, None)) + checksums.append(item.get("checksum", None)) + media.append((fs_filename, fs_filename, None)) - cloud_storage_instance.bulk_download_to_dir(files=files_to_download, upload_dir=tmp_dir) - images = preload_images(images) + cloud_storage_instance.bulk_download_to_dir( + files=files_to_download, upload_dir=tmp_dir + ) + media = preload_images(media) - for checksum, (_, fs_filename, _) in zip(checksums, images): + for checksum, (_, fs_filename, _) in zip(checksums, media): if checksum and not md5_hash(fs_filename) == checksum: - slogger.cloud_storage[db_cloud_storage.id].warning('Hash sums of files {} do not match'.format(file_name)) + slogger.cloud_storage[db_cloud_storage.id].warning( + "Hash sums of files {} do not match".format(file_name) + ) else: for item in reader: - source_path = os.path.join(upload_dir, f"{item['name']}{item['extension']}") - images.append((source_path, source_path, None)) - if dimension == DimensionType.DIM_2D: - images = preload_images(images) - - yield images + source_path = os.path.join( + raw_data_dir, f"{item['name']}{item['extension']}" + ) + media.append((source_path, source_path, None)) + if dimension == models.DimensionType.DIM_2D: + media = preload_images(media) + + yield media finally: - if db_data.storage == StorageChoice.CLOUD_STORAGE and tmp_dir is not None: + if db_data.storage == models.StorageChoice.CLOUD_STORAGE and tmp_dir is not None: shutil.rmtree(tmp_dir) - def _prepare_task_chunk(self, db_data, quality, chunk_number): - FrameProvider = self._get_frame_provider_class() - - writer_classes = { - FrameProvider.Quality.COMPRESSED : Mpeg4CompressedChunkWriter if db_data.compressed_chunk_type == DataChoice.VIDEO else ZipCompressedChunkWriter, - FrameProvider.Quality.ORIGINAL : Mpeg4ChunkWriter if db_data.original_chunk_type == DataChoice.VIDEO else ZipChunkWriter, + def prepare_segment_chunk( + self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality + ) -> DataWithMime: + if db_segment.type == models.SegmentType.RANGE: + return self.prepare_range_segment_chunk(db_segment, chunk_number, quality=quality) + elif db_segment.type == models.SegmentType.SPECIFIC_FRAMES: + return self.prepare_masked_range_segment_chunk( + db_segment, chunk_number, quality=quality + ) + else: + assert False, f"Unknown segment type {db_segment.type}" + + def prepare_range_segment_chunk( + self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality + ) -> DataWithMime: + db_task = db_segment.task + db_data = db_task.data + + chunk_size = db_data.chunk_size + chunk_frames = db_segment.frame_set[ + chunk_size * chunk_number : chunk_number * (chunk_number + 1) + ] + + writer_classes: dict[FrameQuality, Type[IChunkWriter]] = { + FrameQuality.COMPRESSED: ( + Mpeg4CompressedChunkWriter + if db_data.compressed_chunk_type == models.DataChoice.VIDEO + else ZipCompressedChunkWriter + ), + FrameQuality.ORIGINAL: ( + Mpeg4ChunkWriter + if db_data.original_chunk_type == models.DataChoice.VIDEO + else ZipChunkWriter + ), } - image_quality = 100 if writer_classes[quality] in [Mpeg4ChunkWriter, ZipChunkWriter] else db_data.image_quality - mime_type = 'video/mp4' if writer_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] else 'application/zip' + image_quality = ( + 100 + if writer_classes[quality] in [Mpeg4ChunkWriter, ZipChunkWriter] + else db_data.image_quality + ) + mime_type = ( + "video/mp4" + if writer_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] + else "application/zip" + ) kwargs = {} - if self._dimension == DimensionType.DIM_3D: - kwargs["dimension"] = DimensionType.DIM_3D + if db_segment.task.dimension == models.DimensionType.DIM_3D: + kwargs["dimension"] = models.DimensionType.DIM_3D writer = writer_classes[quality](image_quality, **kwargs) - buff = BytesIO() - with self._get_images(db_data, chunk_number, self._dimension) as images: + buff = io.BytesIO() + with self._read_raw_frames(db_task, frames=chunk_frames) as images: writer.save_as_chunk(images, buff) - buff.seek(0) + buff.seek(0) return buff, mime_type - def prepare_selective_job_chunk(self, db_job: Job, quality, chunk_number: int): + def prepare_masked_range_segment_chunk( + self, db_job: models.Job, quality, chunk_number: int + ) -> DataWithMime: db_data = db_job.segment.task.data FrameProvider = self._get_frame_provider_class() @@ -239,10 +320,10 @@ def prepare_selective_job_chunk(self, db_job: Job, quality, chunk_number: int): chunk_frames = [] writer = ZipCompressedChunkWriter(db_data.image_quality, dimension=self._dimension) - dummy_frame = BytesIO() - PIL.Image.new('RGB', (1, 1)).save(dummy_frame, writer.IMAGE_EXT) + dummy_frame = io.BytesIO() + PIL.Image.new("RGB", (1, 1)).save(dummy_frame, writer.IMAGE_EXT) - if hasattr(db_data, 'video'): + if hasattr(db_data, "video"): frame_size = (db_data.video.width, db_data.video.height) else: frame_size = None @@ -266,27 +347,36 @@ def prepare_selective_job_chunk(self, db_job: Job, quality, chunk_number: int): if frame.size != frame_size: frame = frame.resize(frame_size) - frame_bytes = BytesIO() + frame_bytes = io.BytesIO() frame.save(frame_bytes, writer.IMAGE_EXT) frame_bytes.seek(0) else: # Populate skipped frames with placeholder data, # this is required for video chunk decoding implementation in UI - frame_bytes = BytesIO(dummy_frame.getvalue()) + frame_bytes = io.BytesIO(dummy_frame.getvalue()) if frame_bytes is not None: chunk_frames.append((frame_bytes, None, None)) - buff = BytesIO() - writer.save_as_chunk(chunk_frames, buff, compress_frames=False, - zip_compress_level=1 # these are likely to be many skips in SPECIFIC_FRAMES segments + buff = io.BytesIO() + writer.save_as_chunk( + chunk_frames, + buff, + compress_frames=False, + zip_compress_level=1, # there are likely to be many skips in SPECIFIC_FRAMES segments ) buff.seek(0) - return buff, 'application/zip' + return buff, "application/zip" + + def prepare_segment_preview(self, db_segment: models.Segment) -> DataWithMime: + if db_segment.task.data.cloud_storage: + return self._prepare_cloud_segment_preview(db_segment) + else: + return self._prepare_local_segment_preview(db_segment) - def _prepare_local_preview(self, frame_number, db_data): + def _prepare_local_preview(self, db_data: models.Data, frame_number: int) -> DataWithMime: FrameProvider = self._get_frame_provider_class() frame_provider = FrameProvider(db_data, self._dimension) buff, mime_type = frame_provider.get_preview(frame_number) @@ -296,28 +386,31 @@ def _prepare_local_preview(self, frame_number, db_data): def _prepare_cloud_preview(self, db_storage): storage = db_storage_to_storage_instance(db_storage) if not db_storage.manifests.count(): - raise ValidationError('Cannot get the cloud storage preview. There is no manifest file') + raise ValidationError("Cannot get the cloud storage preview. There is no manifest file") preview_path = None for manifest_model in db_storage.manifests.all(): manifest_prefix = os.path.dirname(manifest_model.filename) - full_manifest_path = os.path.join(db_storage.get_storage_dirname(), manifest_model.filename) - if not os.path.exists(full_manifest_path) or \ - datetime.fromtimestamp(os.path.getmtime(full_manifest_path), tz=timezone.utc) < storage.get_file_last_modified(manifest_model.filename): + full_manifest_path = os.path.join( + db_storage.get_storage_dirname(), manifest_model.filename + ) + if not os.path.exists(full_manifest_path) or datetime.fromtimestamp( + os.path.getmtime(full_manifest_path), tz=timezone.utc + ) < storage.get_file_last_modified(manifest_model.filename): storage.download_file(manifest_model.filename, full_manifest_path) manifest = ImageManifestManager( os.path.join(db_storage.get_storage_dirname(), manifest_model.filename), - db_storage.get_storage_dirname() + db_storage.get_storage_dirname(), ) # need to update index manifest.set_index() if not len(manifest): continue preview_info = manifest[0] - preview_filename = ''.join([preview_info['name'], preview_info['extension']]) + preview_filename = "".join([preview_info["name"], preview_info["extension"]]) preview_path = os.path.join(manifest_prefix, preview_filename) break if not preview_path: - msg = 'Cloud storage {} does not contain any images'.format(db_storage.pk) + msg = "Cloud storage {} does not contain any images".format(db_storage.pk) slogger.cloud_storage[db_storage.pk].info(msg) raise NotFound(msg) @@ -326,24 +419,42 @@ def _prepare_cloud_preview(self, db_storage): return buff, mime_type - def _prepare_context_image(self, db_data, frame_number): - zip_buffer = BytesIO() + def prepare_context_images( + self, db_data: models.Data, frame_number: int + ) -> Optional[DataWithMime]: + zip_buffer = io.BytesIO() try: - image = Image.objects.get(data_id=db_data.id, frame=frame_number) - except Image.DoesNotExist: - return None, None - with zipfile.ZipFile(zip_buffer, 'a', zipfile.ZIP_DEFLATED, False) as zip_file: + image = models.Image.objects.get(data_id=db_data.id, frame=frame_number) + except models.Image.DoesNotExist: + return None + + with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED, False) as zip_file: if not image.related_files.count(): return None, None - common_path = os.path.commonpath(list(map(lambda x: str(x.path), image.related_files.all()))) + common_path = os.path.commonpath( + list(map(lambda x: str(x.path), image.related_files.all())) + ) for i in image.related_files.all(): path = os.path.realpath(str(i.path)) name = os.path.relpath(str(i.path), common_path) image = cv2.imread(path) - success, result = cv2.imencode('.JPEG', image) + success, result = cv2.imencode(".JPEG", image) if not success: raise Exception('Failed to encode image to ".jpeg" format') - zip_file.writestr(f'{name}.jpg', result.tobytes()) - mime_type = 'application/zip' + zip_file.writestr(f"{name}.jpg", result.tobytes()) + zip_buffer.seek(0) + mime_type = "application/zip" return zip_buffer, mime_type + + +def prepare_preview_image(image: PIL.Image.Image) -> DataWithMime: + PREVIEW_SIZE = (256, 256) + PREVIEW_MIME = "image/jpeg" + + image = PIL.ImageOps.exif_transpose(image) + image.thumbnail(PREVIEW_SIZE) + + output_buf = io.BytesIO() + image.convert("RGB").save(output_buf, format="JPEG") + return image, PREVIEW_MIME diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 4e2f42ef793..92354723b02 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -3,26 +3,34 @@ # # SPDX-License-Identifier: MIT +from __future__ import annotations + import math -from enum import Enum -from io import BytesIO import os +from dataclasses import dataclass +from enum import Enum, auto +from io import BytesIO +from typing import Any, Callable, Generic, Iterable, Iterator, Optional, Tuple, TypeVar, Union +import av import cv2 import numpy as np -from PIL import Image, ImageOps +from PIL import Image +from rest_framework.exceptions import ValidationError -from cvat.apps.engine.cache import MediaCache -from cvat.apps.engine.media_extractors import VideoReader, ZipReader +from cvat.apps.engine import models +from cvat.apps.engine.cache import DataWithMime, MediaCache, prepare_preview_image +from cvat.apps.engine.media_extractors import FrameQuality, IMediaReader, VideoReader, ZipReader from cvat.apps.engine.mime_types import mimetypes -from cvat.apps.engine.models import DataChoice, StorageMethodChoice, DimensionType -from rest_framework.exceptions import ValidationError -class RandomAccessIterator: - def __init__(self, iterable): - self.iterable = iterable - self.iterator = None - self.pos = -1 +_T = TypeVar("_T") + + +class _RandomAccessIterator(Iterator[_T]): + def __init__(self, iterable: Iterable[_T]): + self.iterable: Iterable[_T] = iterable + self.iterator: Optional[Iterator[_T]] = None + self.pos: int = -1 def __iter__(self): return self @@ -30,7 +38,7 @@ def __iter__(self): def __next__(self): return self[self.pos + 1] - def __getitem__(self, idx): + def __getitem__(self, idx: int) -> Optional[_T]: assert 0 <= idx if self.iterator is None or idx <= self.pos: self.reset() @@ -47,168 +55,241 @@ def reset(self): def close(self): if self.iterator is not None: - if close := getattr(self.iterator, 'close', None): + if close := getattr(self.iterator, "close", None): close() self.iterator = None self.pos = -1 -class FrameProvider: - VIDEO_FRAME_EXT = '.PNG' - VIDEO_FRAME_MIME = 'image/png' - class Quality(Enum): - COMPRESSED = 0 - ORIGINAL = 100 +class _ChunkLoader: + def __init__( + self, reader_class: IMediaReader, path_getter: Callable[[int], DataWithMime] + ) -> None: + self.chunk_id: Optional[int] = None + self.chunk_reader: Optional[_RandomAccessIterator] = None + self.reader_class = reader_class + self.get_chunk_path = path_getter - class Type(Enum): - BUFFER = 0 - PIL = 1 - NUMPY_ARRAY = 2 + def load(self, chunk_id: int) -> _RandomAccessIterator[Tuple[Any, str, int]]: + if self.chunk_id != chunk_id: + self.unload() - class ChunkLoader: - def __init__(self, reader_class, path_getter): - self.chunk_id = None - self.chunk_reader = None - self.reader_class = reader_class - self.get_chunk_path = path_getter - - def load(self, chunk_id): - if self.chunk_id != chunk_id: - self.unload() - - self.chunk_id = chunk_id - self.chunk_reader = RandomAccessIterator( - self.reader_class([self.get_chunk_path(chunk_id)])) - return self.chunk_reader - - def unload(self): - self.chunk_id = None - if self.chunk_reader: - self.chunk_reader.close() - self.chunk_reader = None - - class BuffChunkLoader(ChunkLoader): - def __init__(self, reader_class, path_getter, quality, db_data): - super().__init__(reader_class, path_getter) - self.quality = quality - self.db_data = db_data - - def load(self, chunk_id): - if self.chunk_id != chunk_id: - self.chunk_id = chunk_id - self.chunk_reader = RandomAccessIterator( - self.reader_class([self.get_chunk_path(chunk_id, self.quality, self.db_data)[0]])) - return self.chunk_reader - - def __init__(self, db_data, dimension=DimensionType.DIM_2D): - self._db_data = db_data - self._dimension = dimension - self._loaders = {} - - reader_class = { - DataChoice.IMAGESET: ZipReader, - DataChoice.VIDEO: VideoReader, - } + self.chunk_id = chunk_id + self.chunk_reader = _RandomAccessIterator( + self.reader_class([self.get_chunk_path(chunk_id)]) + ) + return self.chunk_reader - if db_data.storage_method == StorageMethodChoice.CACHE: - cache = MediaCache(dimension=dimension) + def unload(self): + self.chunk_id = None + if self.chunk_reader: + self.chunk_reader.close() + self.chunk_reader = None - self._loaders[self.Quality.COMPRESSED] = self.BuffChunkLoader( - reader_class[db_data.compressed_chunk_type], - cache.get_task_chunk_data_with_mime, - self.Quality.COMPRESSED, - self._db_data) - self._loaders[self.Quality.ORIGINAL] = self.BuffChunkLoader( - reader_class[db_data.original_chunk_type], - cache.get_task_chunk_data_with_mime, - self.Quality.ORIGINAL, - self._db_data) - else: - self._loaders[self.Quality.COMPRESSED] = self.ChunkLoader( - reader_class[db_data.compressed_chunk_type], - db_data.get_compressed_chunk_path) - self._loaders[self.Quality.ORIGINAL] = self.ChunkLoader( - reader_class[db_data.original_chunk_type], - db_data.get_original_chunk_path) - def __len__(self): - return self._db_data.size +class FrameOutputType(Enum): + BUFFER = auto() + PIL = auto() + NUMPY_ARRAY = auto() - def unload(self): - for loader in self._loaders.values(): - loader.unload() - def _validate_frame_number(self, frame_number): - frame_number_ = int(frame_number) - if frame_number_ < 0 or frame_number_ >= self._db_data.size: - raise ValidationError('Incorrect requested frame number: {}'.format(frame_number_)) +Frame2d = Union[BytesIO, np.ndarray, Image.Image] +Frame3d = BytesIO +AnyFrame = Union[Frame2d, Frame3d] - chunk_number = frame_number_ // self._db_data.chunk_size - frame_offset = frame_number_ % self._db_data.chunk_size - return frame_number_, chunk_number, frame_offset +@dataclass +class DataWithMeta(Generic[_T]): + data: _T + mime: str + checksum: int - def get_chunk_number(self, frame_number): - return int(frame_number) // self._db_data.chunk_size - def _validate_chunk_number(self, chunk_number): - chunk_number_ = int(chunk_number) - if chunk_number_ < 0 or chunk_number_ >= math.ceil(self._db_data.size / self._db_data.chunk_size): - raise ValidationError('requested chunk does not exist') +class _FrameProvider: + VIDEO_FRAME_EXT = ".PNG" + VIDEO_FRAME_MIME = "image/png" - return chunk_number_ + def unload(self): + pass @classmethod - def _av_frame_to_png_bytes(cls, av_frame): + def _av_frame_to_png_bytes(cls, av_frame: av.VideoFrame) -> BytesIO: ext = cls.VIDEO_FRAME_EXT - image = av_frame.to_ndarray(format='bgr24') + image = av_frame.to_ndarray(format="bgr24") success, result = cv2.imencode(ext, image) if not success: raise RuntimeError("Failed to encode image to '%s' format" % (ext)) return BytesIO(result.tobytes()) - def _convert_frame(self, frame, reader_class, out_type): - if out_type == self.Type.BUFFER: + def _convert_frame( + self, frame: Any, reader_class: IMediaReader, out_type: FrameOutputType + ) -> AnyFrame: + if out_type == FrameOutputType.BUFFER: return self._av_frame_to_png_bytes(frame) if reader_class is VideoReader else frame - elif out_type == self.Type.PIL: + elif out_type == FrameOutputType.PIL: return frame.to_image() if reader_class is VideoReader else Image.open(frame) - elif out_type == self.Type.NUMPY_ARRAY: + elif out_type == FrameOutputType.NUMPY_ARRAY: if reader_class is VideoReader: - image = frame.to_ndarray(format='bgr24') + image = frame.to_ndarray(format="bgr24") else: image = np.array(Image.open(frame)) if len(image.shape) == 3 and image.shape[2] in {3, 4}: - image[:, :, :3] = image[:, :, 2::-1] # RGB to BGR + image[:, :, :3] = image[:, :, 2::-1] # RGB to BGR return image else: - raise RuntimeError('unsupported output type') + raise RuntimeError("unsupported output type") + + +class TaskFrameProvider(_FrameProvider): + def __init__(self, db_task: models.Task) -> None: + self._db_task = db_task + + def _validate_frame_number(self, frame_number: int) -> int: + if not (0 <= frame_number < self._db_task.data.size): + raise ValidationError(f"Incorrect requested frame number: {frame_number}") + + return frame_number + + def get_preview(self) -> DataWithMeta[BytesIO]: + return self._get_segment_frame_provider(self._db_task.data.start_frame).get_preview() + + def get_chunk( + self, chunk_number: int, *, quality: FrameQuality = FrameQuality.ORIGINAL + ) -> DataWithMeta[BytesIO]: + # TODO: return a joined chunk. Find a solution for segment boundary video chunks + return self._get_segment_frame_provider(frame_number).get_frame( + frame_number, quality=quality, out_type=out_type + ) + + def get_frame( + self, + frame_number: int, + *, + quality: FrameQuality = FrameQuality.ORIGINAL, + out_type: FrameOutputType = FrameOutputType.BUFFER, + ) -> AnyFrame: + return self._get_segment_frame_provider(frame_number).get_frame( + frame_number, quality=quality, out_type=out_type + ) + + def iterate_frames( + self, + *, + start_frame: Optional[int] = None, + stop_frame: Optional[int] = None, + quality: FrameQuality = FrameQuality.ORIGINAL, + out_type: FrameOutputType = FrameOutputType.BUFFER, + ) -> Iterator[AnyFrame]: + # TODO: optimize segment access + for idx in range(start_frame, (stop_frame + 1) if stop_frame else None): + yield self.get_frame(idx, quality=quality, out_type=out_type) - def get_preview(self, frame_number): - PREVIEW_SIZE = (256, 256) - PREVIEW_MIME = 'image/jpeg' + def _get_segment(self, validated_frame_number: int) -> models.Segment: + return next( + s + for s in self._db_task.segments.all() + if s.type == models.SegmentType.RANGE + if validated_frame_number in s.frame_set + ) - if self._dimension == DimensionType.DIM_3D: - # TODO - preview = Image.open(os.path.join(os.path.dirname(__file__), 'assets/3d_preview.jpeg')) + def _get_segment_frame_provider(self, frame_number: int) -> _SegmentFrameProvider: + segment = self._get_segment(self._validate_frame_number(frame_number)) + return _SegmentFrameProvider( + next(job for job in segment.jobs.all() if job.type == models.JobType.ANNOTATION) + ) + + +class _SegmentFrameProvider(_FrameProvider): + def __init__(self, db_segment: models.Segment) -> None: + super().__init__() + self._db_segment = db_segment + + db_data = db_segment.task.data + + reader_class: dict[models.DataChoice, IMediaReader] = { + models.DataChoice.IMAGESET: ZipReader, + models.DataChoice.VIDEO: VideoReader, + } + + self._loaders: dict[FrameQuality, _ChunkLoader] = {} + if db_data.storage_method == models.StorageMethodChoice.CACHE: + cache = MediaCache() + + self._loaders[FrameQuality.COMPRESSED] = _ChunkLoader( + reader_class[db_data.compressed_chunk_type], + lambda chunk_idx: cache.get_segment_chunk( + db_segment, chunk_idx, quality=FrameQuality.COMPRESSED + ), + ) + + self._loaders[FrameQuality.ORIGINAL] = _ChunkLoader( + reader_class[db_data.original_chunk_type], + lambda chunk_idx: cache.get_segment_chunk( + db_segment, chunk_idx, quality=FrameQuality.ORIGINAL + ), + ) else: - preview, _ = self.get_frame(frame_number, self.Quality.COMPRESSED, self.Type.PIL) + self._loaders[FrameQuality.COMPRESSED] = _ChunkLoader( + reader_class[db_data.compressed_chunk_type], db_data.get_compressed_chunk_path + ) - preview = ImageOps.exif_transpose(preview) - preview.thumbnail(PREVIEW_SIZE) + self._loaders[FrameQuality.ORIGINAL] = _ChunkLoader( + reader_class[db_data.original_chunk_type], db_data.get_original_chunk_path + ) - output_buf = BytesIO() - preview.convert('RGB').save(output_buf, format="JPEG") + def unload(self): + for loader in self._loaders.values(): + loader.unload() + + def __len__(self): + return self._db_segment.frame_count + + def _validate_frame_number(self, frame_number: int) -> Tuple[int, int, int]: + # TODO: check for masked range segment + + if frame_number not in self._db_segment.frame_set: + raise ValidationError(f"Incorrect requested frame number: {frame_number}") - return output_buf, PREVIEW_MIME + chunk_number, frame_position = divmod(frame_number, self._db_segment.task.data.chunk_size) + return frame_number, chunk_number, frame_position - def get_chunk(self, chunk_number, quality=Quality.ORIGINAL): - chunk_number = self._validate_chunk_number(chunk_number) - if self._db_data.storage_method == StorageMethodChoice.CACHE: - return self._loaders[quality].get_chunk_path(chunk_number, quality, self._db_data) - return self._loaders[quality].get_chunk_path(chunk_number) + def get_chunk_number(self, frame_number: int) -> int: + return int(frame_number) // self._db_segment.task.data.chunk_size - def get_frame(self, frame_number, quality=Quality.ORIGINAL, - out_type=Type.BUFFER): + def _validate_chunk_number(self, chunk_number: int) -> int: + segment_size = len(self._db_segment.frame_count) + if chunk_number < 0 or chunk_number >= math.ceil( + segment_size / self._db_segment.task.data.chunk_size + ): + raise ValidationError("requested chunk does not exist") + + return chunk_number + + def get_preview(self) -> DataWithMeta[BytesIO]: + if self._db_segment.task.dimension == models.DimensionType.DIM_3D: + # TODO + preview = Image.open(os.path.join(os.path.dirname(__file__), "assets/3d_preview.jpeg")) + else: + preview, _ = self.get_frame( + min(self._db_segment.frame_set), + frame_number=FrameQuality.COMPRESSED, + out_type=FrameOutputType.PIL, + ) + + return prepare_preview_image(preview) + + def get_chunk( + self, chunk_number: int, *, quality: FrameQuality = FrameQuality.ORIGINAL + ) -> DataWithMeta[BytesIO]: + return self._loaders[quality].get_chunk_path(self._validate_chunk_number(chunk_number)) + + def get_frame( + self, + frame_number: int, + *, + quality: FrameQuality = FrameQuality.ORIGINAL, + out_type: FrameOutputType = FrameOutputType.BUFFER, + ) -> AnyFrame: _, chunk_number, frame_offset = self._validate_frame_number(frame_number) loader = self._loaders[quality] chunk_reader = loader.load(chunk_number) @@ -219,10 +300,18 @@ def get_frame(self, frame_number, quality=Quality.ORIGINAL, return (frame, self.VIDEO_FRAME_MIME) return (frame, mimetypes.guess_type(frame_name)[0]) - def get_frames(self, start_frame, stop_frame, quality=Quality.ORIGINAL, out_type=Type.BUFFER): - for idx in range(start_frame, stop_frame): + def iterate_frames( + self, + *, + start_frame: Optional[int] = None, + stop_frame: Optional[int] = None, + quality: FrameQuality = FrameQuality.ORIGINAL, + out_type: FrameOutputType = FrameOutputType.BUFFER, + ) -> Iterator[AnyFrame]: + for idx in range(start_frame, (stop_frame + 1) if stop_frame else None): yield self.get_frame(idx, quality=quality, out_type=out_type) - @property - def data_id(self): - return self._db_data.id + +class JobFrameProvider(_SegmentFrameProvider): + def __init__(self, db_job: models.Job) -> None: + super().__init__(db_job.segment) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 9a352c3b930..7ca6ff0ed54 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -45,6 +45,10 @@ class ORIENTATION(IntEnum): MIRROR_HORIZONTAL_90_ROTATED=7 NORMAL_270_ROTATED=8 +class FrameQuality(IntEnum): + COMPRESSED = 0 + ORIGINAL = 100 + def get_mime(name): for type_name, type_def in MEDIA_TYPES.items(): if type_def['has_mime_type'](name): diff --git a/cvat/apps/engine/pyproject.toml b/cvat/apps/engine/pyproject.toml new file mode 100644 index 00000000000..567b7836258 --- /dev/null +++ b/cvat/apps/engine/pyproject.toml @@ -0,0 +1,12 @@ +[tool.isort] +profile = "black" +forced_separate = ["tests"] +line_length = 100 +skip_gitignore = true # align tool behavior with Black +known_first_party = ["cvat"] + +# Can't just use a pyproject in the root dir, so duplicate +# https://github.com/psf/black/issues/2863 +[tool.black] +line-length = 100 +target-version = ['py38'] From cb4ff9394eb37638de97cbc4365d1032a4d279a8 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 25 Jul 2024 19:39:21 +0300 Subject: [PATCH 002/115] t --- cvat/apps/engine/task.py | 408 +++++++++++++++++++++----------------- cvat/apps/engine/views.py | 21 +- 2 files changed, 237 insertions(+), 192 deletions(-) diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index c44f01e1f35..866e3d0c913 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -1,32 +1,34 @@ # Copyright (C) 2018-2022 Intel Corporation -# Copyright (C) 2022-2023 CVAT.ai Corporation +# Copyright (C) 2022-2024 CVAT.ai Corporation # # SPDX-License-Identifier: MIT import itertools import fnmatch import os -from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Union, Iterable -from rest_framework.serializers import ValidationError import rq import re import shutil +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Sequence, Union, Iterable from urllib import parse as urlparse from urllib import request as urlrequest -import django_rq import concurrent.futures import queue +import django_rq from django.conf import settings from django.db import transaction from django.http import HttpRequest -from datetime import datetime, timezone -from pathlib import Path +from rest_framework.serializers import ValidationError from cvat.apps.engine import models from cvat.apps.engine.log import ServerLogManager -from cvat.apps.engine.media_extractors import (MEDIA_TYPES, ImageListReader, Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter, - ValidateDimension, ZipChunkWriter, ZipCompressedChunkWriter, get_mime, sort) +from cvat.apps.engine.media_extractors import ( + MEDIA_TYPES, IMediaReader, ImageListReader, Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter, + ValidateDimension, ZipChunkWriter, ZipCompressedChunkWriter, get_mime, sort +) from cvat.apps.engine.utils import ( av_scan_paths,get_rq_job_meta, define_dependent_job, get_rq_lock_by_user, preload_images ) @@ -70,6 +72,8 @@ def create( class SegmentParams(NamedTuple): start_frame: int stop_frame: int + type: models.SegmentType = models.SegmentType.RANGE + frames: Optional[Sequence[int]] = [] class SegmentsParams(NamedTuple): segments: Iterator[SegmentParams] @@ -126,10 +130,14 @@ def _segments(): # It is assumed here that files are already saved ordered in the task # Here we just need to create segments by the job sizes start_frame = 0 - for jf in job_file_mapping: - segment_size = len(jf) + for job_files in job_file_mapping: + segment_size = len(job_files) stop_frame = start_frame + segment_size - 1 - yield SegmentParams(start_frame, stop_frame) + yield SegmentParams( + start_frame=start_frame, + stop_frame=stop_frame, + type=models.SegmentType.RANGE, + ) start_frame = stop_frame + 1 @@ -152,31 +160,39 @@ def _segments(): ) segments = ( - SegmentParams(start_frame, min(start_frame + segment_size - 1, data_size - 1)) + SegmentParams( + start_frame=start_frame, + stop_frame=min(start_frame + segment_size - 1, data_size - 1), + type=models.SegmentType.RANGE + ) for start_frame in range(0, data_size - overlap, segment_size - overlap) ) return SegmentsParams(segments, segment_size, overlap) -def _save_task_to_db(db_task: models.Task, *, job_file_mapping: Optional[JobFileMapping] = None): - job = rq.get_current_job() - job.meta['status'] = 'Task is being saved in database' - job.save_meta() +def _save_task_to_db( + db_task: models.Task, + *, + job_file_mapping: Optional[JobFileMapping] = None, +): + rq_job = rq.get_current_job() + rq_job.meta['status'] = 'Task is being saved in database' + rq_job.save_meta() segments, segment_size, overlap = _get_task_segment_data( - db_task=db_task, job_file_mapping=job_file_mapping + db_task=db_task, job_file_mapping=job_file_mapping, ) db_task.segment_size = segment_size db_task.overlap = overlap - for segment_idx, (start_frame, stop_frame) in enumerate(segments): - slogger.glob.info("New segment for task #{}: idx = {}, start_frame = {}, \ - stop_frame = {}".format(db_task.id, segment_idx, start_frame, stop_frame)) + for segment_idx, segment_params in enumerate(segments): + slogger.glob.info( + "New segment for task #{task_id}: idx = {segment_idx}, start_frame = {start_frame}, \ + stop_frame = {stop_frame}".format( + task_id=db_task.id, segment_idx=segment_idx, **segment_params._asdict() + )) - db_segment = models.Segment() - db_segment.task = db_task - db_segment.start_frame = start_frame - db_segment.stop_frame = stop_frame + db_segment = models.Segment(task=db_task, **segment_params._asdict()) db_segment.save() db_job = models.Job(segment=db_segment) @@ -756,7 +772,7 @@ def _update_status(msg: str) -> None: ) # Extract input data - extractor = None + extractor: Optional[IMediaReader] = None manifest_index = _get_manifest_frame_indexer() for media_type, media_files in media.items(): if not media_files: @@ -916,19 +932,6 @@ def _update_status(msg: str) -> None: db_data.compressed_chunk_type = models.DataChoice.VIDEO if task_mode == 'interpolation' and not data['use_zip_chunks'] else models.DataChoice.IMAGESET db_data.original_chunk_type = models.DataChoice.VIDEO if task_mode == 'interpolation' else models.DataChoice.IMAGESET - def update_progress(progress): - progress_animation = '|/-\\' - if not hasattr(update_progress, 'call_counter'): - update_progress.call_counter = 0 - - status_message = 'CVAT is preparing data chunks' - if not progress: - status_message = '{} {}'.format(status_message, progress_animation[update_progress.call_counter]) - job.meta['status'] = status_message - job.meta['task_progress'] = progress or 0. - job.save_meta() - update_progress.call_counter = (update_progress.call_counter + 1) % len(progress_animation) - compressed_chunk_writer_class = Mpeg4CompressedChunkWriter if db_data.compressed_chunk_type == models.DataChoice.VIDEO else ZipCompressedChunkWriter if db_data.original_chunk_type == models.DataChoice.VIDEO: original_chunk_writer_class = Mpeg4ChunkWriter @@ -959,135 +962,220 @@ def update_progress(progress): else: db_data.chunk_size = 36 - video_path = "" - video_size = (0, 0) + # TODO: try to pull up + # replace manifest file (e.g was uploaded 'subdir/manifest.jsonl' or 'some_manifest.jsonl') + if ( + settings.USE_CACHE and db_data.storage_method == models.StorageMethodChoice.CACHE and + manifest_file and not os.path.exists(db_data.get_manifest_path()) + ): + shutil.copyfile(os.path.join(manifest_root, manifest_file), + db_data.get_manifest_path()) + if manifest_root and manifest_root.startswith(db_data.get_upload_dirname()): + os.remove(os.path.join(manifest_root, manifest_file)) + manifest_file = os.path.relpath(db_data.get_manifest_path(), upload_dir) - db_images = [] + video_path: str = "" + video_size: tuple[int, int] = (0, 0) - if settings.USE_CACHE and db_data.storage_method == models.StorageMethodChoice.CACHE: - for media_type, media_files in media.items(): - if not media_files: - continue + images: list[models.Image] = [] - # replace manifest file (e.g was uploaded 'subdir/manifest.jsonl' or 'some_manifest.jsonl') - if manifest_file and not os.path.exists(db_data.get_manifest_path()): - shutil.copyfile(os.path.join(manifest_root, manifest_file), - db_data.get_manifest_path()) - if manifest_root and manifest_root.startswith(db_data.get_upload_dirname()): - os.remove(os.path.join(manifest_root, manifest_file)) - manifest_file = os.path.relpath(db_data.get_manifest_path(), upload_dir) + # Collect media metadata + for media_type, media_files in media.items(): + if not media_files: + continue - if task_mode == MEDIA_TYPES['video']['mode']: + if task_mode == MEDIA_TYPES['video']['mode']: + manifest_is_prepared = False + if manifest_file: try: - manifest_is_prepared = False - if manifest_file: - try: - manifest = VideoManifestValidator(source_path=os.path.join(upload_dir, media_files[0]), - manifest_path=db_data.get_manifest_path()) - manifest.init_index() - manifest.validate_seek_key_frames() - assert len(manifest) > 0, 'No key frames.' - - all_frames = manifest.video_length - video_size = manifest.video_resolution - manifest_is_prepared = True - except Exception as ex: - manifest.remove() - if isinstance(ex, AssertionError): - base_msg = str(ex) - else: - base_msg = 'Invalid manifest file was upload.' - slogger.glob.warning(str(ex)) - _update_status('{} Start prepare a valid manifest file.'.format(base_msg)) - - if not manifest_is_prepared: - _update_status('Start prepare a manifest file') - manifest = VideoManifestManager(db_data.get_manifest_path()) - manifest.link( - media_file=media_files[0], - upload_dir=upload_dir, - chunk_size=db_data.chunk_size - ) - manifest.create() - _update_status('A manifest had been created') + _update_status('Validating the input manifest file') - all_frames = len(manifest.reader) - video_size = manifest.reader.resolution - manifest_is_prepared = True + manifest = VideoManifestValidator( + source_path=os.path.join(upload_dir, media_files[0]), + manifest_path=db_data.get_manifest_path() + ) + manifest.init_index() + manifest.validate_seek_key_frames() - db_data.size = len(range(db_data.start_frame, min(data['stop_frame'] + 1 \ - if data['stop_frame'] else all_frames, all_frames), db_data.get_frame_step())) - video_path = os.path.join(upload_dir, media_files[0]) + if not len(manifest): + raise ValidationError("No key frames found in the manifest") + + all_frames = manifest.video_length + video_size = manifest.video_resolution + manifest_is_prepared = True + except Exception as ex: + manifest.remove() + manifest = None + + slogger.glob.warning(ex, exc_info=True) + if isinstance(ex, (ValidationError, AssertionError)): + _update_status(f'Invalid manifest file was upload: {ex}') + + if ( + settings.USE_CACHE and db_data.storage_method == models.StorageMethodChoice.CACHE + and not manifest_is_prepared + ): + # TODO: check if we can always use video manifest for optimization + try: + _update_status('Preparing a manifest file') + + # TODO: maybe generate manifest in a temp directory + manifest = VideoManifestManager(db_data.get_manifest_path()) + manifest.link( + media_file=media_files[0], + upload_dir=upload_dir, + chunk_size=db_data.chunk_size + ) + manifest.create() + + _update_status('A manifest has been created') + + all_frames = len(manifest.reader) # TODO: check if the field access above and here are equivalent + video_size = manifest.reader.resolution + manifest_is_prepared = True except Exception as ex: - db_data.storage_method = models.StorageMethodChoice.FILE_SYSTEM manifest.remove() - del manifest + manifest = None + + db_data.storage_method = models.StorageMethodChoice.FILE_SYSTEM + base_msg = str(ex) if isinstance(ex, AssertionError) \ else "Uploaded video does not support a quick way of task creating." _update_status("{} The task will be created using the old method".format(base_msg)) - else: # images, archive, pdf - db_data.size = len(extractor) - manifest = ImageManifestManager(db_data.get_manifest_path()) + if not manifest: + all_frames = len(extractor) + video_size = extractor.get_image_size(0) + + db_data.size = len(range( + db_data.start_frame, + min( + data['stop_frame'] + 1 if data['stop_frame'] else all_frames, + all_frames, + ), + db_data.get_frame_step() + )) + video_path = os.path.join(upload_dir, media_files[0]) + else: # images, archive, pdf + db_data.size = len(extractor) + + manifest = None + if settings.USE_CACHE and db_data.storage_method == models.StorageMethodChoice.CACHE: + manifest = ImageManifestManager(db_data.get_manifest_path()) if not manifest.exists: manifest.link( sources=extractor.absolute_source_paths, - meta={ k: {'related_images': related_images[k] } for k in related_images }, + meta={ + k: {'related_images': related_images[k] } + for k in related_images + }, data_dir=upload_dir, DIM_3D=(db_task.dimension == models.DimensionType.DIM_3D), ) manifest.create() else: manifest.init_index() - counter = itertools.count() - for _, chunk_frames in itertools.groupby(extractor.frame_range, lambda x: next(counter) // db_data.chunk_size): - chunk_paths = [(extractor.get_path(i), i) for i in chunk_frames] - img_sizes = [] - - for chunk_path, frame_id in chunk_paths: - properties = manifest[manifest_index(frame_id)] - - # check mapping - if not chunk_path.endswith(f"{properties['name']}{properties['extension']}"): - raise Exception('Incorrect file mapping to manifest content') - - if db_task.dimension == models.DimensionType.DIM_2D and ( - properties.get('width') is not None and - properties.get('height') is not None - ): - resolution = (properties['width'], properties['height']) - elif is_data_in_cloud: - raise Exception( - "Can't find image '{}' width or height info in the manifest" - .format(f"{properties['name']}{properties['extension']}") - ) - else: - resolution = extractor.get_image_size(frame_id) - img_sizes.append(resolution) - - db_images.extend([ - models.Image(data=db_data, - path=os.path.relpath(path, upload_dir), - frame=frame, width=w, height=h) - for (path, frame), (w, h) in zip(chunk_paths, img_sizes) - ]) + + for frame_id in extractor.frame_range: + image_path = extractor.get_path(frame_id) + image_size = None + + if manifest: + image_info = manifest[manifest_index(frame_id)] + + # check mapping + if not image_path.endswith(f"{image_info['name']}{image_info['extension']}"): + raise ValidationError('Incorrect file mapping to manifest content') + + if db_task.dimension == models.DimensionType.DIM_2D and ( + image_info.get('width') is not None and + image_info.get('height') is not None + ): + image_size = (image_info['width'], image_info['height']) + elif is_data_in_cloud: + raise ValidationError( + "Can't find image '{}' width or height info in the manifest" + .format(f"{image_info['name']}{image_info['extension']}") + ) + + if not image_size: + image_size = extractor.get_image_size(frame_id) + + images.append( + models.Image( + data=db_data, + path=os.path.relpath(image_path, upload_dir), + frame=frame_id, + width=image_size[0], + height=image_size[1], + ) + ) + + if db_task.mode == 'annotation': + models.Image.objects.bulk_create(images) + images = models.Image.objects.filter(data_id=db_data.id) + + db_related_files = [ + models.RelatedFile(data=image.data, primary_image=image, path=os.path.join(upload_dir, related_file_path)) + for image in images + for related_file_path in related_images.get(image.path, []) + if not image.is_placeholder # TODO + ] + models.RelatedFile.objects.bulk_create(db_related_files) + else: + models.Video.objects.create( + data=db_data, + path=os.path.relpath(video_path, upload_dir), + width=video_size[0], height=video_size[1] + ) + + # validate stop_frame + if db_data.stop_frame == 0: + db_data.stop_frame = db_data.start_frame + (db_data.size - 1) * db_data.get_frame_step() + else: + db_data.stop_frame = min(db_data.stop_frame, \ + db_data.start_frame + (db_data.size - 1) * db_data.get_frame_step()) + + slogger.glob.info("Found frames {} for Data #{}".format(db_data.size, db_data.id)) + _save_task_to_db(db_task, job_file_mapping=job_file_mapping) # TODO: split into jobs and task saving + + # Save chunks + # TODO: refactor + # TODO: save chunks per job if db_data.storage_method == models.StorageMethodChoice.FILE_SYSTEM or not settings.USE_CACHE: + def update_progress(progress): + # TODO: refactor this function into a class + progress_animation = '|/-\\' + if not hasattr(update_progress, 'call_counter'): + update_progress.call_counter = 0 + + status_message = 'CVAT is preparing data chunks' + if not progress: + status_message = '{} {}'.format(status_message, progress_animation[update_progress.call_counter]) + job.meta['status'] = status_message + job.meta['task_progress'] = progress or 0. + job.save_meta() + update_progress.call_counter = (update_progress.call_counter + 1) % len(progress_animation) + + counter = itertools.count() generator = itertools.groupby(extractor, lambda _: next(counter) // db_data.chunk_size) - generator = ((idx, list(chunk_data)) for idx, chunk_data in generator) + generator = ((chunk_idx, list(chunk_data)) for chunk_idx, chunk_data in generator) def save_chunks( - executor: concurrent.futures.ThreadPoolExecutor, - chunk_idx: int, - chunk_data: Iterable[tuple[str, str, str]]) -> list[tuple[str, int, tuple[int, int]]]: - nonlocal db_data, db_task, extractor, original_chunk_writer, compressed_chunk_writer - if (db_task.dimension == models.DimensionType.DIM_2D and + executor: concurrent.futures.ThreadPoolExecutor, + chunk_idx: int, + chunk_data: Iterable[tuple[str, str, str]] + ) -> list[tuple[str, int, tuple[int, int]]]: + if ( + db_task.dimension == models.DimensionType.DIM_2D and isinstance(extractor, ( MEDIA_TYPES['image']['extractor'], MEDIA_TYPES['zip']['extractor'], MEDIA_TYPES['pdf']['extractor'], MEDIA_TYPES['archive']['extractor'], - ))): + )) + ): chunk_data = preload_images(chunk_data) fs_original = executor.submit( @@ -1100,6 +1188,7 @@ def save_chunks( images=chunk_data, chunk_path=db_data.get_compressed_chunk_path(chunk_idx), ) + # TODO: convert to async for proper concurrency fs_original.result() image_sizes = fs_compressed.result() @@ -1107,58 +1196,17 @@ def save_chunks( return list((i[0][1], i[0][2], i[1]) for i in zip(chunk_data, image_sizes)) def process_results(img_meta: list[tuple[str, int, tuple[int, int]]]): - nonlocal db_images, db_data, video_path, video_size - - if db_task.mode == 'annotation': - db_images.extend( - models.Image( - data=db_data, - path=os.path.relpath(frame_path, upload_dir), - frame=frame_number, - width=frame_size[0], - height=frame_size[1]) - for frame_path, frame_number, frame_size in img_meta) - else: - video_size = img_meta[0][2] - video_path = img_meta[0][0] - - progress = extractor.get_progress(img_meta[-1][1]) + progress = img_meta[-1][1] / db_data.size update_progress(progress) futures = queue.Queue(maxsize=settings.CVAT_CONCURRENT_CHUNK_PROCESSING) - with concurrent.futures.ThreadPoolExecutor(max_workers=2*settings.CVAT_CONCURRENT_CHUNK_PROCESSING) as executor: + with concurrent.futures.ThreadPoolExecutor( + max_workers=2 * settings.CVAT_CONCURRENT_CHUNK_PROCESSING + ) as executor: for chunk_idx, chunk_data in generator: - db_data.size += len(chunk_data) if futures.full(): process_results(futures.get().result()) futures.put(executor.submit(save_chunks, executor, chunk_idx, chunk_data)) while not futures.empty(): process_results(futures.get().result()) - - if db_task.mode == 'annotation': - models.Image.objects.bulk_create(db_images) - created_images = models.Image.objects.filter(data_id=db_data.id) - - db_related_files = [ - models.RelatedFile(data=image.data, primary_image=image, path=os.path.join(upload_dir, related_file_path)) - for image in created_images - for related_file_path in related_images.get(image.path, []) - ] - models.RelatedFile.objects.bulk_create(db_related_files) - db_images = [] - else: - models.Video.objects.create( - data=db_data, - path=os.path.relpath(video_path, upload_dir), - width=video_size[0], height=video_size[1]) - - if db_data.stop_frame == 0: - db_data.stop_frame = db_data.start_frame + (db_data.size - 1) * db_data.get_frame_step() - else: - # validate stop_frame - db_data.stop_frame = min(db_data.stop_frame, \ - db_data.start_frame + (db_data.size - 1) * db_data.get_frame_step()) - - slogger.glob.info("Found frames {} for Data #{}".format(db_data.size, db_data.id)) - _save_task_to_db(db_task, job_file_mapping=job_file_mapping) diff --git a/cvat/apps/engine/views.py b/cvat/apps/engine/views.py index a0edc4ae340..13a7679595e 100644 --- a/cvat/apps/engine/views.py +++ b/cvat/apps/engine/views.py @@ -54,7 +54,7 @@ from cvat.apps.events.handlers import handle_dataset_import from cvat.apps.dataset_manager.bindings import CvatImportError from cvat.apps.dataset_manager.serializers import DatasetFormatsSerializer -from cvat.apps.engine.frame_provider import FrameProvider +from cvat.apps.engine.frame_provider import FrameProvider, JobFrameProvider from cvat.apps.engine.filters import NonModelSimpleFilter, NonModelOrderingFilter, NonModelJsonLogicFilter from cvat.apps.engine.media_extractors import get_mime from cvat.apps.engine.models import ( @@ -625,7 +625,7 @@ def append_backup_chunk(self, request, file_id): def preview(self, request, pk): self._object = self.get_object() # call check_object_permissions as well - first_task = self._object.tasks.order_by('-id').first() + first_task: Optional[models.Task] = self._object.tasks.order_by('-id').first() if not first_task: return HttpResponseNotFound('Project image preview not found') @@ -746,13 +746,13 @@ def __init__(self, job: Job, data_type, data_num, data_quality): self.job = job def _check_frame_range(self, frame: int): - frame_range = self.job.segment.frame_set - if frame not in frame_range: + if frame not in self.job.segment.frame_set: raise ValidationError("The frame number doesn't belong to the job") def __call__(self, request, start, stop, db_data): + # TODO: add segment boundary handling if self.type == 'chunk' and self.job.segment.type == SegmentType.SPECIFIC_FRAMES: - frame_provider = FrameProvider(db_data, self.dimension) + frame_provider = JobFrameProvider(self.job) start_chunk = frame_provider.get_chunk_number(start) stop_chunk = frame_provider.get_chunk_number(stop) @@ -764,12 +764,10 @@ def __call__(self, request, start, stop, db_data): cache = MediaCache() if settings.USE_CACHE and db_data.storage_method == StorageMethodChoice.CACHE: - buf, mime = cache.get_selective_job_chunk_data_with_mime( - chunk_number=self.number, quality=self.quality, job=self.job - ) + buf, mime = cache.get_segment_chunk(self.job, self.number, quality=self.quality) else: - buf, mime = cache.prepare_selective_job_chunk( - chunk_number=self.number, quality=self.quality, db_job=self.job + buf, mime = cache.prepare_masked_range_segment_chunk( + self.job, self.number, quality=self.quality ) return HttpResponse(buf.getvalue(), content_type=mime) @@ -1298,8 +1296,7 @@ def data(self, request, pk): data_num = request.query_params.get('number', None) data_quality = request.query_params.get('quality', 'compressed') - data_getter = DataChunkGetter(data_type, data_num, data_quality, - self._object.dimension) + data_getter = DataChunkGetter(data_type, data_num, data_quality, self._object.dimension) return data_getter(request, self._object.data.start_frame, self._object.data.stop_frame, self._object.data) From d49233c60bb93d3f3fd34bd553b830f7ba18e96f Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 30 Jul 2024 19:18:57 +0300 Subject: [PATCH 003/115] t --- cvat/apps/dataset_manager/bindings.py | 16 +- cvat/apps/dataset_manager/formats/cvat.py | 23 +- cvat/apps/engine/cache.py | 128 +++++---- cvat/apps/engine/frame_provider.py | 313 +++++++++++++++++----- cvat/apps/engine/media_extractors.py | 104 +++---- cvat/apps/engine/models.py | 10 +- cvat/apps/engine/views.py | 189 ++++++------- cvat/apps/lambda_manager/views.py | 8 +- 8 files changed, 452 insertions(+), 339 deletions(-) diff --git a/cvat/apps/dataset_manager/bindings.py b/cvat/apps/dataset_manager/bindings.py index d94d6fd39e3..dfe8d458de6 100644 --- a/cvat/apps/dataset_manager/bindings.py +++ b/cvat/apps/dataset_manager/bindings.py @@ -31,7 +31,7 @@ from cvat.apps.dataset_manager.formats.utils import get_label_color from cvat.apps.dataset_manager.util import add_prefetch_fields -from cvat.apps.engine.frame_provider import FrameProvider +from cvat.apps.engine.frame_provider import TaskFrameProvider, FrameQuality, FrameOutputType from cvat.apps.engine.models import (AttributeSpec, AttributeType, Data, DimensionType, Job, JobType, Label, LabelType, Project, SegmentType, ShapeType, Task) @@ -1348,8 +1348,8 @@ def video_frame_loader(_): # optimization for videos: use numpy arrays instead of bytes # some formats or transforms can require image data return self._frame_provider.get_frame(frame_index, - quality=FrameProvider.Quality.ORIGINAL, - out_type=FrameProvider.Type.NUMPY_ARRAY)[0] + quality=FrameQuality.ORIGINAL, + out_type=FrameOutputType.NUMPY_ARRAY).data return dm.Image(data=video_frame_loader, **image_kwargs) else: def image_loader(_): @@ -1357,8 +1357,8 @@ def image_loader(_): # for images use encoded data to avoid recoding return self._frame_provider.get_frame(frame_index, - quality=FrameProvider.Quality.ORIGINAL, - out_type=FrameProvider.Type.BUFFER)[0].getvalue() + quality=FrameQuality.ORIGINAL, + out_type=FrameOutputType.BUFFER).data.getvalue() return dm.ByteImage(data=image_loader, **image_kwargs) def _load_source(self, source_id: int, source: ImageSource) -> None: @@ -1366,7 +1366,7 @@ def _load_source(self, source_id: int, source: ImageSource) -> None: return self._unload_source() - self._frame_provider = FrameProvider(source.db_data) + self._frame_provider = TaskFrameProvider(next(iter(source.db_data.tasks))) # TODO: refactor self._current_source_id = source_id def _unload_source(self) -> None: @@ -1502,7 +1502,7 @@ def __init__( is_video = instance_meta['mode'] == 'interpolation' ext = '' if is_video: - ext = FrameProvider.VIDEO_FRAME_EXT + ext = TaskFrameProvider.VIDEO_FRAME_EXT if dimension == DimensionType.DIM_3D or include_images: self._image_provider = IMAGE_PROVIDERS_BY_DIMENSION[dimension]( @@ -1593,7 +1593,7 @@ def __init__( ) ext_per_task: Dict[int, str] = { - task.id: FrameProvider.VIDEO_FRAME_EXT if is_video else '' + task.id: TaskFrameProvider.VIDEO_FRAME_EXT if is_video else '' for task in project_data.tasks for is_video in [task.mode == 'interpolation'] } diff --git a/cvat/apps/dataset_manager/formats/cvat.py b/cvat/apps/dataset_manager/formats/cvat.py index 99293fe470d..4d59682f6ee 100644 --- a/cvat/apps/dataset_manager/formats/cvat.py +++ b/cvat/apps/dataset_manager/formats/cvat.py @@ -27,7 +27,7 @@ import_dm_annotations, match_dm_item) from cvat.apps.dataset_manager.util import make_zip_archive -from cvat.apps.engine.frame_provider import FrameProvider +from cvat.apps.engine.frame_provider import FrameQuality, FrameOutputType, make_frame_provider from .registry import dm_env, exporter, importer @@ -1371,16 +1371,19 @@ def dump_project_anno(dst_file: BufferedWriter, project_data: ProjectData, callb dumper.close_document() def dump_media_files(instance_data: CommonData, img_dir: str, project_data: ProjectData = None): + frame_provider = make_frame_provider(instance_data.db_instance) + ext = '' if instance_data.meta[instance_data.META_FIELD]['mode'] == 'interpolation': - ext = FrameProvider.VIDEO_FRAME_EXT - - frame_provider = FrameProvider(instance_data.db_data) - frames = frame_provider.get_frames( - instance_data.start, instance_data.stop, - frame_provider.Quality.ORIGINAL, - frame_provider.Type.BUFFER) - for frame_id, (frame_data, _) in zip(instance_data.rel_range, frames): + ext = frame_provider.VIDEO_FRAME_EXT + + frames = frame_provider.iterate_frames( + start_frame=instance_data.start, + stop_frame=instance_data.stop, + quality=FrameQuality.ORIGINAL, + out_type=FrameOutputType.BUFFER, + ) + for frame_id, frame in zip(instance_data.rel_range, frames): if (project_data is not None and (instance_data.db_instance.id, frame_id) in project_data.deleted_frames) \ or frame_id in instance_data.deleted_frames: continue @@ -1389,7 +1392,7 @@ def dump_media_files(instance_data: CommonData, img_dir: str, project_data: Proj img_path = osp.join(img_dir, frame_name + ext) os.makedirs(osp.dirname(img_path), exist_ok=True) with open(img_path, 'wb') as f: - f.write(frame_data.getvalue()) + f.write(frame.data.getvalue()) def _export_task_or_job(dst_file, temp_dir, instance_data, anno_callback, save_images=False): with open(osp.join(temp_dir, 'annotations.xml'), 'wb') as f: diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 988e7676121..521086296bb 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -8,14 +8,15 @@ import io import os import pickle # nosec -import shutil import tempfile import zipfile import zlib -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from datetime import datetime, timezone -from typing import Any, Callable, Optional, Sequence, Tuple, Type +from itertools import pairwise +from typing import Any, Callable, Iterable, Optional, Sequence, Tuple, Type, Union +import av import cv2 import PIL.Image import PIL.ImageOps @@ -33,10 +34,10 @@ from cvat.apps.engine.media_extractors import ( FrameQuality, IChunkWriter, - ImageDatasetManifestReader, + ImageReaderWithManifest, Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter, - VideoDatasetManifestReader, + VideoReaderWithManifest, ZipChunkWriter, ZipCompressedChunkWriter, ) @@ -127,10 +128,10 @@ def get_selective_job_chunk( ), ) - def get_local_preview(self, db_data: models.Data, frame_number: int) -> DataWithMime: + def get_or_set_segment_preview(self, db_segment: models.Segment) -> DataWithMime: return self._get_or_set_cache_item( - key=f"data_{db_data.id}_{frame_number}_preview", - create_callback=lambda: self._prepare_local_preview(frame_number, db_data), + f"segment_preview_{db_segment.id}", + create_callback=lambda: self._prepare_segment_preview(db_segment), ) def get_cloud_preview(self, db_storage: models.CloudStorage) -> Optional[DataWithMime]: @@ -148,18 +149,17 @@ def get_frame_context_images(self, db_data: models.Data, frame_number: int) -> D create_callback=lambda: self._prepare_context_image(db_data, frame_number), ) - def get_task_preview(self, db_task: models.Task) -> Optional[DataWithMime]: - return self._get(f"task_{db_task.data_id}_preview") - - def get_segment_preview(self, db_segment: models.Segment) -> Optional[DataWithMime]: - return self._get(f"segment_{db_segment.id}_preview") - @contextmanager - def _read_raw_frames(self, db_task: models.Task, frames: Sequence[int]): + def _read_raw_frames( + self, db_task: models.Task, frame_ids: Sequence[int] + ) -> Iterable[Tuple[Union[av.VideoFrame, PIL.Image.Image], str, str]]: + for prev_frame, cur_frame in pairwise(frame_ids): + assert ( + prev_frame <= cur_frame + ), f"Requested frame ids must be sorted, got a ({prev_frame}, {cur_frame}) pair" + db_data = db_task.data - media = [] - tmp_dir = None raw_data_dir = { models.StorageChoice.LOCAL: db_data.get_upload_dirname(), models.StorageChoice.SHARE: settings.SHARE_ROOT, @@ -168,33 +168,20 @@ def _read_raw_frames(self, db_task: models.Task, frames: Sequence[int]): dimension = db_task.dimension - # TODO - try: + media = [] + with ExitStack() as es: if hasattr(db_data, "video"): source_path = os.path.join(raw_data_dir, db_data.video.path) - # TODO: refactor to allow non-manifest videos - reader = VideoDatasetManifestReader( + reader = VideoReaderWithManifest( manifest_path=db_data.get_manifest_path(), source_path=source_path, - chunk_number=chunk_number, - chunk_size=db_data.chunk_size, - start=db_data.start_frame, - stop=db_data.stop_frame, - step=db_data.get_frame_step(), ) - for frame in reader: + for frame in reader.iterate_frames(frame_ids): media.append((frame, source_path, None)) else: - reader = ImageDatasetManifestReader( - manifest_path=db_data.get_manifest_path(), - chunk_number=chunk_number, - chunk_size=db_data.chunk_size, - start=db_data.start_frame, - stop=db_data.stop_frame, - step=db_data.get_frame_step(), - ) - if db_data.storage == StorageChoice.CLOUD_STORAGE: + reader = ImageReaderWithManifest(db_data.get_manifest_path()) + if db_data.storage == models.StorageChoice.CLOUD_STORAGE: db_cloud_storage = db_data.cloud_storage assert db_cloud_storage, "Cloud storage instance was deleted" credentials = Credentials() @@ -213,10 +200,10 @@ def _read_raw_frames(self, db_task: models.Task, frames: Sequence[int]): cloud_provider=db_cloud_storage.provider_type, **details ) - tmp_dir = tempfile.mkdtemp(prefix="cvat") + tmp_dir = es.enter_context(tempfile.TemporaryDirectory(prefix="cvat")) files_to_download = [] checksums = [] - for item in reader: + for item in reader.iterate_frames(frame_ids): file_name = f"{item['name']}{item['extension']}" fs_filename = os.path.join(tmp_dir, file_name) @@ -235,18 +222,16 @@ def _read_raw_frames(self, db_task: models.Task, frames: Sequence[int]): "Hash sums of files {} do not match".format(file_name) ) else: - for item in reader: + for item in reader.iterate_frames(frame_ids): source_path = os.path.join( raw_data_dir, f"{item['name']}{item['extension']}" ) media.append((source_path, source_path, None)) + if dimension == models.DimensionType.DIM_2D: media = preload_images(media) yield media - finally: - if db_data.storage == models.StorageChoice.CLOUD_STORAGE and tmp_dir is not None: - shutil.rmtree(tmp_dir) def prepare_segment_chunk( self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality @@ -267,8 +252,8 @@ def prepare_range_segment_chunk( db_data = db_task.data chunk_size = db_data.chunk_size - chunk_frames = db_segment.frame_set[ - chunk_size * chunk_number : chunk_number * (chunk_number + 1) + chunk_frame_ids = db_segment.frame_set[ + chunk_size * chunk_number : chunk_size * (chunk_number + 1) ] writer_classes: dict[FrameQuality, Type[IChunkWriter]] = { @@ -300,26 +285,29 @@ def prepare_range_segment_chunk( kwargs["dimension"] = models.DimensionType.DIM_3D writer = writer_classes[quality](image_quality, **kwargs) - buff = io.BytesIO() - with self._read_raw_frames(db_task, frames=chunk_frames) as images: - writer.save_as_chunk(images, buff) + buffer = io.BytesIO() + with self._read_raw_frames(db_task, frame_ids=chunk_frame_ids) as images: + writer.save_as_chunk(images, buffer) - buff.seek(0) - return buff, mime_type + buffer.seek(0) + return buffer, mime_type def prepare_masked_range_segment_chunk( - self, db_job: models.Job, quality, chunk_number: int + self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality ) -> DataWithMime: - db_data = db_job.segment.task.data + # TODO: try to refactor into 1 function with prepare_range_segment_chunk() + db_task = db_segment.task + db_data = db_task.data - FrameProvider = self._get_frame_provider_class() - frame_provider = FrameProvider(db_data, self._dimension) + from cvat.apps.engine.frame_provider import TaskFrameProvider - frame_set = db_job.segment.frame_set + frame_provider = TaskFrameProvider(db_task) + + frame_set = db_segment.frame_set frame_step = db_data.get_frame_step() chunk_frames = [] - writer = ZipCompressedChunkWriter(db_data.image_quality, dimension=self._dimension) + writer = ZipCompressedChunkWriter(db_data.image_quality, dimension=db_task.dimension) dummy_frame = io.BytesIO() PIL.Image.new("RGB", (1, 1)).save(dummy_frame, writer.IMAGE_EXT) @@ -370,18 +358,23 @@ def prepare_masked_range_segment_chunk( return buff, "application/zip" - def prepare_segment_preview(self, db_segment: models.Segment) -> DataWithMime: - if db_segment.task.data.cloud_storage: - return self._prepare_cloud_segment_preview(db_segment) + def _prepare_segment_preview(self, db_segment: models.Segment) -> DataWithMime: + if db_segment.task.dimension == models.DimensionType.DIM_3D: + # TODO + preview = PIL.Image.open( + os.path.join(os.path.dirname(__file__), "assets/3d_preview.jpeg") + ) else: - return self._prepare_local_segment_preview(db_segment) + from cvat.apps.engine.frame_provider import FrameOutputType, SegmentFrameProvider - def _prepare_local_preview(self, db_data: models.Data, frame_number: int) -> DataWithMime: - FrameProvider = self._get_frame_provider_class() - frame_provider = FrameProvider(db_data, self._dimension) - buff, mime_type = frame_provider.get_preview(frame_number) + segment_frame_provider = SegmentFrameProvider(db_segment) + preview = segment_frame_provider.get_frame( + min(db_segment.frame_set), + quality=FrameQuality.COMPRESSED, + out_type=FrameOutputType.PIL, + ).data - return buff, mime_type + return prepare_preview_image(preview) def _prepare_cloud_preview(self, db_storage): storage = db_storage_to_storage_instance(db_storage) @@ -428,9 +421,10 @@ def prepare_context_images( except models.Image.DoesNotExist: return None + if not image.related_files.count(): + return None, None + with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED, False) as zip_file: - if not image.related_files.count(): - return None, None common_path = os.path.commonpath( list(map(lambda x: str(x.path), image.related_files.all())) ) @@ -457,4 +451,4 @@ def prepare_preview_image(image: PIL.Image.Image) -> DataWithMime: output_buf = io.BytesIO() image.convert("RGB").save(output_buf, format="JPEG") - return image, PREVIEW_MIME + return output_buf, PREVIEW_MIME diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 92354723b02..77794d60ff7 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -5,8 +5,9 @@ from __future__ import annotations +import io import math -import os +from abc import ABCMeta, abstractmethod from dataclasses import dataclass from enum import Enum, auto from io import BytesIO @@ -19,8 +20,14 @@ from rest_framework.exceptions import ValidationError from cvat.apps.engine import models -from cvat.apps.engine.cache import DataWithMime, MediaCache, prepare_preview_image -from cvat.apps.engine.media_extractors import FrameQuality, IMediaReader, VideoReader, ZipReader +from cvat.apps.engine.cache import DataWithMime, MediaCache +from cvat.apps.engine.media_extractors import ( + FrameQuality, + IMediaReader, + VideoReader, + ZipCompressedChunkWriter, + ZipReader, +) from cvat.apps.engine.mime_types import mimetypes _T = TypeVar("_T") @@ -61,14 +68,11 @@ def close(self): self.pos = -1 -class _ChunkLoader: - def __init__( - self, reader_class: IMediaReader, path_getter: Callable[[int], DataWithMime] - ) -> None: +class _ChunkLoader(metaclass=ABCMeta): + def __init__(self, reader_class: IMediaReader) -> None: self.chunk_id: Optional[int] = None self.chunk_reader: Optional[_RandomAccessIterator] = None self.reader_class = reader_class - self.get_chunk_path = path_getter def load(self, chunk_id: int) -> _RandomAccessIterator[Tuple[Any, str, int]]: if self.chunk_id != chunk_id: @@ -76,7 +80,7 @@ def load(self, chunk_id: int) -> _RandomAccessIterator[Tuple[Any, str, int]]: self.chunk_id = chunk_id self.chunk_reader = _RandomAccessIterator( - self.reader_class([self.get_chunk_path(chunk_id)]) + self.reader_class([self.read_chunk(chunk_id)[0]]) ) return self.chunk_reader @@ -86,6 +90,36 @@ def unload(self): self.chunk_reader.close() self.chunk_reader = None + @abstractmethod + def read_chunk(self, chunk_id: int) -> DataWithMime: ... + + +class _FileChunkLoader(_ChunkLoader): + def __init__( + self, reader_class: IMediaReader, get_chunk_path_callback: Callable[[int], str] + ) -> None: + super().__init__(reader_class) + self.get_chunk_path = get_chunk_path_callback + + def read_chunk(self, chunk_id: int) -> DataWithMime: + chunk_path = self.get_chunk_path(chunk_id) + with open(chunk_path, "r") as f: + return ( + io.BytesIO(f.read()), + mimetypes.guess_type(chunk_path)[0], + ) + + +class _BufferChunkLoader(_ChunkLoader): + def __init__( + self, reader_class: IMediaReader, get_chunk_callback: Callable[[int], DataWithMime] + ) -> None: + super().__init__(reader_class) + self.get_chunk = get_chunk_callback + + def read_chunk(self, chunk_id: int) -> DataWithMime: + return self.get_chunk(chunk_id) + class FrameOutputType(Enum): BUFFER = auto() @@ -105,7 +139,7 @@ class DataWithMeta(Generic[_T]): checksum: int -class _FrameProvider: +class IFrameProvider(metaclass=ABCMeta): VIDEO_FRAME_EXT = ".PNG" VIDEO_FRAME_MIME = "image/png" @@ -118,7 +152,7 @@ def _av_frame_to_png_bytes(cls, av_frame: av.VideoFrame) -> BytesIO: image = av_frame.to_ndarray(format="bgr24") success, result = cv2.imencode(ext, image) if not success: - raise RuntimeError("Failed to encode image to '%s' format" % (ext)) + raise RuntimeError(f"Failed to encode image to '{ext}' format") return BytesIO(result.tobytes()) def _convert_frame( @@ -139,27 +173,145 @@ def _convert_frame( else: raise RuntimeError("unsupported output type") + @abstractmethod + def validate_frame_number(self, frame_number: int) -> int: ... + + @abstractmethod + def validate_chunk_number(self, chunk_number: int) -> int: ... + + @abstractmethod + def get_chunk_number(self, frame_number: int) -> int: ... + + @abstractmethod + def get_preview(self) -> DataWithMeta[BytesIO]: ... + + @abstractmethod + def get_chunk( + self, chunk_number: int, *, quality: FrameQuality = FrameQuality.ORIGINAL + ) -> DataWithMeta[BytesIO]: ... + + @abstractmethod + def get_frame( + self, + frame_number: int, + *, + quality: FrameQuality = FrameQuality.ORIGINAL, + out_type: FrameOutputType = FrameOutputType.BUFFER, + ) -> DataWithMeta[AnyFrame]: ... + + @abstractmethod + def get_frame_context_images( + self, + frame_number: int, + ) -> Optional[DataWithMeta[BytesIO]]: ... + + @abstractmethod + def iterate_frames( + self, + *, + start_frame: Optional[int] = None, + stop_frame: Optional[int] = None, + quality: FrameQuality = FrameQuality.ORIGINAL, + out_type: FrameOutputType = FrameOutputType.BUFFER, + ) -> Iterator[DataWithMeta[AnyFrame]]: ... + -class TaskFrameProvider(_FrameProvider): +class TaskFrameProvider(IFrameProvider): def __init__(self, db_task: models.Task) -> None: self._db_task = db_task - def _validate_frame_number(self, frame_number: int) -> int: - if not (0 <= frame_number < self._db_task.data.size): - raise ValidationError(f"Incorrect requested frame number: {frame_number}") + def validate_frame_number(self, frame_number: int) -> int: + start = self._db_task.data.start_frame + stop = self._db_task.data.stop_frame + if frame_number not in range(start, stop + 1, self._db_task.data.get_frame_step()): + raise ValidationError( + f"Invalid frame '{frame_number}'. " + f"The frame number should be in the [{start}, {stop}] range" + ) return frame_number + def validate_chunk_number(self, chunk_number: int) -> int: + start_chunk = 0 + stop_chunk = math.ceil(self._db_task.data.size / self._db_task.data.chunk_size) + if not (start_chunk <= chunk_number <= stop_chunk): + raise ValidationError( + f"Invalid chunk number '{chunk_number}'. " + f"The chunk number should be in the [{start_chunk}, {stop_chunk}] range" + ) + + return chunk_number + + def get_chunk_number(self, frame_number: int) -> int: + return int(frame_number) // self._db_task.data.chunk_size + def get_preview(self) -> DataWithMeta[BytesIO]: return self._get_segment_frame_provider(self._db_task.data.start_frame).get_preview() def get_chunk( self, chunk_number: int, *, quality: FrameQuality = FrameQuality.ORIGINAL ) -> DataWithMeta[BytesIO]: + return_type = DataWithMeta[BytesIO] + chunk_number = self.validate_chunk_number(chunk_number) + # TODO: return a joined chunk. Find a solution for segment boundary video chunks - return self._get_segment_frame_provider(frame_number).get_frame( - frame_number, quality=quality, out_type=out_type + db_data = self._db_task.data + step = db_data.get_frame_step() + task_chunk_start_frame = chunk_number * db_data.chunk_size + task_chunk_stop_frame = (chunk_number + 1) * db_data.chunk_size - 1 + task_chunk_frame_set = set( + range( + db_data.start_frame + task_chunk_start_frame * step, + min(db_data.start_frame + task_chunk_stop_frame * step, db_data.stop_frame) + step, + step, + ) + ) + + matching_segments = sorted( + [ + s + for s in self._db_task.segment_set + if s.type == models.SegmentType.RANGE + if not task_chunk_frame_set.isdisjoint(s.frame_set) + ], + key=lambda s: s.start_frame, + ) + assert matching_segments + + if len(matching_segments) == 1: + segment_frame_provider = SegmentFrameProvider(matching_segments[0]) + return segment_frame_provider.get_chunk( + segment_frame_provider.get_chunk_number(task_chunk_start_frame), quality=quality + ) + + task_chunk_frames = [] + for db_segment in matching_segments: + segment_frame_provider = SegmentFrameProvider(db_segment) + segment_frame_set = db_segment.frame_set + + for task_chunk_frame_id in task_chunk_frame_set: + if task_chunk_frame_id not in segment_frame_set: + continue + + frame = segment_frame_provider.get_frame( + task_chunk_frame_id, quality=quality, out_type=FrameOutputType.BUFFER + ) + task_chunk_frames.append((frame, None, None)) + + merged_chunk_writer = ZipCompressedChunkWriter( + db_data.image_quality, dimension=self._db_task.dimension + ) + + buffer = io.BytesIO() + merged_chunk_writer.save_as_chunk( + task_chunk_frames, + buffer, + compress_frames=False, + zip_compress_level=1, ) + buffer.seek(0) + + return return_type(data=buffer, mime="application/zip", checksum=None) def get_frame( self, @@ -167,11 +319,17 @@ def get_frame( *, quality: FrameQuality = FrameQuality.ORIGINAL, out_type: FrameOutputType = FrameOutputType.BUFFER, - ) -> AnyFrame: + ) -> DataWithMeta[AnyFrame]: return self._get_segment_frame_provider(frame_number).get_frame( frame_number, quality=quality, out_type=out_type ) + def get_frame_context_images( + self, + frame_number: int, + ) -> Optional[DataWithMeta[BytesIO]]: + return self._get_segment_frame_provider(frame_number).get_frame_context_images(frame_number) + def iterate_frames( self, *, @@ -179,7 +337,7 @@ def iterate_frames( stop_frame: Optional[int] = None, quality: FrameQuality = FrameQuality.ORIGINAL, out_type: FrameOutputType = FrameOutputType.BUFFER, - ) -> Iterator[AnyFrame]: + ) -> Iterator[DataWithMeta[AnyFrame]]: # TODO: optimize segment access for idx in range(start_frame, (stop_frame + 1) if stop_frame else None): yield self.get_frame(idx, quality=quality, out_type=out_type) @@ -187,19 +345,16 @@ def iterate_frames( def _get_segment(self, validated_frame_number: int) -> models.Segment: return next( s - for s in self._db_task.segments.all() + for s in self._db_task.segment_set.all() if s.type == models.SegmentType.RANGE if validated_frame_number in s.frame_set ) - def _get_segment_frame_provider(self, frame_number: int) -> _SegmentFrameProvider: - segment = self._get_segment(self._validate_frame_number(frame_number)) - return _SegmentFrameProvider( - next(job for job in segment.jobs.all() if job.type == models.JobType.ANNOTATION) - ) + def _get_segment_frame_provider(self, frame_number: int) -> SegmentFrameProvider: + return SegmentFrameProvider(self._get_segment(self.validate_frame_number(frame_number))) -class _SegmentFrameProvider(_FrameProvider): +class SegmentFrameProvider(IFrameProvider): def __init__(self, db_segment: models.Segment) -> None: super().__init__() self._db_segment = db_segment @@ -215,26 +370,32 @@ def __init__(self, db_segment: models.Segment) -> None: if db_data.storage_method == models.StorageMethodChoice.CACHE: cache = MediaCache() - self._loaders[FrameQuality.COMPRESSED] = _ChunkLoader( - reader_class[db_data.compressed_chunk_type], - lambda chunk_idx: cache.get_segment_chunk( + self._loaders[FrameQuality.COMPRESSED] = _BufferChunkLoader( + reader_class=reader_class[db_data.compressed_chunk_type], + get_chunk_callback=lambda chunk_idx: cache.get_segment_chunk( db_segment, chunk_idx, quality=FrameQuality.COMPRESSED ), ) - self._loaders[FrameQuality.ORIGINAL] = _ChunkLoader( - reader_class[db_data.original_chunk_type], - lambda chunk_idx: cache.get_segment_chunk( + self._loaders[FrameQuality.ORIGINAL] = _BufferChunkLoader( + reader_class=reader_class[db_data.original_chunk_type], + get_chunk_callback=lambda chunk_idx: cache.get_segment_chunk( db_segment, chunk_idx, quality=FrameQuality.ORIGINAL ), ) else: - self._loaders[FrameQuality.COMPRESSED] = _ChunkLoader( - reader_class[db_data.compressed_chunk_type], db_data.get_compressed_chunk_path + self._loaders[FrameQuality.COMPRESSED] = _FileChunkLoader( + reader_class=reader_class[db_data.compressed_chunk_type], + get_chunk_path_callback=lambda chunk_idx: db_data.get_compressed_segment_chunk_path( + chunk_idx, segment=db_segment.id + ), ) - self._loaders[FrameQuality.ORIGINAL] = _ChunkLoader( - reader_class[db_data.original_chunk_type], db_data.get_original_chunk_path + self._loaders[FrameQuality.ORIGINAL] = _FileChunkLoader( + reader_class=reader_class[db_data.original_chunk_type], + get_chunk_path_callback=lambda chunk_idx: db_data.get_original_segment_chunk_path( + chunk_idx, segment=db_segment.id + ), ) def unload(self): @@ -244,9 +405,7 @@ def unload(self): def __len__(self): return self._db_segment.frame_count - def _validate_frame_number(self, frame_number: int) -> Tuple[int, int, int]: - # TODO: check for masked range segment - + def validate_frame_number(self, frame_number: int) -> Tuple[int, int, int]: if frame_number not in self._db_segment.frame_set: raise ValidationError(f"Incorrect requested frame number: {frame_number}") @@ -256,32 +415,29 @@ def _validate_frame_number(self, frame_number: int) -> Tuple[int, int, int]: def get_chunk_number(self, frame_number: int) -> int: return int(frame_number) // self._db_segment.task.data.chunk_size - def _validate_chunk_number(self, chunk_number: int) -> int: - segment_size = len(self._db_segment.frame_count) - if chunk_number < 0 or chunk_number >= math.ceil( - segment_size / self._db_segment.task.data.chunk_size - ): - raise ValidationError("requested chunk does not exist") + def validate_chunk_number(self, chunk_number: int) -> int: + segment_size = self._db_segment.frame_count + start_chunk = 0 + stop_chunk = math.ceil(segment_size / self._db_segment.task.data.chunk_size) + if not (start_chunk <= chunk_number <= stop_chunk): + raise ValidationError( + f"Invalid chunk number '{chunk_number}'. " + f"The chunk number should be in the [{start_chunk}, {stop_chunk}] range" + ) return chunk_number def get_preview(self) -> DataWithMeta[BytesIO]: - if self._db_segment.task.dimension == models.DimensionType.DIM_3D: - # TODO - preview = Image.open(os.path.join(os.path.dirname(__file__), "assets/3d_preview.jpeg")) - else: - preview, _ = self.get_frame( - min(self._db_segment.frame_set), - frame_number=FrameQuality.COMPRESSED, - out_type=FrameOutputType.PIL, - ) - - return prepare_preview_image(preview) + cache = MediaCache() + preview, mime = cache.get_or_set_segment_preview(self._db_segment) + return DataWithMeta[BytesIO](preview, mime=mime, checksum=None) def get_chunk( self, chunk_number: int, *, quality: FrameQuality = FrameQuality.ORIGINAL ) -> DataWithMeta[BytesIO]: - return self._loaders[quality].get_chunk_path(self._validate_chunk_number(chunk_number)) + chunk_number = self.validate_chunk_number(chunk_number) + chunk_data, mime = self._loaders[quality].read_chunk(chunk_number) + return DataWithMeta[BytesIO](chunk_data, mime=mime, checksum=None) def get_frame( self, @@ -289,16 +445,36 @@ def get_frame( *, quality: FrameQuality = FrameQuality.ORIGINAL, out_type: FrameOutputType = FrameOutputType.BUFFER, - ) -> AnyFrame: - _, chunk_number, frame_offset = self._validate_frame_number(frame_number) + ) -> DataWithMeta[AnyFrame]: + return_type = DataWithMeta[AnyFrame] + + _, chunk_number, frame_offset = self.validate_frame_number(frame_number) loader = self._loaders[quality] chunk_reader = loader.load(chunk_number) frame, frame_name, _ = chunk_reader[frame_offset] frame = self._convert_frame(frame, loader.reader_class, out_type) if loader.reader_class is VideoReader: - return (frame, self.VIDEO_FRAME_MIME) - return (frame, mimetypes.guess_type(frame_name)[0]) + return return_type(frame, mime=self.VIDEO_FRAME_MIME, checksum=None) + + return return_type(frame, mime=mimetypes.guess_type(frame_name)[0], checksum=None) + + def get_frame_context_images( + self, + frame_number: int, + ) -> Optional[DataWithMeta[BytesIO]]: + # TODO: refactor, optimize + cache = MediaCache() + + if self._db_segment.task.data.storage_method == models.StorageMethodChoice.CACHE: + data, mime = cache.get_frame_context_images(self._db_segment.task.data, frame_number) + else: + data, mime = cache.prepare_context_images(self._db_segment.task.data, frame_number) + + if not data: + return None + + return DataWithMeta[BytesIO](data, mime=mime, checksum=None) def iterate_frames( self, @@ -307,11 +483,22 @@ def iterate_frames( stop_frame: Optional[int] = None, quality: FrameQuality = FrameQuality.ORIGINAL, out_type: FrameOutputType = FrameOutputType.BUFFER, - ) -> Iterator[AnyFrame]: + ) -> Iterator[DataWithMeta[AnyFrame]]: for idx in range(start_frame, (stop_frame + 1) if stop_frame else None): yield self.get_frame(idx, quality=quality, out_type=out_type) -class JobFrameProvider(_SegmentFrameProvider): +class JobFrameProvider(SegmentFrameProvider): def __init__(self, db_job: models.Job) -> None: super().__init__(db_job.segment) + + +def make_frame_provider(data_source: Union[models.Job, models.Task, Any]) -> IFrameProvider: + if isinstance(data_source, models.Task): + frame_provider = TaskFrameProvider(data_source) + elif isinstance(data_source, models.Job): + frame_provider = JobFrameProvider(data_source) + else: + raise TypeError(f"Unexpected data source type {type(data_source)}") + + return frame_provider diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 7ca6ff0ed54..7e67d45b686 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -11,9 +11,10 @@ import io import itertools import struct -from enum import IntEnum from abc import ABC, abstractmethod +from bisect import bisect from contextlib import closing +from enum import IntEnum from typing import Iterable import av @@ -493,79 +494,42 @@ def get_image_size(self, i): image = (next(iter(self)))[0] return image.width, image.height -class FragmentMediaReader: - def __init__(self, chunk_number, chunk_size, start, stop, step=1): - self._start = start - self._stop = stop + 1 # up to the last inclusive - self._step = step - self._chunk_number = chunk_number - self._chunk_size = chunk_size - self._start_chunk_frame_number = \ - self._start + self._chunk_number * self._chunk_size * self._step - self._end_chunk_frame_number = min(self._start_chunk_frame_number \ - + (self._chunk_size - 1) * self._step + 1, self._stop) - self._frame_range = self._get_frame_range() - - @property - def frame_range(self): - return self._frame_range - - def _get_frame_range(self): - frame_range = [] - for idx in range(self._start, self._stop, self._step): - if idx < self._start_chunk_frame_number: - continue - elif idx < self._end_chunk_frame_number and \ - not (idx - self._start_chunk_frame_number) % self._step: - frame_range.append(idx) - elif (idx - self._start_chunk_frame_number) % self._step: - continue - else: - break - return frame_range - -class ImageDatasetManifestReader(FragmentMediaReader): - def __init__(self, manifest_path, **kwargs): - super().__init__(**kwargs) +class ImageReaderWithManifest: + def __init__(self, manifest_path: str): self._manifest = ImageManifestManager(manifest_path) self._manifest.init_index() - def __iter__(self): - for idx in self._frame_range: + def iterate_frames(self, frame_ids: Iterable[int]): + for idx in frame_ids: yield self._manifest[idx] -class VideoDatasetManifestReader(FragmentMediaReader): - def __init__(self, manifest_path, **kwargs): - self.source_path = kwargs.pop('source_path') - super().__init__(**kwargs) +class VideoReaderWithManifest: + def __init__(self, manifest_path: str, source_path: str): + self._source_path = source_path self._manifest = VideoManifestManager(manifest_path) self._manifest.init_index() - def _get_nearest_left_key_frame(self): - if self._start_chunk_frame_number >= \ - self._manifest[len(self._manifest) - 1].get('number'): - left_border = len(self._manifest) - 1 - else: - left_border = 0 - delta = len(self._manifest) - while delta: - step = delta // 2 - cur_position = left_border + step - if self._manifest[cur_position].get('number') < self._start_chunk_frame_number: - cur_position += 1 - left_border = cur_position - delta -= step + 1 - else: - delta = step - if self._manifest[cur_position].get('number') > self._start_chunk_frame_number: - left_border -= 1 - frame_number = self._manifest[left_border].get('number') - timestamp = self._manifest[left_border].get('pts') + def _get_nearest_left_key_frame(self, frame_id: int) -> tuple[int, int]: + nearest_left_keyframe_pos = bisect( + self._manifest, frame_id, key=lambda entry: entry.get('number') + ) + frame_number = self._manifest[nearest_left_keyframe_pos].get('number') + timestamp = self._manifest[nearest_left_keyframe_pos].get('pts') return frame_number, timestamp - def __iter__(self): - start_decode_frame_number, start_decode_timestamp = self._get_nearest_left_key_frame() - with closing(av.open(self.source_path, mode='r')) as container: + def iterate_frames(self, frame_ids: Iterable[int]) -> Iterable[av.VideoFrame]: + "frame_ids must be an ordered sequence in the ascending order" + + frame_ids_iter = iter(frame_ids) + frame_ids_frame = next(frame_ids_iter, None) + if frame_ids_frame is None: + return + + start_decode_frame_number, start_decode_timestamp = self._get_nearest_left_key_frame( + frame_ids_frame + ) + + with closing(av.open(self._source_path, mode='r')) as container: video_stream = next(stream for stream in container.streams if stream.type == 'video') video_stream.thread_type = 'AUTO' @@ -575,7 +539,10 @@ def __iter__(self): for packet in container.demux(video_stream): for frame in packet.decode(): frame_number += 1 - if frame_number in self._frame_range: + + if frame_number < frame_ids_frame: + continue + elif frame_number == frame_ids_frame: if video_stream.metadata.get('rotate'): frame = av.VideoFrame().from_ndarray( rotate_image( @@ -584,11 +551,12 @@ def __iter__(self): ), format ='bgr24' ) + yield frame - elif frame_number < self._frame_range[-1]: - continue else: - return + frame_ids_frame = next(frame_ids_iter, None) + if frame_ids_frame is None: + return class IChunkWriter(ABC): def __init__(self, quality, dimension=DimensionType.DIM_2D): diff --git a/cvat/apps/engine/models.py b/cvat/apps/engine/models.py index aab67ac13af..8be47acecb9 100644 --- a/cvat/apps/engine/models.py +++ b/cvat/apps/engine/models.py @@ -272,12 +272,12 @@ def _get_compressed_chunk_name(self, chunk_number): def _get_original_chunk_name(self, chunk_number): return self._get_chunk_name(chunk_number, self.original_chunk_type) - def get_original_chunk_path(self, chunk_number): - return os.path.join(self.get_original_cache_dirname(), + def get_original_segment_chunk_path(self, chunk_number: int, segment: int) -> str: + return os.path.join(self.get_original_cache_dirname(), f'segment_{segment}', self._get_original_chunk_name(chunk_number)) - def get_compressed_chunk_path(self, chunk_number): - return os.path.join(self.get_compressed_cache_dirname(), + def get_compressed_segment_chunk_path(self, chunk_number: int, segment: int) -> str: + return os.path.join(self.get_compressed_cache_dirname(), f'segment_{segment}', self._get_compressed_chunk_name(chunk_number)) def get_manifest_path(self): @@ -558,7 +558,7 @@ def __str__(self): class Segment(models.Model): # Common fields - task = models.ForeignKey(Task, on_delete=models.CASCADE) + task = models.ForeignKey(Task, on_delete=models.CASCADE) # TODO: add related name start_frame = models.IntegerField() stop_frame = models.IntegerField() type = models.CharField(choices=SegmentType.choices(), default=SegmentType.RANGE, max_length=32) diff --git a/cvat/apps/engine/views.py b/cvat/apps/engine/views.py index 13a7679595e..0fdaca30a74 100644 --- a/cvat/apps/engine/views.py +++ b/cvat/apps/engine/views.py @@ -3,12 +3,13 @@ # # SPDX-License-Identifier: MIT +from abc import ABCMeta, abstractmethod import os import os.path as osp import functools from PIL import Image from types import SimpleNamespace -from typing import Optional, Any, Dict, List, cast, Callable, Mapping, Iterable +from typing import Optional, Any, Dict, List, Union, cast, Callable, Mapping, Iterable import traceback import textwrap from collections import namedtuple @@ -54,7 +55,9 @@ from cvat.apps.events.handlers import handle_dataset_import from cvat.apps.dataset_manager.bindings import CvatImportError from cvat.apps.dataset_manager.serializers import DatasetFormatsSerializer -from cvat.apps.engine.frame_provider import FrameProvider, JobFrameProvider +from cvat.apps.engine.frame_provider import ( + IFrameProvider, TaskFrameProvider, JobFrameProvider, FrameQuality, FrameOutputType +) from cvat.apps.engine.filters import NonModelSimpleFilter, NonModelOrderingFilter, NonModelJsonLogicFilter from cvat.apps.engine.media_extractors import get_mime from cvat.apps.engine.models import ( @@ -629,15 +632,13 @@ def preview(self, request, pk): if not first_task: return HttpResponseNotFound('Project image preview not found') - data_getter = DataChunkGetter( + data_getter = _TaskDataGetter( + db_task=first_task, data_type='preview', data_quality='compressed', - data_num=first_task.data.start_frame, - task_dim=first_task.dimension ) - return data_getter(request, first_task.data.start_frame, - first_task.data.stop_frame, first_task.data) + return data_getter(request) @staticmethod def _get_rq_response(queue, job_id): @@ -657,80 +658,51 @@ def _get_rq_response(queue, job_id): return response -class DataChunkGetter: - def __init__(self, data_type, data_num, data_quality, task_dim): +class _DataGetter(metaclass=ABCMeta): + def __init__( + self, data_type: str, data_num: Optional[Union[str, int]], data_quality: str + ) -> None: possible_data_type_values = ('chunk', 'frame', 'preview', 'context_image') possible_quality_values = ('compressed', 'original') if not data_type or data_type not in possible_data_type_values: raise ValidationError('Data type not specified or has wrong value') elif data_type == 'chunk' or data_type == 'frame' or data_type == 'preview': - if data_num is None: + if data_num is None and data_type != 'preview': raise ValidationError('Number is not specified') elif data_quality not in possible_quality_values: raise ValidationError('Wrong quality value') self.type = data_type self.number = int(data_num) if data_num is not None else None - self.quality = FrameProvider.Quality.COMPRESSED \ - if data_quality == 'compressed' else FrameProvider.Quality.ORIGINAL - - self.dimension = task_dim + self.quality = FrameQuality.COMPRESSED \ + if data_quality == 'compressed' else FrameQuality.ORIGINAL - def _check_frame_range(self, frame: int): - frame_range = range(self._start, self._stop + 1, self._db_data.get_frame_step()) - if frame not in frame_range: - raise ValidationError( - f'The frame number should be in the [{self._start}, {self._stop}] range' - ) + @abstractmethod + def _get_frame_provider(self) -> IFrameProvider: + ... - def __call__(self, request, start: int, stop: int, db_data: Optional[Data]): - if not db_data: - raise NotFound(detail='Cannot find requested data') - - self._start = start - self._stop = stop - self._db_data = db_data - - frame_provider = FrameProvider(db_data, self.dimension) + def __call__(self): + frame_provider = self._get_frame_provider() try: if self.type == 'chunk': - start_chunk = frame_provider.get_chunk_number(start) - stop_chunk = frame_provider.get_chunk_number(stop) - # pylint: disable=superfluous-parens - if not (start_chunk <= self.number <= stop_chunk): - raise ValidationError('The chunk number should be in the ' + - f'[{start_chunk}, {stop_chunk}] range') - - # TODO: av.FFmpegError processing - if settings.USE_CACHE and db_data.storage_method == StorageMethodChoice.CACHE: - buff, mime_type = frame_provider.get_chunk(self.number, self.quality) - return HttpResponse(buff.getvalue(), content_type=mime_type) - - # Follow symbol links if the chunk is a link on a real image otherwise - # mimetype detection inside sendfile will work incorrectly. - path = os.path.realpath(frame_provider.get_chunk(self.number, self.quality)) - return sendfile(request, path) + data = frame_provider.get_chunk(self.number, quality=self.quality) + return HttpResponse(data.data.getvalue(), content_type=data.mime) # TODO: add new headers elif self.type == 'frame' or self.type == 'preview': - self._check_frame_range(self.number) - if self.type == 'preview': - cache = MediaCache(self.dimension) - buf, mime = cache.get_local_preview_with_mime(self.number, db_data) + data = frame_provider.get_preview() else: - buf, mime = frame_provider.get_frame(self.number, self.quality) + data = frame_provider.get_frame(self.number, quality=self.quality) - return HttpResponse(buf.getvalue(), content_type=mime) + return HttpResponse(data.data.getvalue(), content_type=data.mime) elif self.type == 'context_image': - self._check_frame_range(self.number) - - cache = MediaCache(self.dimension) - buff, mime = cache.get_frame_context_images(db_data, self.number) - if not buff: + data = frame_provider.get_frame_context_images(self.number) + if not data: return HttpResponseNotFound() - return HttpResponse(buff, content_type=mime) + + return HttpResponse(data.data, content_type=data.mime) else: return Response(data='unknown data type {}.'.format(self.type), status=status.HTTP_400_BAD_REQUEST) @@ -739,41 +711,36 @@ def __call__(self, request, start: int, stop: int, db_data: Optional[Data]): '\n'.join([str(d) for d in ex.detail]) return Response(data=msg, status=ex.status_code) - -class JobDataGetter(DataChunkGetter): - def __init__(self, job: Job, data_type, data_num, data_quality): - super().__init__(data_type, data_num, data_quality, task_dim=job.segment.task.dimension) - self.job = job - - def _check_frame_range(self, frame: int): - if frame not in self.job.segment.frame_set: - raise ValidationError("The frame number doesn't belong to the job") - - def __call__(self, request, start, stop, db_data): - # TODO: add segment boundary handling - if self.type == 'chunk' and self.job.segment.type == SegmentType.SPECIFIC_FRAMES: - frame_provider = JobFrameProvider(self.job) - - start_chunk = frame_provider.get_chunk_number(start) - stop_chunk = frame_provider.get_chunk_number(stop) - # pylint: disable=superfluous-parens - if not (start_chunk <= self.number <= stop_chunk): - raise ValidationError('The chunk number should be in the ' + - f'[{start_chunk}, {stop_chunk}] range') - - cache = MediaCache() - - if settings.USE_CACHE and db_data.storage_method == StorageMethodChoice.CACHE: - buf, mime = cache.get_segment_chunk(self.job, self.number, quality=self.quality) - else: - buf, mime = cache.prepare_masked_range_segment_chunk( - self.job, self.number, quality=self.quality - ) - - return HttpResponse(buf.getvalue(), content_type=mime) - - else: - return super().__call__(request, start, stop, db_data) +class _TaskDataGetter(_DataGetter): + def __init__( + self, + db_task: models.Task, + *, + data_type: str, + data_quality: str, + data_num: Optional[Union[str, int]] = None, + ) -> None: + super().__init__(data_type=data_type, data_num=data_num, data_quality=data_quality) + self._db_task = db_task + + def _get_frame_provider(self) -> IFrameProvider: + return TaskFrameProvider(self._db_task) + + +class _JobDataGetter(_DataGetter): + def __init__( + self, + db_job: models.Job, + *, + data_type: str, + data_quality: str, + data_num: Optional[Union[str, int]] = None, + ) -> None: + super().__init__(data_type=data_type, data_num=data_num, data_quality=data_quality) + self._db_job = db_job + + def _get_frame_provider(self) -> IFrameProvider: + return JobFrameProvider(self._db_job) @extend_schema(tags=['tasks']) @@ -1296,10 +1263,10 @@ def data(self, request, pk): data_num = request.query_params.get('number', None) data_quality = request.query_params.get('quality', 'compressed') - data_getter = DataChunkGetter(data_type, data_num, data_quality, self._object.dimension) - - return data_getter(request, self._object.data.start_frame, - self._object.data.stop_frame, self._object.data) + data_getter = _TaskDataGetter( + self._object, data_type=data_type, data_num=data_num, data_quality=data_quality + ) + return data_getter() @tus_chunk_action(detail=True, suffix_base="data") def append_data_chunk(self, request, pk, file_id): @@ -1640,15 +1607,12 @@ def preview(self, request, pk): if not self._object.data: return HttpResponseNotFound('Task image preview not found') - data_getter = DataChunkGetter( + data_getter = _TaskDataGetter( + db_task=self._object, data_type='preview', data_quality='compressed', - data_num=self._object.data.start_frame, - task_dim=self._object.dimension ) - - return data_getter(request, self._object.data.start_frame, - self._object.data.stop_frame, self._object.data) + return data_getter() @extend_schema(tags=['jobs']) @@ -2030,10 +1994,10 @@ def data(self, request, pk): data_num = request.query_params.get('number', None) data_quality = request.query_params.get('quality', 'compressed') - data_getter = JobDataGetter(db_job, data_type, data_num, data_quality) - - return data_getter(request, db_job.segment.start_frame, - db_job.segment.stop_frame, db_job.segment.task.data) + data_getter = _JobDataGetter( + db_job, data_type=data_type, data_num=data_num, data_quality=data_quality + ) + return data_getter() @extend_schema(methods=['GET'], summary='Get metainformation for media files in a job', @@ -2126,15 +2090,12 @@ def metadata(self, request, pk): def preview(self, request, pk): self._object = self.get_object() # call check_object_permissions as well - data_getter = DataChunkGetter( + data_getter = _JobDataGetter( + db_job=self._object, data_type='preview', data_quality='compressed', - data_num=self._object.segment.start_frame, - task_dim=self._object.segment.task.dimension ) - - return data_getter(request, self._object.segment.start_frame, - self._object.segment.stop_frame, self._object.segment.task.data) + return data_getter() @extend_schema(tags=['issues']) @@ -2704,12 +2665,12 @@ def preview(self, request, pk): # The idea is try to define real manifest preview only for the storages that have related manifests # because otherwise it can lead to extra calls to a bucket, that are usually not free. if not db_storage.has_at_least_one_manifest: - result = cache.get_cloud_preview_with_mime(db_storage) + result = cache.get_cloud_preview(db_storage) if not result: return HttpResponseNotFound('Cloud storage preview not found') return HttpResponse(result[0], result[1]) - preview, mime = cache.get_or_set_cloud_preview_with_mime(db_storage) + preview, mime = cache.get_or_set_cloud_preview(db_storage) return HttpResponse(preview, mime) except CloudStorageModel.DoesNotExist: message = f"Storage {pk} does not exist" diff --git a/cvat/apps/lambda_manager/views.py b/cvat/apps/lambda_manager/views.py index a6e37933f32..3e1f4691bba 100644 --- a/cvat/apps/lambda_manager/views.py +++ b/cvat/apps/lambda_manager/views.py @@ -32,7 +32,7 @@ from rest_framework.request import Request import cvat.apps.dataset_manager as dm -from cvat.apps.engine.frame_provider import FrameProvider +from cvat.apps.engine.frame_provider import FrameQuality, TaskFrameProvider from cvat.apps.engine.models import Job, ShapeType, SourceType, Task, Label from cvat.apps.engine.serializers import LabeledDataSerializer from cvat.apps.lambda_manager.permissions import LambdaPermission @@ -480,16 +480,16 @@ def transform_attributes(input_attributes, attr_mapping, db_attributes): def _get_image(self, db_task, frame, quality): if quality is None or quality == "original": - quality = FrameProvider.Quality.ORIGINAL + quality = FrameQuality.ORIGINAL elif quality == "compressed": - quality = FrameProvider.Quality.COMPRESSED + quality = FrameQuality.COMPRESSED else: raise ValidationError( '`{}` lambda function was run '.format(self.id) + 'with wrong arguments (quality={})'.format(quality), code=status.HTTP_400_BAD_REQUEST) - frame_provider = FrameProvider(db_task.data) + frame_provider = TaskFrameProvider(db_task) image = frame_provider.get_frame(frame, quality=quality) return base64.b64encode(image[0].getvalue()).decode('utf-8') From 146a896f6d33a4153e98de95e75dc01a53b657c9 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 1 Aug 2024 19:37:47 +0300 Subject: [PATCH 004/115] Support static chunk building, fix av memory leak, add caching media iterator --- cvat/apps/engine/frame_provider.py | 98 +++++------ cvat/apps/engine/media_extractors.py | 248 +++++++++++++++++++++++---- cvat/apps/engine/models.py | 20 +-- cvat/apps/engine/task.py | 123 ++++++++----- 4 files changed, 350 insertions(+), 139 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 77794d60ff7..f47e3073203 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -11,7 +11,7 @@ from dataclasses import dataclass from enum import Enum, auto from io import BytesIO -from typing import Any, Callable, Generic, Iterable, Iterator, Optional, Tuple, TypeVar, Union +from typing import Any, Callable, Generic, Iterator, Optional, Tuple, Type, TypeVar, Union import av import cv2 @@ -23,8 +23,13 @@ from cvat.apps.engine.cache import DataWithMime, MediaCache from cvat.apps.engine.media_extractors import ( FrameQuality, + IChunkWriter, IMediaReader, + Mpeg4ChunkWriter, + Mpeg4CompressedChunkWriter, + RandomAccessIterator, VideoReader, + ZipChunkWriter, ZipCompressedChunkWriter, ZipReader, ) @@ -33,53 +38,18 @@ _T = TypeVar("_T") -class _RandomAccessIterator(Iterator[_T]): - def __init__(self, iterable: Iterable[_T]): - self.iterable: Iterable[_T] = iterable - self.iterator: Optional[Iterator[_T]] = None - self.pos: int = -1 - - def __iter__(self): - return self - - def __next__(self): - return self[self.pos + 1] - - def __getitem__(self, idx: int) -> Optional[_T]: - assert 0 <= idx - if self.iterator is None or idx <= self.pos: - self.reset() - v = None - while self.pos < idx: - # NOTE: don't keep the last item in self, it can be expensive - v = next(self.iterator) - self.pos += 1 - return v - - def reset(self): - self.close() - self.iterator = iter(self.iterable) - - def close(self): - if self.iterator is not None: - if close := getattr(self.iterator, "close", None): - close() - self.iterator = None - self.pos = -1 - - class _ChunkLoader(metaclass=ABCMeta): def __init__(self, reader_class: IMediaReader) -> None: self.chunk_id: Optional[int] = None - self.chunk_reader: Optional[_RandomAccessIterator] = None + self.chunk_reader: Optional[RandomAccessIterator] = None self.reader_class = reader_class - def load(self, chunk_id: int) -> _RandomAccessIterator[Tuple[Any, str, int]]: + def load(self, chunk_id: int) -> RandomAccessIterator[Tuple[Any, str, int]]: if self.chunk_id != chunk_id: self.unload() self.chunk_id = chunk_id - self.chunk_reader = _RandomAccessIterator( + self.chunk_reader = RandomAccessIterator( self.reader_class([self.read_chunk(chunk_id)[0]]) ) return self.chunk_reader @@ -103,7 +73,7 @@ def __init__( def read_chunk(self, chunk_id: int) -> DataWithMime: chunk_path = self.get_chunk_path(chunk_id) - with open(chunk_path, "r") as f: + with open(chunk_path, "rb") as f: return ( io.BytesIO(f.read()), mimetypes.guess_type(chunk_path)[0], @@ -254,7 +224,6 @@ def get_chunk( return_type = DataWithMeta[BytesIO] chunk_number = self.validate_chunk_number(chunk_number) - # TODO: return a joined chunk. Find a solution for segment boundary video chunks db_data = self._db_task.data step = db_data.get_frame_step() task_chunk_start_frame = chunk_number * db_data.chunk_size @@ -270,7 +239,7 @@ def get_chunk( matching_segments = sorted( [ s - for s in self._db_task.segment_set + for s in self._db_task.segment_set.all() if s.type == models.SegmentType.RANGE if not task_chunk_frame_set.isdisjoint(s.frame_set) ], @@ -284,6 +253,8 @@ def get_chunk( segment_frame_provider.get_chunk_number(task_chunk_start_frame), quality=quality ) + # Create and return a joined chunk + # TODO: refactor into another class, optimize (don't visit frames twice) task_chunk_frames = [] for db_segment in matching_segments: segment_frame_provider = SegmentFrameProvider(db_segment) @@ -295,13 +266,38 @@ def get_chunk( frame = segment_frame_provider.get_frame( task_chunk_frame_id, quality=quality, out_type=FrameOutputType.BUFFER - ) + ).data task_chunk_frames.append((frame, None, None)) - merged_chunk_writer = ZipCompressedChunkWriter( - db_data.image_quality, dimension=self._db_task.dimension + writer_classes: dict[FrameQuality, Type[IChunkWriter]] = { + FrameQuality.COMPRESSED: ( + Mpeg4CompressedChunkWriter + if db_data.compressed_chunk_type == models.DataChoice.VIDEO + else ZipCompressedChunkWriter + ), + FrameQuality.ORIGINAL: ( + Mpeg4ChunkWriter + if db_data.original_chunk_type == models.DataChoice.VIDEO + else ZipChunkWriter + ), + } + + image_quality = ( + 100 + if writer_classes[quality] in [Mpeg4ChunkWriter, ZipChunkWriter] + else db_data.image_quality + ) + mime_type = ( + "video/mp4" + if writer_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] + else "application/zip" ) + kwargs = {} + if self._db_task.dimension == models.DimensionType.DIM_3D: + kwargs["dimension"] = models.DimensionType.DIM_3D + merged_chunk_writer = writer_classes[quality](image_quality, **kwargs) + buffer = io.BytesIO() merged_chunk_writer.save_as_chunk( task_chunk_frames, @@ -311,7 +307,9 @@ def get_chunk( ) buffer.seek(0) - return return_type(data=buffer, mime="application/zip", checksum=None) + # TODO: add caching + + return return_type(data=buffer, mime=mime_type, checksum=None) def get_frame( self, @@ -406,10 +404,14 @@ def __len__(self): return self._db_segment.frame_count def validate_frame_number(self, frame_number: int) -> Tuple[int, int, int]: - if frame_number not in self._db_segment.frame_set: + frame_sequence = list(self._db_segment.frame_set) + if frame_number not in frame_sequence: raise ValidationError(f"Incorrect requested frame number: {frame_number}") - chunk_number, frame_position = divmod(frame_number, self._db_segment.task.data.chunk_size) + # TODO: maybe optimize search + chunk_number, frame_position = divmod( + frame_sequence.index(frame_number), self._db_segment.task.data.chunk_size + ) return frame_number, chunk_number, frame_position def get_chunk_number(self, frame_number: int) -> int: diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 7e67d45b686..80d8dfc1dba 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -3,6 +3,8 @@ # # SPDX-License-Identifier: MIT +from __future__ import annotations + import os import sysconfig import tempfile @@ -13,11 +15,15 @@ import struct from abc import ABC, abstractmethod from bisect import bisect -from contextlib import closing +from contextlib import ExitStack, closing +from dataclasses import dataclass from enum import IntEnum -from typing import Iterable +from typing import Callable, Iterable, Iterator, Optional, Protocol, Tuple, TypeVar import av +import av.codec +import av.container +import av.video.stream import numpy as np from natsort import os_sorted from pyunpack import Archive @@ -92,6 +98,97 @@ def image_size_within_orientation(img: Image): def has_exif_rotation(img: Image): return img.getexif().get(ORIENTATION_EXIF_TAG, ORIENTATION.NORMAL_HORIZONTAL) != ORIENTATION.NORMAL_HORIZONTAL +_T = TypeVar("_T") + + +class RandomAccessIterator(Iterator[_T]): + def __init__(self, iterable: Iterable[_T]): + self.iterable: Iterable[_T] = iterable + self.iterator: Optional[Iterator[_T]] = None + self.pos: int = -1 + + def __iter__(self): + return self + + def __next__(self): + return self[self.pos + 1] + + def __getitem__(self, idx: int) -> Optional[_T]: + assert 0 <= idx + if self.iterator is None or idx <= self.pos: + self.reset() + v = None + while self.pos < idx: + # NOTE: don't keep the last item in self, it can be expensive + v = next(self.iterator) + self.pos += 1 + return v + + def reset(self): + self.close() + self.iterator = iter(self.iterable) + + def close(self): + if self.iterator is not None: + if close := getattr(self.iterator, "close", None): + close() + self.iterator = None + self.pos = -1 + + +class Sized(Protocol): + def get_size(self) -> int: ... + +_MediaT = TypeVar("_MediaT", bound=Sized) + +class CachingMediaIterator(RandomAccessIterator[_MediaT]): + @dataclass + class _CacheItem: + value: _MediaT + size: int + + def __init__( + self, + iterable: Iterable, + *, + max_cache_memory: int, + max_cache_entries: int, + object_size_callback: Optional[Callable[[_MediaT], int]] = None, + ): + super().__init__(iterable) + self.max_cache_entries = max_cache_entries + self.max_cache_memory = max_cache_memory + self._get_object_size_callback = object_size_callback + self.used_cache_memory = 0 + self._cache: dict[int, self._CacheItem] = {} + + def _get_object_size(self, obj: _MediaT) -> int: + if self._get_object_size_callback: + return self._get_object_size_callback(obj) + + return obj.get_size() + + def __getitem__(self, idx: int): + cache_item = self._cache.get(idx) + if cache_item: + return cache_item.value + + value = super().__getitem__(idx) + value_size = self._get_object_size(value) + + while ( + len(self._cache) + 1 > self.max_cache_entries or + self.used_cache_memory + value_size > self.max_cache_memory + ): + min_key = min(self._cache.keys()) + self._cache.pop(min_key) + + if self.used_cache_memory + value_size <= self.max_cache_memory: + self._cache[idx] = self._CacheItem(value, value_size) + + return value + + class IMediaReader(ABC): def __init__(self, source_path, step, start, stop, dimension): self._source_path = source_path @@ -409,7 +506,10 @@ def extract(self): os.remove(self._zip_source.filename) class VideoReader(IMediaReader): - def __init__(self, source_path, step=1, start=0, stop=None, dimension=DimensionType.DIM_2D): + def __init__( + self, source_path, step=1, start=0, stop=None, + dimension=DimensionType.DIM_2D, *, allow_threading: bool = True + ): super().__init__( source_path=source_path, step=step, @@ -418,6 +518,10 @@ def __init__(self, source_path, step=1, start=0, stop=None, dimension=DimensionT dimension=dimension, ) + self.allow_threading = allow_threading + self._frame_count: Optional[int] = None + self._frame_size: Optional[tuple[int, int]] = None # (w, h) + def _has_frame(self, i): if i >= self._start: if (i - self._start) % self._step == 0: @@ -426,26 +530,59 @@ def _has_frame(self, i): return False - def __iter__(self): - with self._get_av_container() as container: - stream = container.streams.video[0] - stream.thread_type = 'AUTO' + def _make_frame_iterator( + self, + *, + apply_filter: bool = True, + stream: Optional[av.video.stream.VideoStream] = None, + ) -> Iterator[Tuple[av.VideoFrame, str, int]]: + es = ExitStack() + + need_init = stream is None + if need_init: + container = es.enter_context(self._get_av_container()) + else: + container = stream.container + + with es: + if need_init: + stream = container.streams.video[0] + + if self.allow_threading: + stream.thread_type = 'AUTO' + + es.enter_context(closing(stream.codec_context)) + frame_num = 0 + for packet in container.demux(stream): for image in packet.decode(): frame_num += 1 - if self._has_frame(frame_num - 1): - if packet.stream.metadata.get('rotate'): - pts = image.pts - image = av.VideoFrame().from_ndarray( - rotate_image( - image.to_ndarray(format='bgr24'), - 360 - int(stream.metadata.get('rotate')) - ), - format ='bgr24' - ) - image.pts = pts - yield (image, self._source_path[0], image.pts) + + if apply_filter and not self._has_frame(frame_num - 1): + continue + + if stream.metadata.get('rotate'): + pts = image.pts + image = av.VideoFrame().from_ndarray( + rotate_image( + image.to_ndarray(format='bgr24'), + 360 - int(stream.metadata.get('rotate')) + ), + format ='bgr24' + ) + image.pts = pts + + if self._frame_size is None: + self._frame_size = (image.width, image.height) + + yield (image, self._source_path[0], image.pts) + + if self._frame_count is None: + self._frame_count = frame_num + + def __iter__(self): + return self._make_frame_iterator() def get_progress(self, pos): duration = self._get_duration() @@ -457,8 +594,11 @@ def _get_av_container(self): return av.open(self._source_path[0]) def _get_duration(self): - with self._get_av_container() as container: + with ExitStack() as es: + container = es.enter_context(self._get_av_container()) stream = container.streams.video[0] + es.enter_context(closing(stream.codec_context)) + duration = None if stream.duration: duration = stream.duration @@ -473,26 +613,52 @@ def _get_duration(self): return duration def get_preview(self, frame): - with self._get_av_container() as container: + with ExitStack() as es: + container = es.enter_context(self._get_av_container()) stream = container.streams.video[0] + es.enter_context(closing(stream.codec_context)) + tb_denominator = stream.time_base.denominator needed_time = int((frame / stream.guessed_rate) * tb_denominator) container.seek(offset=needed_time, stream=stream) - for packet in container.demux(stream): - for frame in packet.decode(): - return self._get_preview(frame.to_image() if not stream.metadata.get('rotate') \ - else av.VideoFrame().from_ndarray( - rotate_image( - frame.to_ndarray(format='bgr24'), - 360 - int(container.streams.video[0].metadata.get('rotate')) - ), - format ='bgr24' - ).to_image() - ) + + with closing(self._make_frame_iterator(stream=stream)) as frame_iter: + return self._get_preview(next(frame_iter)) def get_image_size(self, i): - image = (next(iter(self)))[0] - return image.width, image.height + if self._frame_size is not None: + return self._frame_size + + with closing(iter(self)) as frame_iter: + image = next(frame_iter)[0] + self._frame_size = (image.width, image.height) + + return self._frame_size + + def get_frame_count(self) -> int: + """ + Returns total frame count in the video + + Note that not all videos provide length / duration metainfo, so the + result may require full video decoding. + + The total count is NOT affected by the frame filtering options of the object, + i.e. start frame, end frame and frame step. + """ + # It's possible to retrieve frame count from the stream.frames, + # but the number may be incorrect. + # https://superuser.com/questions/1512575/why-total-frame-count-is-different-in-ffmpeg-than-ffprobe + if self._frame_count is not None: + return self._frame_count + + frame_count = 0 + for _ in self._make_frame_iterator(apply_filter=False): + frame_count += 1 + + self._frame_count = frame_count + + return frame_count + class ImageReaderWithManifest: def __init__(self, manifest_path: str): @@ -723,7 +889,7 @@ def __init__(self, quality=67): "preset": "ultrafast", } - def _add_video_stream(self, container, w, h, rate, options): + def _add_video_stream(self, container: av.container.OutputContainer, w, h, rate, options): # x264 requires width and height must be divisible by 2 for yuv420p if h % 2: h += 1 @@ -760,11 +926,15 @@ def save_as_chunk(self, images, chunk_path): options=self._codec_opts, ) - self._encode_images(images, output_container, output_v_stream) + with closing(output_v_stream): + self._encode_images(images, output_container, output_v_stream) + return [(input_w, input_h)] @staticmethod - def _encode_images(images, container, stream): + def _encode_images( + images, container: av.container.OutputContainer, stream: av.video.stream.VideoStream + ): for frame, _, _ in images: # let libav set the correct pts and time_base frame.pts = None @@ -812,7 +982,9 @@ def save_as_chunk(self, images, chunk_path): options=self._codec_opts, ) - self._encode_images(images, output_container, output_v_stream) + with closing(output_v_stream): + self._encode_images(images, output_container, output_v_stream) + return [(input_w, input_h)] def _is_archive(path): diff --git a/cvat/apps/engine/models.py b/cvat/apps/engine/models.py index 8be47acecb9..ec0870a0e22 100644 --- a/cvat/apps/engine/models.py +++ b/cvat/apps/engine/models.py @@ -256,7 +256,7 @@ def get_original_cache_dirname(self): return os.path.join(self.get_data_dirname(), "original") @staticmethod - def _get_chunk_name(chunk_number, chunk_type): + def _get_chunk_name(segment_id: int, chunk_number: int, chunk_type: DataChoice | str) -> str: if chunk_type == DataChoice.VIDEO: ext = 'mp4' elif chunk_type == DataChoice.IMAGESET: @@ -264,21 +264,21 @@ def _get_chunk_name(chunk_number, chunk_type): else: ext = 'list' - return '{}.{}'.format(chunk_number, ext) + return 'segment_{}-{}.{}'.format(segment_id, chunk_number, ext) - def _get_compressed_chunk_name(self, chunk_number): - return self._get_chunk_name(chunk_number, self.compressed_chunk_type) + def _get_compressed_chunk_name(self, segment_id: int, chunk_number: int) -> str: + return self._get_chunk_name(segment_id, chunk_number, self.compressed_chunk_type) - def _get_original_chunk_name(self, chunk_number): - return self._get_chunk_name(chunk_number, self.original_chunk_type) + def _get_original_chunk_name(self, segment_id: int, chunk_number: int) -> str: + return self._get_chunk_name(segment_id, chunk_number, self.original_chunk_type) def get_original_segment_chunk_path(self, chunk_number: int, segment: int) -> str: - return os.path.join(self.get_original_cache_dirname(), f'segment_{segment}', - self._get_original_chunk_name(chunk_number)) + return os.path.join(self.get_original_cache_dirname(), + self._get_original_chunk_name(segment, chunk_number)) def get_compressed_segment_chunk_path(self, chunk_number: int, segment: int) -> str: - return os.path.join(self.get_compressed_cache_dirname(), f'segment_{segment}', - self._get_compressed_chunk_name(chunk_number)) + return os.path.join(self.get_compressed_cache_dirname(), + self._get_compressed_chunk_name(segment, chunk_number)) def get_manifest_path(self): return os.path.join(self.get_upload_dirname(), 'manifest.jsonl') diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index 866e3d0c913..274cb3c7206 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -3,15 +3,17 @@ # # SPDX-License-Identifier: MIT +from contextlib import closing import itertools import fnmatch import os +import av import rq import re import shutil from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Sequence, Union, Iterable +from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple, Union, Iterable from urllib import parse as urlparse from urllib import request as urlrequest import concurrent.futures @@ -26,7 +28,8 @@ from cvat.apps.engine import models from cvat.apps.engine.log import ServerLogManager from cvat.apps.engine.media_extractors import ( - MEDIA_TYPES, IMediaReader, ImageListReader, Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter, + MEDIA_TYPES, CachingMediaIterator, IMediaReader, ImageListReader, + Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter, RandomAccessIterator, ValidateDimension, ZipChunkWriter, ZipCompressedChunkWriter, get_mime, sort ) from cvat.apps.engine.utils import ( @@ -119,7 +122,7 @@ def _copy_data_from_share_point( os.makedirs(target_dir) shutil.copyfile(source_path, target_path) -def _get_task_segment_data( +def _generate_segment_params( db_task: models.Task, *, data_size: Optional[int] = None, @@ -170,7 +173,7 @@ def _segments(): return SegmentsParams(segments, segment_size, overlap) -def _save_task_to_db( +def _create_segments_and_jobs( db_task: models.Task, *, job_file_mapping: Optional[JobFileMapping] = None, @@ -179,7 +182,7 @@ def _save_task_to_db( rq_job.meta['status'] = 'Task is being saved in database' rq_job.save_meta() - segments, segment_size, overlap = _get_task_segment_data( + segments, segment_size, overlap = _generate_segment_params( db_task=db_task, job_file_mapping=job_file_mapping, ) db_task.segment_size = segment_size @@ -975,7 +978,7 @@ def _update_status(msg: str) -> None: manifest_file = os.path.relpath(db_data.get_manifest_path(), upload_dir) video_path: str = "" - video_size: tuple[int, int] = (0, 0) + video_frame_size: tuple[int, int] = (0, 0) images: list[models.Image] = [] @@ -1000,8 +1003,8 @@ def _update_status(msg: str) -> None: if not len(manifest): raise ValidationError("No key frames found in the manifest") - all_frames = manifest.video_length - video_size = manifest.video_resolution + video_frame_count = manifest.video_length + video_frame_size = manifest.video_resolution manifest_is_prepared = True except Exception as ex: manifest.remove() @@ -1030,8 +1033,8 @@ def _update_status(msg: str) -> None: _update_status('A manifest has been created') - all_frames = len(manifest.reader) # TODO: check if the field access above and here are equivalent - video_size = manifest.reader.resolution + video_frame_count = len(manifest.reader) # TODO: check if the field access above and here are equivalent + video_frame_size = manifest.reader.resolution manifest_is_prepared = True except Exception as ex: manifest.remove() @@ -1044,14 +1047,14 @@ def _update_status(msg: str) -> None: _update_status("{} The task will be created using the old method".format(base_msg)) if not manifest: - all_frames = len(extractor) - video_size = extractor.get_image_size(0) + video_frame_count = extractor.get_frame_count() + video_frame_size = extractor.get_image_size(0) db_data.size = len(range( db_data.start_frame, min( - data['stop_frame'] + 1 if data['stop_frame'] else all_frames, - all_frames, + data['stop_frame'] + 1 if data['stop_frame'] else video_frame_count, + video_frame_count, ), db_data.get_frame_step() )) @@ -1126,7 +1129,7 @@ def _update_status(msg: str) -> None: models.Video.objects.create( data=db_data, path=os.path.relpath(video_path, upload_dir), - width=video_size[0], height=video_size[1] + width=video_frame_size[0], height=video_frame_size[1] ) # validate stop_frame @@ -1137,13 +1140,12 @@ def _update_status(msg: str) -> None: db_data.start_frame + (db_data.size - 1) * db_data.get_frame_step()) slogger.glob.info("Found frames {} for Data #{}".format(db_data.size, db_data.id)) - _save_task_to_db(db_task, job_file_mapping=job_file_mapping) # TODO: split into jobs and task saving + _create_segments_and_jobs(db_task, job_file_mapping=job_file_mapping) # Save chunks - # TODO: refactor - # TODO: save chunks per job + # TODO: refactor into a separate class / function for chunk creation if db_data.storage_method == models.StorageMethodChoice.FILE_SYSTEM or not settings.USE_CACHE: - def update_progress(progress): + def update_progress(progress: float): # TODO: refactor this function into a class progress_animation = '|/-\\' if not hasattr(update_progress, 'call_counter'): @@ -1157,16 +1159,48 @@ def update_progress(progress): job.save_meta() update_progress.call_counter = (update_progress.call_counter + 1) % len(progress_animation) + db_segments = db_task.segment_set.all() + + def generate_chunk_params() -> Iterator[Tuple[models.Segment, int, Sequence[Any]]]: + if isinstance(extractor, MEDIA_TYPES['video']['extractor']): + def _get_frame_size(frame_tuple: Tuple[av.VideoFrame, Any, Any]) -> int: + # There is no need to be absolutely precise here, + # just need to provide the reasonable upper boundary. + # Return bytes needed for 1 frame + frame = frame_tuple[0] + return frame.width * frame.height * (frame.format.padded_bits_per_pixel // 8) + + media_iterator = CachingMediaIterator( + extractor, + max_cache_memory=2 ** 30, max_cache_entries=db_task.overlap, + object_size_callback=_get_frame_size + ) + else: + media_iterator = RandomAccessIterator(extractor) + + with closing(media_iterator): + for db_segment in db_segments: + counter = itertools.count() + generator = ( + media_iterator[frame_idx] + for frame_idx in db_segment.frame_set # TODO: check absolute vs relative ids + ) # probably won't work for GT job segments with gaps + generator = itertools.groupby( + generator, lambda _: next(counter) // db_data.chunk_size + ) + generator = ( + (chunk_idx, list(chunk_data)) + for chunk_idx, chunk_data in generator + ) - counter = itertools.count() - generator = itertools.groupby(extractor, lambda _: next(counter) // db_data.chunk_size) - generator = ((chunk_idx, list(chunk_data)) for chunk_idx, chunk_data in generator) + yield (db_segment, generator) def save_chunks( executor: concurrent.futures.ThreadPoolExecutor, + db_segment: models.Segment, chunk_idx: int, - chunk_data: Iterable[tuple[str, str, str]] - ) -> list[tuple[str, int, tuple[int, int]]]: + chunk_data: Iterable[tuple[Any, str, str]] + ): if ( db_task.dimension == models.DimensionType.DIM_2D and isinstance(extractor, ( @@ -1181,32 +1215,35 @@ def save_chunks( fs_original = executor.submit( original_chunk_writer.save_as_chunk, images=chunk_data, - chunk_path=db_data.get_original_chunk_path(chunk_idx) + chunk_path=db_data.get_original_segment_chunk_path( + chunk_idx, segment=db_segment.id + ), ) - fs_compressed = executor.submit( - compressed_chunk_writer.save_as_chunk, + compressed_chunk_writer.save_as_chunk( images=chunk_data, - chunk_path=db_data.get_compressed_chunk_path(chunk_idx), + chunk_path=db_data.get_compressed_segment_chunk_path( + chunk_idx, segment=db_segment.id + ), ) - # TODO: convert to async for proper concurrency - fs_original.result() - image_sizes = fs_compressed.result() - - # (path, frame, size) - return list((i[0][1], i[0][2], i[1]) for i in zip(chunk_data, image_sizes)) - def process_results(img_meta: list[tuple[str, int, tuple[int, int]]]): - progress = img_meta[-1][1] / db_data.size - update_progress(progress) + fs_original.result() futures = queue.Queue(maxsize=settings.CVAT_CONCURRENT_CHUNK_PROCESSING) with concurrent.futures.ThreadPoolExecutor( - max_workers=2 * settings.CVAT_CONCURRENT_CHUNK_PROCESSING + max_workers=2 * settings.CVAT_CONCURRENT_CHUNK_PROCESSING # TODO: remove 2 * or configuration ) as executor: - for chunk_idx, chunk_data in generator: - if futures.full(): - process_results(futures.get().result()) - futures.put(executor.submit(save_chunks, executor, chunk_idx, chunk_data)) + # TODO: maybe make real multithreading support, currently the code is limited by 1 + # segment chunk, even if more threads are available + for segment_idx, (segment, segment_chunk_params) in enumerate(generate_chunk_params()): + for chunk_idx, chunk_data in segment_chunk_params: + if futures.full(): + futures.get().result() + + futures.put(executor.submit( + save_chunks, executor, segment, chunk_idx, chunk_data + )) + + update_progress(segment_idx / len(db_segments)) while not futures.empty(): - process_results(futures.get().result()) + futures.get().result() From 52d1bacce9d258595b6caaa499426c7ac3ea9bed Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 2 Aug 2024 17:26:13 +0300 Subject: [PATCH 005/115] Refactor static chunk generation - extract function, revise threading --- cvat/apps/engine/task.py | 244 +++++++++++++++++++++------------------ 1 file changed, 131 insertions(+), 113 deletions(-) diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index 274cb3c7206..b6c9a0aae38 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -3,22 +3,22 @@ # # SPDX-License-Identifier: MIT -from contextlib import closing +import concurrent.futures import itertools import fnmatch import os -import av -import rq import re +import rq import shutil +from contextlib import closing from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple, Union, Iterable from urllib import parse as urlparse from urllib import request as urlrequest -import concurrent.futures -import queue +import av +import attrs import django_rq from django.conf import settings from django.db import transaction @@ -190,8 +190,8 @@ def _create_segments_and_jobs( for segment_idx, segment_params in enumerate(segments): slogger.glob.info( - "New segment for task #{task_id}: idx = {segment_idx}, start_frame = {start_frame}, \ - stop_frame = {stop_frame}".format( + "New segment for task #{task_id}: idx = {segment_idx}, start_frame = {start_frame}, " + "stop_frame = {stop_frame}".format( task_id=db_task.id, segment_idx=segment_idx, **segment_params._asdict() )) @@ -936,24 +936,10 @@ def _update_status(msg: str) -> None: db_data.original_chunk_type = models.DataChoice.VIDEO if task_mode == 'interpolation' else models.DataChoice.IMAGESET compressed_chunk_writer_class = Mpeg4CompressedChunkWriter if db_data.compressed_chunk_type == models.DataChoice.VIDEO else ZipCompressedChunkWriter - if db_data.original_chunk_type == models.DataChoice.VIDEO: - original_chunk_writer_class = Mpeg4ChunkWriter - # Let's use QP=17 (that is 67 for 0-100 range) for the original chunks, which should be visually lossless or nearly so. - # A lower value will significantly increase the chunk size with a slight increase of quality. - original_quality = 67 - else: - original_chunk_writer_class = ZipChunkWriter - original_quality = 100 - - kwargs = {} - if validate_dimension.dimension == models.DimensionType.DIM_3D: - kwargs["dimension"] = validate_dimension.dimension - compressed_chunk_writer = compressed_chunk_writer_class(db_data.image_quality, **kwargs) - original_chunk_writer = original_chunk_writer_class(original_quality, **kwargs) # calculate chunk size if it isn't specified if db_data.chunk_size is None: - if isinstance(compressed_chunk_writer, ZipCompressedChunkWriter): + if issubclass(compressed_chunk_writer_class, ZipCompressedChunkWriter): first_image_idx = db_data.start_frame if not is_data_in_cloud: w, h = extractor.get_image_size(first_image_idx) @@ -1027,7 +1013,7 @@ def _update_status(msg: str) -> None: manifest.link( media_file=media_files[0], upload_dir=upload_dir, - chunk_size=db_data.chunk_size + chunk_size=db_data.chunk_size # TODO: why it's needed here? ) manifest.create() @@ -1142,108 +1128,140 @@ def _update_status(msg: str) -> None: slogger.glob.info("Found frames {} for Data #{}".format(db_data.size, db_data.id)) _create_segments_and_jobs(db_task, job_file_mapping=job_file_mapping) - # Save chunks - # TODO: refactor into a separate class / function for chunk creation if db_data.storage_method == models.StorageMethodChoice.FILE_SYSTEM or not settings.USE_CACHE: - def update_progress(progress: float): - # TODO: refactor this function into a class + _create_static_chunks(db_task, media_extractor=extractor) + +def _create_static_chunks(db_task: models.Task, *, media_extractor: IMediaReader): + @attrs.define + class _ChunkProgressUpdater: + _call_counter: int = attrs.field(default=0, init=False) + _rq_job: rq.job.Job = attrs.field(factory=rq.get_current_job) + + def update_progress(self, progress: float): progress_animation = '|/-\\' - if not hasattr(update_progress, 'call_counter'): - update_progress.call_counter = 0 status_message = 'CVAT is preparing data chunks' if not progress: - status_message = '{} {}'.format(status_message, progress_animation[update_progress.call_counter]) - job.meta['status'] = status_message - job.meta['task_progress'] = progress or 0. - job.save_meta() - update_progress.call_counter = (update_progress.call_counter + 1) % len(progress_animation) - - db_segments = db_task.segment_set.all() - - def generate_chunk_params() -> Iterator[Tuple[models.Segment, int, Sequence[Any]]]: - if isinstance(extractor, MEDIA_TYPES['video']['extractor']): - def _get_frame_size(frame_tuple: Tuple[av.VideoFrame, Any, Any]) -> int: - # There is no need to be absolutely precise here, - # just need to provide the reasonable upper boundary. - # Return bytes needed for 1 frame - frame = frame_tuple[0] - return frame.width * frame.height * (frame.format.padded_bits_per_pixel // 8) - - media_iterator = CachingMediaIterator( - extractor, - max_cache_memory=2 ** 30, max_cache_entries=db_task.overlap, - object_size_callback=_get_frame_size + status_message = '{} {}'.format( + status_message, progress_animation[self._call_counter] ) - else: - media_iterator = RandomAccessIterator(extractor) - - with closing(media_iterator): - for db_segment in db_segments: - counter = itertools.count() - generator = ( - media_iterator[frame_idx] - for frame_idx in db_segment.frame_set # TODO: check absolute vs relative ids - ) # probably won't work for GT job segments with gaps - generator = itertools.groupby( - generator, lambda _: next(counter) // db_data.chunk_size - ) - generator = ( - (chunk_idx, list(chunk_data)) - for chunk_idx, chunk_data in generator - ) - yield (db_segment, generator) + self._rq_job.meta['status'] = status_message + self._rq_job.meta['task_progress'] = progress or 0. + self._rq_job.save_meta() + + self._call_counter = (self._call_counter + 1) % len(progress_animation) - def save_chunks( - executor: concurrent.futures.ThreadPoolExecutor, - db_segment: models.Segment, - chunk_idx: int, - chunk_data: Iterable[tuple[Any, str, str]] + def save_chunks( + executor: concurrent.futures.ThreadPoolExecutor, + db_segment: models.Segment, + chunk_idx: int, + chunk_frame_ids: Sequence[int] + ): + chunk_data = [media_iterator[frame_idx] for frame_idx in chunk_frame_ids] + + if ( + db_task.dimension == models.DimensionType.DIM_2D and + isinstance(media_extractor, ( + MEDIA_TYPES['image']['extractor'], + MEDIA_TYPES['zip']['extractor'], + MEDIA_TYPES['pdf']['extractor'], + MEDIA_TYPES['archive']['extractor'], + )) ): - if ( - db_task.dimension == models.DimensionType.DIM_2D and - isinstance(extractor, ( - MEDIA_TYPES['image']['extractor'], - MEDIA_TYPES['zip']['extractor'], - MEDIA_TYPES['pdf']['extractor'], - MEDIA_TYPES['archive']['extractor'], - )) - ): - chunk_data = preload_images(chunk_data) + chunk_data = preload_images(chunk_data) - fs_original = executor.submit( - original_chunk_writer.save_as_chunk, - images=chunk_data, - chunk_path=db_data.get_original_segment_chunk_path( - chunk_idx, segment=db_segment.id - ), - ) - compressed_chunk_writer.save_as_chunk( - images=chunk_data, - chunk_path=db_data.get_compressed_segment_chunk_path( - chunk_idx, segment=db_segment.id - ), - ) + # TODO: extract into a class - fs_original.result() + fs_original = executor.submit( + original_chunk_writer.save_as_chunk, + images=chunk_data, + chunk_path=db_data.get_original_segment_chunk_path( + chunk_idx, segment=db_segment.id + ), + ) + compressed_chunk_writer.save_as_chunk( + images=chunk_data, + chunk_path=db_data.get_compressed_segment_chunk_path( + chunk_idx, segment=db_segment.id + ), + ) + + fs_original.result() - futures = queue.Queue(maxsize=settings.CVAT_CONCURRENT_CHUNK_PROCESSING) - with concurrent.futures.ThreadPoolExecutor( - max_workers=2 * settings.CVAT_CONCURRENT_CHUNK_PROCESSING # TODO: remove 2 * or configuration - ) as executor: - # TODO: maybe make real multithreading support, currently the code is limited by 1 - # segment chunk, even if more threads are available - for segment_idx, (segment, segment_chunk_params) in enumerate(generate_chunk_params()): - for chunk_idx, chunk_data in segment_chunk_params: - if futures.full(): - futures.get().result() + db_data = db_task.data - futures.put(executor.submit( - save_chunks, executor, segment, chunk_idx, chunk_data - )) + if db_data.compressed_chunk_type == models.DataChoice.VIDEO: + compressed_chunk_writer_class = Mpeg4CompressedChunkWriter + else: + compressed_chunk_writer_class = ZipCompressedChunkWriter - update_progress(segment_idx / len(db_segments)) + if db_data.original_chunk_type == models.DataChoice.VIDEO: + original_chunk_writer_class = Mpeg4ChunkWriter + + # Let's use QP=17 (that is 67 for 0-100 range) for the original chunks, + # which should be visually lossless or nearly so. + # A lower value will significantly increase the chunk size with a slight increase of quality. + original_quality = 67 + else: + original_chunk_writer_class = ZipChunkWriter + original_quality = 100 + + chunk_writer_kwargs = {} + if db_task.dimension == models.DimensionType.DIM_3D: + chunk_writer_kwargs["dimension"] = db_task.dimension + compressed_chunk_writer = compressed_chunk_writer_class( + db_data.image_quality, **chunk_writer_kwargs + ) + original_chunk_writer = original_chunk_writer_class(original_quality, **chunk_writer_kwargs) + + db_segments = db_task.segment_set.all() + + if isinstance(media_extractor, MEDIA_TYPES['video']['extractor']): + def _get_frame_size(frame_tuple: Tuple[av.VideoFrame, Any, Any]) -> int: + # There is no need to be absolutely precise here, + # just need to provide the reasonable upper boundary. + # Return bytes needed for 1 frame + frame = frame_tuple[0] + return frame.width * frame.height * (frame.format.padded_bits_per_pixel // 8) + + # Currently, we only optimize video creation for sequential + # chunks with potential overlap, so parallel processing is likely to + # help only for image datasets + media_iterator = CachingMediaIterator( + media_extractor, + max_cache_memory=2 ** 30, max_cache_entries=db_task.overlap, + object_size_callback=_get_frame_size + ) + else: + media_iterator = RandomAccessIterator(media_extractor) + + with closing(media_iterator): + progress_updater = _ChunkProgressUpdater() + + # TODO: remove 2 * or the configuration option + # TODO: maybe make real multithreading support, currently the code is limited by 1 + # video segment chunk, even if more threads are available + max_concurrency = 2 * settings.CVAT_CONCURRENT_CHUNK_PROCESSING if not isinstance( + media_extractor, MEDIA_TYPES['video']['extractor'] + ) else 2 + with concurrent.futures.ThreadPoolExecutor(max_workers=max_concurrency) as executor: + frame_step = db_data.get_frame_step() + for segment_idx, db_segment in enumerate(db_segments): + frame_counter = itertools.count() + for chunk_idx, chunk_frame_ids in ( + (chunk_idx, list(chunk_frame_ids)) + for chunk_idx, chunk_frame_ids in itertools.groupby( + ( + # Convert absolute to relative ids (extractor output positions) + # Extractor will skip frames outside requested + (abs_frame_id - db_data.start_frame) // frame_step + for abs_frame_id in db_segment.frame_set + # TODO: is start frame different for video and images? + ), + lambda _: next(frame_counter) // db_data.chunk_size + ) + ): + save_chunks(executor, db_segment, chunk_idx, chunk_frame_ids) - while not futures.empty(): - futures.get().result() + progress_updater.update_progress(segment_idx / len(db_segments)) From 0c5343683898d2025d213160f0afc180a6dd2fa6 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 2 Aug 2024 19:07:39 +0300 Subject: [PATCH 006/115] Refactor and fix task chunk creation from segment chunks, any storage --- cvat/apps/engine/frame_provider.py | 58 ++++++++++-------------------- cvat/apps/engine/task.py | 1 - 2 files changed, 19 insertions(+), 40 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index f47e3073203..0c7c5d4b49f 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -204,7 +204,7 @@ def validate_frame_number(self, frame_number: int) -> int: def validate_chunk_number(self, chunk_number: int) -> int: start_chunk = 0 stop_chunk = math.ceil(self._db_task.data.size / self._db_task.data.chunk_size) - if not (start_chunk <= chunk_number <= stop_chunk): + if not (start_chunk <= chunk_number < stop_chunk): raise ValidationError( f"Invalid chunk number '{chunk_number}'. " f"The chunk number should be in the [{start_chunk}, {stop_chunk}] range" @@ -247,69 +247,49 @@ def get_chunk( ) assert matching_segments - if len(matching_segments) == 1: + if len(matching_segments) == 1 and task_chunk_frame_set == set( + matching_segments[0].frame_set + ): segment_frame_provider = SegmentFrameProvider(matching_segments[0]) return segment_frame_provider.get_chunk( segment_frame_provider.get_chunk_number(task_chunk_start_frame), quality=quality ) - # Create and return a joined chunk - # TODO: refactor into another class, optimize (don't visit frames twice) - task_chunk_frames = [] + # Create and return a joined / cleaned chunk + # TODO: refactor into another class, maybe optimize + task_chunk_frames = {} for db_segment in matching_segments: segment_frame_provider = SegmentFrameProvider(db_segment) segment_frame_set = db_segment.frame_set - for task_chunk_frame_id in task_chunk_frame_set: - if task_chunk_frame_id not in segment_frame_set: + for task_chunk_frame_id in sorted(task_chunk_frame_set): + if ( + task_chunk_frame_id not in segment_frame_set + or task_chunk_frame_id in task_chunk_frames + ): continue frame = segment_frame_provider.get_frame( task_chunk_frame_id, quality=quality, out_type=FrameOutputType.BUFFER ).data - task_chunk_frames.append((frame, None, None)) - - writer_classes: dict[FrameQuality, Type[IChunkWriter]] = { - FrameQuality.COMPRESSED: ( - Mpeg4CompressedChunkWriter - if db_data.compressed_chunk_type == models.DataChoice.VIDEO - else ZipCompressedChunkWriter - ), - FrameQuality.ORIGINAL: ( - Mpeg4ChunkWriter - if db_data.original_chunk_type == models.DataChoice.VIDEO - else ZipChunkWriter - ), - } - - image_quality = ( - 100 - if writer_classes[quality] in [Mpeg4ChunkWriter, ZipChunkWriter] - else db_data.image_quality - ) - mime_type = ( - "video/mp4" - if writer_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] - else "application/zip" - ) + task_chunk_frames[task_chunk_frame_id] = (frame, None, None) kwargs = {} if self._db_task.dimension == models.DimensionType.DIM_3D: kwargs["dimension"] = models.DimensionType.DIM_3D - merged_chunk_writer = writer_classes[quality](image_quality, **kwargs) + merged_chunk_writer = ZipCompressedChunkWriter( + 100 if quality == FrameQuality.ORIGINAL else db_data.image_quality, **kwargs + ) buffer = io.BytesIO() merged_chunk_writer.save_as_chunk( - task_chunk_frames, - buffer, - compress_frames=False, - zip_compress_level=1, + task_chunk_frames.values(), buffer, compress_frames=False, zip_compress_level=1 ) buffer.seek(0) - # TODO: add caching + # TODO: add caching in media cache for the resulting chunk - return return_type(data=buffer, mime=mime_type, checksum=None) + return return_type(data=buffer, mime="application/zip", checksum=None) def get_frame( self, diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index b6c9a0aae38..0416704b913 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -1257,7 +1257,6 @@ def _get_frame_size(frame_tuple: Tuple[av.VideoFrame, Any, Any]) -> int: # Extractor will skip frames outside requested (abs_frame_id - db_data.start_frame) // frame_step for abs_frame_id in db_segment.frame_set - # TODO: is start frame different for video and images? ), lambda _: next(frame_counter) // db_data.chunk_size ) From c1661237b1eea5b79393e823a33dbb136eef57ee Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 5 Aug 2024 16:19:02 +0300 Subject: [PATCH 007/115] Fix chunk number validation --- cvat/apps/engine/frame_provider.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 0c7c5d4b49f..0da5c5f1f86 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -202,12 +202,11 @@ def validate_frame_number(self, frame_number: int) -> int: return frame_number def validate_chunk_number(self, chunk_number: int) -> int: - start_chunk = 0 - stop_chunk = math.ceil(self._db_task.data.size / self._db_task.data.chunk_size) - if not (start_chunk <= chunk_number < stop_chunk): + last_chunk = math.ceil(self._db_task.data.size / self._db_task.data.chunk_size) - 1 + if not (0 <= chunk_number <= last_chunk): raise ValidationError( f"Invalid chunk number '{chunk_number}'. " - f"The chunk number should be in the [{start_chunk}, {stop_chunk}] range" + f"The chunk number should be in the [0, {last_chunk}] range" ) return chunk_number @@ -399,12 +398,11 @@ def get_chunk_number(self, frame_number: int) -> int: def validate_chunk_number(self, chunk_number: int) -> int: segment_size = self._db_segment.frame_count - start_chunk = 0 - stop_chunk = math.ceil(segment_size / self._db_segment.task.data.chunk_size) - if not (start_chunk <= chunk_number <= stop_chunk): + last_chunk = math.ceil(segment_size / self._db_segment.task.data.chunk_size) - 1 + if not (0 <= chunk_number <= last_chunk): raise ValidationError( f"Invalid chunk number '{chunk_number}'. " - f"The chunk number should be in the [{start_chunk}, {stop_chunk}] range" + f"The chunk number should be in the [0, {last_chunk}] range" ) return chunk_number From 630c97edc78f2f7a5def7fb187c7183863c05f91 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 5 Aug 2024 16:26:15 +0300 Subject: [PATCH 008/115] Enable formatting for updated components --- dev/format_python_code.sh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dev/format_python_code.sh b/dev/format_python_code.sh index a67bf08572e..10469178f90 100755 --- a/dev/format_python_code.sh +++ b/dev/format_python_code.sh @@ -23,6 +23,8 @@ for paths in \ "tests/python/" \ "cvat/apps/quality_control" \ "cvat/apps/analytics_report" \ + "cvat/apps/engine/frame_provider.py" \ + "cvat/apps/engine/cache.py" \ ; do ${BLACK} -- ${paths} ${ISORT} -- ${paths} From 8d710e701013ec2b0542d601175cf88bc06f1ea1 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 5 Aug 2024 16:49:00 +0300 Subject: [PATCH 009/115] Remove the checksum field --- cvat/apps/engine/frame_provider.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 0da5c5f1f86..d3235781b49 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -106,7 +106,6 @@ class FrameOutputType(Enum): class DataWithMeta(Generic[_T]): data: _T mime: str - checksum: int class IFrameProvider(metaclass=ABCMeta): @@ -288,7 +287,7 @@ def get_chunk( # TODO: add caching in media cache for the resulting chunk - return return_type(data=buffer, mime="application/zip", checksum=None) + return return_type(data=buffer, mime="application/zip") def get_frame( self, @@ -410,14 +409,14 @@ def validate_chunk_number(self, chunk_number: int) -> int: def get_preview(self) -> DataWithMeta[BytesIO]: cache = MediaCache() preview, mime = cache.get_or_set_segment_preview(self._db_segment) - return DataWithMeta[BytesIO](preview, mime=mime, checksum=None) + return DataWithMeta[BytesIO](preview, mime=mime) def get_chunk( self, chunk_number: int, *, quality: FrameQuality = FrameQuality.ORIGINAL ) -> DataWithMeta[BytesIO]: chunk_number = self.validate_chunk_number(chunk_number) chunk_data, mime = self._loaders[quality].read_chunk(chunk_number) - return DataWithMeta[BytesIO](chunk_data, mime=mime, checksum=None) + return DataWithMeta[BytesIO](chunk_data, mime=mime) def get_frame( self, @@ -435,9 +434,9 @@ def get_frame( frame = self._convert_frame(frame, loader.reader_class, out_type) if loader.reader_class is VideoReader: - return return_type(frame, mime=self.VIDEO_FRAME_MIME, checksum=None) + return return_type(frame, mime=self.VIDEO_FRAME_MIME) - return return_type(frame, mime=mimetypes.guess_type(frame_name)[0], checksum=None) + return return_type(frame, mime=mimetypes.guess_type(frame_name)[0]) def get_frame_context_images( self, @@ -454,7 +453,7 @@ def get_frame_context_images( if not data: return None - return DataWithMeta[BytesIO](data, mime=mime, checksum=None) + return DataWithMeta[BytesIO](data, mime=mime) def iterate_frames( self, From 654a827ce56f6debe51a875d40e5d7d0079c140c Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 6 Aug 2024 12:19:07 +0300 Subject: [PATCH 010/115] Be consistent about returned task chunk types (allow video chunks) --- cvat/apps/engine/frame_provider.py | 75 ++++++++++++++++++++-------- cvat/apps/engine/media_extractors.py | 6 ++- 2 files changed, 57 insertions(+), 24 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index d3235781b49..a299ba29a0d 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -125,14 +125,14 @@ def _av_frame_to_png_bytes(cls, av_frame: av.VideoFrame) -> BytesIO: return BytesIO(result.tobytes()) def _convert_frame( - self, frame: Any, reader_class: IMediaReader, out_type: FrameOutputType + self, frame: Any, reader: IMediaReader, out_type: FrameOutputType ) -> AnyFrame: if out_type == FrameOutputType.BUFFER: - return self._av_frame_to_png_bytes(frame) if reader_class is VideoReader else frame + return self._av_frame_to_png_bytes(frame) if isinstance(reader, VideoReader) else frame elif out_type == FrameOutputType.PIL: - return frame.to_image() if reader_class is VideoReader else Image.open(frame) + return frame.to_image() if isinstance(reader, VideoReader) else Image.open(frame) elif out_type == FrameOutputType.NUMPY_ARRAY: - if reader_class is VideoReader: + if isinstance(reader, VideoReader): image = frame.to_ndarray(format="bgr24") else: image = np.array(Image.open(frame)) @@ -255,6 +255,28 @@ def get_chunk( # Create and return a joined / cleaned chunk # TODO: refactor into another class, maybe optimize + + writer_classes: dict[FrameQuality, Type[IChunkWriter]] = { + FrameQuality.COMPRESSED: ( + Mpeg4CompressedChunkWriter + if db_data.compressed_chunk_type == models.DataChoice.VIDEO + else ZipCompressedChunkWriter + ), + FrameQuality.ORIGINAL: ( + Mpeg4ChunkWriter + if db_data.original_chunk_type == models.DataChoice.VIDEO + else ZipChunkWriter + ), + } + + writer_class = writer_classes[quality] + + mime_type = ( + "video/mp4" + if writer_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] + else "application/zip" + ) + task_chunk_frames = {} for db_segment in matching_segments: segment_frame_provider = SegmentFrameProvider(db_segment) @@ -267,27 +289,27 @@ def get_chunk( ): continue - frame = segment_frame_provider.get_frame( - task_chunk_frame_id, quality=quality, out_type=FrameOutputType.BUFFER - ).data + frame, _, _ = segment_frame_provider._get_raw_frame( + task_chunk_frame_id, quality=quality + ) task_chunk_frames[task_chunk_frame_id] = (frame, None, None) - kwargs = {} + writer_kwargs = {} if self._db_task.dimension == models.DimensionType.DIM_3D: - kwargs["dimension"] = models.DimensionType.DIM_3D - merged_chunk_writer = ZipCompressedChunkWriter( - 100 if quality == FrameQuality.ORIGINAL else db_data.image_quality, **kwargs - ) + writer_kwargs["dimension"] = models.DimensionType.DIM_3D + merged_chunk_writer = writer_class(int(quality), **writer_kwargs) + + writer_kwargs = {} + if isinstance(merged_chunk_writer, ZipCompressedChunkWriter): + writer_kwargs = dict(compress_frames=False, zip_compress_level=1) buffer = io.BytesIO() - merged_chunk_writer.save_as_chunk( - task_chunk_frames.values(), buffer, compress_frames=False, zip_compress_level=1 - ) + merged_chunk_writer.save_as_chunk(list(task_chunk_frames.values()), buffer, **writer_kwargs) buffer.seek(0) # TODO: add caching in media cache for the resulting chunk - return return_type(data=buffer, mime="application/zip") + return return_type(data=buffer, mime=mime_type) def get_frame( self, @@ -418,6 +440,18 @@ def get_chunk( chunk_data, mime = self._loaders[quality].read_chunk(chunk_number) return DataWithMeta[BytesIO](chunk_data, mime=mime) + def _get_raw_frame( + self, + frame_number: int, + *, + quality: FrameQuality = FrameQuality.ORIGINAL, + ) -> Tuple[Any, str, IMediaReader]: + _, chunk_number, frame_offset = self.validate_frame_number(frame_number) + loader = self._loaders[quality] + chunk_reader = loader.load(chunk_number) + frame, frame_name, _ = chunk_reader[frame_offset] + return frame, frame_name, chunk_reader + def get_frame( self, frame_number: int, @@ -427,13 +461,10 @@ def get_frame( ) -> DataWithMeta[AnyFrame]: return_type = DataWithMeta[AnyFrame] - _, chunk_number, frame_offset = self.validate_frame_number(frame_number) - loader = self._loaders[quality] - chunk_reader = loader.load(chunk_number) - frame, frame_name, _ = chunk_reader[frame_offset] + frame, frame_name, reader = self._get_raw_frame(frame_number, quality=quality) - frame = self._convert_frame(frame, loader.reader_class, out_type) - if loader.reader_class is VideoReader: + frame = self._convert_frame(frame, reader, out_type) + if isinstance(reader, VideoReader): return return_type(frame, mime=self.VIDEO_FRAME_MIME) return return_type(frame, mime=mimetypes.guess_type(frame_name)[0]) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 80d8dfc1dba..883de77c96a 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -18,7 +18,7 @@ from contextlib import ExitStack, closing from dataclasses import dataclass from enum import IntEnum -from typing import Callable, Iterable, Iterator, Optional, Protocol, Tuple, TypeVar +from typing import Callable, Iterable, Iterator, Optional, Protocol, Sequence, Tuple, TypeVar import av import av.codec @@ -910,7 +910,9 @@ def _add_video_stream(self, container: av.container.OutputContainer, w, h, rate, return video_stream - def save_as_chunk(self, images, chunk_path): + def save_as_chunk( + self, images: Sequence[av.VideoFrame], chunk_path: str + ) -> Sequence[Tuple[int, int]]: if not images: raise Exception('no images to save') From 12e5f2ac7dfc46b059810ffa38116b4b7b5f8661 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 6 Aug 2024 13:01:02 +0300 Subject: [PATCH 011/115] Support iterator input in video chunk writing --- cvat/apps/engine/frame_provider.py | 2 +- cvat/apps/engine/media_extractors.py | 31 +++++++++++++++++++++------- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index a299ba29a0d..2ad8854d520 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -304,7 +304,7 @@ def get_chunk( writer_kwargs = dict(compress_frames=False, zip_compress_level=1) buffer = io.BytesIO() - merged_chunk_writer.save_as_chunk(list(task_chunk_frames.values()), buffer, **writer_kwargs) + merged_chunk_writer.save_as_chunk(task_chunk_frames.values(), buffer, **writer_kwargs) buffer.seek(0) # TODO: add caching in media cache for the resulting chunk diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 883de77c96a..33e247a5a7c 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -18,7 +18,7 @@ from contextlib import ExitStack, closing from dataclasses import dataclass from enum import IntEnum -from typing import Callable, Iterable, Iterator, Optional, Protocol, Sequence, Tuple, TypeVar +from typing import Any, Callable, Iterable, Iterator, Optional, Protocol, Sequence, Tuple, TypeVar import av import av.codec @@ -910,14 +910,28 @@ def _add_video_stream(self, container: av.container.OutputContainer, w, h, rate, return video_stream + FrameDescriptor = Tuple[av.VideoFrame, Any, Any] + + def _peek_first_frame( + self, frame_iter: Iterator[FrameDescriptor] + ) -> Tuple[Optional[FrameDescriptor], Iterator[FrameDescriptor]]: + "Gets the first frame and returns the same full iterator" + + if not hasattr(frame_iter, '__next__'): + frame_iter = iter(frame_iter) + + first_frame = next(frame_iter, None) + return first_frame, itertools.chain((first_frame, ), frame_iter) + def save_as_chunk( - self, images: Sequence[av.VideoFrame], chunk_path: str + self, images: Iterator[FrameDescriptor], chunk_path: str ) -> Sequence[Tuple[int, int]]: - if not images: + first_frame, images = self._peek_first_frame(images) + if not first_frame: raise Exception('no images to save') - input_w = images[0][0].width - input_h = images[0][0].height + input_w = first_frame[0].width + input_h = first_frame[0].height with av.open(chunk_path, 'w', format=self.FORMAT) as output_container: output_v_stream = self._add_video_stream( @@ -962,11 +976,12 @@ def __init__(self, quality): } def save_as_chunk(self, images, chunk_path): - if not images: + first_frame, images = self._peek_first_frame(images) + if not first_frame: raise Exception('no images to save') - input_w = images[0][0].width - input_h = images[0][0].height + input_w = first_frame[0].width + input_h = first_frame[0].height downscale_factor = 1 while input_h / downscale_factor >= 1080: From a79a681c8d369a0670496b61c053f7a2279f1e04 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 6 Aug 2024 13:05:29 +0300 Subject: [PATCH 012/115] Fix type annotation --- cvat/apps/engine/frame_provider.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 2ad8854d520..e9b655aa5c2 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -39,7 +39,7 @@ class _ChunkLoader(metaclass=ABCMeta): - def __init__(self, reader_class: IMediaReader) -> None: + def __init__(self, reader_class: Type[IMediaReader]) -> None: self.chunk_id: Optional[int] = None self.chunk_reader: Optional[RandomAccessIterator] = None self.reader_class = reader_class @@ -66,7 +66,7 @@ def read_chunk(self, chunk_id: int) -> DataWithMime: ... class _FileChunkLoader(_ChunkLoader): def __init__( - self, reader_class: IMediaReader, get_chunk_path_callback: Callable[[int], str] + self, reader_class: Type[IMediaReader], get_chunk_path_callback: Callable[[int], str] ) -> None: super().__init__(reader_class) self.get_chunk_path = get_chunk_path_callback @@ -82,7 +82,7 @@ def read_chunk(self, chunk_id: int) -> DataWithMime: class _BufferChunkLoader(_ChunkLoader): def __init__( - self, reader_class: IMediaReader, get_chunk_callback: Callable[[int], DataWithMime] + self, reader_class: Type[IMediaReader], get_chunk_callback: Callable[[int], DataWithMime] ) -> None: super().__init__(reader_class) self.get_chunk = get_chunk_callback @@ -359,7 +359,7 @@ def __init__(self, db_segment: models.Segment) -> None: db_data = db_segment.task.data - reader_class: dict[models.DataChoice, IMediaReader] = { + reader_class: dict[models.DataChoice, Type[IMediaReader]] = { models.DataChoice.IMAGESET: ZipReader, models.DataChoice.VIDEO: VideoReader, } From d5118a2d253e5eb9388bfcc9f98286c083c8d61c Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 6 Aug 2024 13:43:03 +0300 Subject: [PATCH 013/115] Refactor video reader memory leak fix, add to reader with manifest --- cvat/apps/engine/media_extractors.py | 69 +++++++++++++++++++--------- 1 file changed, 47 insertions(+), 22 deletions(-) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 33e247a5a7c..c7314bcf4b6 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -15,7 +15,7 @@ import struct from abc import ABC, abstractmethod from bisect import bisect -from contextlib import ExitStack, closing +from contextlib import ExitStack, closing, contextmanager from dataclasses import dataclass from enum import IntEnum from typing import Any, Callable, Iterable, Iterator, Optional, Protocol, Sequence, Tuple, TypeVar @@ -538,21 +538,19 @@ def _make_frame_iterator( ) -> Iterator[Tuple[av.VideoFrame, str, int]]: es = ExitStack() - need_init = stream is None - if need_init: - container = es.enter_context(self._get_av_container()) + needs_init = stream is None + if needs_init: + container = es.enter_context(self._read_av_container()) else: container = stream.container with es: - if need_init: + if needs_init: stream = container.streams.video[0] if self.allow_threading: stream.thread_type = 'AUTO' - es.enter_context(closing(stream.codec_context)) - frame_num = 0 for packet in container.demux(stream): @@ -588,16 +586,27 @@ def get_progress(self, pos): duration = self._get_duration() return pos / duration if duration else None - def _get_av_container(self): + @contextmanager + def _read_av_container(self): if isinstance(self._source_path[0], io.BytesIO): self._source_path[0].seek(0) # required for re-reading - return av.open(self._source_path[0]) + + container = av.open(self._source_path[0]) + try: + yield container + finally: + # fixes a memory leak in input container closing + # https://github.com/PyAV-Org/PyAV/issues/1117 + for stream in container.streams: + context = stream.codec_context + if context and context.is_open: + context.close() + + container.close() def _get_duration(self): - with ExitStack() as es: - container = es.enter_context(self._get_av_container()) + with self._read_av_container() as container: stream = container.streams.video[0] - es.enter_context(closing(stream.codec_context)) duration = None if stream.duration: @@ -613,10 +622,8 @@ def _get_duration(self): return duration def get_preview(self, frame): - with ExitStack() as es: - container = es.enter_context(self._get_av_container()) + with self._read_av_container() as container: stream = container.streams.video[0] - es.enter_context(closing(stream.codec_context)) tb_denominator = stream.time_base.denominator needed_time = int((frame / stream.guessed_rate) * tb_denominator) @@ -670,11 +677,28 @@ def iterate_frames(self, frame_ids: Iterable[int]): yield self._manifest[idx] class VideoReaderWithManifest: - def __init__(self, manifest_path: str, source_path: str): + def __init__(self, manifest_path: str, source_path: str, *, allow_threading: bool = False): self._source_path = source_path self._manifest = VideoManifestManager(manifest_path) self._manifest.init_index() + self.allow_threading = allow_threading + + @contextmanager + def _read_av_container(self): + container = av.open(self._source_path[0]) + try: + yield container + finally: + # fixes a memory leak in input container closing + # https://github.com/PyAV-Org/PyAV/issues/1117 + for stream in container.streams: + context = stream.codec_context + if context and context.is_open: + context.close() + + container.close() + def _get_nearest_left_key_frame(self, frame_id: int) -> tuple[int, int]: nearest_left_keyframe_pos = bisect( self._manifest, frame_id, key=lambda entry: entry.get('number') @@ -695,21 +719,22 @@ def iterate_frames(self, frame_ids: Iterable[int]) -> Iterable[av.VideoFrame]: frame_ids_frame ) - with closing(av.open(self._source_path, mode='r')) as container: - video_stream = next(stream for stream in container.streams if stream.type == 'video') - video_stream.thread_type = 'AUTO' + with self._read_av_container() as container: + stream = container.streams.video[0] + if self.allow_threading: + stream.thread_type = 'AUTO' - container.seek(offset=start_decode_timestamp, stream=video_stream) + container.seek(offset=start_decode_timestamp, stream=stream) frame_number = start_decode_frame_number - 1 - for packet in container.demux(video_stream): + for packet in container.demux(stream): for frame in packet.decode(): frame_number += 1 if frame_number < frame_ids_frame: continue elif frame_number == frame_ids_frame: - if video_stream.metadata.get('rotate'): + if stream.metadata.get('rotate'): frame = av.VideoFrame().from_ndarray( rotate_image( frame.to_ndarray(format='bgr24'), From 1b429cffb4e4a9649cfc78b777426098567cadc4 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 6 Aug 2024 13:43:34 +0300 Subject: [PATCH 014/115] Disable threading in video reading in frame provider --- cvat/apps/engine/frame_provider.py | 52 +++++++++++++++++++++------- cvat/apps/engine/media_extractors.py | 2 +- 2 files changed, 40 insertions(+), 14 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index e9b655aa5c2..e06be42d64a 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -39,10 +39,16 @@ class _ChunkLoader(metaclass=ABCMeta): - def __init__(self, reader_class: Type[IMediaReader]) -> None: + def __init__( + self, + reader_class: Type[IMediaReader], + *, + reader_params: Optional[dict] = None, + ) -> None: self.chunk_id: Optional[int] = None self.chunk_reader: Optional[RandomAccessIterator] = None self.reader_class = reader_class + self.reader_params = reader_params def load(self, chunk_id: int) -> RandomAccessIterator[Tuple[Any, str, int]]: if self.chunk_id != chunk_id: @@ -50,7 +56,10 @@ def load(self, chunk_id: int) -> RandomAccessIterator[Tuple[Any, str, int]]: self.chunk_id = chunk_id self.chunk_reader = RandomAccessIterator( - self.reader_class([self.read_chunk(chunk_id)[0]]) + self.reader_class( + [self.read_chunk(chunk_id)[0]], + **(self.reader_params or {}), + ) ) return self.chunk_reader @@ -66,9 +75,13 @@ def read_chunk(self, chunk_id: int) -> DataWithMime: ... class _FileChunkLoader(_ChunkLoader): def __init__( - self, reader_class: Type[IMediaReader], get_chunk_path_callback: Callable[[int], str] + self, + reader_class: Type[IMediaReader], + get_chunk_path_callback: Callable[[int], str], + *, + reader_params: Optional[dict] = None, ) -> None: - super().__init__(reader_class) + super().__init__(reader_class, reader_params=reader_params) self.get_chunk_path = get_chunk_path_callback def read_chunk(self, chunk_id: int) -> DataWithMime: @@ -82,9 +95,13 @@ def read_chunk(self, chunk_id: int) -> DataWithMime: class _BufferChunkLoader(_ChunkLoader): def __init__( - self, reader_class: Type[IMediaReader], get_chunk_callback: Callable[[int], DataWithMime] + self, + reader_class: Type[IMediaReader], + get_chunk_callback: Callable[[int], DataWithMime], + *, + reader_params: Optional[dict] = None, ) -> None: - super().__init__(reader_class) + super().__init__(reader_class, reader_params=reader_params) self.get_chunk = get_chunk_callback def read_chunk(self, chunk_id: int) -> DataWithMime: @@ -359,9 +376,14 @@ def __init__(self, db_segment: models.Segment) -> None: db_data = db_segment.task.data - reader_class: dict[models.DataChoice, Type[IMediaReader]] = { - models.DataChoice.IMAGESET: ZipReader, - models.DataChoice.VIDEO: VideoReader, + reader_class: dict[models.DataChoice, Tuple[Type[IMediaReader], Optional[dict]]] = { + models.DataChoice.IMAGESET: (ZipReader, None), + models.DataChoice.VIDEO: (VideoReader, { + "allow_threading": False + # disable threading to avoid unpredictable server + # resource consumption during reading in endpoints + # can be enabled for other clients + }), } self._loaders: dict[FrameQuality, _ChunkLoader] = {} @@ -369,28 +391,32 @@ def __init__(self, db_segment: models.Segment) -> None: cache = MediaCache() self._loaders[FrameQuality.COMPRESSED] = _BufferChunkLoader( - reader_class=reader_class[db_data.compressed_chunk_type], + reader_class=reader_class[db_data.compressed_chunk_type][0], + reader_params=reader_class[db_data.compressed_chunk_type][1], get_chunk_callback=lambda chunk_idx: cache.get_segment_chunk( db_segment, chunk_idx, quality=FrameQuality.COMPRESSED ), ) self._loaders[FrameQuality.ORIGINAL] = _BufferChunkLoader( - reader_class=reader_class[db_data.original_chunk_type], + reader_class=reader_class[db_data.original_chunk_type][0], + reader_params=reader_class[db_data.original_chunk_type][1], get_chunk_callback=lambda chunk_idx: cache.get_segment_chunk( db_segment, chunk_idx, quality=FrameQuality.ORIGINAL ), ) else: self._loaders[FrameQuality.COMPRESSED] = _FileChunkLoader( - reader_class=reader_class[db_data.compressed_chunk_type], + reader_class=reader_class[db_data.compressed_chunk_type][0], + reader_params=reader_class[db_data.compressed_chunk_type][1], get_chunk_path_callback=lambda chunk_idx: db_data.get_compressed_segment_chunk_path( chunk_idx, segment=db_segment.id ), ) self._loaders[FrameQuality.ORIGINAL] = _FileChunkLoader( - reader_class=reader_class[db_data.original_chunk_type], + reader_class=reader_class[db_data.original_chunk_type][0], + reader_params=reader_class[db_data.original_chunk_type][1], get_chunk_path_callback=lambda chunk_idx: db_data.get_original_segment_chunk_path( chunk_idx, segment=db_segment.id ), diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index c7314bcf4b6..ef2bd8e0e9e 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -686,7 +686,7 @@ def __init__(self, manifest_path: str, source_path: str, *, allow_threading: boo @contextmanager def _read_av_container(self): - container = av.open(self._source_path[0]) + container = av.open(self._source_path) try: yield container finally: From d512312508c7ef5915d302e01c4bd50e0bd8b3b8 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 6 Aug 2024 14:11:43 +0300 Subject: [PATCH 015/115] Fix keyframe search --- cvat/apps/engine/media_extractors.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index ef2bd8e0e9e..41170736055 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -703,8 +703,12 @@ def _get_nearest_left_key_frame(self, frame_id: int) -> tuple[int, int]: nearest_left_keyframe_pos = bisect( self._manifest, frame_id, key=lambda entry: entry.get('number') ) - frame_number = self._manifest[nearest_left_keyframe_pos].get('number') - timestamp = self._manifest[nearest_left_keyframe_pos].get('pts') + if nearest_left_keyframe_pos: + frame_number = self._manifest[nearest_left_keyframe_pos - 1].get('number') + timestamp = self._manifest[nearest_left_keyframe_pos - 1].get('pts') + else: + frame_number = 0 + timestamp = 0 return frame_number, timestamp def iterate_frames(self, frame_ids: Iterable[int]) -> Iterable[av.VideoFrame]: @@ -744,10 +748,12 @@ def iterate_frames(self, frame_ids: Iterable[int]) -> Iterable[av.VideoFrame]: ) yield frame - else: + frame_ids_frame = next(frame_ids_iter, None) - if frame_ids_frame is None: - return + + if frame_ids_frame is None: + return + class IChunkWriter(ABC): def __init__(self, quality, dimension=DimensionType.DIM_2D): From 167ee12b43471486f5db2856a55aadcc4f09a89d Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 6 Aug 2024 14:56:20 +0300 Subject: [PATCH 016/115] Return frames as generator in dynamic chunk creation --- cvat/apps/engine/cache.py | 63 +++++++++++++--------------- cvat/apps/engine/frame_provider.py | 15 ++++--- cvat/apps/engine/media_extractors.py | 4 +- 3 files changed, 39 insertions(+), 43 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 521086296bb..66c08385340 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -11,10 +11,10 @@ import tempfile import zipfile import zlib -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack, closing from datetime import datetime, timezone from itertools import pairwise -from typing import Any, Callable, Iterable, Optional, Sequence, Tuple, Type, Union +from typing import Any, Callable, Iterator, Optional, Sequence, Tuple, Type, Union import av import cv2 @@ -149,10 +149,9 @@ def get_frame_context_images(self, db_data: models.Data, frame_number: int) -> D create_callback=lambda: self._prepare_context_image(db_data, frame_number), ) - @contextmanager def _read_raw_frames( self, db_task: models.Task, frame_ids: Sequence[int] - ) -> Iterable[Tuple[Union[av.VideoFrame, PIL.Image.Image], str, str]]: + ) -> Iterator[Tuple[Union[av.VideoFrame, PIL.Image.Image], str, str]]: for prev_frame, cur_frame in pairwise(frame_ids): assert ( prev_frame <= cur_frame @@ -168,20 +167,19 @@ def _read_raw_frames( dimension = db_task.dimension - media = [] - with ExitStack() as es: - if hasattr(db_data, "video"): - source_path = os.path.join(raw_data_dir, db_data.video.path) - - reader = VideoReaderWithManifest( - manifest_path=db_data.get_manifest_path(), - source_path=source_path, - ) - for frame in reader.iterate_frames(frame_ids): - media.append((frame, source_path, None)) - else: - reader = ImageReaderWithManifest(db_data.get_manifest_path()) - if db_data.storage == models.StorageChoice.CLOUD_STORAGE: + if hasattr(db_data, "video"): + source_path = os.path.join(raw_data_dir, db_data.video.path) + + reader = VideoReaderWithManifest( + manifest_path=db_data.get_manifest_path(), + source_path=source_path, + ) + for frame in reader.iterate_frames(frame_ids): + yield (frame, source_path, None) + else: + reader = ImageReaderWithManifest(db_data.get_manifest_path()) + if db_data.storage == models.StorageChoice.CLOUD_STORAGE: + with ExitStack() as es: db_cloud_storage = db_data.cloud_storage assert db_cloud_storage, "Cloud storage instance was deleted" credentials = Credentials() @@ -221,17 +219,17 @@ def _read_raw_frames( slogger.cloud_storage[db_cloud_storage.id].warning( "Hash sums of files {} do not match".format(file_name) ) - else: - for item in reader.iterate_frames(frame_ids): - source_path = os.path.join( - raw_data_dir, f"{item['name']}{item['extension']}" - ) - media.append((source_path, source_path, None)) - if dimension == models.DimensionType.DIM_2D: - media = preload_images(media) + yield from media + else: + for item in reader.iterate_frames(frame_ids): + source_path = os.path.join(raw_data_dir, f"{item['name']}{item['extension']}") + media.append((source_path, source_path, None)) - yield media + if dimension == models.DimensionType.DIM_2D: + media = preload_images(media) + + yield from media def prepare_segment_chunk( self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality @@ -269,11 +267,6 @@ def prepare_range_segment_chunk( ), } - image_quality = ( - 100 - if writer_classes[quality] in [Mpeg4ChunkWriter, ZipChunkWriter] - else db_data.image_quality - ) mime_type = ( "video/mp4" if writer_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] @@ -283,11 +276,11 @@ def prepare_range_segment_chunk( kwargs = {} if db_segment.task.dimension == models.DimensionType.DIM_3D: kwargs["dimension"] = models.DimensionType.DIM_3D - writer = writer_classes[quality](image_quality, **kwargs) + writer = writer_classes[quality](int(quality), **kwargs) buffer = io.BytesIO() - with self._read_raw_frames(db_task, frame_ids=chunk_frame_ids) as images: - writer.save_as_chunk(images, buffer) + with closing(self._read_raw_frames(db_task, frame_ids=chunk_frame_ids)) as frame_iter: + writer.save_as_chunk(frame_iter, buffer) buffer.seek(0) return buffer, mime_type diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index e06be42d64a..6e5980527bb 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -378,12 +378,15 @@ def __init__(self, db_segment: models.Segment) -> None: reader_class: dict[models.DataChoice, Tuple[Type[IMediaReader], Optional[dict]]] = { models.DataChoice.IMAGESET: (ZipReader, None), - models.DataChoice.VIDEO: (VideoReader, { - "allow_threading": False - # disable threading to avoid unpredictable server - # resource consumption during reading in endpoints - # can be enabled for other clients - }), + models.DataChoice.VIDEO: ( + VideoReader, + { + "allow_threading": False + # disable threading to avoid unpredictable server + # resource consumption during reading in endpoints + # can be enabled for other clients + }, + ), } self._loaders: dict[FrameQuality, _ChunkLoader] = {} diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 41170736055..b1049cbdd37 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -827,7 +827,7 @@ def _write_pcd_file(self, image: str|io.BytesIO) -> tuple[io.BytesIO, str, int, if isinstance(image, str): image_buf.close() - def save_as_chunk(self, images: Iterable[tuple[Image.Image|io.IOBase|str, str, str]], chunk_path: str): + def save_as_chunk(self, images: Iterator[tuple[Image.Image|io.IOBase|str, str, str]], chunk_path: str): with zipfile.ZipFile(chunk_path, 'x') as zip_chunk: for idx, (image, path, _) in enumerate(images): ext = os.path.splitext(path)[1].replace('.', '') @@ -872,7 +872,7 @@ def save_as_chunk(self, images: Iterable[tuple[Image.Image|io.IOBase|str, str, s class ZipCompressedChunkWriter(ZipChunkWriter): def save_as_chunk( self, - images: Iterable[tuple[Image.Image|io.IOBase|str, str, str]], + images: Iterator[tuple[Image.Image|io.IOBase|str, str, str]], chunk_path: str, *, compress_frames: bool = True, zip_compress_level: int = 0 ): image_sizes = [] From 88a9cb258e067efea269bccba1a510bd88cf2437 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 7 Aug 2024 12:29:56 +0300 Subject: [PATCH 017/115] Update chunk requests in UI --- cvat-core/src/frames.ts | 39 +++++++++++++++++++++------------------ 1 file changed, 21 insertions(+), 18 deletions(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index 8aab9f7f0eb..7dd2ccc06fa 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -25,7 +25,7 @@ const frameDataCache: Record | null; activeContextRequest: Promise> | null; @@ -208,10 +208,12 @@ export class FrameData { class PrefetchAnalyzer { #chunkSize: number; #requestedFrames: number[]; + #startFrame: number; - constructor(chunkSize) { + constructor(chunkSize, startFrame) { this.#chunkSize = chunkSize; this.#requestedFrames = []; + this.#startFrame = startFrame; } shouldPrefetchNext(current: number, isPlaying: boolean, isChunkCached: (chunk) => boolean): boolean { @@ -219,13 +221,13 @@ class PrefetchAnalyzer { return true; } - const currentChunk = Math.floor(current / this.#chunkSize); + const currentChunk = Math.floor((current - this.#startFrame) / this.#chunkSize); const { length } = this.#requestedFrames; const isIncreasingOrder = this.#requestedFrames .every((val, index) => index === 0 || val > this.#requestedFrames[index - 1]); if ( length && (isIncreasingOrder && current > this.#requestedFrames[length - 1]) && - (current % this.#chunkSize) >= Math.ceil(this.#chunkSize / 2) && + ((current - this.#startFrame) % this.#chunkSize) >= Math.ceil(this.#chunkSize / 2) && !isChunkCached(currentChunk + 1) ) { // is increasing order including the current frame @@ -262,19 +264,20 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { imageData: ImageBitmap | Blob; } | Blob>((resolve, reject) => { const { - provider, prefetchAnalizer, chunkSize, stopFrame, decodeForward, forwardStep, decodedBlocksCacheSize, + provider, prefetchAnalyzer, chunkSize, startFrame, stopFrame, + decodeForward, forwardStep, decodedBlocksCacheSize, } = frameDataCache[this.jobID]; const requestId = +_.uniqueId(); - const chunkNumber = Math.floor(this.number / chunkSize); + const chunkNumber = Math.floor((this.number - startFrame) / chunkSize); const frame = provider.frame(this.number); function findTheNextNotDecodedChunk(searchFrom: number): number { let firstFrameInNextChunk = searchFrom + forwardStep; - let nextChunkNumber = Math.floor(firstFrameInNextChunk / chunkSize); + let nextChunkNumber = Math.floor((firstFrameInNextChunk - startFrame) / chunkSize); while (nextChunkNumber === chunkNumber) { firstFrameInNextChunk += forwardStep; - nextChunkNumber = Math.floor(firstFrameInNextChunk / chunkSize); + nextChunkNumber = Math.floor((firstFrameInNextChunk - startFrame) / chunkSize); } if (provider.isChunkCached(nextChunkNumber)) { @@ -286,7 +289,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { if (frame) { if ( - prefetchAnalizer.shouldPrefetchNext( + prefetchAnalyzer.shouldPrefetchNext( this.number, decodeForward, (chunk) => provider.isChunkCached(chunk), @@ -294,7 +297,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { ) { const nextChunkNumber = findTheNextNotDecodedChunk(this.number); const predecodeChunksMax = Math.floor(decodedBlocksCacheSize / 2); - if (nextChunkNumber * chunkSize <= stopFrame && + if (startFrame + nextChunkNumber * chunkSize <= stopFrame && nextChunkNumber <= chunkNumber + predecodeChunksMax ) { frameDataCache[this.jobID].activeChunkRequest = new Promise((resolveForward) => { @@ -316,8 +319,8 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { provider.cleanup(1); provider.requestDecodeBlock( chunk, - nextChunkNumber * chunkSize, - Math.min(stopFrame, (nextChunkNumber + 1) * chunkSize - 1), + startFrame + nextChunkNumber * chunkSize, + Math.min(stopFrame, startFrame + (nextChunkNumber + 1) * chunkSize - 1), () => {}, releasePromise, releasePromise, @@ -334,7 +337,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { renderHeight: this.height, imageData: frame, }); - prefetchAnalizer.addRequested(this.number); + prefetchAnalyzer.addRequested(this.number); return; } @@ -355,7 +358,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { renderHeight: this.height, imageData: currentFrame, }); - prefetchAnalizer.addRequested(this.number); + prefetchAnalyzer.addRequested(this.number); return; } @@ -378,8 +381,8 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { provider .requestDecodeBlock( chunk, - chunkNumber * chunkSize, - Math.min(stopFrame, (chunkNumber + 1) * chunkSize - 1), + startFrame + chunkNumber * chunkSize, + Math.min(stopFrame, startFrame + (chunkNumber + 1) * chunkSize - 1), (_frame: number, bitmap: ImageBitmap | Blob) => { if (decodeForward) { // resolve immediately only if is not playing @@ -395,7 +398,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { renderHeight: this.height, imageData: bitmap, }); - prefetchAnalizer.addRequested(this.number); + prefetchAnalyzer.addRequested(this.number); } }, () => { frameDataCache[this.jobID].activeChunkRequest = null; @@ -614,7 +617,7 @@ export async function getFrame( decodedBlocksCacheSize, dimension, ), - prefetchAnalizer: new PrefetchAnalyzer(chunkSize), + prefetchAnalyzer: new PrefetchAnalyzer(chunkSize, startFrame), decodedBlocksCacheSize, activeChunkRequest: null, activeContextRequest: null, From 30bf8fd8c3bb41dc8ffaca0fbf690d59dd538cb0 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 7 Aug 2024 14:27:13 +0300 Subject: [PATCH 018/115] Update cache indices in FrameDecoder, enable video play --- cvat-core/src/frames.ts | 1 + cvat-data/src/ts/cvat-data.ts | 7 +++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index 7dd2ccc06fa..ff6200ce91b 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -615,6 +615,7 @@ export async function getFrame( blockType, chunkSize, decodedBlocksCacheSize, + startFrame, dimension, ), prefetchAnalyzer: new PrefetchAnalyzer(chunkSize, startFrame), diff --git a/cvat-data/src/ts/cvat-data.ts b/cvat-data/src/ts/cvat-data.ts index 2f832ac9d3f..ec9fc5cccc5 100644 --- a/cvat-data/src/ts/cvat-data.ts +++ b/cvat-data/src/ts/cvat-data.ts @@ -100,11 +100,13 @@ export class FrameDecoder { private renderHeight: number; private zipWorker: Worker | null; private videoWorker: Worker | null; + private startFrame: number; constructor( blockType: BlockType, chunkSize: number, cachedBlockCount: number, + startFrame: number, dimension: DimensionType = DimensionType.DIMENSION_2D, ) { this.mutex = new Mutex(); @@ -118,6 +120,7 @@ export class FrameDecoder { this.renderWidth = 1920; this.renderHeight = 1080; this.chunkSize = chunkSize; + this.startFrame = startFrame; this.blockType = blockType; this.decodedChunks = {}; @@ -203,7 +206,7 @@ export class FrameDecoder { } frame(frameNumber: number): ImageBitmap | Blob | null { - const chunkNumber = Math.floor(frameNumber / this.chunkSize); + const chunkNumber = Math.floor((frameNumber - this.startFrame) / this.chunkSize); if (chunkNumber in this.decodedChunks) { return this.decodedChunks[chunkNumber][frameNumber]; } @@ -262,7 +265,7 @@ export class FrameDecoder { throw new RequestOutdatedError(); } - const chunkNumber = Math.floor(start / this.chunkSize); + const chunkNumber = Math.floor((start - this.startFrame) / this.chunkSize); this.orderedStack = [chunkNumber, ...this.orderedStack]; this.cleanup(); const decodedFrames: Record = {}; From ee3c905debc918fc5514ef3049d231486cced833 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 7 Aug 2024 14:27:35 +0300 Subject: [PATCH 019/115] Fix frame retrieval for video --- cvat/apps/engine/frame_provider.py | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 6e5980527bb..078f626cd56 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -142,14 +142,18 @@ def _av_frame_to_png_bytes(cls, av_frame: av.VideoFrame) -> BytesIO: return BytesIO(result.tobytes()) def _convert_frame( - self, frame: Any, reader: IMediaReader, out_type: FrameOutputType + self, frame: Any, reader_class: Type[IMediaReader], out_type: FrameOutputType ) -> AnyFrame: if out_type == FrameOutputType.BUFFER: - return self._av_frame_to_png_bytes(frame) if isinstance(reader, VideoReader) else frame + return ( + self._av_frame_to_png_bytes(frame) + if issubclass(reader_class, VideoReader) + else frame + ) elif out_type == FrameOutputType.PIL: - return frame.to_image() if isinstance(reader, VideoReader) else Image.open(frame) + return frame.to_image() if issubclass(reader_class, VideoReader) else Image.open(frame) elif out_type == FrameOutputType.NUMPY_ARRAY: - if isinstance(reader, VideoReader): + if issubclass(reader_class, VideoReader): image = frame.to_ndarray(format="bgr24") else: image = np.array(Image.open(frame)) @@ -358,6 +362,9 @@ def iterate_frames( yield self.get_frame(idx, quality=quality, out_type=out_type) def _get_segment(self, validated_frame_number: int) -> models.Segment: + if not self._db_task.data or not self._db_task.data.size: + raise ValidationError("Task has no data") + return next( s for s in self._db_task.segment_set.all() @@ -474,12 +481,12 @@ def _get_raw_frame( frame_number: int, *, quality: FrameQuality = FrameQuality.ORIGINAL, - ) -> Tuple[Any, str, IMediaReader]: + ) -> Tuple[Any, str, Type[IMediaReader]]: _, chunk_number, frame_offset = self.validate_frame_number(frame_number) loader = self._loaders[quality] chunk_reader = loader.load(chunk_number) frame, frame_name, _ = chunk_reader[frame_offset] - return frame, frame_name, chunk_reader + return frame, frame_name, loader.reader_class def get_frame( self, @@ -490,10 +497,10 @@ def get_frame( ) -> DataWithMeta[AnyFrame]: return_type = DataWithMeta[AnyFrame] - frame, frame_name, reader = self._get_raw_frame(frame_number, quality=quality) + frame, frame_name, reader_class = self._get_raw_frame(frame_number, quality=quality) - frame = self._convert_frame(frame, reader, out_type) - if isinstance(reader, VideoReader): + frame = self._convert_frame(frame, reader_class, out_type) + if issubclass(reader_class, VideoReader): return return_type(frame, mime=self.VIDEO_FRAME_MIME) return return_type(frame, mime=mimetypes.guess_type(frame_name)[0]) From dc03220c02cdee7997549c3be8e470b519cfbd16 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 7 Aug 2024 16:45:17 +0300 Subject: [PATCH 020/115] Fix frame reading in updated dynamic cache building --- cvat/apps/engine/cache.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 66c08385340..b374213cb8e 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -201,6 +201,7 @@ def _read_raw_frames( tmp_dir = es.enter_context(tempfile.TemporaryDirectory(prefix="cvat")) files_to_download = [] checksums = [] + media = [] for item in reader.iterate_frames(frame_ids): file_name = f"{item['name']}{item['extension']}" fs_filename = os.path.join(tmp_dir, file_name) @@ -222,6 +223,7 @@ def _read_raw_frames( yield from media else: + media = [] for item in reader.iterate_frames(frame_ids): source_path = os.path.join(raw_data_dir, f"{item['name']}{item['extension']}") media.append((source_path, source_path, None)) From 4bb8a74e24b0f9dc759ecdfdc987818440db0009 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 9 Aug 2024 17:29:28 +0300 Subject: [PATCH 021/115] Fix invalid frame quality --- cvat/apps/engine/cache.py | 4 +++- cvat/apps/engine/frame_provider.py | 4 +++- cvat/apps/engine/task.py | 2 +- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index b374213cb8e..204fdc251e3 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -269,6 +269,8 @@ def prepare_range_segment_chunk( ), } + image_quality = 100 if quality == FrameQuality.ORIGINAL else db_data.image_quality + mime_type = ( "video/mp4" if writer_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] @@ -278,7 +280,7 @@ def prepare_range_segment_chunk( kwargs = {} if db_segment.task.dimension == models.DimensionType.DIM_3D: kwargs["dimension"] = models.DimensionType.DIM_3D - writer = writer_classes[quality](int(quality), **kwargs) + writer = writer_classes[quality](image_quality, **kwargs) buffer = io.BytesIO() with closing(self._read_raw_frames(db_task, frame_ids=chunk_frame_ids)) as frame_iter: diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 078f626cd56..93a4e747c51 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -292,6 +292,8 @@ def get_chunk( writer_class = writer_classes[quality] + image_quality = 100 if quality == FrameQuality.ORIGINAL else db_data.image_quality + mime_type = ( "video/mp4" if writer_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] @@ -318,7 +320,7 @@ def get_chunk( writer_kwargs = {} if self._db_task.dimension == models.DimensionType.DIM_3D: writer_kwargs["dimension"] = models.DimensionType.DIM_3D - merged_chunk_writer = writer_class(int(quality), **writer_kwargs) + merged_chunk_writer = writer_class(image_quality, **writer_kwargs) writer_kwargs = {} if isinstance(merged_chunk_writer, ZipCompressedChunkWriter): diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index 0416704b913..a0422ee9c29 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -1202,7 +1202,7 @@ def save_chunks( # Let's use QP=17 (that is 67 for 0-100 range) for the original chunks, # which should be visually lossless or nearly so. # A lower value will significantly increase the chunk size with a slight increase of quality. - original_quality = 67 + original_quality = 67 # TODO: fix discrepancy in values in different parts of code else: original_chunk_writer_class = ZipChunkWriter original_quality = 100 From f7d2c4c2c12e302c0af8e7618dfe64b15ad951be Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 9 Aug 2024 17:32:21 +0300 Subject: [PATCH 022/115] Fix video reading in media_extractors - exception handling, frame mismatches --- cvat/apps/engine/media_extractors.py | 187 ++++++++++++++++----------- 1 file changed, 115 insertions(+), 72 deletions(-) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index b1049cbdd37..3cb613fadbe 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -18,7 +18,7 @@ from contextlib import ExitStack, closing, contextmanager from dataclasses import dataclass from enum import IntEnum -from typing import Any, Callable, Iterable, Iterator, Optional, Protocol, Sequence, Tuple, TypeVar +from typing import Any, Callable, Iterable, Iterator, Optional, Protocol, Sequence, Tuple, TypeVar, Union import av import av.codec @@ -507,8 +507,14 @@ def extract(self): class VideoReader(IMediaReader): def __init__( - self, source_path, step=1, start=0, stop=None, - dimension=DimensionType.DIM_2D, *, allow_threading: bool = True + self, + source_path: Union[str, io.BytesIO], + step: int = 1, + start: int = 0, + stop: Optional[int] = None, + dimension: DimensionType = DimensionType.DIM_2D, + *, + allow_threading: bool = True, ): super().__init__( source_path=source_path, @@ -522,65 +528,88 @@ def __init__( self._frame_count: Optional[int] = None self._frame_size: Optional[tuple[int, int]] = None # (w, h) - def _has_frame(self, i): - if i >= self._start: - if (i - self._start) % self._step == 0: - if self._stop is None or i < self._stop: - return True - - return False - - def _make_frame_iterator( + def iterate_frames( self, *, - apply_filter: bool = True, - stream: Optional[av.video.stream.VideoStream] = None, + frame_filter: Union[bool, Iterable[int]] = True, + video_stream: Optional[av.video.stream.VideoStream] = None, ) -> Iterator[Tuple[av.VideoFrame, str, int]]: + """ + If provided, frame_filter must be an ordered sequence in the ascending order. + 'True' means using the frames configured in the reader object. + 'False' or 'None' means returning all the video frames. + """ + + if frame_filter is True: + frame_filter = itertools.count(self._start, self._step) + if self._stop: + frame_filter = itertools.takewhile(lambda x: x <= self._stop, frame_filter) + elif not frame_filter: + frame_filter = itertools.count() + + frame_filter_iter = iter(frame_filter) + next_frame_filter_frame = next(frame_filter_iter, None) + if next_frame_filter_frame is None: + return + es = ExitStack() - needs_init = stream is None + needs_init = video_stream is None if needs_init: container = es.enter_context(self._read_av_container()) else: - container = stream.container + container = video_stream.container with es: if needs_init: - stream = container.streams.video[0] + video_stream = container.streams.video[0] if self.allow_threading: - stream.thread_type = 'AUTO' + video_stream.thread_type = 'AUTO' - frame_num = 0 + exception = None + frame_number = 0 + for packet in container.demux(video_stream): + try: + for frame in packet.decode(): + if frame_number == next_frame_filter_frame: + if video_stream.metadata.get('rotate'): + pts = frame.pts + frame = av.VideoFrame().from_ndarray( + rotate_image( + frame.to_ndarray(format='bgr24'), + 360 - int(video_stream.metadata.get('rotate')) + ), + format ='bgr24' + ) + frame.pts = pts - for packet in container.demux(stream): - for image in packet.decode(): - frame_num += 1 + if self._frame_size is None: + self._frame_size = (frame.width, frame.height) - if apply_filter and not self._has_frame(frame_num - 1): - continue + yield (frame, self._source_path[0], frame.pts) - if stream.metadata.get('rotate'): - pts = image.pts - image = av.VideoFrame().from_ndarray( - rotate_image( - image.to_ndarray(format='bgr24'), - 360 - int(stream.metadata.get('rotate')) - ), - format ='bgr24' - ) - image.pts = pts + next_frame_filter_frame = next(frame_filter_iter, None) - if self._frame_size is None: - self._frame_size = (image.width, image.height) + if next_frame_filter_frame is None: + return - yield (image, self._source_path[0], image.pts) + frame_number += 1 + except Exception as e: + if av.__version__ == "9.2.0": + # av v9.2.0 seems to have a memory corruption + # in exception handling for demux() in the multithreaded mode. + # Instead of breaking the iteration, we iterate over packets till the end. + # Fixed in v12.2.0. + exception = e + if video_stream.thread_type != 'AUTO': + break - if self._frame_count is None: - self._frame_count = frame_num + if exception: + raise exception def __iter__(self): - return self._make_frame_iterator() + return self.iterate_frames() def get_progress(self, pos): duration = self._get_duration() @@ -602,7 +631,8 @@ def _read_av_container(self): if context and context.is_open: context.close() - container.close() + if container.open_files: + container.close() def _get_duration(self): with self._read_av_container() as container: @@ -629,7 +659,7 @@ def get_preview(self, frame): needed_time = int((frame / stream.guessed_rate) * tb_denominator) container.seek(offset=needed_time, stream=stream) - with closing(self._make_frame_iterator(stream=stream)) as frame_iter: + with closing(self.iterate_frames(video_stream=stream)) as frame_iter: return self._get_preview(next(frame_iter)) def get_image_size(self, i): @@ -637,8 +667,8 @@ def get_image_size(self, i): return self._frame_size with closing(iter(self)) as frame_iter: - image = next(frame_iter)[0] - self._frame_size = (image.width, image.height) + frame = next(frame_iter)[0] + self._frame_size = (frame.width, frame.height) return self._frame_size @@ -659,7 +689,7 @@ def get_frame_count(self) -> int: return self._frame_count frame_count = 0 - for _ in self._make_frame_iterator(apply_filter=False): + for _ in self.iterate_frames(frame_filter=False): frame_count += 1 self._frame_count = frame_count @@ -677,10 +707,13 @@ def iterate_frames(self, frame_ids: Iterable[int]): yield self._manifest[idx] class VideoReaderWithManifest: + # TODO: merge this class with VideoReader + def __init__(self, manifest_path: str, source_path: str, *, allow_threading: bool = False): self._source_path = source_path self._manifest = VideoManifestManager(manifest_path) - self._manifest.init_index() + if self._manifest.exists: + self._manifest.init_index() self.allow_threading = allow_threading @@ -711,49 +744,59 @@ def _get_nearest_left_key_frame(self, frame_id: int) -> tuple[int, int]: timestamp = 0 return frame_number, timestamp - def iterate_frames(self, frame_ids: Iterable[int]) -> Iterable[av.VideoFrame]: + def iterate_frames(self, *, frame_filter: Iterable[int]) -> Iterable[av.VideoFrame]: "frame_ids must be an ordered sequence in the ascending order" - frame_ids_iter = iter(frame_ids) - frame_ids_frame = next(frame_ids_iter, None) - if frame_ids_frame is None: + frame_filter_iter = iter(frame_filter) + next_frame_filter_frame = next(frame_filter_iter, None) + if next_frame_filter_frame is None: return start_decode_frame_number, start_decode_timestamp = self._get_nearest_left_key_frame( - frame_ids_frame + next_frame_filter_frame ) with self._read_av_container() as container: - stream = container.streams.video[0] + video_stream = container.streams.video[0] if self.allow_threading: - stream.thread_type = 'AUTO' + video_stream.thread_type = 'AUTO' - container.seek(offset=start_decode_timestamp, stream=stream) + container.seek(offset=start_decode_timestamp, stream=video_stream) frame_number = start_decode_frame_number - 1 - for packet in container.demux(stream): - for frame in packet.decode(): - frame_number += 1 + for packet in container.demux(video_stream): + try: + for frame in packet.decode(): + if frame_number == next_frame_filter_frame: + if video_stream.metadata.get('rotate'): + frame = av.VideoFrame().from_ndarray( + rotate_image( + frame.to_ndarray(format='bgr24'), + 360 - int(video_stream.metadata.get('rotate')) + ), + format ='bgr24' + ) - if frame_number < frame_ids_frame: - continue - elif frame_number == frame_ids_frame: - if stream.metadata.get('rotate'): - frame = av.VideoFrame().from_ndarray( - rotate_image( - frame.to_ndarray(format='bgr24'), - 360 - int(container.streams.video[0].metadata.get('rotate')) - ), - format ='bgr24' - ) + yield frame - yield frame + next_frame_filter_frame = next(frame_filter_iter, None) - frame_ids_frame = next(frame_ids_iter, None) + if next_frame_filter_frame is None: + return - if frame_ids_frame is None: - return + frame_number += 1 + except Exception as e: + if av.__version__ == "9.2.0": + # av v9.2.0 seems to have a memory corruption + # in exception handling for demux() in the multithreaded mode. + # Instead of breaking the iteration, we iterate over packets till the end. + # Fixed in v12.2.0. + exception = e + if video_stream.thread_type != 'AUTO': + break + if exception: + raise exception class IChunkWriter(ABC): def __init__(self, quality, dimension=DimensionType.DIM_2D): From 34d9ca0be23a1e4b9ce4b2b467ef0017e937db74 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 9 Aug 2024 17:32:47 +0300 Subject: [PATCH 023/115] Allow disabling static chunks, add seamless switching --- cvat/apps/engine/apps.py | 8 +++ cvat/apps/engine/cache.py | 35 ++++++++--- cvat/apps/engine/default_settings.py | 16 ++++++ cvat/apps/engine/frame_provider.py | 7 ++- cvat/apps/engine/task.py | 86 ++++++++++++++-------------- 5 files changed, 102 insertions(+), 50 deletions(-) create mode 100644 cvat/apps/engine/default_settings.py diff --git a/cvat/apps/engine/apps.py b/cvat/apps/engine/apps.py index 326920e8b49..bcad84510f5 100644 --- a/cvat/apps/engine/apps.py +++ b/cvat/apps/engine/apps.py @@ -10,6 +10,14 @@ class EngineConfig(AppConfig): name = 'cvat.apps.engine' def ready(self): + from django.conf import settings + + from . import default_settings + + for key in dir(default_settings): + if key.isupper() and not hasattr(settings, key): + setattr(settings, key, getattr(default_settings, key)) + # Required to define signals in application import cvat.apps.engine.signals # Required in order to silent "unused-import" in pyflake diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 204fdc251e3..6295cbb0e36 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -7,6 +7,7 @@ import io import os +import os.path import pickle # nosec import tempfile import zipfile @@ -37,6 +38,7 @@ ImageReaderWithManifest, Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter, + VideoReader, VideoReaderWithManifest, ZipChunkWriter, ZipCompressedChunkWriter, @@ -166,19 +168,38 @@ def _read_raw_frames( }[db_data.storage] dimension = db_task.dimension + manifest_path = db_data.get_manifest_path() if hasattr(db_data, "video"): source_path = os.path.join(raw_data_dir, db_data.video.path) reader = VideoReaderWithManifest( - manifest_path=db_data.get_manifest_path(), + manifest_path=manifest_path, source_path=source_path, + allow_threading=False, ) - for frame in reader.iterate_frames(frame_ids): - yield (frame, source_path, None) + if not os.path.isfile(manifest_path): + try: + reader._manifest.link(source_path, force=True) + reader._manifest.create() + except Exception as e: + # TODO: improve logging + reader = None + + if reader: + for frame in reader.iterate_frames(frame_filter=frame_ids): + yield (frame, source_path, None) + else: + reader = VideoReader([source_path], allow_threading=False) + + for frame_tuple in reader.iterate_frames(frame_filter=frame_ids): + yield frame_tuple else: - reader = ImageReaderWithManifest(db_data.get_manifest_path()) - if db_data.storage == models.StorageChoice.CLOUD_STORAGE: + if ( + os.path.isfile(manifest_path) + and db_data.storage == models.StorageChoice.CLOUD_STORAGE + ): + reader = ImageReaderWithManifest() with ExitStack() as es: db_cloud_storage = db_data.cloud_storage assert db_cloud_storage, "Cloud storage instance was deleted" @@ -224,8 +245,8 @@ def _read_raw_frames( yield from media else: media = [] - for item in reader.iterate_frames(frame_ids): - source_path = os.path.join(raw_data_dir, f"{item['name']}{item['extension']}") + for image in sorted(db_data.images.all(), key=lambda image: image.frame): + source_path = os.path.join(raw_data_dir, image.path) media.append((source_path, source_path, None)) if dimension == models.DimensionType.DIM_2D: diff --git a/cvat/apps/engine/default_settings.py b/cvat/apps/engine/default_settings.py new file mode 100644 index 00000000000..5e74759f6b4 --- /dev/null +++ b/cvat/apps/engine/default_settings.py @@ -0,0 +1,16 @@ +# Copyright (C) 2024 CVAT.ai Corporation +# +# SPDX-License-Identifier: MIT + +import os + +from attrs.converters import to_bool + +MEDIA_CACHE_ALLOW_STATIC_CHUNKS = to_bool(os.getenv("CVAT_ALLOW_STATIC_CHUNKS", False)) +""" +Allow or disallow static media chunks. +If disabled, CVAT will only use the dynamic media cache. New tasks requesting static media cache +will be automatically switched to the dynamic cache. +When enabled, this option can increase data access speed and reduce server load, +but significantly increase disk space occupied by tasks. +""" diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 93a4e747c51..b51a61a512f 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -16,6 +16,7 @@ import av import cv2 import numpy as np +from django.conf import settings from PIL import Image from rest_framework.exceptions import ValidationError @@ -399,7 +400,11 @@ def __init__(self, db_segment: models.Segment) -> None: } self._loaders: dict[FrameQuality, _ChunkLoader] = {} - if db_data.storage_method == models.StorageMethodChoice.CACHE: + if ( + db_data.storage_method == models.StorageMethodChoice.CACHE + or not settings.MEDIA_CACHE_ALLOW_STATIC_CHUNKS + # TODO: separate handling, extract cache creation logic from media cache + ): cache = MediaCache() self._loaders[FrameQuality.COMPRESSED] = _BufferChunkLoader( diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index a0422ee9c29..4ff62dc1f29 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -953,28 +953,24 @@ def _update_status(msg: str) -> None: # TODO: try to pull up # replace manifest file (e.g was uploaded 'subdir/manifest.jsonl' or 'some_manifest.jsonl') - if ( - settings.USE_CACHE and db_data.storage_method == models.StorageMethodChoice.CACHE and - manifest_file and not os.path.exists(db_data.get_manifest_path()) - ): + if (manifest_file and not os.path.exists(db_data.get_manifest_path())): shutil.copyfile(os.path.join(manifest_root, manifest_file), db_data.get_manifest_path()) if manifest_root and manifest_root.startswith(db_data.get_upload_dirname()): os.remove(os.path.join(manifest_root, manifest_file)) manifest_file = os.path.relpath(db_data.get_manifest_path(), upload_dir) + # Create task frames from the metadata collected video_path: str = "" video_frame_size: tuple[int, int] = (0, 0) images: list[models.Image] = [] - # Collect media metadata for media_type, media_files in media.items(): if not media_files: continue if task_mode == MEDIA_TYPES['video']['mode']: - manifest_is_prepared = False if manifest_file: try: _update_status('Validating the input manifest file') @@ -989,22 +985,21 @@ def _update_status(msg: str) -> None: if not len(manifest): raise ValidationError("No key frames found in the manifest") - video_frame_count = manifest.video_length - video_frame_size = manifest.video_resolution - manifest_is_prepared = True except Exception as ex: manifest.remove() manifest = None - slogger.glob.warning(ex, exc_info=True) if isinstance(ex, (ValidationError, AssertionError)): - _update_status(f'Invalid manifest file was upload: {ex}') + base_msg = f"Invalid manifest file was uploaded: {ex}" + else: + base_msg = "Failed to parse the uploaded manifest file" + slogger.glob.warning(ex, exc_info=True) - if ( - settings.USE_CACHE and db_data.storage_method == models.StorageMethodChoice.CACHE - and not manifest_is_prepared - ): - # TODO: check if we can always use video manifest for optimization + _update_status(base_msg) + else: + manifest = None + + if not manifest: try: _update_status('Preparing a manifest file') @@ -1013,26 +1008,32 @@ def _update_status(msg: str) -> None: manifest.link( media_file=media_files[0], upload_dir=upload_dir, - chunk_size=db_data.chunk_size # TODO: why it's needed here? + chunk_size=db_data.chunk_size, # TODO: why it's needed here? + force=True ) manifest.create() _update_status('A manifest has been created') - video_frame_count = len(manifest.reader) # TODO: check if the field access above and here are equivalent - video_frame_size = manifest.reader.resolution - manifest_is_prepared = True except Exception as ex: manifest.remove() manifest = None - db_data.storage_method = models.StorageMethodChoice.FILE_SYSTEM + if isinstance(ex, AssertionError): + base_msg = f": {ex}" + else: + base_msg = "" + slogger.glob.warning(ex, exc_info=True) - base_msg = str(ex) if isinstance(ex, AssertionError) \ - else "Uploaded video does not support a quick way of task creating." - _update_status("{} The task will be created using the old method".format(base_msg)) + _update_status( + f"Failed to create manifest for the uploaded video{base_msg}. " + "A manifest will not be used in this task" + ) - if not manifest: + if manifest: + video_frame_count = manifest.video_length + video_frame_size = manifest.video_resolution + else: video_frame_count = extractor.get_frame_count() video_frame_size = extractor.get_image_size(0) @@ -1048,22 +1049,20 @@ def _update_status(msg: str) -> None: else: # images, archive, pdf db_data.size = len(extractor) - manifest = None - if settings.USE_CACHE and db_data.storage_method == models.StorageMethodChoice.CACHE: - manifest = ImageManifestManager(db_data.get_manifest_path()) - if not manifest.exists: - manifest.link( - sources=extractor.absolute_source_paths, - meta={ - k: {'related_images': related_images[k] } - for k in related_images - }, - data_dir=upload_dir, - DIM_3D=(db_task.dimension == models.DimensionType.DIM_3D), - ) - manifest.create() - else: - manifest.init_index() + manifest = ImageManifestManager(db_data.get_manifest_path()) + if not manifest.exists: + manifest.link( + sources=extractor.absolute_source_paths, + meta={ + k: {'related_images': related_images[k] } + for k in related_images + }, + data_dir=upload_dir, + DIM_3D=(db_task.dimension == models.DimensionType.DIM_3D), + ) + manifest.create() + else: + manifest.init_index() for frame_id in extractor.frame_range: image_path = extractor.get_path(frame_id) @@ -1128,7 +1127,10 @@ def _update_status(msg: str) -> None: slogger.glob.info("Found frames {} for Data #{}".format(db_data.size, db_data.id)) _create_segments_and_jobs(db_task, job_file_mapping=job_file_mapping) - if db_data.storage_method == models.StorageMethodChoice.FILE_SYSTEM or not settings.USE_CACHE: + if ( + settings.MEDIA_CACHE_ALLOW_STATIC_CHUNKS and + db_data.storage_method == models.StorageMethodChoice.FILE_SYSTEM + ): _create_static_chunks(db_task, media_extractor=extractor) def _create_static_chunks(db_task: models.Task, *, media_extractor: IMediaReader): From 8c97967bb0b32f72ad97cb594f3392ef7d049c3b Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 9 Aug 2024 17:33:55 +0300 Subject: [PATCH 024/115] Extend code formatting --- dev/format_python_code.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/dev/format_python_code.sh b/dev/format_python_code.sh index 10469178f90..8c224c3caaf 100755 --- a/dev/format_python_code.sh +++ b/dev/format_python_code.sh @@ -25,6 +25,7 @@ for paths in \ "cvat/apps/analytics_report" \ "cvat/apps/engine/frame_provider.py" \ "cvat/apps/engine/cache.py" \ + "cvat/apps/engine/default_settings.py" \ ; do ${BLACK} -- ${paths} ${ISORT} -- ${paths} From a0fd0ba00c89cb1fe61253330ed7add0e7ae8bfa Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 9 Aug 2024 17:34:45 +0300 Subject: [PATCH 025/115] Rename function argument --- cvat/apps/engine/frame_provider.py | 4 ++-- cvat/apps/engine/models.py | 8 ++++---- cvat/apps/engine/task.py | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index b51a61a512f..fe55a5df5dc 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -427,7 +427,7 @@ def __init__(self, db_segment: models.Segment) -> None: reader_class=reader_class[db_data.compressed_chunk_type][0], reader_params=reader_class[db_data.compressed_chunk_type][1], get_chunk_path_callback=lambda chunk_idx: db_data.get_compressed_segment_chunk_path( - chunk_idx, segment=db_segment.id + chunk_idx, segment_id=db_segment.id ), ) @@ -435,7 +435,7 @@ def __init__(self, db_segment: models.Segment) -> None: reader_class=reader_class[db_data.original_chunk_type][0], reader_params=reader_class[db_data.original_chunk_type][1], get_chunk_path_callback=lambda chunk_idx: db_data.get_original_segment_chunk_path( - chunk_idx, segment=db_segment.id + chunk_idx, segment_id=db_segment.id ), ) diff --git a/cvat/apps/engine/models.py b/cvat/apps/engine/models.py index ec0870a0e22..1cc26af3876 100644 --- a/cvat/apps/engine/models.py +++ b/cvat/apps/engine/models.py @@ -272,13 +272,13 @@ def _get_compressed_chunk_name(self, segment_id: int, chunk_number: int) -> str: def _get_original_chunk_name(self, segment_id: int, chunk_number: int) -> str: return self._get_chunk_name(segment_id, chunk_number, self.original_chunk_type) - def get_original_segment_chunk_path(self, chunk_number: int, segment: int) -> str: + def get_original_segment_chunk_path(self, chunk_number: int, segment_id: int) -> str: return os.path.join(self.get_original_cache_dirname(), - self._get_original_chunk_name(segment, chunk_number)) + self._get_original_chunk_name(segment_id, chunk_number)) - def get_compressed_segment_chunk_path(self, chunk_number: int, segment: int) -> str: + def get_compressed_segment_chunk_path(self, chunk_number: int, segment_id: int) -> str: return os.path.join(self.get_compressed_cache_dirname(), - self._get_compressed_chunk_name(segment, chunk_number)) + self._get_compressed_chunk_name(segment_id, chunk_number)) def get_manifest_path(self): return os.path.join(self.get_upload_dirname(), 'manifest.jsonl') diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index 4ff62dc1f29..7ea9b7d2aa9 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -1179,13 +1179,13 @@ def save_chunks( original_chunk_writer.save_as_chunk, images=chunk_data, chunk_path=db_data.get_original_segment_chunk_path( - chunk_idx, segment=db_segment.id + chunk_idx, segment_id=db_segment.id ), ) compressed_chunk_writer.save_as_chunk( images=chunk_data, chunk_path=db_data.get_compressed_segment_chunk_path( - chunk_idx, segment=db_segment.id + chunk_idx, segment_id=db_segment.id ), ) From c0480c9065d6ff8f55e67929ff2d48308ae19184 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 9 Aug 2024 17:44:31 +0300 Subject: [PATCH 026/115] Rename configuration parameter --- cvat/apps/engine/default_settings.py | 4 ++-- cvat/apps/engine/frame_provider.py | 2 +- cvat/apps/engine/task.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/cvat/apps/engine/default_settings.py b/cvat/apps/engine/default_settings.py index 5e74759f6b4..826fe1c9bef 100644 --- a/cvat/apps/engine/default_settings.py +++ b/cvat/apps/engine/default_settings.py @@ -6,9 +6,9 @@ from attrs.converters import to_bool -MEDIA_CACHE_ALLOW_STATIC_CHUNKS = to_bool(os.getenv("CVAT_ALLOW_STATIC_CHUNKS", False)) +MEDIA_CACHE_ALLOW_STATIC_CACHE = to_bool(os.getenv("CVAT_ALLOW_STATIC_CACHE", False)) """ -Allow or disallow static media chunks. +Allow or disallow static media cache. If disabled, CVAT will only use the dynamic media cache. New tasks requesting static media cache will be automatically switched to the dynamic cache. When enabled, this option can increase data access speed and reduce server load, diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index fe55a5df5dc..ac764f00c2a 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -402,7 +402,7 @@ def __init__(self, db_segment: models.Segment) -> None: self._loaders: dict[FrameQuality, _ChunkLoader] = {} if ( db_data.storage_method == models.StorageMethodChoice.CACHE - or not settings.MEDIA_CACHE_ALLOW_STATIC_CHUNKS + or not settings.MEDIA_CACHE_ALLOW_STATIC_CACHE # TODO: separate handling, extract cache creation logic from media cache ): cache = MediaCache() diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index 7ea9b7d2aa9..e021761b5b0 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -1128,7 +1128,7 @@ def _update_status(msg: str) -> None: _create_segments_and_jobs(db_task, job_file_mapping=job_file_mapping) if ( - settings.MEDIA_CACHE_ALLOW_STATIC_CHUNKS and + settings.MEDIA_CACHE_ALLOW_STATIC_CACHE and db_data.storage_method == models.StorageMethodChoice.FILE_SYSTEM ): _create_static_chunks(db_task, media_extractor=extractor) From 5caf283824580c279fe805dc136487081243ca3a Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 10:56:08 +0300 Subject: [PATCH 027/115] Add av version comment --- cvat/requirements/base.in | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cvat/requirements/base.in b/cvat/requirements/base.in index 6cbfb10321a..44c5c9b8c87 100644 --- a/cvat/requirements/base.in +++ b/cvat/requirements/base.in @@ -1,7 +1,13 @@ -r ../../utils/dataset_manifest/requirements.in attrs==21.4.0 + +# This is the last version of av that supports ffmpeg we depend on. +# Changing ffmpeg is undesirable, as there might be video decoding differences +# between versions. +# TODO: try to move to the newer version av==9.2.0 + azure-storage-blob==12.13.0 boto3==1.17.61 clickhouse-connect==0.6.8 From efbe3a00b69ffbe0f99bf121e06f122750c359c4 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 12:29:34 +0300 Subject: [PATCH 028/115] Refactor av video reading --- cvat/apps/engine/media_extractors.py | 195 +++++++++++++-------------- 1 file changed, 93 insertions(+), 102 deletions(-) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 3cb613fadbe..3b843ea3c9c 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -18,7 +18,10 @@ from contextlib import ExitStack, closing, contextmanager from dataclasses import dataclass from enum import IntEnum -from typing import Any, Callable, Iterable, Iterator, Optional, Protocol, Sequence, Tuple, TypeVar, Union +from typing import ( + Any, Callable, ContextManager, Generator, Iterable, Iterator, Optional, Protocol, + Sequence, Tuple, TypeVar, Union +) import av import av.codec @@ -505,6 +508,43 @@ def extract(self): if not self.extract_dir: os.remove(self._zip_source.filename) +class _AvVideoReading: + @contextmanager + def read_av_container(self, source: Union[str, io.BytesIO]) -> av.container.InputContainer: + if isinstance(source, io.BytesIO): + source.seek(0) # required for re-reading + + container = av.open(source) + try: + yield container + finally: + # fixes a memory leak in input container closing + # https://github.com/PyAV-Org/PyAV/issues/1117 + for stream in container.streams: + context = stream.codec_context + if context and context.is_open: + context.close() + + if container.open_files: + container.close() + + def decode_stream( + self, container: av.container.Container, video_stream: av.video.stream.VideoStream + ) -> Generator[av.VideoFrame, None, None]: + demux_iter = container.demux(video_stream) + try: + for packet in demux_iter: + yield from packet.decode() + finally: + # av v9.2.0 seems to have a memory corruption or a deadlock + # in exception handling for demux() in the multithreaded mode. + # Instead of breaking the iteration, we iterate over packets till the end. + # Fixed in av v12.2.0. + if av.__version__ == "9.2.0" and video_stream.thread_type == 'AUTO': + exhausted = object() + while next(demux_iter, exhausted) is not exhausted: + pass + class VideoReader(IMediaReader): def __init__( self, @@ -567,72 +607,45 @@ def iterate_frames( if self.allow_threading: video_stream.thread_type = 'AUTO' - exception = None - frame_number = 0 - for packet in container.demux(video_stream): - try: - for frame in packet.decode(): - if frame_number == next_frame_filter_frame: - if video_stream.metadata.get('rotate'): - pts = frame.pts - frame = av.VideoFrame().from_ndarray( - rotate_image( - frame.to_ndarray(format='bgr24'), - 360 - int(video_stream.metadata.get('rotate')) - ), - format ='bgr24' - ) - frame.pts = pts + frame_counter = itertools.count() + with closing(self._decode_stream(container, video_stream)) as stream_decoder: + for frame, frame_number in zip(stream_decoder, frame_counter): + if frame_number == next_frame_filter_frame: + if video_stream.metadata.get('rotate'): + pts = frame.pts + frame = av.VideoFrame().from_ndarray( + rotate_image( + frame.to_ndarray(format='bgr24'), + 360 - int(video_stream.metadata.get('rotate')) + ), + format ='bgr24' + ) + frame.pts = pts - if self._frame_size is None: - self._frame_size = (frame.width, frame.height) + if self._frame_size is None: + self._frame_size = (frame.width, frame.height) - yield (frame, self._source_path[0], frame.pts) + yield (frame, self._source_path[0], frame.pts) - next_frame_filter_frame = next(frame_filter_iter, None) + next_frame_filter_frame = next(frame_filter_iter, None) - if next_frame_filter_frame is None: - return + if next_frame_filter_frame is None: + return - frame_number += 1 - except Exception as e: - if av.__version__ == "9.2.0": - # av v9.2.0 seems to have a memory corruption - # in exception handling for demux() in the multithreaded mode. - # Instead of breaking the iteration, we iterate over packets till the end. - # Fixed in v12.2.0. - exception = e - if video_stream.thread_type != 'AUTO': - break - - if exception: - raise exception - - def __iter__(self): + def __iter__(self) -> Iterator[Tuple[av.VideoFrame, str, int]]: return self.iterate_frames() def get_progress(self, pos): duration = self._get_duration() return pos / duration if duration else None - @contextmanager - def _read_av_container(self): - if isinstance(self._source_path[0], io.BytesIO): - self._source_path[0].seek(0) # required for re-reading - - container = av.open(self._source_path[0]) - try: - yield container - finally: - # fixes a memory leak in input container closing - # https://github.com/PyAV-Org/PyAV/issues/1117 - for stream in container.streams: - context = stream.codec_context - if context and context.is_open: - context.close() + def _read_av_container(self) -> ContextManager[av.container.InputContainer]: + return _AvVideoReading().read_av_container(self._source_path[0]) - if container.open_files: - container.close() + def _decode_stream( + self, container: av.container.Container, video_stream: av.video.stream.VideoStream + ) -> Generator[av.VideoFrame, None, None]: + return _AvVideoReading().decode_stream(container, video_stream) def _get_duration(self): with self._read_av_container() as container: @@ -717,20 +730,13 @@ def __init__(self, manifest_path: str, source_path: str, *, allow_threading: boo self.allow_threading = allow_threading - @contextmanager - def _read_av_container(self): - container = av.open(self._source_path) - try: - yield container - finally: - # fixes a memory leak in input container closing - # https://github.com/PyAV-Org/PyAV/issues/1117 - for stream in container.streams: - context = stream.codec_context - if context and context.is_open: - context.close() + def _read_av_container(self) -> ContextManager[av.container.InputContainer]: + return _AvVideoReading().read_av_container(self._source_path) - container.close() + def _decode_stream( + self, container: av.container.Container, video_stream: av.video.stream.VideoStream + ) -> Generator[av.VideoFrame, None, None]: + return _AvVideoReading().decode_stream(container, video_stream) def _get_nearest_left_key_frame(self, frame_id: int) -> tuple[int, int]: nearest_left_keyframe_pos = bisect( @@ -763,40 +769,25 @@ def iterate_frames(self, *, frame_filter: Iterable[int]) -> Iterable[av.VideoFra container.seek(offset=start_decode_timestamp, stream=video_stream) - frame_number = start_decode_frame_number - 1 - for packet in container.demux(video_stream): - try: - for frame in packet.decode(): - if frame_number == next_frame_filter_frame: - if video_stream.metadata.get('rotate'): - frame = av.VideoFrame().from_ndarray( - rotate_image( - frame.to_ndarray(format='bgr24'), - 360 - int(video_stream.metadata.get('rotate')) - ), - format ='bgr24' - ) - - yield frame - - next_frame_filter_frame = next(frame_filter_iter, None) - - if next_frame_filter_frame is None: - return - - frame_number += 1 - except Exception as e: - if av.__version__ == "9.2.0": - # av v9.2.0 seems to have a memory corruption - # in exception handling for demux() in the multithreaded mode. - # Instead of breaking the iteration, we iterate over packets till the end. - # Fixed in v12.2.0. - exception = e - if video_stream.thread_type != 'AUTO': - break - - if exception: - raise exception + frame_counter = itertools.count(start_decode_frame_number - 1) + with closing(self._decode_stream(container, video_stream)) as stream_decoder: + for frame, frame_number in zip(stream_decoder, frame_counter): + if frame_number == next_frame_filter_frame: + if video_stream.metadata.get('rotate'): + frame = av.VideoFrame().from_ndarray( + rotate_image( + frame.to_ndarray(format='bgr24'), + 360 - int(video_stream.metadata.get('rotate')) + ), + format ='bgr24' + ) + + yield frame + + next_frame_filter_frame = next(frame_filter_iter, None) + + if next_frame_filter_frame is None: + return class IChunkWriter(ABC): def __init__(self, quality, dimension=DimensionType.DIM_2D): From fb1284d3c6b8781fda4633a4495b8392fb2f0fd6 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 14:34:14 +0300 Subject: [PATCH 029/115] Fix manifest access --- cvat/apps/engine/cache.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 6295cbb0e36..be5c2d799e8 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -199,7 +199,7 @@ def _read_raw_frames( os.path.isfile(manifest_path) and db_data.storage == models.StorageChoice.CLOUD_STORAGE ): - reader = ImageReaderWithManifest() + reader = ImageReaderWithManifest(manifest_path) with ExitStack() as es: db_cloud_storage = db_data.cloud_storage assert db_cloud_storage, "Cloud storage instance was deleted" From 8edcfc53991286da203bc8ae00b2d30f53316988 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 15:27:35 +0300 Subject: [PATCH 030/115] Add migration --- .../migrations/0082_move_to_segment_chunks.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) create mode 100644 cvat/apps/engine/migrations/0082_move_to_segment_chunks.py diff --git a/cvat/apps/engine/migrations/0082_move_to_segment_chunks.py b/cvat/apps/engine/migrations/0082_move_to_segment_chunks.py new file mode 100644 index 00000000000..42ee41dd657 --- /dev/null +++ b/cvat/apps/engine/migrations/0082_move_to_segment_chunks.py @@ -0,0 +1,52 @@ +# Generated by Django 4.2.13 on 2024-08-12 09:49 + +import os +from django.db import migrations +from cvat.apps.engine.log import get_migration_logger + +def switch_tasks_with_static_chunks_to_dynamic_chunks(apps, schema_editor): + migration_name = os.path.splitext(os.path.basename(__file__))[0] + with get_migration_logger(migration_name) as common_logger: + Data = apps.get_model("engine", "Data") + + data_with_static_cache_query = ( + Data.objects + .filter(storage_method="file_system") + ) + + data_with_static_cache_ids = list( + v[0] + for v in ( + data_with_static_cache_query + .order_by('id') + .values_list('id') + .iterator(chunk_size=100000) + ) + ) + + data_with_static_cache_query.update(storage_method="cache") + + updated_data_ids_filename = migration_name + "-data_ids.log" + with open(updated_data_ids_filename, "w") as data_ids_file: + print( + "The following Data ids have been switched from using \"filesystem\" chunk storage " + "to \"cache\":", + file=data_ids_file + ) + for data_id in data_with_static_cache_ids: + print(data_id, file=data_ids_file) + + common_logger.info( + "Information about migrated tasks is available in the migration log file: " + f"{updated_data_ids_filename}. You will need to remove data manually for these tasks." + ) + +class Migration(migrations.Migration): + + dependencies = [ + ("engine", "0081_job_assignee_updated_date_and_more"), + ] + + operations = [ + migrations.RunPython(switch_tasks_with_static_chunks_to_dynamic_chunks) + ] From 51a7f83473b286e4ae5048e40b935b7d6ff87a50 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 15:28:20 +0300 Subject: [PATCH 031/115] Update downloading from cloud storage for packed data in task creation --- cvat/apps/engine/task.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index e021761b5b0..8701c979537 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -686,14 +686,10 @@ def _update_status(msg: str) -> None: is_media_sorted = False if is_data_in_cloud: - # first we need to filter files and keep only supported ones - if any([v for k, v in media.items() if k != 'image']) and db_data.storage_method == models.StorageMethodChoice.CACHE: - # FUTURE-FIXME: This is a temporary workaround for creating tasks - # with unsupported cloud storage data (video, archive, pdf) when use_cache is enabled - db_data.storage_method = models.StorageMethodChoice.FILE_SYSTEM - _update_status("The 'use cache' option is ignored") - - if db_data.storage_method == models.StorageMethodChoice.FILE_SYSTEM or not settings.USE_CACHE: + # Packed media must be downloaded for task creation + if any(v for k, v in media.items() if k != 'image'): + _update_status("The input media is packed - downloading it for further processing") + filtered_data = [] for files in (i for i in media.values() if i): filtered_data.extend(files) @@ -708,9 +704,11 @@ def _update_status(msg: str) -> None: step = db_data.get_frame_step() if start_frame or step != 1 or stop_frame != len(filtered_data) - 1: media_to_download = filtered_data[start_frame : stop_frame + 1: step] + _download_data_from_cloud_storage(db_data.cloud_storage, media_to_download, upload_dir) del media_to_download del filtered_data + is_data_in_cloud = False db_data.storage = models.StorageChoice.LOCAL else: From 65e417495bf10c6992df0ad680f5db19765b1b8d Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 16:22:39 +0300 Subject: [PATCH 032/115] Update changelog --- changelog.d/20240812_161617_mzhiltso_job_chunks.md | 4 ++++ changelog.d/20240812_161734_mzhiltso_job_chunks.md | 4 ++++ changelog.d/20240812_161912_mzhiltso_job_chunks.md | 6 ++++++ 3 files changed, 14 insertions(+) create mode 100644 changelog.d/20240812_161617_mzhiltso_job_chunks.md create mode 100644 changelog.d/20240812_161734_mzhiltso_job_chunks.md create mode 100644 changelog.d/20240812_161912_mzhiltso_job_chunks.md diff --git a/changelog.d/20240812_161617_mzhiltso_job_chunks.md b/changelog.d/20240812_161617_mzhiltso_job_chunks.md new file mode 100644 index 00000000000..f78376d9443 --- /dev/null +++ b/changelog.d/20240812_161617_mzhiltso_job_chunks.md @@ -0,0 +1,4 @@ +### Added + +- A server setting to disable media chunks on the local filesystem + () diff --git a/changelog.d/20240812_161734_mzhiltso_job_chunks.md b/changelog.d/20240812_161734_mzhiltso_job_chunks.md new file mode 100644 index 00000000000..2a587593b4f --- /dev/null +++ b/changelog.d/20240812_161734_mzhiltso_job_chunks.md @@ -0,0 +1,4 @@ +### Changed + +- Jobs now have separate chunk ids starting from 0, instead of using ones from the task + () diff --git a/changelog.d/20240812_161912_mzhiltso_job_chunks.md b/changelog.d/20240812_161912_mzhiltso_job_chunks.md new file mode 100644 index 00000000000..2a0198dd8f7 --- /dev/null +++ b/changelog.d/20240812_161912_mzhiltso_job_chunks.md @@ -0,0 +1,6 @@ +### Fixed + +- Various memory leaks in video reading on the server + () +- Job assignees will not receive frames from adjacent jobs in the boundary chunks + () From 34f972fc01031f1457090581d91da4182490b1f8 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 16:24:17 +0300 Subject: [PATCH 033/115] Update migration name --- ...move_to_segment_chunks.py => 0083_move_to_segment_chunks.py} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename cvat/apps/engine/migrations/{0082_move_to_segment_chunks.py => 0083_move_to_segment_chunks.py} (96%) diff --git a/cvat/apps/engine/migrations/0082_move_to_segment_chunks.py b/cvat/apps/engine/migrations/0083_move_to_segment_chunks.py similarity index 96% rename from cvat/apps/engine/migrations/0082_move_to_segment_chunks.py rename to cvat/apps/engine/migrations/0083_move_to_segment_chunks.py index 42ee41dd657..d21a7f669b2 100644 --- a/cvat/apps/engine/migrations/0082_move_to_segment_chunks.py +++ b/cvat/apps/engine/migrations/0083_move_to_segment_chunks.py @@ -44,7 +44,7 @@ def switch_tasks_with_static_chunks_to_dynamic_chunks(apps, schema_editor): class Migration(migrations.Migration): dependencies = [ - ("engine", "0081_job_assignee_updated_date_and_more"), + ("engine", "0082_alter_labeledimage_job_and_more"), ] operations = [ From 2bb2b17d78bdcfc5f3b17f6e1b78c45269ea8b14 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 17:32:28 +0300 Subject: [PATCH 034/115] Polish some code --- cvat/apps/engine/cache.py | 12 ++++++------ cvat/apps/engine/frame_provider.py | 4 ++-- cvat/apps/engine/media_extractors.py | 16 ++++++++-------- cvat/apps/engine/task.py | 2 +- cvat/apps/engine/views.py | 2 +- 5 files changed, 18 insertions(+), 18 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index be5c2d799e8..83cdecec13e 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -15,7 +15,7 @@ from contextlib import ExitStack, closing from datetime import datetime, timezone from itertools import pairwise -from typing import Any, Callable, Iterator, Optional, Sequence, Tuple, Type, Union +from typing import Callable, Iterator, Optional, Sequence, Tuple, Type, Union import av import cv2 @@ -58,8 +58,6 @@ class MediaCache: def __init__(self) -> None: self._cache = caches["media"] - # TODO migrate keys (check if they will be removed) - def get_checksum(self, value: bytes) -> int: return zlib.crc32(value) @@ -180,10 +178,12 @@ def _read_raw_frames( ) if not os.path.isfile(manifest_path): try: - reader._manifest.link(source_path, force=True) - reader._manifest.create() + reader.manifest.link(source_path, force=True) + reader.manifest.create() except Exception as e: - # TODO: improve logging + slogger.task[db_task.id].warning( + f"Failed to create video manifest: {e}", exc_info=True + ) reader = None if reader: diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index ac764f00c2a..1a0e713eef1 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -224,7 +224,7 @@ def validate_frame_number(self, frame_number: int) -> int: def validate_chunk_number(self, chunk_number: int) -> int: last_chunk = math.ceil(self._db_task.data.size / self._db_task.data.chunk_size) - 1 - if not (0 <= chunk_number <= last_chunk): + if not 0 <= chunk_number <= last_chunk: raise ValidationError( f"Invalid chunk number '{chunk_number}'. " f"The chunk number should be in the [0, {last_chunk}] range" @@ -463,7 +463,7 @@ def get_chunk_number(self, frame_number: int) -> int: def validate_chunk_number(self, chunk_number: int) -> int: segment_size = self._db_segment.frame_count last_chunk = math.ceil(segment_size / self._db_segment.task.data.chunk_size) - 1 - if not (0 <= chunk_number <= last_chunk): + if not 0 <= chunk_number <= last_chunk: raise ValidationError( f"Invalid chunk number '{chunk_number}'. " f"The chunk number should be in the [0, {last_chunk}] range" diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 3b843ea3c9c..296023c8f93 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -723,15 +723,15 @@ class VideoReaderWithManifest: # TODO: merge this class with VideoReader def __init__(self, manifest_path: str, source_path: str, *, allow_threading: bool = False): - self._source_path = source_path - self._manifest = VideoManifestManager(manifest_path) - if self._manifest.exists: - self._manifest.init_index() + self.source_path = source_path + self.manifest = VideoManifestManager(manifest_path) + if self.manifest.exists: + self.manifest.init_index() self.allow_threading = allow_threading def _read_av_container(self) -> ContextManager[av.container.InputContainer]: - return _AvVideoReading().read_av_container(self._source_path) + return _AvVideoReading().read_av_container(self.source_path) def _decode_stream( self, container: av.container.Container, video_stream: av.video.stream.VideoStream @@ -740,11 +740,11 @@ def _decode_stream( def _get_nearest_left_key_frame(self, frame_id: int) -> tuple[int, int]: nearest_left_keyframe_pos = bisect( - self._manifest, frame_id, key=lambda entry: entry.get('number') + self.manifest, frame_id, key=lambda entry: entry.get('number') ) if nearest_left_keyframe_pos: - frame_number = self._manifest[nearest_left_keyframe_pos - 1].get('number') - timestamp = self._manifest[nearest_left_keyframe_pos - 1].get('pts') + frame_number = self.manifest[nearest_left_keyframe_pos - 1].get('number') + timestamp = self.manifest[nearest_left_keyframe_pos - 1].get('pts') else: frame_number = 0 timestamp = 0 diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index 8701c979537..db093165b85 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -13,7 +13,7 @@ from contextlib import closing from datetime import datetime, timezone from pathlib import Path -from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple, Union, Iterable +from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple, Union from urllib import parse as urlparse from urllib import request as urlrequest diff --git a/cvat/apps/engine/views.py b/cvat/apps/engine/views.py index d7448472c8a..f0670b40bb2 100644 --- a/cvat/apps/engine/views.py +++ b/cvat/apps/engine/views.py @@ -60,7 +60,7 @@ from cvat.apps.dataset_manager.bindings import CvatImportError from cvat.apps.dataset_manager.serializers import DatasetFormatsSerializer from cvat.apps.engine.frame_provider import ( - IFrameProvider, TaskFrameProvider, JobFrameProvider, FrameQuality, FrameOutputType + IFrameProvider, TaskFrameProvider, JobFrameProvider, FrameQuality ) from cvat.apps.engine.filters import NonModelSimpleFilter, NonModelOrderingFilter, NonModelJsonLogicFilter from cvat.apps.engine.media_extractors import get_mime From 3788917573fdb820c6474960638370db836692b8 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 17:57:37 +0300 Subject: [PATCH 035/115] Fix frame retrieval by id --- cvat/apps/engine/cache.py | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 83cdecec13e..ff92614763b 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -244,10 +244,31 @@ def _read_raw_frames( yield from media else: + requested_frame_iter = iter(frame_ids) + next_requested_frame_id = next(requested_frame_iter, None) + if next_requested_frame_id is None: + return + + # TODO: find a way to use prefetched results, if provided + db_images = ( + db_data.images.order_by("frame") + .filter(frame__gte=frame_ids[0], frame__lte=frame_ids[-1]) + .values_list("frame", "path") + .all() + ) + media = [] - for image in sorted(db_data.images.all(), key=lambda image: image.frame): - source_path = os.path.join(raw_data_dir, image.path) - media.append((source_path, source_path, None)) + for frame_id, frame_path in db_images: + if frame_id == next_requested_frame_id: + source_path = os.path.join(raw_data_dir, frame_path) + media.append((source_path, source_path, None)) + + next_requested_frame_id = next(requested_frame_iter, None) + + if next_requested_frame_id is None: + break + + assert next_requested_frame_id is None if dimension == models.DimensionType.DIM_2D: media = preload_images(media) From f695ae1ef91d42ce480f7acef98e23dda4ed757a Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 17:58:43 +0300 Subject: [PATCH 036/115] Remove extra import --- cvat/apps/engine/views.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/apps/engine/views.py b/cvat/apps/engine/views.py index f0670b40bb2..f5e3827e919 100644 --- a/cvat/apps/engine/views.py +++ b/cvat/apps/engine/views.py @@ -66,7 +66,7 @@ from cvat.apps.engine.media_extractors import get_mime from cvat.apps.engine.permissions import AnnotationGuidePermission, get_iam_context from cvat.apps.engine.models import ( - ClientFile, Job, JobType, Label, SegmentType, Task, Project, Issue, Data, + ClientFile, Job, JobType, Label, Task, Project, Issue, Data, Comment, StorageMethodChoice, StorageChoice, CloudProviderChoice, Location, CloudStorage as CloudStorageModel, Asset, AnnotationGuide) From 14a9033baba979cbaad87bf9daad1a525ad8dc86 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 18:25:50 +0300 Subject: [PATCH 037/115] Fix frame access in gt jobs --- cvat/apps/engine/cache.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index ff92614763b..9a5ae932bef 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -365,7 +365,7 @@ def prepare_masked_range_segment_chunk( frame_bytes = None if frame_idx in frame_set: - frame_bytes = frame_provider.get_frame(frame_idx, quality=quality)[0] + frame_bytes = frame_provider.get_frame(frame_idx, quality=quality).data if frame_size is not None: # Decoded video frames can have different size, restore the original one From e8bebe9e9b5bddfde2cf8e4f9e17bc8fee5fd5e3 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 18:53:32 +0300 Subject: [PATCH 038/115] Fix frame access in export --- cvat/apps/dataset_manager/bindings.py | 27 +++++++++++++++++++-------- 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/cvat/apps/dataset_manager/bindings.py b/cvat/apps/dataset_manager/bindings.py index 0ab7d9d6a65..eb648ba2acc 100644 --- a/cvat/apps/dataset_manager/bindings.py +++ b/cvat/apps/dataset_manager/bindings.py @@ -1333,7 +1333,7 @@ def add_task(self, task, files): @attrs(frozen=True, auto_attribs=True) class ImageSource: - db_data: Data + db_task: Task is_video: bool = attrib(kw_only=True) class ImageProvider: @@ -1363,7 +1363,9 @@ def video_frame_loader(_): # some formats or transforms can require image data return self._frame_provider.get_frame(frame_index, quality=FrameQuality.ORIGINAL, - out_type=FrameOutputType.NUMPY_ARRAY).data + out_type=FrameOutputType.NUMPY_ARRAY + ).data + return dm.Image(data=video_frame_loader, **image_kwargs) else: def image_loader(_): @@ -1372,7 +1374,9 @@ def image_loader(_): # for images use encoded data to avoid recoding return self._frame_provider.get_frame(frame_index, quality=FrameQuality.ORIGINAL, - out_type=FrameOutputType.BUFFER).data.getvalue() + out_type=FrameOutputType.BUFFER + ).data.getvalue() + return dm.ByteImage(data=image_loader, **image_kwargs) def _load_source(self, source_id: int, source: ImageSource) -> None: @@ -1380,7 +1384,7 @@ def _load_source(self, source_id: int, source: ImageSource) -> None: return self._unload_source() - self._frame_provider = TaskFrameProvider(next(iter(source.db_data.tasks))) # TODO: refactor + self._frame_provider = TaskFrameProvider(source.db_task) self._current_source_id = source_id def _unload_source(self) -> None: @@ -1396,7 +1400,7 @@ def __init__(self, sources: Dict[int, ImageSource]) -> None: self._images_per_source = { source_id: { image.id: image - for image in source.db_data.images.prefetch_related('related_files') + for image in source.db_task.data.images.prefetch_related('related_files') } for source_id, source in sources.items() } @@ -1405,7 +1409,7 @@ def get_image_for_frame(self, source_id: int, frame_id: int, **image_kwargs): source = self._sources[source_id] point_cloud_path = osp.join( - source.db_data.get_upload_dirname(), image_kwargs['path'], + source.db_task.data.get_upload_dirname(), image_kwargs['path'], ) image = self._images_per_source[source_id][frame_id] @@ -1521,8 +1525,15 @@ def __init__( ext = TaskFrameProvider.VIDEO_FRAME_EXT if dimension == DimensionType.DIM_3D or include_images: + if isinstance(instance_data, TaskData): + db_task = instance_data.db_instance + elif isinstance(instance_data, JobData): + db_task = instance_data.db_instance.segment.task + else: + assert False + self._image_provider = IMAGE_PROVIDERS_BY_DIMENSION[dimension]( - {0: ImageSource(instance_data.db_data, is_video=is_video)} + {0: ImageSource(db_task, is_video=is_video)} ) for frame_data in instance_data.group_by_frame(include_empty=True): @@ -1604,7 +1615,7 @@ def __init__( if self._dimension == DimensionType.DIM_3D or include_images: self._image_provider = IMAGE_PROVIDERS_BY_DIMENSION[self._dimension]( { - task.id: ImageSource(task.data, is_video=task.mode == 'interpolation') + task.id: ImageSource(task, is_video=task.mode == 'interpolation') for task in project_data.tasks } ) From bbef52f1f270757505593407b9449956d2578532 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 12 Aug 2024 19:51:53 +0300 Subject: [PATCH 039/115] Fix frame iteration for frame step and excluded frames, fix export in cvat format --- cvat/apps/dataset_manager/bindings.py | 8 +++--- cvat/apps/dataset_manager/formats/cvat.py | 4 +-- cvat/apps/engine/frame_provider.py | 33 +++++++++++++++++++---- 3 files changed, 34 insertions(+), 11 deletions(-) diff --git a/cvat/apps/dataset_manager/bindings.py b/cvat/apps/dataset_manager/bindings.py index eb648ba2acc..ec717d205fe 100644 --- a/cvat/apps/dataset_manager/bindings.py +++ b/cvat/apps/dataset_manager/bindings.py @@ -240,7 +240,7 @@ def start(self) -> int: @property def stop(self) -> int: - return len(self) + return max(0, len(self) - 1) def _get_queryset(self): raise NotImplementedError() @@ -376,7 +376,7 @@ def _export_tag(self, tag): def _export_track(self, track, idx): track['shapes'] = list(filter(lambda x: not self._is_frame_deleted(x['frame']), track['shapes'])) tracked_shapes = TrackManager.get_interpolated_shapes( - track, 0, self.stop, self._annotation_ir.dimension) + track, 0, self.stop + 1, self._annotation_ir.dimension) for tracked_shape in tracked_shapes: tracked_shape["attributes"] += track["attributes"] tracked_shape["track_id"] = track["track_id"] if self._use_server_track_ids else idx @@ -432,7 +432,7 @@ def get_frame(idx): anno_manager = AnnotationManager(self._annotation_ir) for shape in sorted( - anno_manager.to_shapes(self.stop, self._annotation_ir.dimension, + anno_manager.to_shapes(self.stop + 1, self._annotation_ir.dimension, # Skip outside, deleted and excluded frames included_frames=included_frames, include_outside=False, @@ -763,7 +763,7 @@ def start(self) -> int: @property def stop(self) -> int: segment = self._db_job.segment - return segment.stop_frame + 1 + return segment.stop_frame @property def db_instance(self): diff --git a/cvat/apps/dataset_manager/formats/cvat.py b/cvat/apps/dataset_manager/formats/cvat.py index 4651fd39845..299982f2dcf 100644 --- a/cvat/apps/dataset_manager/formats/cvat.py +++ b/cvat/apps/dataset_manager/formats/cvat.py @@ -1378,8 +1378,8 @@ def dump_media_files(instance_data: CommonData, img_dir: str, project_data: Proj ext = frame_provider.VIDEO_FRAME_EXT frames = frame_provider.iterate_frames( - start_frame=instance_data.start, - stop_frame=instance_data.stop, + start_frame=instance_data.abs_frame_id(instance_data.start), + stop_frame=instance_data.abs_frame_id(instance_data.stop), quality=FrameQuality.ORIGINAL, out_type=FrameOutputType.BUFFER, ) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 1a0e713eef1..c25c961ce10 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -6,6 +6,7 @@ from __future__ import annotations import io +import itertools import math from abc import ABCMeta, abstractmethod from dataclasses import dataclass @@ -360,9 +361,25 @@ def iterate_frames( quality: FrameQuality = FrameQuality.ORIGINAL, out_type: FrameOutputType = FrameOutputType.BUFFER, ) -> Iterator[DataWithMeta[AnyFrame]]: - # TODO: optimize segment access - for idx in range(start_frame, (stop_frame + 1) if stop_frame else None): - yield self.get_frame(idx, quality=quality, out_type=out_type) + frame_range = itertools.count(start_frame, self._db_task.data.get_frame_step()) + if stop_frame: + frame_range = itertools.takewhile(lambda x: x <= stop_frame, frame_range) + + db_segment = None + db_segment_frame_set = None + db_segment_frame_provider = None + for idx in frame_range: + if db_segment and idx not in db_segment_frame_set: + db_segment = None + db_segment_frame_set = None + db_segment_frame_provider = None + + if not db_segment: + db_segment = self._get_segment(idx) + db_segment_frame_set = set(db_segment.frame_set) + db_segment_frame_provider = SegmentFrameProvider(db_segment) + + yield db_segment_frame_provider.get_frame(idx, quality=quality, out_type=out_type) def _get_segment(self, validated_frame_number: int) -> models.Segment: if not self._db_task.data or not self._db_task.data.size: @@ -537,8 +554,14 @@ def iterate_frames( quality: FrameQuality = FrameQuality.ORIGINAL, out_type: FrameOutputType = FrameOutputType.BUFFER, ) -> Iterator[DataWithMeta[AnyFrame]]: - for idx in range(start_frame, (stop_frame + 1) if stop_frame else None): - yield self.get_frame(idx, quality=quality, out_type=out_type) + frame_range = itertools.count(start_frame, self._db_segment.task.data.get_frame_step()) + if stop_frame: + frame_range = itertools.takewhile(lambda x: x <= stop_frame, frame_range) + + segment_frame_set = set(self._db_segment.frame_set) + for idx in frame_range: + if idx in segment_frame_set: + yield self.get_frame(idx, quality=quality, out_type=out_type) class JobFrameProvider(SegmentFrameProvider): From 3d5bb5203d76ec8f0b6292e19d1abcd3b5209b74 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 13 Aug 2024 11:57:36 +0300 Subject: [PATCH 040/115] Remove unused import --- cvat/apps/dataset_manager/bindings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/apps/dataset_manager/bindings.py b/cvat/apps/dataset_manager/bindings.py index ec717d205fe..f8dd470b64a 100644 --- a/cvat/apps/dataset_manager/bindings.py +++ b/cvat/apps/dataset_manager/bindings.py @@ -32,7 +32,7 @@ from cvat.apps.dataset_manager.formats.utils import get_label_color from cvat.apps.dataset_manager.util import add_prefetch_fields from cvat.apps.engine.frame_provider import TaskFrameProvider, FrameQuality, FrameOutputType -from cvat.apps.engine.models import (AttributeSpec, AttributeType, Data, DimensionType, Job, +from cvat.apps.engine.models import (AttributeSpec, AttributeType, DimensionType, Job, JobType, Label, LabelType, Project, SegmentType, ShapeType, Task) from cvat.apps.engine.rq_job_handler import RQJobMetaField From 0e9c5c8b391a7f89e6eca127aa69f381a046147d Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 13 Aug 2024 12:05:16 +0300 Subject: [PATCH 041/115] Fix error check in test --- tests/python/rest_api/test_jobs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/rest_api/test_jobs.py b/tests/python/rest_api/test_jobs.py index cd344506336..f873e862a6a 100644 --- a/tests/python/rest_api/test_jobs.py +++ b/tests/python/rest_api/test_jobs.py @@ -823,7 +823,7 @@ def test_can_get_gt_job_frame(self, admin_user, tasks, jobs, task_mode, quality, _check_status=False, ) assert response.status == HTTPStatus.BAD_REQUEST - assert b"The frame number doesn't belong to the job" in response.data + assert b"Incorrect requested frame number" in response.data (_, response) = api_client.jobs_api.retrieve_data( gt_job.id, number=included_frames[0], quality=quality, type="frame" From 351bdc87bbf4d8eb0f41abecb309f39723d5176b Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 13 Aug 2024 12:05:29 +0300 Subject: [PATCH 042/115] Fix cleanup in test --- tests/python/rest_api/test_jobs.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/python/rest_api/test_jobs.py b/tests/python/rest_api/test_jobs.py index f873e862a6a..3f49e6a35b6 100644 --- a/tests/python/rest_api/test_jobs.py +++ b/tests/python/rest_api/test_jobs.py @@ -629,12 +629,11 @@ def test_can_get_gt_job_meta(self, admin_user, tasks, jobs, task_mode, request): :job_frame_count ] gt_job = self._create_gt_job(admin_user, task_id, job_frame_ids) + request.addfinalizer(lambda: self._delete_gt_job(user, gt_job.id)) with make_api_client(user) as api_client: (gt_job_meta, _) = api_client.jobs_api.retrieve_data_meta(gt_job.id) - request.addfinalizer(lambda: self._delete_gt_job(user, gt_job.id)) - # These values are relative to the resulting task frames, unlike meta values assert 0 == gt_job.start_frame assert task_meta.size - 1 == gt_job.stop_frame @@ -684,12 +683,11 @@ def test_can_get_gt_job_meta_with_complex_frame_setup(self, admin_user, request) task_frame_ids = range(start_frame, stop_frame, frame_step) job_frame_ids = list(task_frame_ids[::3]) gt_job = self._create_gt_job(admin_user, task_id, job_frame_ids) + request.addfinalizer(lambda: self._delete_gt_job(admin_user, gt_job.id)) with make_api_client(admin_user) as api_client: (gt_job_meta, _) = api_client.jobs_api.retrieve_data_meta(gt_job.id) - request.addfinalizer(lambda: self._delete_gt_job(admin_user, gt_job.id)) - # These values are relative to the resulting task frames, unlike meta values assert 0 == gt_job.start_frame assert len(task_frame_ids) - 1 == gt_job.stop_frame @@ -731,6 +729,7 @@ def test_can_get_gt_job_chunk(self, admin_user, tasks, jobs, task_mode, quality, :job_frame_count ] gt_job = self._create_gt_job(admin_user, task_id, job_frame_ids) + request.addfinalizer(lambda: self._delete_gt_job(admin_user, gt_job.id)) with make_api_client(admin_user) as api_client: (chunk_file, response) = api_client.jobs_api.retrieve_data( @@ -738,8 +737,6 @@ def test_can_get_gt_job_chunk(self, admin_user, tasks, jobs, task_mode, quality, ) assert response.status == HTTPStatus.OK - request.addfinalizer(lambda: self._delete_gt_job(admin_user, gt_job.id)) - frame_range = range( task_meta.start_frame, min(task_meta.stop_frame + 1, task_meta.chunk_size), frame_step ) @@ -806,6 +803,7 @@ def test_can_get_gt_job_frame(self, admin_user, tasks, jobs, task_mode, quality, :job_frame_count ] gt_job = self._create_gt_job(admin_user, task_id, job_frame_ids) + request.addfinalizer(lambda: self._delete_gt_job(admin_user, gt_job.id)) frame_range = range( task_meta.start_frame, min(task_meta.stop_frame + 1, task_meta.chunk_size), frame_step @@ -830,8 +828,6 @@ def test_can_get_gt_job_frame(self, admin_user, tasks, jobs, task_mode, quality, ) assert response.status == HTTPStatus.OK - request.addfinalizer(lambda: self._delete_gt_job(admin_user, gt_job.id)) - @pytest.mark.usefixtures("restore_db_per_class") class TestListJobs: From a71852c5f8d777f8dca1942d3803d2ca2be5f27c Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 13 Aug 2024 16:35:32 +0300 Subject: [PATCH 043/115] Add handling for disabled static cache during task creation --- cvat/apps/engine/task.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index db093165b85..c233aa47fa0 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -569,6 +569,12 @@ def _update_status(msg: str) -> None: else: assert False, f"Unknown file storage {db_data.storage}" + if ( + db_data.storage_method == models.StorageMethodChoice.FILE_SYSTEM and + not settings.MEDIA_CACHE_ALLOW_STATIC_CACHE + ): + db_data.storage_method = models.StorageMethodChoice.CACHE + manifest_file = _validate_manifest( manifest_files, manifest_root, From d90ca0df014c95c750ee8903c56307ac165b4bf5 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 13 Aug 2024 16:35:48 +0300 Subject: [PATCH 044/115] Refactor some code --- cvat/apps/engine/task.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index c233aa47fa0..efe4e892292 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -540,18 +540,18 @@ def _create_thread( slogger.glob.info("create task #{}".format(db_task.id)) - job_file_mapping = _validate_job_file_mapping(db_task, data) - - db_data = db_task.data - upload_dir = db_data.get_upload_dirname() if db_data.storage != models.StorageChoice.SHARE else settings.SHARE_ROOT - is_data_in_cloud = db_data.storage == models.StorageChoice.CLOUD_STORAGE - job = rq.get_current_job() def _update_status(msg: str) -> None: job.meta['status'] = msg job.save_meta() + job_file_mapping = _validate_job_file_mapping(db_task, data) + + db_data = db_task.data + upload_dir = db_data.get_upload_dirname() if db_data.storage != models.StorageChoice.SHARE else settings.SHARE_ROOT + is_data_in_cloud = db_data.storage == models.StorageChoice.CLOUD_STORAGE + if data['remote_files'] and not isDatasetImport: data['remote_files'] = _download_data(data['remote_files'], upload_dir) From 03e749a07a46dbe031e3f5b15ce22d1a8deb4a69 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 13 Aug 2024 16:36:30 +0300 Subject: [PATCH 045/115] Fix downloading for cloud data in task creation --- cvat/apps/engine/task.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index efe4e892292..31eb2db3c65 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -692,9 +692,15 @@ def _update_status(msg: str) -> None: is_media_sorted = False if is_data_in_cloud: - # Packed media must be downloaded for task creation - if any(v for k, v in media.items() if k != 'image'): - _update_status("The input media is packed - downloading it for further processing") + if ( + # Download remote data if local storage is requested + # TODO: maybe move into cache building to fail faster on invalid task configurations + db_data.storage_method == models.StorageMethodChoice.FILE_SYSTEM or + + # Packed media must be downloaded for task creation + any(v for k, v in media.items() if k != 'image') + ): + _update_status("Downloading input media") filtered_data = [] for files in (i for i in media.values() if i): From c0822a0020212edb18e6a2019ff80decd45b4c25 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 13 Aug 2024 18:17:54 +0300 Subject: [PATCH 046/115] Fix preview reading for projects --- cvat/apps/engine/views.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/apps/engine/views.py b/cvat/apps/engine/views.py index f5e3827e919..9dd42dad9c6 100644 --- a/cvat/apps/engine/views.py +++ b/cvat/apps/engine/views.py @@ -647,7 +647,7 @@ def preview(self, request, pk): data_quality='compressed', ) - return data_getter(request) + return data_getter() @staticmethod def _get_rq_response(queue, job_id): From 56d413fc939e60874737dee1e5583aa2fb97a095 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 13 Aug 2024 19:24:20 +0300 Subject: [PATCH 047/115] Fix failing sdk tests --- cvat/apps/engine/frame_provider.py | 4 ++-- cvat/apps/engine/media_extractors.py | 29 +++++++++++++++--------- tests/python/sdk/test_auto_annotation.py | 1 + tests/python/sdk/test_datasets.py | 1 + tests/python/sdk/test_jobs.py | 1 + tests/python/sdk/test_projects.py | 1 + tests/python/sdk/test_tasks.py | 1 + 7 files changed, 25 insertions(+), 13 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index c25c961ce10..6a6ea1c82c7 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -314,10 +314,10 @@ def get_chunk( ): continue - frame, _, _ = segment_frame_provider._get_raw_frame( + frame, frame_name, _ = segment_frame_provider._get_raw_frame( task_chunk_frame_id, quality=quality ) - task_chunk_frames[task_chunk_frame_id] = (frame, None, None) + task_chunk_frames[task_chunk_frame_id] = (frame, frame_name, None) writer_kwargs = {} if self._db_task.dimension == models.DimensionType.DIM_3D: diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 296023c8f93..0fb3f53658c 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -92,13 +92,13 @@ def sort(images, sorting_method=SortingMethod.LEXICOGRAPHICAL, func=None): else: raise NotImplementedError() -def image_size_within_orientation(img: Image): +def image_size_within_orientation(img: Image.Image): orientation = img.getexif().get(ORIENTATION_EXIF_TAG, ORIENTATION.NORMAL_HORIZONTAL) if orientation > 4: return img.height, img.width return img.width, img.height -def has_exif_rotation(img: Image): +def has_exif_rotation(img: Image.Image): return img.getexif().get(ORIENTATION_EXIF_TAG, ORIENTATION.NORMAL_HORIZONTAL) != ORIENTATION.NORMAL_HORIZONTAL _T = TypeVar("_T") @@ -851,33 +851,37 @@ class ZipChunkWriter(IChunkWriter): POINT_CLOUD_EXT = 'pcd' def _write_pcd_file(self, image: str|io.BytesIO) -> tuple[io.BytesIO, str, int, int]: - image_buf = open(image, "rb") if isinstance(image, str) else image - try: + with ExitStack() as es: + if isinstance(image, str): + image_buf = es.enter_context(open(image, "rb")) + else: + image_buf = image + properties = ValidateDimension.get_pcd_properties(image_buf) w, h = int(properties["WIDTH"]), int(properties["HEIGHT"]) image_buf.seek(0, 0) return io.BytesIO(image_buf.read()), self.POINT_CLOUD_EXT, w, h - finally: - if isinstance(image, str): - image_buf.close() def save_as_chunk(self, images: Iterator[tuple[Image.Image|io.IOBase|str, str, str]], chunk_path: str): with zipfile.ZipFile(chunk_path, 'x') as zip_chunk: for idx, (image, path, _) in enumerate(images): ext = os.path.splitext(path)[1].replace('.', '') - output = io.BytesIO() + if self._dimension == DimensionType.DIM_2D: # current version of Pillow applies exif rotation immediately when TIFF image opened # and it removes rotation tag after that # so, has_exif_rotation(image) will return False for TIFF images even if they were actually rotated # and original files will be added to the archive (without applied rotation) # that is why we need the second part of the condition - if has_exif_rotation(image) or image.format == 'TIFF': + if isinstance(image, Image.Image) and ( + has_exif_rotation(image) or image.format == 'TIFF' + ): + output = io.BytesIO() rot_image = ImageOps.exif_transpose(image) try: if image.format == 'TIFF': # https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html - # use loseless lzw compression for tiff images + # use lossless lzw compression for tiff images rot_image.save(output, format='TIFF', compression='tiff_lzw') else: rot_image.save( @@ -889,16 +893,19 @@ def save_as_chunk(self, images: Iterator[tuple[Image.Image|io.IOBase|str, str, s ) finally: rot_image.close() + elif isinstance(image, io.IOBase): + output = image else: output = path else: output, ext = self._write_pcd_file(path)[0:2] - arcname = '{:06d}.{}'.format(idx, ext) + arcname = '{:06d}.{}'.format(idx, ext) if isinstance(output, io.BytesIO): zip_chunk.writestr(arcname, output.getvalue()) else: zip_chunk.write(filename=output, arcname=arcname) + # return empty list because ZipChunkWriter write files as is # and does not decode it to know img size. return [] diff --git a/tests/python/sdk/test_auto_annotation.py b/tests/python/sdk/test_auto_annotation.py index 142c4354c4d..e7ac8418b69 100644 --- a/tests/python/sdk/test_auto_annotation.py +++ b/tests/python/sdk/test_auto_annotation.py @@ -29,6 +29,7 @@ def _common_setup( tmp_path: Path, fxt_login: Tuple[Client, str], fxt_logger: Tuple[Logger, io.StringIO], + restore_redis_ondisk_per_function, ): logger = fxt_logger[0] client = fxt_login[0] diff --git a/tests/python/sdk/test_datasets.py b/tests/python/sdk/test_datasets.py index d5fbc0957eb..542ad9a1e80 100644 --- a/tests/python/sdk/test_datasets.py +++ b/tests/python/sdk/test_datasets.py @@ -23,6 +23,7 @@ def _common_setup( tmp_path: Path, fxt_login: Tuple[Client, str], fxt_logger: Tuple[Logger, io.StringIO], + restore_redis_ondisk_per_function, ): logger = fxt_logger[0] client = fxt_login[0] diff --git a/tests/python/sdk/test_jobs.py b/tests/python/sdk/test_jobs.py index a3e3b9d8516..f73b9cae7af 100644 --- a/tests/python/sdk/test_jobs.py +++ b/tests/python/sdk/test_jobs.py @@ -24,6 +24,7 @@ def setup( fxt_login: Tuple[Client, str], fxt_logger: Tuple[Logger, io.StringIO], fxt_stdout: io.StringIO, + restore_redis_ondisk_per_function, ): self.tmp_path = tmp_path logger, self.logger_stream = fxt_logger diff --git a/tests/python/sdk/test_projects.py b/tests/python/sdk/test_projects.py index 852a286fc28..09592dddbac 100644 --- a/tests/python/sdk/test_projects.py +++ b/tests/python/sdk/test_projects.py @@ -26,6 +26,7 @@ def setup( fxt_login: Tuple[Client, str], fxt_logger: Tuple[Logger, io.StringIO], fxt_stdout: io.StringIO, + restore_redis_ondisk_per_function, ): self.tmp_path = tmp_path logger, self.logger_stream = fxt_logger diff --git a/tests/python/sdk/test_tasks.py b/tests/python/sdk/test_tasks.py index f97e88a2924..3cdbf69f75e 100644 --- a/tests/python/sdk/test_tasks.py +++ b/tests/python/sdk/test_tasks.py @@ -29,6 +29,7 @@ def setup( fxt_login: Tuple[Client, str], fxt_logger: Tuple[Logger, io.StringIO], fxt_stdout: io.StringIO, + restore_redis_ondisk_per_function, ): self.tmp_path = tmp_path logger, self.logger_stream = fxt_logger From 48f479413cd77b3e37aeff105eaa6af4ac1ba440 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 13 Aug 2024 20:44:29 +0300 Subject: [PATCH 048/115] Fix other failing sdk tests --- tests/python/sdk/test_pytorch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/python/sdk/test_pytorch.py b/tests/python/sdk/test_pytorch.py index 722cb37ab00..2bcbd122abf 100644 --- a/tests/python/sdk/test_pytorch.py +++ b/tests/python/sdk/test_pytorch.py @@ -36,6 +36,7 @@ def _common_setup( tmp_path: Path, fxt_login: Tuple[Client, str], fxt_logger: Tuple[Logger, io.StringIO], + restore_redis_ondisk_per_function, ): logger = fxt_logger[0] client = fxt_login[0] From 5c0cc1ae10ca5385faec1fe2cbe9a3afc95fdff9 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 14 Aug 2024 18:13:33 +0300 Subject: [PATCH 049/115] Improve logging for migration --- cvat/apps/engine/log.py | 39 +++++++++++-------- .../migrations/0083_move_to_segment_chunks.py | 11 ++++-- 2 files changed, 30 insertions(+), 20 deletions(-) diff --git a/cvat/apps/engine/log.py b/cvat/apps/engine/log.py index 5f123d33eef..6f1740e74fd 100644 --- a/cvat/apps/engine/log.py +++ b/cvat/apps/engine/log.py @@ -59,24 +59,31 @@ def get_logger(logger_name, log_file): vlogger = logging.getLogger('vector') + +def get_migration_log_dir() -> str: + return settings.MIGRATIONS_LOGS_ROOT + +def get_migration_log_file_path(migration_name: str) -> str: + return osp.join(get_migration_log_dir(), f'{migration_name}.log') + @contextmanager def get_migration_logger(migration_name): - migration_log_file = '{}.log'.format(migration_name) + migration_log_file_path = get_migration_log_file_path(migration_name) stdout = sys.stdout stderr = sys.stderr + # redirect all stdout to the file - log_file_object = open(osp.join(settings.MIGRATIONS_LOGS_ROOT, migration_log_file), 'w') - sys.stdout = log_file_object - sys.stderr = log_file_object - - log = logging.getLogger(migration_name) - log.addHandler(logging.StreamHandler(stdout)) - log.addHandler(logging.StreamHandler(log_file_object)) - log.setLevel(logging.INFO) - - try: - yield log - finally: - log_file_object.close() - sys.stdout = stdout - sys.stderr = stderr + with open(migration_log_file_path, 'w') as log_file_object: + sys.stdout = log_file_object + sys.stderr = log_file_object + + log = logging.getLogger(migration_name) + log.addHandler(logging.StreamHandler(stdout)) + log.addHandler(logging.StreamHandler(log_file_object)) + log.setLevel(logging.INFO) + + try: + yield log + finally: + sys.stdout = stdout + sys.stderr = stderr diff --git a/cvat/apps/engine/migrations/0083_move_to_segment_chunks.py b/cvat/apps/engine/migrations/0083_move_to_segment_chunks.py index d21a7f669b2..c9f59593d23 100644 --- a/cvat/apps/engine/migrations/0083_move_to_segment_chunks.py +++ b/cvat/apps/engine/migrations/0083_move_to_segment_chunks.py @@ -2,10 +2,11 @@ import os from django.db import migrations -from cvat.apps.engine.log import get_migration_logger +from cvat.apps.engine.log import get_migration_logger, get_migration_log_dir def switch_tasks_with_static_chunks_to_dynamic_chunks(apps, schema_editor): migration_name = os.path.splitext(os.path.basename(__file__))[0] + migration_log_dir = get_migration_log_dir() with get_migration_logger(migration_name) as common_logger: Data = apps.get_model("engine", "Data") @@ -26,8 +27,8 @@ def switch_tasks_with_static_chunks_to_dynamic_chunks(apps, schema_editor): data_with_static_cache_query.update(storage_method="cache") - updated_data_ids_filename = migration_name + "-data_ids.log" - with open(updated_data_ids_filename, "w") as data_ids_file: + updated_ids_filename = os.path.join(migration_log_dir, migration_name + "-data_ids.log") + with open(updated_ids_filename, "w") as data_ids_file: print( "The following Data ids have been switched from using \"filesystem\" chunk storage " "to \"cache\":", @@ -38,7 +39,9 @@ def switch_tasks_with_static_chunks_to_dynamic_chunks(apps, schema_editor): common_logger.info( "Information about migrated tasks is available in the migration log file: " - f"{updated_data_ids_filename}. You will need to remove data manually for these tasks." + "{}. You will need to remove data manually for these tasks.".format( + updated_ids_filename + ) ) class Migration(migrations.Migration): From 5abd89137cd133ccf27332b66b1054a3f6e2dc2f Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 14 Aug 2024 18:55:40 +0300 Subject: [PATCH 050/115] Fix invalid starting index --- cvat/apps/engine/media_extractors.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 0fb3f53658c..40fae1176a9 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -769,7 +769,7 @@ def iterate_frames(self, *, frame_filter: Iterable[int]) -> Iterable[av.VideoFra container.seek(offset=start_decode_timestamp, stream=video_stream) - frame_counter = itertools.count(start_decode_frame_number - 1) + frame_counter = itertools.count(start_decode_frame_number) with closing(self._decode_stream(container, video_stream)) as stream_decoder: for frame, frame_number in zip(stream_decoder, frame_counter): if frame_number == next_frame_filter_frame: From 749b970d206352e4fe75210f1549f5c61701a3b5 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 14 Aug 2024 18:55:56 +0300 Subject: [PATCH 051/115] Fix frame reading in lambda functions --- cvat/apps/lambda_manager/views.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/apps/lambda_manager/views.py b/cvat/apps/lambda_manager/views.py index 776911eac3f..fdb8aaa869e 100644 --- a/cvat/apps/lambda_manager/views.py +++ b/cvat/apps/lambda_manager/views.py @@ -499,7 +499,7 @@ def _get_image(self, db_task, frame, quality): frame_provider = TaskFrameProvider(db_task) image = frame_provider.get_frame(frame, quality=quality) - return base64.b64encode(image[0].getvalue()).decode('utf-8') + return base64.b64encode(image.data.getvalue()).decode('utf-8') class LambdaQueue: RESULT_TTL = timedelta(minutes=30) From 9105cd3b3d9869803a67b2be36dc38b57c75242f Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 14 Aug 2024 19:58:08 +0300 Subject: [PATCH 052/115] Fix unintended frame indexing changes --- cvat/apps/dataset_manager/formats/cvat.py | 4 +-- cvat/apps/engine/cache.py | 7 +++-- cvat/apps/engine/frame_provider.py | 38 +++++++++++++++++------ 3 files changed, 35 insertions(+), 14 deletions(-) diff --git a/cvat/apps/dataset_manager/formats/cvat.py b/cvat/apps/dataset_manager/formats/cvat.py index 299982f2dcf..4651fd39845 100644 --- a/cvat/apps/dataset_manager/formats/cvat.py +++ b/cvat/apps/dataset_manager/formats/cvat.py @@ -1378,8 +1378,8 @@ def dump_media_files(instance_data: CommonData, img_dir: str, project_data: Proj ext = frame_provider.VIDEO_FRAME_EXT frames = frame_provider.iterate_frames( - start_frame=instance_data.abs_frame_id(instance_data.start), - stop_frame=instance_data.abs_frame_id(instance_data.stop), + start_frame=instance_data.start, + stop_frame=instance_data.stop, quality=FrameQuality.ORIGINAL, out_type=FrameOutputType.BUFFER, ) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 9a5ae932bef..d7af8b86153 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -404,11 +404,14 @@ def _prepare_segment_preview(self, db_segment: models.Segment) -> DataWithMime: os.path.join(os.path.dirname(__file__), "assets/3d_preview.jpeg") ) else: - from cvat.apps.engine.frame_provider import FrameOutputType, SegmentFrameProvider + from cvat.apps.engine.frame_provider import ( + FrameOutputType, SegmentFrameProvider, TaskFrameProvider + ) + task_frame_provider = TaskFrameProvider(db_segment.task) segment_frame_provider = SegmentFrameProvider(db_segment) preview = segment_frame_provider.get_frame( - min(db_segment.frame_set), + task_frame_provider.get_rel_frame_number(min(db_segment.frame_set)), quality=FrameQuality.COMPRESSED, out_type=FrameOutputType.PIL, ).data diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 6a6ea1c82c7..8ce5be00f5a 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -207,18 +207,22 @@ def iterate_frames( out_type: FrameOutputType = FrameOutputType.BUFFER, ) -> Iterator[DataWithMeta[AnyFrame]]: ... + def _get_abs_frame_number(self, db_data: models.Data, rel_frame_number: int) -> int: + return db_data.start_frame + rel_frame_number * db_data.get_frame_step() + + def _get_rel_frame_number(self, db_data: models.Data, abs_frame_number: int) -> int: + return (abs_frame_number - db_data.start_frame) // db_data.get_frame_step() + class TaskFrameProvider(IFrameProvider): def __init__(self, db_task: models.Task) -> None: self._db_task = db_task def validate_frame_number(self, frame_number: int) -> int: - start = self._db_task.data.start_frame - stop = self._db_task.data.stop_frame - if frame_number not in range(start, stop + 1, self._db_task.data.get_frame_step()): + if frame_number not in range(0, self._db_task.data.size): raise ValidationError( f"Invalid frame '{frame_number}'. " - f"The frame number should be in the [{start}, {stop}] range" + f"The frame number should be in the [0, {self._db_task.data.size}] range" ) return frame_number @@ -236,6 +240,17 @@ def validate_chunk_number(self, chunk_number: int) -> int: def get_chunk_number(self, frame_number: int) -> int: return int(frame_number) // self._db_task.data.chunk_size + def get_abs_frame_number(self, rel_frame_number: int) -> int: + "Returns absolute frame number in the task (in the range [start, stop, step])" + return super()._get_abs_frame_number(self._db_task.data, rel_frame_number) + + def get_rel_frame_number(self, abs_frame_number: int) -> int: + """ + Returns relative frame number in the task (in the range [0, task_size - 1]). + This is the "normal" frame number, expected in other methods. + """ + return super()._get_rel_frame_number(self._db_task.data, abs_frame_number) + def get_preview(self) -> DataWithMeta[BytesIO]: return self._get_segment_frame_provider(self._db_task.data.start_frame).get_preview() @@ -315,7 +330,7 @@ def get_chunk( continue frame, frame_name, _ = segment_frame_provider._get_raw_frame( - task_chunk_frame_id, quality=quality + self.get_rel_frame_number(task_chunk_frame_id), quality=quality ) task_chunk_frames[task_chunk_frame_id] = (frame, frame_name, None) @@ -385,11 +400,13 @@ def _get_segment(self, validated_frame_number: int) -> models.Segment: if not self._db_task.data or not self._db_task.data.size: raise ValidationError("Task has no data") + abs_frame_number = self.get_abs_frame_number(validated_frame_number) + return next( s for s in self._db_task.segment_set.all() if s.type == models.SegmentType.RANGE - if validated_frame_number in s.frame_set + if abs_frame_number in s.frame_set ) def _get_segment_frame_provider(self, frame_number: int) -> SegmentFrameProvider: @@ -465,12 +482,13 @@ def __len__(self): def validate_frame_number(self, frame_number: int) -> Tuple[int, int, int]: frame_sequence = list(self._db_segment.frame_set) - if frame_number not in frame_sequence: + abs_frame_number = self._get_abs_frame_number(self._db_segment.task.data, frame_number) + if abs_frame_number not in frame_sequence: raise ValidationError(f"Incorrect requested frame number: {frame_number}") # TODO: maybe optimize search chunk_number, frame_position = divmod( - frame_sequence.index(frame_number), self._db_segment.task.data.chunk_size + frame_sequence.index(abs_frame_number), self._db_segment.task.data.chunk_size ) return frame_number, chunk_number, frame_position @@ -554,13 +572,13 @@ def iterate_frames( quality: FrameQuality = FrameQuality.ORIGINAL, out_type: FrameOutputType = FrameOutputType.BUFFER, ) -> Iterator[DataWithMeta[AnyFrame]]: - frame_range = itertools.count(start_frame, self._db_segment.task.data.get_frame_step()) + frame_range = itertools.count(start_frame) if stop_frame: frame_range = itertools.takewhile(lambda x: x <= stop_frame, frame_range) segment_frame_set = set(self._db_segment.frame_set) for idx in frame_range: - if idx in segment_frame_set: + if self._get_abs_frame_number(self._db_segment.task.data, idx) in segment_frame_set: yield self.get_frame(idx, quality=quality, out_type=out_type) From 8dafcbe4ebf41054a2c676bba43511b170cddb76 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 14 Aug 2024 20:46:58 +0300 Subject: [PATCH 053/115] Fix various indexing errors in media extractors --- cvat/apps/engine/media_extractors.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 40fae1176a9..17c81569504 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -250,7 +250,7 @@ def __len__(self): @property def frame_range(self): - return range(self._start, self._stop, self._step) + return range(self._start, self._stop + 1, self._step) class ImageListReader(IMediaReader): def __init__(self, @@ -264,9 +264,9 @@ def __init__(self, raise Exception('No image found') if not stop: - stop = len(source_path) + stop = max(0, len(source_path) - 1) else: - stop = min(len(source_path), stop + 1) + stop = max(0, min(len(source_path), stop + 1) - 1) step = max(step, 1) assert stop > start @@ -281,7 +281,7 @@ def __init__(self, self._sorting_method = sorting_method def __iter__(self): - for i in range(self._start, self._stop, self._step): + for i in range(self._start, self._stop + 1, self._step): yield (self.get_image(i), self.get_path(i), i) def __contains__(self, media_file): @@ -294,7 +294,7 @@ def filter(self, callback): source_path, step=self._step, start=self._start, - stop=self._stop - 1, + stop=self._stop, dimension=self._dimension, sorting_method=self._sorting_method ) @@ -306,7 +306,7 @@ def get_image(self, i): return self._source_path[i] def get_progress(self, pos): - return (pos - self._start + 1) / (self._stop - self._start) + return (pos + 1) / (len(self.frame_range) or 1) def get_preview(self, frame): if self._dimension == DimensionType.DIM_3D: @@ -560,7 +560,7 @@ def __init__( source_path=source_path, step=step, start=start, - stop=stop + 1 if stop is not None else stop, + stop=stop if stop is not None else stop, dimension=dimension, ) From 4cbf82f2d8315b389c0ca63beb8ad48488a8897a Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 14 Aug 2024 20:47:29 +0300 Subject: [PATCH 054/115] Fix temp resource cleanup in server tests --- cvat/apps/engine/tests/utils.py | 18 ++++----- cvat/apps/lambda_manager/tests/test_lambda.py | 38 +++---------------- 2 files changed, 14 insertions(+), 42 deletions(-) diff --git a/cvat/apps/engine/tests/utils.py b/cvat/apps/engine/tests/utils.py index b884b3e9b4c..87c5911e12b 100644 --- a/cvat/apps/engine/tests/utils.py +++ b/cvat/apps/engine/tests/utils.py @@ -92,14 +92,7 @@ def clear_rq_jobs(): class ApiTestBase(APITestCase): - def _clear_rq_jobs(self): - clear_rq_jobs() - - def setUp(self): - super().setUp() - self.client = APIClient() - - def tearDown(self): + def _clear_temp_data(self): # Clear server frame/chunk cache. # The parent class clears DB changes, and it can lead to under-cleaned task data, # which can affect other tests. @@ -112,7 +105,14 @@ def tearDown(self): # Clear any remaining RQ jobs produced by the tests executed self._clear_rq_jobs() - return super().tearDown() + def _clear_rq_jobs(self): + clear_rq_jobs() + + def setUp(self): + self._clear_temp_data() + + super().setUp() + self.client = APIClient() def generate_image_file(filename, size=(100, 100)): diff --git a/cvat/apps/lambda_manager/tests/test_lambda.py b/cvat/apps/lambda_manager/tests/test_lambda.py index e360ab1996d..51973fdbd4a 100644 --- a/cvat/apps/lambda_manager/tests/test_lambda.py +++ b/cvat/apps/lambda_manager/tests/test_lambda.py @@ -1,11 +1,10 @@ # Copyright (C) 2021-2022 Intel Corporation -# Copyright (C) 2023 CVAT.ai Corporation +# Copyright (C) 2023-2024 CVAT.ai Corporation # # SPDX-License-Identifier: MIT from collections import OrderedDict from itertools import groupby -from io import BytesIO from typing import Dict, Optional from unittest import mock, skip import json @@ -14,11 +13,11 @@ import requests from django.contrib.auth.models import Group, User from django.http import HttpResponseNotFound, HttpResponseServerError -from PIL import Image from rest_framework import status -from rest_framework.test import APIClient, APITestCase -from cvat.apps.engine.tests.utils import filter_dict, get_paginated_collection +from cvat.apps.engine.tests.utils import ( + ApiTestBase, filter_dict, ForceLogin, generate_image_file, get_paginated_collection +) LAMBDA_ROOT_PATH = '/api/lambda' LAMBDA_FUNCTIONS_PATH = f'{LAMBDA_ROOT_PATH}/functions' @@ -49,35 +48,8 @@ with open(path) as f: functions = json.load(f) - -def generate_image_file(filename, size=(100, 100)): - f = BytesIO() - image = Image.new('RGB', size=size) - image.save(f, 'jpeg') - f.name = filename - f.seek(0) - return f - - -class ForceLogin: - def __init__(self, user, client): - self.user = user - self.client = client - - def __enter__(self): - if self.user: - self.client.force_login(self.user, backend='django.contrib.auth.backends.ModelBackend') - - return self - - def __exit__(self, exception_type, exception_value, traceback): - if self.user: - self.client.logout() - -class _LambdaTestCaseBase(APITestCase): +class _LambdaTestCaseBase(ApiTestBase): def setUp(self): - self.client = APIClient() - http_patcher = mock.patch('cvat.apps.lambda_manager.views.LambdaGateway._http', side_effect = self._get_data_from_lambda_manager_http) self.addCleanup(http_patcher.stop) http_patcher.start() From 88c34a3445bfc02bdc1c27fc00dcabe52af032bb Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 15 Aug 2024 18:46:16 +0300 Subject: [PATCH 055/115] Refactor some code --- cvat/apps/engine/media_extractors.py | 50 +++++++++++++++++++--------- 1 file changed, 35 insertions(+), 15 deletions(-) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 17c81569504..1567ff0574c 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -193,11 +193,25 @@ def __getitem__(self, idx: int): class IMediaReader(ABC): - def __init__(self, source_path, step, start, stop, dimension): + def __init__( + self, + source_path, + *, + start: int = 0, + stop: Optional[int] = None, + step: int = 1, + dimension: DimensionType = DimensionType.DIM_2D + ): self._source_path = source_path + self._step = step + self._start = start + "The first included index" + self._stop = stop + "The last included index" + self._dimension = dimension @abstractmethod @@ -245,28 +259,27 @@ def _get_preview(obj): def get_image_size(self, i): pass + @abstractmethod def __len__(self): - return len(self.frame_range) - - @property - def frame_range(self): - return range(self._start, self._stop + 1, self._step) + pass class ImageListReader(IMediaReader): def __init__(self, - source_path, - step=1, - start=0, - stop=None, - dimension=DimensionType.DIM_2D, - sorting_method=SortingMethod.LEXICOGRAPHICAL): + source_path, + step: int = 1, + start: int = 0, + stop: Optional[int] = None, + dimension: DimensionType = DimensionType.DIM_2D, + sorting_method: SortingMethod = SortingMethod.LEXICOGRAPHICAL, + ): if not source_path: raise Exception('No image found') if not stop: - stop = max(0, len(source_path) - 1) + stop = len(source_path) - 1 else: - stop = max(0, min(len(source_path), stop + 1) - 1) + stop = min(len(source_path) - 1, stop) + step = max(step, 1) assert stop > start @@ -281,7 +294,7 @@ def __init__(self, self._sorting_method = sorting_method def __iter__(self): - for i in range(self._start, self._stop + 1, self._step): + for i in self.frame_range: yield (self.get_image(i), self.get_path(i), i) def __contains__(self, media_file): @@ -338,6 +351,13 @@ def reconcile(self, source_files, step=1, start=0, stop=None, dimension=Dimensio def absolute_source_paths(self): return [self.get_path(idx) for idx, _ in enumerate(self._source_path)] + def __len__(self): + return len(self.frame_range) + + @property + def frame_range(self): + return range(self._start, self._stop + 1, self._step) + class DirectoryReader(ImageListReader): def __init__(self, source_path, From b0fd0064fcee6dd8f59bfaefd57c5156d358c623 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 15 Aug 2024 21:11:47 +0300 Subject: [PATCH 056/115] Remove duplicated tests --- cvat/apps/engine/tests/test_rest_api.py | 42 ------------------------- 1 file changed, 42 deletions(-) diff --git a/cvat/apps/engine/tests/test_rest_api.py b/cvat/apps/engine/tests/test_rest_api.py index 47758be11d1..2f8a00b4111 100644 --- a/cvat/apps/engine/tests/test_rest_api.py +++ b/cvat/apps/engine/tests/test_rest_api.py @@ -4115,48 +4115,6 @@ def _test_api_v2_tasks_id_data_create_can_use_server_images_and_manifest(self, u for i, fn in enumerate(images + [manifest_name]) }) - for copy_data in [True, False]: - with self.subTest(current_function_name(), copy=copy_data): - task_spec = task_spec_common.copy() - task_spec['name'] = task_spec['name'] + f' copy={copy_data}' - task_data['copy_data'] = copy_data - self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data, - self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, - image_sizes, StorageMethodChoice.CACHE, - StorageChoice.LOCAL if copy_data else StorageChoice.SHARE) - - with self.subTest(current_function_name() + ' file order mismatch'): - task_spec = task_spec_common.copy() - task_spec['name'] = task_spec['name'] + f' mismatching file order' - task_data_copy = task_data.copy() - task_data_copy[f'server_files[{len(images)}]'] = "images_manifest.jsonl" - self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data_copy, - self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, - image_sizes, StorageMethodChoice.CACHE, StorageChoice.SHARE, - expected_task_creation_status_state='Failed', - expected_task_creation_status_reason='Incorrect file mapping to manifest content') - - for copy_data in [True, False]: - with self.subTest(current_function_name(), copy=copy_data): - task_spec = task_spec_common.copy() - task_spec['name'] = task_spec['name'] + f' copy={copy_data}' - task_data['copy_data'] = copy_data - self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data, - self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, - image_sizes, StorageMethodChoice.CACHE, - StorageChoice.LOCAL if copy_data else StorageChoice.SHARE) - - with self.subTest(current_function_name() + ' file order mismatch'): - task_spec = task_spec_common.copy() - task_spec['name'] = task_spec['name'] + f' mismatching file order' - task_data_copy = task_data.copy() - task_data_copy[f'server_files[{len(images)}]'] = "images_manifest.jsonl" - self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data_copy, - self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, - image_sizes, StorageMethodChoice.CACHE, StorageChoice.SHARE, - expected_task_creation_status_state='Failed', - expected_task_creation_status_reason='Incorrect file mapping to manifest content') - for copy_data in [True, False]: with self.subTest(current_function_name(), copy=copy_data): task_spec = task_spec_common.copy() From 2eac04a2d0469eeae0c2b7b167a0108b32a5c317 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 15 Aug 2024 21:12:15 +0300 Subject: [PATCH 057/115] Remove extra change --- cvat/apps/engine/task.py | 1 - 1 file changed, 1 deletion(-) diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index 31eb2db3c65..8024ed4c376 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -1117,7 +1117,6 @@ def _update_status(msg: str) -> None: models.RelatedFile(data=image.data, primary_image=image, path=os.path.join(upload_dir, related_file_path)) for image in images for related_file_path in related_images.get(image.path, []) - if not image.is_placeholder # TODO ] models.RelatedFile.objects.bulk_create(db_related_files) else: From 640518ced993904647593889606d9203eb8634f6 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 15 Aug 2024 21:12:49 +0300 Subject: [PATCH 058/115] Fix method name, remove extra method --- cvat/apps/engine/cache.py | 2 +- cvat/apps/engine/media_extractors.py | 6 +----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index d7af8b86153..ff1851b2f93 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -146,7 +146,7 @@ def get_or_set_cloud_preview(self, db_storage: models.CloudStorage) -> DataWithM def get_frame_context_images(self, db_data: models.Data, frame_number: int) -> DataWithMime: return self._get_or_set_cache_item( key=f"context_image_{db_data.id}_{frame_number}", - create_callback=lambda: self._prepare_context_image(db_data, frame_number), + create_callback=lambda: self.prepare_context_images(db_data, frame_number), ) def _read_raw_frames( diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 1567ff0574c..afdeaba9efa 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -259,10 +259,6 @@ def _get_preview(obj): def get_image_size(self, i): pass - @abstractmethod - def __len__(self): - pass - class ImageListReader(IMediaReader): def __init__(self, source_path, @@ -281,7 +277,7 @@ def __init__(self, stop = min(len(source_path) - 1, stop) step = max(step, 1) - assert stop > start + assert stop >= start super().__init__( source_path=sort(source_path, sorting_method), From 3a246b3f896fb809665d158963ce4f97a45a4cad Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 15 Aug 2024 21:15:46 +0300 Subject: [PATCH 059/115] Remove some shared code in tests, add temp data cleanup --- .../dataset_manager/tests/test_formats.py | 37 +++---------------- .../tests/test_rest_api_formats.py | 12 +++--- cvat/apps/engine/tests/test_rest_api_3D.py | 4 +- cvat/apps/lambda_manager/tests/test_lambda.py | 2 + 4 files changed, 16 insertions(+), 39 deletions(-) diff --git a/cvat/apps/dataset_manager/tests/test_formats.py b/cvat/apps/dataset_manager/tests/test_formats.py index 1c7db60814d..ea5021d7d57 100644 --- a/cvat/apps/dataset_manager/tests/test_formats.py +++ b/cvat/apps/dataset_manager/tests/test_formats.py @@ -1,6 +1,6 @@ # Copyright (C) 2020-2022 Intel Corporation -# Copyright (C) 2022 CVAT.ai Corporation +# Copyright (C) 2022-2024 CVAT.ai Corporation # # SPDX-License-Identifier: MIT @@ -14,10 +14,8 @@ from datumaro.components.dataset import Dataset, DatasetItem from datumaro.components.annotation import Mask from django.contrib.auth.models import Group, User -from PIL import Image from rest_framework import status -from rest_framework.test import APIClient, APITestCase import cvat.apps.dataset_manager as dm from cvat.apps.dataset_manager.annotation import AnnotationIR @@ -26,36 +24,13 @@ from cvat.apps.dataset_manager.task import TaskAnnotation from cvat.apps.dataset_manager.util import make_zip_archive from cvat.apps.engine.models import Task -from cvat.apps.engine.tests.utils import get_paginated_collection +from cvat.apps.engine.tests.utils import ( + get_paginated_collection, ForceLogin, generate_image_file, ApiTestBase +) - -def generate_image_file(filename, size=(100, 100)): - f = BytesIO() - image = Image.new('RGB', size=size) - image.save(f, 'jpeg') - f.name = filename - f.seek(0) - return f - -class ForceLogin: - def __init__(self, user, client): - self.user = user - self.client = client - - def __enter__(self): - if self.user: - self.client.force_login(self.user, - backend='django.contrib.auth.backends.ModelBackend') - - return self - - def __exit__(self, exception_type, exception_value, traceback): - if self.user: - self.client.logout() - -class _DbTestBase(APITestCase): +class _DbTestBase(ApiTestBase): def setUp(self): - self.client = APIClient() + super().setUp() @classmethod def setUpTestData(cls): diff --git a/cvat/apps/dataset_manager/tests/test_rest_api_formats.py b/cvat/apps/dataset_manager/tests/test_rest_api_formats.py index 767fb07fe96..6ed6bbee3d5 100644 --- a/cvat/apps/dataset_manager/tests/test_rest_api_formats.py +++ b/cvat/apps/dataset_manager/tests/test_rest_api_formats.py @@ -419,7 +419,7 @@ def test_api_v2_dump_and_upload_annotations_with_objects_type_is_shape(self): url = self._generate_url_dump_tasks_annotations(task_id) for user, edata in list(expected.items()): - self._clear_rq_jobs() # clean up from previous tests and iterations + self._clear_temp_data() # clean up from previous tests and iterations user_name = edata['name'] file_zip_name = osp.join(test_dir, f'{test_name}_{user_name}_{dump_format_name}.zip') @@ -526,7 +526,7 @@ def test_api_v2_dump_annotations_with_objects_type_is_track(self): url = self._generate_url_dump_tasks_annotations(task_id) for user, edata in list(expected.items()): - self._clear_rq_jobs() # clean up from previous tests and iterations + self._clear_temp_data() # clean up from previous tests and iterations user_name = edata['name'] file_zip_name = osp.join(test_dir, f'{test_name}_{user_name}_{dump_format_name}.zip') @@ -613,7 +613,7 @@ def test_api_v2_dump_tag_annotations(self): for user, edata in list(expected.items()): with self.subTest(format=f"{edata['name']}"): with TestDir() as test_dir: - self._clear_rq_jobs() # clean up from previous tests and iterations + self._clear_temp_data() # clean up from previous tests and iterations user_name = edata['name'] url = self._generate_url_dump_tasks_annotations(task_id) @@ -857,7 +857,7 @@ def test_api_v2_export_dataset(self): # dump annotations url = self._generate_url_dump_task_dataset(task_id) for user, edata in list(expected.items()): - self._clear_rq_jobs() # clean up from previous tests and iterations + self._clear_temp_data() # clean up from previous tests and iterations user_name = edata['name'] file_zip_name = osp.join(test_dir, f'{test_name}_{user_name}_{dump_format_name}.zip') @@ -2058,7 +2058,7 @@ def test_api_v2_export_import_dataset(self): self._create_annotations(task, dump_format_name, "random") for user, edata in list(expected.items()): - self._clear_rq_jobs() # clean up from previous tests and iterations + self._clear_temp_data() # clean up from previous tests and iterations user_name = edata['name'] file_zip_name = osp.join(test_dir, f'{test_name}_{user_name}_{dump_format_name}.zip') @@ -2140,7 +2140,7 @@ def test_api_v2_export_annotations(self): url = self._generate_url_dump_project_annotations(project['id'], dump_format_name) for user, edata in list(expected.items()): - self._clear_rq_jobs() # clean up from previous tests and iterations + self._clear_temp_data() # clean up from previous tests and iterations user_name = edata['name'] file_zip_name = osp.join(test_dir, f'{test_name}_{user_name}_{dump_format_name}.zip') diff --git a/cvat/apps/engine/tests/test_rest_api_3D.py b/cvat/apps/engine/tests/test_rest_api_3D.py index a67a79109f3..0bdd3a00105 100644 --- a/cvat/apps/engine/tests/test_rest_api_3D.py +++ b/cvat/apps/engine/tests/test_rest_api_3D.py @@ -527,7 +527,7 @@ def test_api_v2_dump_and_upload_annotation(self): for user, edata in list(self.expected_dump_upload.items()): with self.subTest(format=f"{format_name}_{edata['name']}_dump"): - self._clear_rq_jobs() # clean up from previous tests and iterations + self._clear_temp_data() # clean up from previous tests and iterations url = self._generate_url_dump_tasks_annotations(task_id) file_name = osp.join(test_dir, f"{format_name}_{edata['name']}.zip") @@ -718,7 +718,7 @@ def test_api_v2_export_dataset(self): for user, edata in list(self.expected_dump_upload.items()): with self.subTest(format=f"{format_name}_{edata['name']}_export"): - self._clear_rq_jobs() # clean up from previous tests and iterations + self._clear_temp_data() # clean up from previous tests and iterations url = self._generate_url_dump_dataset(task_id) file_name = osp.join(test_dir, f"{format_name}_{edata['name']}.zip") diff --git a/cvat/apps/lambda_manager/tests/test_lambda.py b/cvat/apps/lambda_manager/tests/test_lambda.py index 51973fdbd4a..7d57332c330 100644 --- a/cvat/apps/lambda_manager/tests/test_lambda.py +++ b/cvat/apps/lambda_manager/tests/test_lambda.py @@ -50,6 +50,8 @@ class _LambdaTestCaseBase(ApiTestBase): def setUp(self): + super().setUp() + http_patcher = mock.patch('cvat.apps.lambda_manager.views.LambdaGateway._http', side_effect = self._get_data_from_lambda_manager_http) self.addCleanup(http_patcher.stop) http_patcher.start() From a0704f46d190b00f9604e8a51e77c0363f0213d5 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 15 Aug 2024 21:18:52 +0300 Subject: [PATCH 060/115] Add checks for successful task creation in tests --- .../dataset_manager/tests/test_formats.py | 5 +++++ .../tests/test_rest_api_formats.py | 5 +++++ cvat/apps/engine/tests/test_rest_api.py | 20 ++++++++++++++++++- cvat/apps/engine/tests/test_rest_api_3D.py | 8 ++++++-- cvat/apps/lambda_manager/tests/test_lambda.py | 5 +++++ 5 files changed, 40 insertions(+), 3 deletions(-) diff --git a/cvat/apps/dataset_manager/tests/test_formats.py b/cvat/apps/dataset_manager/tests/test_formats.py index ea5021d7d57..f83672f1475 100644 --- a/cvat/apps/dataset_manager/tests/test_formats.py +++ b/cvat/apps/dataset_manager/tests/test_formats.py @@ -69,6 +69,11 @@ def _create_task(self, data, image_data): response = self.client.post("/api/tasks/%s/data" % tid, data=image_data) assert response.status_code == status.HTTP_202_ACCEPTED, response.status_code + rq_id = response.json()["rq_id"] + + response = self.client.get(f"/api/requests/{rq_id}") + assert response.status_code == status.HTTP_200_OK, response.status_code + assert response.json()["status"] == "finished", response.json().get("status") response = self.client.get("/api/tasks/%s" % tid) diff --git a/cvat/apps/dataset_manager/tests/test_rest_api_formats.py b/cvat/apps/dataset_manager/tests/test_rest_api_formats.py index 6ed6bbee3d5..bc23a253ee1 100644 --- a/cvat/apps/dataset_manager/tests/test_rest_api_formats.py +++ b/cvat/apps/dataset_manager/tests/test_rest_api_formats.py @@ -151,6 +151,11 @@ def _create_task(self, data, image_data): response = self.client.post("/api/tasks/%s/data" % tid, data=image_data) assert response.status_code == status.HTTP_202_ACCEPTED, response.status_code + rq_id = response.json()["rq_id"] + + response = self.client.get(f"/api/requests/{rq_id}") + assert response.status_code == status.HTTP_200_OK, response.status_code + assert response.json()["status"] == "finished", response.json().get("status") response = self.client.get("/api/tasks/%s" % tid) diff --git a/cvat/apps/engine/tests/test_rest_api.py b/cvat/apps/engine/tests/test_rest_api.py index 2f8a00b4111..16356ef8a0a 100644 --- a/cvat/apps/engine/tests/test_rest_api.py +++ b/cvat/apps/engine/tests/test_rest_api.py @@ -1422,7 +1422,13 @@ def _create_task(task_data, media_data): if isinstance(media, io.BytesIO): media.seek(0) response = cls.client.post("/api/tasks/{}/data".format(tid), data=media_data) - assert response.status_code == status.HTTP_202_ACCEPTED + assert response.status_code == status.HTTP_202_ACCEPTED, response.status_code + rq_id = response.json()["rq_id"] + + response = cls.client.get(f"/api/requests/{rq_id}") + assert response.status_code == status.HTTP_200_OK, response.status_code + assert response.json()["status"] == "finished", response.json().get("status") + response = cls.client.get("/api/tasks/{}".format(tid)) data_id = response.data["data"] cls.tasks.append({ @@ -1766,6 +1772,12 @@ def _create_task(task_data, media_data): media.seek(0) response = self.client.post("/api/tasks/{}/data".format(tid), data=media_data) assert response.status_code == status.HTTP_202_ACCEPTED + rq_id = response.json()["rq_id"] + + response = self.client.get(f"/api/requests/{rq_id}") + assert response.status_code == status.HTTP_200_OK, response.status_code + assert response.json()["status"] == "finished", response.json().get("status") + response = self.client.get("/api/tasks/{}".format(tid)) data_id = response.data["data"] self.tasks.append({ @@ -2882,6 +2894,12 @@ def _create_task(task_data, media_data): media.seek(0) response = self.client.post("/api/tasks/{}/data".format(tid), data=media_data) assert response.status_code == status.HTTP_202_ACCEPTED + rq_id = response.json()["rq_id"] + + response = self.client.get(f"/api/requests/{rq_id}") + assert response.status_code == status.HTTP_200_OK, response.status_code + assert response.json()["status"] == "finished", response.json().get("status") + response = self.client.get("/api/tasks/{}".format(tid)) data_id = response.data["data"] self.tasks.append({ diff --git a/cvat/apps/engine/tests/test_rest_api_3D.py b/cvat/apps/engine/tests/test_rest_api_3D.py index 0bdd3a00105..785eb875579 100644 --- a/cvat/apps/engine/tests/test_rest_api_3D.py +++ b/cvat/apps/engine/tests/test_rest_api_3D.py @@ -86,9 +86,13 @@ def _create_task(self, data, image_data): assert response.status_code == status.HTTP_201_CREATED, response.status_code tid = response.data["id"] - response = self.client.post("/api/tasks/%s/data" % tid, - data=image_data) + response = self.client.post("/api/tasks/%s/data" % tid, data=image_data) assert response.status_code == status.HTTP_202_ACCEPTED, response.status_code + rq_id = response.json()["rq_id"] + + response = self.client.get(f"/api/requests/{rq_id}") + assert response.status_code == status.HTTP_200_OK, response.status_code + assert response.json()["status"] == "finished", response.json().get("status") response = self.client.get("/api/tasks/%s" % tid) diff --git a/cvat/apps/lambda_manager/tests/test_lambda.py b/cvat/apps/lambda_manager/tests/test_lambda.py index 7d57332c330..92b51862e2a 100644 --- a/cvat/apps/lambda_manager/tests/test_lambda.py +++ b/cvat/apps/lambda_manager/tests/test_lambda.py @@ -155,6 +155,11 @@ def _create_task(self, task_spec, data, *, owner=None, org_id=None): data=data, QUERY_STRING=f'org_id={org_id}' if org_id is not None else None) assert response.status_code == status.HTTP_202_ACCEPTED, response.status_code + rq_id = response.json()["rq_id"] + + response = self.client.get(f"/api/requests/{rq_id}") + assert response.status_code == status.HTTP_200_OK, response.status_code + assert response.json()["status"] == "finished", response.json().get("status") response = self.client.get("/api/tasks/%s" % tid, QUERY_STRING=f'org_id={org_id}' if org_id is not None else None) From cf026ef343565cd18e8753fe6990257496a4e656 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 15 Aug 2024 21:19:13 +0300 Subject: [PATCH 061/115] Fix invalid variable access in test --- cvat/apps/engine/tests/test_rest_api_3D.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/cvat/apps/engine/tests/test_rest_api_3D.py b/cvat/apps/engine/tests/test_rest_api_3D.py index 785eb875579..9f000be5d21 100644 --- a/cvat/apps/engine/tests/test_rest_api_3D.py +++ b/cvat/apps/engine/tests/test_rest_api_3D.py @@ -744,6 +744,8 @@ def test_api_v2_export_dataset(self): content = io.BytesIO(b"".join(response.streaming_content)) with open(file_name, "wb") as f: f.write(content.getvalue()) - self.assertEqual(osp.exists(file_name), edata['file_exists']) - self._check_dump_content(content, task_ann_prev.data, format_name,related_files=False) + self.assertEqual(osp.exists(file_name), edata['file_exists']) + self._check_dump_content( + content, task_ann_prev.data, format_name, related_files=False + ) From f73cef30db9f070d2d4743426878922c16754e44 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 15 Aug 2024 21:40:01 +0300 Subject: [PATCH 062/115] Update default cache location in test checks --- cvat/apps/engine/tests/test_rest_api.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/cvat/apps/engine/tests/test_rest_api.py b/cvat/apps/engine/tests/test_rest_api.py index 16356ef8a0a..8c08948166a 100644 --- a/cvat/apps/engine/tests/test_rest_api.py +++ b/cvat/apps/engine/tests/test_rest_api.py @@ -3451,7 +3451,7 @@ def _test_api_v2_tasks_id_data_spec(self, user, spec, data, expected_compressed_type, expected_original_type, expected_image_sizes, - expected_storage_method=StorageMethodChoice.FILE_SYSTEM, + expected_storage_method=None, expected_uploaded_data_location=StorageChoice.LOCAL, dimension=DimensionType.DIM_2D, expected_task_creation_status_state='Finished', @@ -3466,6 +3466,12 @@ def _test_api_v2_tasks_id_data_spec(self, user, spec, data, if get_status_callback is None: get_status_callback = self._get_task_creation_status + if expected_storage_method is None: + if settings.MEDIA_CACHE_ALLOW_STATIC_CACHE: + expected_storage_method = StorageMethodChoice.FILE_SYSTEM + else: + expected_storage_method = StorageMethodChoice.CACHE + # create task response = self._create_task(user, spec) self.assertEqual(response.status_code, status.HTTP_201_CREATED) @@ -4025,7 +4031,7 @@ def _test_api_v2_tasks_id_data_create_can_use_chunked_local_video(self, user): image_sizes = self._share_image_sizes['test_rotated_90_video.mp4'] self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data, self.ChunkType.IMAGESET, - self.ChunkType.VIDEO, image_sizes, StorageMethodChoice.FILE_SYSTEM) + self.ChunkType.VIDEO, image_sizes, StorageMethodChoice.CACHE) def _test_api_v2_tasks_id_data_create_can_use_chunked_cached_local_video(self, user): task_spec = { @@ -4195,7 +4201,7 @@ def _test_api_v2_tasks_id_data_create_can_use_server_images_with_predefined_sort task_data = task_data_common.copy() task_data["use_cache"] = caching_enabled - if caching_enabled: + if caching_enabled or not settings.MEDIA_CACHE_ALLOW_STATIC_CACHE: storage_method = StorageMethodChoice.CACHE else: storage_method = StorageMethodChoice.FILE_SYSTEM @@ -4254,7 +4260,7 @@ def _test_api_v2_tasks_id_data_create_can_use_local_images_with_predefined_sorti sorting_method=SortingMethod.PREDEFINED) task_data_common["use_cache"] = caching_enabled - if caching_enabled: + if caching_enabled or not settings.MEDIA_CACHE_ALLOW_STATIC_CACHE: storage_method = StorageMethodChoice.CACHE else: storage_method = StorageMethodChoice.FILE_SYSTEM @@ -4315,7 +4321,7 @@ def _test_api_v2_tasks_id_data_create_can_use_server_archive_with_predefined_sor task_data = task_data_common.copy() task_data["use_cache"] = caching_enabled - if caching_enabled: + if caching_enabled or not settings.MEDIA_CACHE_ALLOW_STATIC_CACHE: storage_method = StorageMethodChoice.CACHE else: storage_method = StorageMethodChoice.FILE_SYSTEM @@ -4388,7 +4394,7 @@ def _test_api_v2_tasks_id_data_create_can_use_local_archive_with_predefined_sort sorting_method=SortingMethod.PREDEFINED) task_data["use_cache"] = caching_enabled - if caching_enabled: + if caching_enabled or not settings.MEDIA_CACHE_ALLOW_STATIC_CACHE: storage_method = StorageMethodChoice.CACHE else: storage_method = StorageMethodChoice.FILE_SYSTEM @@ -4566,7 +4572,7 @@ def _send_data_and_fail(*args, **kwargs): self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data, self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, - image_sizes, StorageMethodChoice.FILE_SYSTEM, StorageChoice.LOCAL, + image_sizes, expected_uploaded_data_location=StorageChoice.LOCAL, send_data_callback=_send_data) with self.subTest(current_function_name() + ' mismatching file sets - extra files'): @@ -4580,7 +4586,7 @@ def _send_data_and_fail(*args, **kwargs): with self.assertRaisesMessage(Exception, "(extra)"): self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data, self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, - image_sizes, StorageMethodChoice.FILE_SYSTEM, StorageChoice.LOCAL, + image_sizes, expected_uploaded_data_location=StorageChoice.LOCAL, send_data_callback=_send_data_and_fail) with self.subTest(current_function_name() + ' mismatching file sets - missing files'): @@ -4594,7 +4600,7 @@ def _send_data_and_fail(*args, **kwargs): with self.assertRaisesMessage(Exception, "(missing)"): self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data, self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, - image_sizes, StorageMethodChoice.FILE_SYSTEM, StorageChoice.LOCAL, + image_sizes, expected_uploaded_data_location=StorageChoice.LOCAL, send_data_callback=_send_data_and_fail) def _test_api_v2_tasks_id_data_create_can_use_server_rar(self, user): From 258c800d9101e03bd99a9160508a73317bbfefdc Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 16 Aug 2024 11:46:25 +0300 Subject: [PATCH 063/115] Update manifest validation logic, allow manifest input in any task data source --- cvat/apps/engine/task.py | 57 ++++++++----------------- cvat/apps/engine/tests/test_rest_api.py | 56 +++++++++++------------- 2 files changed, 43 insertions(+), 70 deletions(-) diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index 8024ed4c376..eb2da924509 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -340,48 +340,28 @@ def _validate_manifest( *, is_in_cloud: bool, db_cloud_storage: Optional[Any], - data_storage_method: str, - data_sorting_method: str, - isBackupRestore: bool, ) -> Optional[str]: - if manifests: - if len(manifests) != 1: - raise ValidationError('Only one manifest file can be attached to data') - manifest_file = manifests[0] - full_manifest_path = os.path.join(root_dir, manifests[0]) - - if is_in_cloud: - cloud_storage_instance = db_storage_to_storage_instance(db_cloud_storage) - # check that cloud storage manifest file exists and is up to date - if not os.path.exists(full_manifest_path) or \ - datetime.fromtimestamp(os.path.getmtime(full_manifest_path), tz=timezone.utc) \ - < cloud_storage_instance.get_file_last_modified(manifest_file): - cloud_storage_instance.download_file(manifest_file, full_manifest_path) - - if is_manifest(full_manifest_path): - if not ( - data_sorting_method == models.SortingMethod.PREDEFINED or - (settings.USE_CACHE and data_storage_method == models.StorageMethodChoice.CACHE) or - isBackupRestore or is_in_cloud - ): - cache_disabled_message = "" - if data_storage_method == models.StorageMethodChoice.CACHE and not settings.USE_CACHE: - cache_disabled_message = ( - "This server doesn't allow to use cache for data. " - "Please turn 'use cache' off and try to recreate the task" - ) - slogger.glob.warning(cache_disabled_message) + if not manifests: + return None - raise ValidationError( - "A manifest file can only be used with the 'use cache' option " - "or when 'sorting_method' is 'predefined'" + \ - (". " + cache_disabled_message if cache_disabled_message else "") - ) - return manifest_file + if len(manifests) != 1: + raise ValidationError('Only one manifest file can be attached to data') + manifest_file = manifests[0] + full_manifest_path = os.path.join(root_dir, manifests[0]) + + if is_in_cloud: + cloud_storage_instance = db_storage_to_storage_instance(db_cloud_storage) + # check that cloud storage manifest file exists and is up to date + if not os.path.exists(full_manifest_path) or ( + datetime.fromtimestamp(os.path.getmtime(full_manifest_path), tz=timezone.utc) \ + < cloud_storage_instance.get_file_last_modified(manifest_file) + ): + cloud_storage_instance.download_file(manifest_file, full_manifest_path) + if not is_manifest(full_manifest_path): raise ValidationError('Invalid manifest was uploaded') - return None + return manifest_file def _validate_scheme(url): ALLOWED_SCHEMES = ['http', 'https'] @@ -580,9 +560,6 @@ def _update_status(msg: str) -> None: manifest_root, is_in_cloud=is_data_in_cloud, db_cloud_storage=db_data.cloud_storage if is_data_in_cloud else None, - data_storage_method=db_data.storage_method, - data_sorting_method=data['sorting_method'], - isBackupRestore=isBackupRestore, ) manifest = None diff --git a/cvat/apps/engine/tests/test_rest_api.py b/cvat/apps/engine/tests/test_rest_api.py index 8c08948166a..36fd9ed00d2 100644 --- a/cvat/apps/engine/tests/test_rest_api.py +++ b/cvat/apps/engine/tests/test_rest_api.py @@ -4128,7 +4128,6 @@ def _test_api_v2_tasks_id_data_create_can_use_server_images_and_manifest(self, u task_data = { "image_quality": 70, - "use_cache": True } manifest_name = "images_manifest_sorted.jsonl" @@ -4139,37 +4138,34 @@ def _test_api_v2_tasks_id_data_create_can_use_server_images_and_manifest(self, u for i, fn in enumerate(images + [manifest_name]) }) - for copy_data in [True, False]: - with self.subTest(current_function_name(), copy=copy_data): + for use_cache in [True, False]: + task_data['use_cache'] = use_cache + + for copy_data in [True, False]: + with self.subTest(current_function_name(), copy=copy_data, use_cache=use_cache): + task_spec = task_spec_common.copy() + task_spec['name'] = task_spec['name'] + f' copy={copy_data}' + task_data_copy = task_data.copy() + task_data_copy['copy_data'] = copy_data + self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data_copy, + self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, + image_sizes, + expected_uploaded_data_location=( + StorageChoice.LOCAL if copy_data else StorageChoice.SHARE + ) + ) + + with self.subTest(current_function_name() + ' file order mismatch', use_cache=use_cache): task_spec = task_spec_common.copy() - task_spec['name'] = task_spec['name'] + f' copy={copy_data}' - task_data['copy_data'] = copy_data - self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data, + task_spec['name'] = task_spec['name'] + f' mismatching file order' + task_data_copy = task_data.copy() + task_data_copy[f'server_files[{len(images)}]'] = "images_manifest.jsonl" + self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data_copy, self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, - image_sizes, StorageMethodChoice.CACHE, - StorageChoice.LOCAL if copy_data else StorageChoice.SHARE) - - with self.subTest(current_function_name() + ' file order mismatch'): - task_spec = task_spec_common.copy() - task_spec['name'] = task_spec['name'] + f' mismatching file order' - task_data_copy = task_data.copy() - task_data_copy[f'server_files[{len(images)}]'] = "images_manifest.jsonl" - self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data_copy, - self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, - image_sizes, StorageMethodChoice.CACHE, StorageChoice.SHARE, - expected_task_creation_status_state='Failed', - expected_task_creation_status_reason='Incorrect file mapping to manifest content') - - with self.subTest(current_function_name() + ' without use cache'): - task_spec = task_spec_common.copy() - task_spec['name'] = task_spec['name'] + f' manifest without cache' - task_data_copy = task_data.copy() - task_data_copy['use_cache'] = False - self._test_api_v2_tasks_id_data_spec(user, task_spec, task_data_copy, - self.ChunkType.IMAGESET, self.ChunkType.IMAGESET, - image_sizes, StorageMethodChoice.CACHE, StorageChoice.SHARE, - expected_task_creation_status_state='Failed', - expected_task_creation_status_reason="A manifest file can only be used with the 'use cache' option") + image_sizes, + expected_uploaded_data_location=StorageChoice.SHARE, + expected_task_creation_status_state='Failed', + expected_task_creation_status_reason='Incorrect file mapping to manifest content') def _test_api_v2_tasks_id_data_create_can_use_server_images_with_predefined_sorting(self, user): task_spec = { From 5e89ef49074715706b2f96ba2b936fc103d1d547 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 16 Aug 2024 14:05:55 +0300 Subject: [PATCH 064/115] Add task chunk caching, refactor chunk building --- cvat/apps/engine/cache.py | 119 ++++++++++++++++++++--------- cvat/apps/engine/frame_provider.py | 93 ++++++++-------------- 2 files changed, 115 insertions(+), 97 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index ff1851b2f93..a16071d446e 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -15,7 +15,7 @@ from contextlib import ExitStack, closing from datetime import datetime, timezone from itertools import pairwise -from typing import Callable, Iterator, Optional, Sequence, Tuple, Type, Union +from typing import Any, Callable, Iterator, Optional, Sequence, Tuple, Type, Union import av import cv2 @@ -118,6 +118,29 @@ def get_segment_chunk( ), ) + def _make_task_chunk_key( + self, db_task: models.Task, chunk_number: int, *, quality: FrameQuality + ) -> str: + return f"task_{db_task.id}_{chunk_number}_{quality}" + + def get_task_chunk( + self, db_task: models.Task, chunk_number: int, *, quality: FrameQuality + ) -> Optional[DataWithMime]: + return self._get(key=self._make_task_chunk_key(db_task, chunk_number, quality=quality)) + + def get_or_set_task_chunk( + self, + db_task: models.Task, + chunk_number: int, + *, + quality: FrameQuality, + set_callback: Callable[[], DataWithMime], + ) -> DataWithMime: + return self._get_or_set_cache_item( + key=self._make_task_chunk_key(db_task, chunk_number, quality=quality), + create_callback=lambda: set_callback(db_task, chunk_number, quality=quality), + ) + def get_selective_job_chunk( self, db_job: models.Job, chunk_number: int, *, quality: FrameQuality ) -> DataWithMime: @@ -298,43 +321,12 @@ def prepare_range_segment_chunk( chunk_size * chunk_number : chunk_size * (chunk_number + 1) ] - writer_classes: dict[FrameQuality, Type[IChunkWriter]] = { - FrameQuality.COMPRESSED: ( - Mpeg4CompressedChunkWriter - if db_data.compressed_chunk_type == models.DataChoice.VIDEO - else ZipCompressedChunkWriter - ), - FrameQuality.ORIGINAL: ( - Mpeg4ChunkWriter - if db_data.original_chunk_type == models.DataChoice.VIDEO - else ZipChunkWriter - ), - } - - image_quality = 100 if quality == FrameQuality.ORIGINAL else db_data.image_quality - - mime_type = ( - "video/mp4" - if writer_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] - else "application/zip" - ) - - kwargs = {} - if db_segment.task.dimension == models.DimensionType.DIM_3D: - kwargs["dimension"] = models.DimensionType.DIM_3D - writer = writer_classes[quality](image_quality, **kwargs) - - buffer = io.BytesIO() with closing(self._read_raw_frames(db_task, frame_ids=chunk_frame_ids)) as frame_iter: - writer.save_as_chunk(frame_iter, buffer) - - buffer.seek(0) - return buffer, mime_type + return prepare_chunk(frame_iter, quality=quality, db_task=db_task) def prepare_masked_range_segment_chunk( self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality ) -> DataWithMime: - # TODO: try to refactor into 1 function with prepare_range_segment_chunk() db_task = db_segment.task db_data = db_task.data @@ -395,7 +387,7 @@ def prepare_masked_range_segment_chunk( ) buff.seek(0) - return buff, "application/zip" + return buff, get_chunk_mime_type_for_writer(writer) def _prepare_segment_preview(self, db_segment: models.Segment) -> DataWithMime: if db_segment.task.dimension == models.DimensionType.DIM_3D: @@ -405,7 +397,9 @@ def _prepare_segment_preview(self, db_segment: models.Segment) -> DataWithMime: ) else: from cvat.apps.engine.frame_provider import ( - FrameOutputType, SegmentFrameProvider, TaskFrameProvider + FrameOutputType, + SegmentFrameProvider, + TaskFrameProvider, ) task_frame_provider = TaskFrameProvider(db_segment.task) @@ -494,3 +488,58 @@ def prepare_preview_image(image: PIL.Image.Image) -> DataWithMime: output_buf = io.BytesIO() image.convert("RGB").save(output_buf, format="JPEG") return output_buf, PREVIEW_MIME + + +def prepare_chunk( + task_chunk_frames: Iterator[Tuple[Any, str, int]], + *, + quality: FrameQuality, + db_task: models.Task, +) -> DataWithMime: + # TODO: refactor all chunk building into another class + + db_data = db_task.data + + writer_classes: dict[FrameQuality, Type[IChunkWriter]] = { + FrameQuality.COMPRESSED: ( + Mpeg4CompressedChunkWriter + if db_data.compressed_chunk_type == models.DataChoice.VIDEO + else ZipCompressedChunkWriter + ), + FrameQuality.ORIGINAL: ( + Mpeg4ChunkWriter + if db_data.original_chunk_type == models.DataChoice.VIDEO + else ZipChunkWriter + ), + } + + writer_class = writer_classes[quality] + + image_quality = 100 if quality == FrameQuality.ORIGINAL else db_data.image_quality + + writer_kwargs = {} + if db_task.dimension == models.DimensionType.DIM_3D: + writer_kwargs["dimension"] = models.DimensionType.DIM_3D + merged_chunk_writer = writer_class(image_quality, **writer_kwargs) + + writer_kwargs = {} + if isinstance(merged_chunk_writer, ZipCompressedChunkWriter): + writer_kwargs = dict(compress_frames=False, zip_compress_level=1) + + buffer = io.BytesIO() + merged_chunk_writer.save_as_chunk(task_chunk_frames, buffer, **writer_kwargs) + + buffer.seek(0) + return buffer, get_chunk_mime_type_for_writer(writer_class) + + +def get_chunk_mime_type_for_writer(writer: Union[IChunkWriter, Type[IChunkWriter]]) -> str: + if isinstance(writer, IChunkWriter): + writer_class = type(writer) + + if issubclass(writer_class, ZipChunkWriter): + return "application/zip" + elif issubclass(writer_class, Mpeg4ChunkWriter): + return "video/mp4" + else: + assert False, f"Unknown chunk writer class {writer_class}" diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 8ce5be00f5a..28e6c396a76 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -22,17 +22,12 @@ from rest_framework.exceptions import ValidationError from cvat.apps.engine import models -from cvat.apps.engine.cache import DataWithMime, MediaCache +from cvat.apps.engine.cache import DataWithMime, MediaCache, prepare_chunk from cvat.apps.engine.media_extractors import ( FrameQuality, - IChunkWriter, IMediaReader, - Mpeg4ChunkWriter, - Mpeg4CompressedChunkWriter, RandomAccessIterator, VideoReader, - ZipChunkWriter, - ZipCompressedChunkWriter, ZipReader, ) from cvat.apps.engine.mime_types import mimetypes @@ -261,6 +256,12 @@ def get_chunk( chunk_number = self.validate_chunk_number(chunk_number) db_data = self._db_task.data + + cache = MediaCache() + cached_chunk = cache.get_task_chunk(self._db_task, chunk_number, quality=quality) + if cached_chunk: + return return_type(cached_chunk[0], cached_chunk[1]) + step = db_data.get_frame_step() task_chunk_start_frame = chunk_number * db_data.chunk_size task_chunk_stop_frame = (chunk_number + 1) * db_data.chunk_size - 1 @@ -283,6 +284,7 @@ def get_chunk( ) assert matching_segments + # Don't put this into set_callback to avoid data duplication in the cache if len(matching_segments) == 1 and task_chunk_frame_set == set( matching_segments[0].frame_set ): @@ -291,64 +293,31 @@ def get_chunk( segment_frame_provider.get_chunk_number(task_chunk_start_frame), quality=quality ) - # Create and return a joined / cleaned chunk - # TODO: refactor into another class, maybe optimize - - writer_classes: dict[FrameQuality, Type[IChunkWriter]] = { - FrameQuality.COMPRESSED: ( - Mpeg4CompressedChunkWriter - if db_data.compressed_chunk_type == models.DataChoice.VIDEO - else ZipCompressedChunkWriter - ), - FrameQuality.ORIGINAL: ( - Mpeg4ChunkWriter - if db_data.original_chunk_type == models.DataChoice.VIDEO - else ZipChunkWriter - ), - } - - writer_class = writer_classes[quality] - - image_quality = 100 if quality == FrameQuality.ORIGINAL else db_data.image_quality - - mime_type = ( - "video/mp4" - if writer_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] - else "application/zip" + def _set_callback() -> DataWithMime: + # Create and return a joined / cleaned chunk + task_chunk_frames = {} + for db_segment in matching_segments: + segment_frame_provider = SegmentFrameProvider(db_segment) + segment_frame_set = db_segment.frame_set + + for task_chunk_frame_id in sorted(task_chunk_frame_set): + if ( + task_chunk_frame_id not in segment_frame_set + or task_chunk_frame_id in task_chunk_frames + ): + continue + + frame, frame_name, _ = segment_frame_provider._get_raw_frame( + self.get_rel_frame_number(task_chunk_frame_id), quality=quality + ) + task_chunk_frames[task_chunk_frame_id] = (frame, frame_name, None) + + return prepare_chunk(task_chunk_frames.values(), quality=quality, db_task=self._db_task) + + buffer, mime_type = cache.get_or_set_task_chunk( + self._db_task, chunk_number, quality=quality, set_callback=_set_callback ) - task_chunk_frames = {} - for db_segment in matching_segments: - segment_frame_provider = SegmentFrameProvider(db_segment) - segment_frame_set = db_segment.frame_set - - for task_chunk_frame_id in sorted(task_chunk_frame_set): - if ( - task_chunk_frame_id not in segment_frame_set - or task_chunk_frame_id in task_chunk_frames - ): - continue - - frame, frame_name, _ = segment_frame_provider._get_raw_frame( - self.get_rel_frame_number(task_chunk_frame_id), quality=quality - ) - task_chunk_frames[task_chunk_frame_id] = (frame, frame_name, None) - - writer_kwargs = {} - if self._db_task.dimension == models.DimensionType.DIM_3D: - writer_kwargs["dimension"] = models.DimensionType.DIM_3D - merged_chunk_writer = writer_class(image_quality, **writer_kwargs) - - writer_kwargs = {} - if isinstance(merged_chunk_writer, ZipCompressedChunkWriter): - writer_kwargs = dict(compress_frames=False, zip_compress_level=1) - - buffer = io.BytesIO() - merged_chunk_writer.save_as_chunk(task_chunk_frames.values(), buffer, **writer_kwargs) - buffer.seek(0) - - # TODO: add caching in media cache for the resulting chunk - return return_type(data=buffer, mime=mime_type) def get_frame( From c5edcda519c20f07c3f62233e1ff903dd0302952 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 16 Aug 2024 14:24:49 +0300 Subject: [PATCH 065/115] Refactor some code --- cvat/apps/engine/cache.py | 175 ++++++++++++++++++++------------------ 1 file changed, 93 insertions(+), 82 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index a16071d446e..b2b0be9a151 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -15,7 +15,7 @@ from contextlib import ExitStack, closing from datetime import datetime, timezone from itertools import pairwise -from typing import Any, Callable, Iterator, Optional, Sequence, Tuple, Type, Union +from typing import Any, Callable, Generator, Iterator, Optional, Sequence, Tuple, Type, Union import av import cv2 @@ -172,9 +172,97 @@ def get_frame_context_images(self, db_data: models.Data, frame_number: int) -> D create_callback=lambda: self.prepare_context_images(db_data, frame_number), ) + def _read_raw_images( + self, + db_task: models.Task, + frame_ids: Sequence[int], + *, + raw_data_dir: str, + manifest_path: str, + ): + db_data = db_task.data + dimension = db_task.dimension + + if os.path.isfile(manifest_path) and db_data.storage == models.StorageChoice.CLOUD_STORAGE: + reader = ImageReaderWithManifest(manifest_path) + with ExitStack() as es: + db_cloud_storage = db_data.cloud_storage + assert db_cloud_storage, "Cloud storage instance was deleted" + credentials = Credentials() + credentials.convert_from_db( + { + "type": db_cloud_storage.credentials_type, + "value": db_cloud_storage.credentials, + } + ) + details = { + "resource": db_cloud_storage.resource, + "credentials": credentials, + "specific_attributes": db_cloud_storage.get_specific_attributes(), + } + cloud_storage_instance = get_cloud_storage_instance( + cloud_provider=db_cloud_storage.provider_type, **details + ) + + tmp_dir = es.enter_context(tempfile.TemporaryDirectory(prefix="cvat")) + files_to_download = [] + checksums = [] + media = [] + for item in reader.iterate_frames(frame_ids): + file_name = f"{item['name']}{item['extension']}" + fs_filename = os.path.join(tmp_dir, file_name) + + files_to_download.append(file_name) + checksums.append(item.get("checksum", None)) + media.append((fs_filename, fs_filename, None)) + + cloud_storage_instance.bulk_download_to_dir( + files=files_to_download, upload_dir=tmp_dir + ) + media = preload_images(media) + + for checksum, (_, fs_filename, _) in zip(checksums, media): + if checksum and not md5_hash(fs_filename) == checksum: + slogger.cloud_storage[db_cloud_storage.id].warning( + "Hash sums of files {} do not match".format(file_name) + ) + + yield from media + else: + requested_frame_iter = iter(frame_ids) + next_requested_frame_id = next(requested_frame_iter, None) + if next_requested_frame_id is None: + return + + # TODO: find a way to use prefetched results, if provided + db_images = ( + db_data.images.order_by("frame") + .filter(frame__gte=frame_ids[0], frame__lte=frame_ids[-1]) + .values_list("frame", "path") + .all() + ) + + media = [] + for frame_id, frame_path in db_images: + if frame_id == next_requested_frame_id: + source_path = os.path.join(raw_data_dir, frame_path) + media.append((source_path, source_path, None)) + + next_requested_frame_id = next(requested_frame_iter, None) + + if next_requested_frame_id is None: + break + + assert next_requested_frame_id is None + + if dimension == models.DimensionType.DIM_2D: + media = preload_images(media) + + yield from media + def _read_raw_frames( self, db_task: models.Task, frame_ids: Sequence[int] - ) -> Iterator[Tuple[Union[av.VideoFrame, PIL.Image.Image], str, str]]: + ) -> Generator[Tuple[Union[av.VideoFrame, PIL.Image.Image], str, str], None, None]: for prev_frame, cur_frame in pairwise(frame_ids): assert ( prev_frame <= cur_frame @@ -188,7 +276,6 @@ def _read_raw_frames( models.StorageChoice.CLOUD_STORAGE: db_data.get_upload_dirname(), }[db_data.storage] - dimension = db_task.dimension manifest_path = db_data.get_manifest_path() if hasattr(db_data, "video"): @@ -218,85 +305,9 @@ def _read_raw_frames( for frame_tuple in reader.iterate_frames(frame_filter=frame_ids): yield frame_tuple else: - if ( - os.path.isfile(manifest_path) - and db_data.storage == models.StorageChoice.CLOUD_STORAGE - ): - reader = ImageReaderWithManifest(manifest_path) - with ExitStack() as es: - db_cloud_storage = db_data.cloud_storage - assert db_cloud_storage, "Cloud storage instance was deleted" - credentials = Credentials() - credentials.convert_from_db( - { - "type": db_cloud_storage.credentials_type, - "value": db_cloud_storage.credentials, - } - ) - details = { - "resource": db_cloud_storage.resource, - "credentials": credentials, - "specific_attributes": db_cloud_storage.get_specific_attributes(), - } - cloud_storage_instance = get_cloud_storage_instance( - cloud_provider=db_cloud_storage.provider_type, **details - ) - - tmp_dir = es.enter_context(tempfile.TemporaryDirectory(prefix="cvat")) - files_to_download = [] - checksums = [] - media = [] - for item in reader.iterate_frames(frame_ids): - file_name = f"{item['name']}{item['extension']}" - fs_filename = os.path.join(tmp_dir, file_name) - - files_to_download.append(file_name) - checksums.append(item.get("checksum", None)) - media.append((fs_filename, fs_filename, None)) - - cloud_storage_instance.bulk_download_to_dir( - files=files_to_download, upload_dir=tmp_dir - ) - media = preload_images(media) - - for checksum, (_, fs_filename, _) in zip(checksums, media): - if checksum and not md5_hash(fs_filename) == checksum: - slogger.cloud_storage[db_cloud_storage.id].warning( - "Hash sums of files {} do not match".format(file_name) - ) - - yield from media - else: - requested_frame_iter = iter(frame_ids) - next_requested_frame_id = next(requested_frame_iter, None) - if next_requested_frame_id is None: - return - - # TODO: find a way to use prefetched results, if provided - db_images = ( - db_data.images.order_by("frame") - .filter(frame__gte=frame_ids[0], frame__lte=frame_ids[-1]) - .values_list("frame", "path") - .all() - ) - - media = [] - for frame_id, frame_path in db_images: - if frame_id == next_requested_frame_id: - source_path = os.path.join(raw_data_dir, frame_path) - media.append((source_path, source_path, None)) - - next_requested_frame_id = next(requested_frame_iter, None) - - if next_requested_frame_id is None: - break - - assert next_requested_frame_id is None - - if dimension == models.DimensionType.DIM_2D: - media = preload_images(media) - - yield from media + return self._read_raw_images( + db_task, frame_ids, raw_data_dir=raw_data_dir, manifest_path=manifest_path + ) def prepare_segment_chunk( self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality From 7f5c7229981a98bc5be7f8e61029c8a69052d132 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 16 Aug 2024 14:57:00 +0300 Subject: [PATCH 066/115] Refactor some code --- cvat/apps/engine/cache.py | 7 +++++-- cvat/apps/engine/frame_provider.py | 4 +++- cvat/apps/engine/media_extractors.py | 2 +- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index b2b0be9a151..a15dd538528 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -305,7 +305,7 @@ def _read_raw_frames( for frame_tuple in reader.iterate_frames(frame_filter=frame_ids): yield frame_tuple else: - return self._read_raw_images( + yield from self._read_raw_images( db_task, frame_ids, raw_data_dir=raw_data_dir, manifest_path=manifest_path ) @@ -506,6 +506,7 @@ def prepare_chunk( *, quality: FrameQuality, db_task: models.Task, + dump: bool = False, ) -> DataWithMime: # TODO: refactor all chunk building into another class @@ -534,7 +535,7 @@ def prepare_chunk( merged_chunk_writer = writer_class(image_quality, **writer_kwargs) writer_kwargs = {} - if isinstance(merged_chunk_writer, ZipCompressedChunkWriter): + if dump and isinstance(merged_chunk_writer, ZipCompressedChunkWriter): writer_kwargs = dict(compress_frames=False, zip_compress_level=1) buffer = io.BytesIO() @@ -547,6 +548,8 @@ def prepare_chunk( def get_chunk_mime_type_for_writer(writer: Union[IChunkWriter, Type[IChunkWriter]]) -> str: if isinstance(writer, IChunkWriter): writer_class = type(writer) + else: + writer_class = writer if issubclass(writer_class, ZipChunkWriter): return "application/zip" diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 28e6c396a76..c37489dda57 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -312,7 +312,9 @@ def _set_callback() -> DataWithMime: ) task_chunk_frames[task_chunk_frame_id] = (frame, frame_name, None) - return prepare_chunk(task_chunk_frames.values(), quality=quality, db_task=self._db_task) + return prepare_chunk( + task_chunk_frames.values(), quality=quality, db_task=self._db_task, dump=True + ) buffer, mime_type = cache.get_or_set_task_chunk( self._db_task, chunk_number, quality=quality, set_callback=_set_callback diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index afdeaba9efa..1d5c6fde76e 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -576,7 +576,7 @@ def __init__( source_path=source_path, step=step, start=start, - stop=stop if stop is not None else stop, + stop=stop, dimension=dimension, ) From daf4035e4a0456401e3433a38aa2e7d5a6e1e733 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 16 Aug 2024 14:59:42 +0300 Subject: [PATCH 067/115] Improve parameter name --- cvat/apps/engine/cache.py | 4 ++-- cvat/apps/engine/frame_provider.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index a15dd538528..02c4ec4c5fb 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -506,7 +506,7 @@ def prepare_chunk( *, quality: FrameQuality, db_task: models.Task, - dump: bool = False, + dump_unchanged: bool = False, ) -> DataWithMime: # TODO: refactor all chunk building into another class @@ -535,7 +535,7 @@ def prepare_chunk( merged_chunk_writer = writer_class(image_quality, **writer_kwargs) writer_kwargs = {} - if dump and isinstance(merged_chunk_writer, ZipCompressedChunkWriter): + if dump_unchanged and isinstance(merged_chunk_writer, ZipCompressedChunkWriter): writer_kwargs = dict(compress_frames=False, zip_compress_level=1) buffer = io.BytesIO() diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index c37489dda57..77c92b6a040 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -313,7 +313,7 @@ def _set_callback() -> DataWithMime: task_chunk_frames[task_chunk_frame_id] = (frame, frame_name, None) return prepare_chunk( - task_chunk_frames.values(), quality=quality, db_task=self._db_task, dump=True + task_chunk_frames.values(), quality=quality, db_task=self._db_task, dump_unchanged=True ) buffer, mime_type = cache.get_or_set_task_chunk( From 8c1b82c783931f3096af5a9c9cc85c3d4629f7b5 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 16 Aug 2024 15:05:13 +0300 Subject: [PATCH 068/115] Fix function call --- cvat/apps/engine/cache.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 02c4ec4c5fb..4995c046664 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -138,7 +138,7 @@ def get_or_set_task_chunk( ) -> DataWithMime: return self._get_or_set_cache_item( key=self._make_task_chunk_key(db_task, chunk_number, quality=quality), - create_callback=lambda: set_callback(db_task, chunk_number, quality=quality), + create_callback=set_callback, ) def get_selective_job_chunk( From f172865af2188a7a807a09a641adcf4fcb3e1921 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 16 Aug 2024 21:44:34 +0300 Subject: [PATCH 069/115] Add basic test set for meta, frames, and chunks reading in tasks --- cvat/apps/engine/frame_provider.py | 5 +- tests/python/rest_api/test_tasks.py | 294 ++++++++++++++++++++++++++- tests/python/shared/fixtures/init.py | 18 ++ tests/python/shared/utils/helpers.py | 24 ++- 4 files changed, 335 insertions(+), 6 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 77c92b6a040..11ea5539e29 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -313,7 +313,10 @@ def _set_callback() -> DataWithMime: task_chunk_frames[task_chunk_frame_id] = (frame, frame_name, None) return prepare_chunk( - task_chunk_frames.values(), quality=quality, db_task=self._db_task, dump_unchanged=True + task_chunk_frames.values(), + quality=quality, + db_task=self._db_task, + dump_unchanged=True, ) buffer, mime_type = cache.get_or_set_task_chunk( diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py index f8f8510fdc7..f8d0f6951d7 100644 --- a/tests/python/rest_api/test_tasks.py +++ b/tests/python/rest_api/test_tasks.py @@ -6,10 +6,14 @@ import io import itertools import json +import math import os import os.path as osp import zipfile +from abc import ABCMeta, abstractmethod +from contextlib import closing from copy import deepcopy +from enum import Enum from functools import partial from http import HTTPStatus from itertools import chain, product @@ -18,8 +22,10 @@ from pathlib import Path from tempfile import NamedTemporaryFile, TemporaryDirectory from time import sleep, time -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Tuple +import attrs +import numpy as np import pytest from cvat_sdk import Client, Config, exceptions from cvat_sdk.api_client import models @@ -30,6 +36,7 @@ from cvat_sdk.core.uploading import Uploader from deepdiff import DeepDiff from PIL import Image +from pytest_cases import fixture_ref, parametrize import shared.utils.s3 as s3 from shared.fixtures.init import docker_exec_cvat, kube_exec_cvat @@ -48,6 +55,7 @@ generate_image_files, generate_manifest, generate_video_file, + read_video_file, ) from .utils import ( @@ -1968,6 +1976,290 @@ def test_create_task_with_cloud_storage_directories_and_default_bucket_prefix( assert task.size == expected_task_size +class _SourceDataType(str, Enum): + images = "images" + video = "video" + + +@pytest.mark.usefixtures("restore_db_per_class") +@pytest.mark.usefixtures("restore_redis_ondisk_per_class") +@pytest.mark.usefixtures("restore_cvat_data") +class TestTaskData: + _USERNAME = "admin1" + + class _TaskSpec(models.ITaskWriteRequest, models.IDataRequest, metaclass=ABCMeta): + size: int + frame_step: int + source_data_type: _SourceDataType + + @abstractmethod + def read_frame(self, i: int) -> Image.Image: ... + + @attrs.define + class _TaskSpecBase(_TaskSpec): + _params: dict | models.TaskWriteRequest + _data_params: dict | models.DataRequest + size: int = attrs.field(kw_only=True) + + @property + def frame_step(self) -> int: + v = getattr(self, "frame_filter", "step=1") + return int(v.split("=")[-1]) + + def __getattr__(self, k: str) -> Any: + notfound = object() + + for params in [self._params, self._data_params]: + if isinstance(params, dict): + v = params.get(k, notfound) + else: + v = getattr(params, k, notfound) + + if v is not notfound: + return v + + raise AttributeError(k) + + @attrs.define + class _ImagesTaskSpec(_TaskSpecBase): + source_data_type: ClassVar[_SourceDataType] = _SourceDataType.images + + _get_frame: Callable[[int], bytes] = attrs.field(kw_only=True) + + def read_frame(self, i: int) -> Image.Image: + return Image.open(io.BytesIO(self._get_frame(i))) + + @attrs.define + class _VideoTaskSpec(_TaskSpecBase): + source_data_type: ClassVar[_SourceDataType] = _SourceDataType.video + + _get_video_file: Callable[[], io.IOBase] = attrs.field(kw_only=True) + + def read_frame(self, i: int) -> Image.Image: + with closing(read_video_file(self._get_video_file())) as reader: + for _ in range(i + 1): + frame = next(reader) + + return frame + + def _uploaded_images_task_fxt_base( + self, + request: pytest.FixtureRequest, + *, + frame_count: int = 10, + segment_size: Optional[int] = None, + ) -> Generator[tuple[_TaskSpec, int], None, None]: + task_params = { + "name": request.node.name, + "labels": [{"name": "a"}], + } + if segment_size: + task_params["segment_size"] = segment_size + + image_files = generate_image_files(frame_count) + images_data = [f.getvalue() for f in image_files] + data_params = { + "image_quality": 70, + "client_files": image_files, + } + + def get_frame(i: int) -> bytes: + return images_data[i] + + task_id, _ = create_task(self._USERNAME, spec=task_params, data=data_params) + yield self._ImagesTaskSpec( + models.TaskWriteRequest._from_openapi_data(**task_params), + models.DataRequest._from_openapi_data(**data_params), + get_frame=get_frame, + size=len(images_data), + ), task_id + + @pytest.fixture(scope="class") + def fxt_uploaded_images_task( + self, request: pytest.FixtureRequest + ) -> Generator[tuple[_TaskSpec, int], None, None]: + yield from self._uploaded_images_task_fxt_base(request=request) + + @pytest.fixture(scope="class") + def fxt_uploaded_images_task_with_segments( + self, request: pytest.FixtureRequest + ) -> Generator[tuple[_TaskSpec, int], None, None]: + yield from self._uploaded_images_task_fxt_base(request=request, segment_size=4) + + def _uploaded_video_task_fxt_base( + self, + request: pytest.FixtureRequest, + *, + frame_count: int = 10, + segment_size: Optional[int] = None, + ) -> Generator[tuple[_TaskSpec, int], None, None]: + task_params = { + "name": request.node.name, + "labels": [{"name": "a"}], + } + if segment_size: + task_params["segment_size"] = segment_size + + video_file = generate_video_file(frame_count) + video_data = video_file.getvalue() + data_params = { + "image_quality": 70, + "client_files": [video_file], + } + + def get_video_file() -> io.BytesIO: + return io.BytesIO(video_data) + + task_id, _ = create_task(self._USERNAME, spec=task_params, data=data_params) + yield self._VideoTaskSpec( + models.TaskWriteRequest._from_openapi_data(**task_params), + models.DataRequest._from_openapi_data(**data_params), + get_video_file=get_video_file, + size=frame_count, + ), task_id + + @pytest.fixture(scope="class") + def fxt_uploaded_video_task( + self, + request: pytest.FixtureRequest, + ) -> Generator[tuple[_TaskSpec, 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[_TaskSpec, int], None, None]: + yield from self._uploaded_video_task_fxt_base(request=request, segment_size=4) + + _default_task_cases = [ + fixture_ref("fxt_uploaded_images_task"), + fixture_ref("fxt_uploaded_images_task_with_segments"), + fixture_ref("fxt_uploaded_video_task"), + fixture_ref("fxt_uploaded_video_task_with_segments"), + ] + + @parametrize("task_spec, task_id", _default_task_cases) + def test_can_get_task_meta(self, task_spec: _TaskSpec, task_id: int): + with make_api_client(self._USERNAME) as api_client: + (task_meta, _) = api_client.tasks_api.retrieve_data_meta(task_id) + + assert task_meta.size == task_spec.size + + assert task_meta.start_frame == getattr(task_spec, "start_frame", 0) + + if getattr(task_spec, "stop_frame", None): + assert task_meta.stop_frame == task_spec.stop_frame + + assert task_meta.frame_filter == getattr(task_spec, "frame_filter", "") + + task_frame_set = set( + range(task_meta.start_frame, task_meta.stop_frame + 1, task_spec.frame_step) + ) + assert len(task_frame_set) == task_meta.size + + if getattr(task_spec, "chunk_size", None): + assert task_meta.chunk_size == task_spec.chunk_size + + if task_spec.source_data_type == _SourceDataType.video: + assert len(task_meta.frames) == 1 + else: + assert len(task_meta.frames) == task_meta.size + + @parametrize("task_spec, task_id", _default_task_cases) + def test_can_get_task_frames(self, task_spec: _TaskSpec, task_id: int): + with make_api_client(self._USERNAME) as api_client: + (task_meta, _) = api_client.tasks_api.retrieve_data_meta(task_id) + + for quality, frame_id in product(["original", "compressed"], range(task_meta.size)): + (_, response) = api_client.tasks_api.retrieve_data( + task_id, + type="frame", + quality=quality, + number=frame_id, + _parse_response=False, + ) + + if task_spec.source_data_type == _SourceDataType.video: + frame_size = (task_meta.frames[0].height, task_meta.frames[0].width) + else: + frame_size = ( + task_meta.frames[frame_id].height, + task_meta.frames[frame_id].width, + ) + + expected_pixels = np.array(task_spec.read_frame(frame_id)) + assert frame_size == expected_pixels.shape[:2] + + frame_pixels = np.array(Image.open(io.BytesIO(response.data))) + assert frame_size == frame_pixels.shape[:2] + + if ( + quality == "original" + and task_spec.source_data_type == _SourceDataType.images + # video chunks can have slightly changed colors, due to codec specifics + ): + assert np.array_equal(frame_pixels, expected_pixels) + + @parametrize("task_spec, task_id", _default_task_cases) + def test_can_get_task_chunks(self, task_spec: _TaskSpec, task_id: int): + with make_api_client(self._USERNAME) as api_client: + (task, _) = api_client.tasks_api.retrieve(task_id) + (task_meta, _) = api_client.tasks_api.retrieve_data_meta(task_id) + + if task_spec.source_data_type == _SourceDataType.images: + assert task.data_original_chunk_type == "imageset" + assert task.data_compressed_chunk_type == "imageset" + elif task_spec.source_data_type == _SourceDataType.video: + assert task.data_original_chunk_type == "video" + + if getattr(task_spec, "use_zip_chunks", False): + assert task.data_compressed_chunk_type == "imageset" + else: + assert task.data_compressed_chunk_type == "video" + else: + assert False + + chunk_count = math.ceil(task_meta.size / task_meta.chunk_size) + for quality, chunk_id in product(["original", "compressed"], range(chunk_count)): + expected_chunk_frame_ids = range( + chunk_id * task_meta.chunk_size, + min((chunk_id + 1) * task_meta.chunk_size, task_meta.size), + ) + + (_, response) = api_client.tasks_api.retrieve_data( + task_id, type="chunk", quality=quality, number=chunk_id, _parse_response=False + ) + + chunk_file = io.BytesIO(response.data) + if zipfile.is_zipfile(chunk_file): + with zipfile.ZipFile(chunk_file, "r") as chunk_archive: + chunk_images = { + int(os.path.splitext(name)[0]): np.array( + Image.open(io.BytesIO(chunk_archive.read(name))) + ) + for name in chunk_archive.namelist() + } + + else: + chunk_images = dict(enumerate(read_video_file(chunk_file))) + + assert sorted(chunk_images.keys()) == list( + v - chunk_id * task_meta.chunk_size for v in expected_chunk_frame_ids + ) + + for chunk_frame, frame_id in zip(chunk_images, expected_chunk_frame_ids): + expected_pixels = np.array(task_spec.read_frame(frame_id)) + chunk_frame_pixels = np.array(chunk_images[chunk_frame]) + assert expected_pixels.shape == chunk_frame_pixels.shape + + if ( + quality == "original" + and task_spec.source_data_type == _SourceDataType.images + # video chunks can have slightly changed colors, due to codec specifics + ): + assert np.array_equal(chunk_frame_pixels, expected_pixels) + + @pytest.mark.usefixtures("restore_db_per_function") class TestPatchTaskLabel: def _get_task_labels(self, pid, user, **kwargs) -> List[models.Label]: diff --git a/tests/python/shared/fixtures/init.py b/tests/python/shared/fixtures/init.py index 8e9d334f7a4..0b4a13f5ec9 100644 --- a/tests/python/shared/fixtures/init.py +++ b/tests/python/shared/fixtures/init.py @@ -592,6 +592,15 @@ def restore_redis_inmem_per_function(request): kube_restore_redis_inmem() +@pytest.fixture(scope="class") +def restore_redis_inmem_per_class(request): + platform = request.config.getoption("--platform") + if platform == "local": + docker_restore_redis_inmem() + else: + kube_restore_redis_inmem() + + @pytest.fixture(scope="function") def restore_redis_ondisk_per_function(request): platform = request.config.getoption("--platform") @@ -599,3 +608,12 @@ def restore_redis_ondisk_per_function(request): docker_restore_redis_ondisk() else: kube_restore_redis_ondisk() + + +@pytest.fixture(scope="class") +def restore_redis_ondisk_per_class(request): + platform = request.config.getoption("--platform") + if platform == "local": + docker_restore_redis_ondisk() + else: + kube_restore_redis_ondisk() diff --git a/tests/python/shared/utils/helpers.py b/tests/python/shared/utils/helpers.py index f336cb3f911..2200fd7f4b2 100644 --- a/tests/python/shared/utils/helpers.py +++ b/tests/python/shared/utils/helpers.py @@ -1,10 +1,11 @@ -# Copyright (C) 2022 CVAT.ai Corporation +# Copyright (C) 2022-2024 CVAT.ai Corporation # # SPDX-License-Identifier: MIT import subprocess +from contextlib import closing from io import BytesIO -from typing import List, Optional +from typing import Generator, List, Optional import av import av.video.reformatter @@ -13,7 +14,7 @@ from shared.fixtures.init import get_server_image_tag -def generate_image_file(filename="image.png", size=(50, 50), color=(0, 0, 0)): +def generate_image_file(filename="image.png", size=(100, 50), color=(0, 0, 0)): f = BytesIO() f.name = filename image = Image.new("RGB", size=size, color=color) @@ -40,7 +41,7 @@ def generate_image_files( return images -def generate_video_file(num_frames: int, size=(50, 50)) -> BytesIO: +def generate_video_file(num_frames: int, size=(100, 50)) -> BytesIO: f = BytesIO() f.name = "video.avi" @@ -60,6 +61,21 @@ def generate_video_file(num_frames: int, size=(50, 50)) -> BytesIO: return f +def read_video_file(file: BytesIO) -> Generator[Image.Image, None, None]: + file.seek(0) + + with av.open(file) as container: + video_stream = container.streams.video[0] + + with ( + closing(video_stream.codec_context), # pyav has a memory leak in stream.close() + closing(container.demux(video_stream)) as demux_iter, + ): + for packet in demux_iter: + for frame in packet.decode(): + yield frame.to_image() + + def generate_manifest(path: str) -> None: command = [ "docker", From aacceeed1661b9e9fab9cb38382a708140f5e5bb Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 16 Aug 2024 21:53:00 +0300 Subject: [PATCH 070/115] Move class declaration for pylint compatibility --- tests/python/rest_api/test_tasks.py | 98 +++++++++++++++-------------- 1 file changed, 51 insertions(+), 47 deletions(-) diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py index f8d0f6951d7..e91562ab6c9 100644 --- a/tests/python/rest_api/test_tasks.py +++ b/tests/python/rest_api/test_tasks.py @@ -1981,66 +1981,70 @@ class _SourceDataType(str, Enum): video = "video" -@pytest.mark.usefixtures("restore_db_per_class") -@pytest.mark.usefixtures("restore_redis_ondisk_per_class") -@pytest.mark.usefixtures("restore_cvat_data") -class TestTaskData: - _USERNAME = "admin1" +class _TaskSpec(models.ITaskWriteRequest, models.IDataRequest, metaclass=ABCMeta): + size: int + frame_step: int + source_data_type: _SourceDataType - class _TaskSpec(models.ITaskWriteRequest, models.IDataRequest, metaclass=ABCMeta): - size: int - frame_step: int - source_data_type: _SourceDataType + @abstractmethod + def read_frame(self, i: int) -> Image.Image: ... - @abstractmethod - def read_frame(self, i: int) -> Image.Image: ... - @attrs.define - class _TaskSpecBase(_TaskSpec): - _params: dict | models.TaskWriteRequest - _data_params: dict | models.DataRequest - size: int = attrs.field(kw_only=True) +@attrs.define +class _TaskSpecBase(_TaskSpec): + _params: dict | models.TaskWriteRequest + _data_params: dict | models.DataRequest + size: int = attrs.field(kw_only=True) - @property - def frame_step(self) -> int: - v = getattr(self, "frame_filter", "step=1") - return int(v.split("=")[-1]) + @property + def frame_step(self) -> int: + v = getattr(self, "frame_filter", "step=1") + return int(v.split("=")[-1]) - def __getattr__(self, k: str) -> Any: - notfound = object() + def __getattr__(self, k: str) -> Any: + notfound = object() - for params in [self._params, self._data_params]: - if isinstance(params, dict): - v = params.get(k, notfound) - else: - v = getattr(params, k, notfound) + for params in [self._params, self._data_params]: + if isinstance(params, dict): + v = params.get(k, notfound) + else: + v = getattr(params, k, notfound) + + if v is not notfound: + return v - if v is not notfound: - return v + raise AttributeError(k) - raise AttributeError(k) - @attrs.define - class _ImagesTaskSpec(_TaskSpecBase): - source_data_type: ClassVar[_SourceDataType] = _SourceDataType.images +@attrs.define +class _ImagesTaskSpec(_TaskSpecBase): + source_data_type: ClassVar[_SourceDataType] = _SourceDataType.images - _get_frame: Callable[[int], bytes] = attrs.field(kw_only=True) + _get_frame: Callable[[int], bytes] = attrs.field(kw_only=True) - def read_frame(self, i: int) -> Image.Image: - return Image.open(io.BytesIO(self._get_frame(i))) + def read_frame(self, i: int) -> Image.Image: + return Image.open(io.BytesIO(self._get_frame(i))) - @attrs.define - class _VideoTaskSpec(_TaskSpecBase): - source_data_type: ClassVar[_SourceDataType] = _SourceDataType.video - _get_video_file: Callable[[], io.IOBase] = attrs.field(kw_only=True) +@attrs.define +class _VideoTaskSpec(_TaskSpecBase): + source_data_type: ClassVar[_SourceDataType] = _SourceDataType.video - def read_frame(self, i: int) -> Image.Image: - with closing(read_video_file(self._get_video_file())) as reader: - for _ in range(i + 1): - frame = next(reader) + _get_video_file: Callable[[], io.IOBase] = attrs.field(kw_only=True) - return frame + def read_frame(self, i: int) -> Image.Image: + with closing(read_video_file(self._get_video_file())) as reader: + for _ in range(i + 1): + frame = next(reader) + + return frame + + +@pytest.mark.usefixtures("restore_db_per_class") +@pytest.mark.usefixtures("restore_redis_ondisk_per_class") +@pytest.mark.usefixtures("restore_cvat_data") +class TestTaskData: + _USERNAME = "admin1" def _uploaded_images_task_fxt_base( self, @@ -2067,7 +2071,7 @@ def get_frame(i: int) -> bytes: return images_data[i] task_id, _ = create_task(self._USERNAME, spec=task_params, data=data_params) - yield self._ImagesTaskSpec( + yield _ImagesTaskSpec( models.TaskWriteRequest._from_openapi_data(**task_params), models.DataRequest._from_openapi_data(**data_params), get_frame=get_frame, @@ -2111,7 +2115,7 @@ def get_video_file() -> io.BytesIO: return io.BytesIO(video_data) task_id, _ = create_task(self._USERNAME, spec=task_params, data=data_params) - yield self._VideoTaskSpec( + yield _VideoTaskSpec( models.TaskWriteRequest._from_openapi_data(**task_params), models.DataRequest._from_openapi_data(**data_params), get_video_file=get_video_file, From c8dbb7c938adf227f2d3678e055225c1e6b21fed Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Sat, 17 Aug 2024 00:46:27 +0300 Subject: [PATCH 071/115] Add missing original chunk type field in job responses --- cvat/apps/engine/serializers.py | 4 +++- cvat/schema.yml | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/cvat/apps/engine/serializers.py b/cvat/apps/engine/serializers.py index d0177140e45..0db2e45ada8 100644 --- a/cvat/apps/engine/serializers.py +++ b/cvat/apps/engine/serializers.py @@ -594,6 +594,7 @@ class JobReadSerializer(serializers.ModelSerializer): dimension = serializers.CharField(max_length=2, source='segment.task.dimension', read_only=True) data_chunk_size = serializers.ReadOnlyField(source='segment.task.data.chunk_size') organization = serializers.ReadOnlyField(source='segment.task.organization.id', allow_null=True) + data_original_chunk_type = serializers.ReadOnlyField(source='segment.task.data.original_chunk_type') data_compressed_chunk_type = serializers.ReadOnlyField(source='segment.task.data.compressed_chunk_type') mode = serializers.ReadOnlyField(source='segment.task.mode') bug_tracker = serializers.CharField(max_length=2000, source='get_bug_tracker', @@ -607,7 +608,8 @@ class Meta: model = models.Job fields = ('url', 'id', 'task_id', 'project_id', 'assignee', 'guide_id', 'dimension', 'bug_tracker', 'status', 'stage', 'state', 'mode', 'frame_count', - 'start_frame', 'stop_frame', 'data_chunk_size', 'data_compressed_chunk_type', + 'start_frame', 'stop_frame', + 'data_chunk_size', 'data_compressed_chunk_type', 'data_original_chunk_type', 'created_date', 'updated_date', 'issues', 'labels', 'type', 'organization', 'target_storage', 'source_storage', 'assignee_updated_date') read_only_fields = fields diff --git a/cvat/schema.yml b/cvat/schema.yml index d4f01b9642a..590dab4bc1e 100644 --- a/cvat/schema.yml +++ b/cvat/schema.yml @@ -8064,6 +8064,10 @@ components: allOf: - $ref: '#/components/schemas/ChunkType' readOnly: true + data_original_chunk_type: + allOf: + - $ref: '#/components/schemas/ChunkType' + readOnly: true created_date: type: string format: date-time From 6b9a3e9f664831f2c37f49f5c9ced86e07d534ae Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Sat, 17 Aug 2024 01:53:24 +0300 Subject: [PATCH 072/115] Add tests for job data access --- tests/python/rest_api/test_tasks.py | 250 ++++++++++++++++++++++++---- 1 file changed, 216 insertions(+), 34 deletions(-) diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py index e91562ab6c9..a4c247ba93b 100644 --- a/tests/python/rest_api/test_tasks.py +++ b/tests/python/rest_api/test_tasks.py @@ -22,7 +22,7 @@ from pathlib import Path from tempfile import NamedTemporaryFile, TemporaryDirectory from time import sleep, time -from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Tuple +from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Tuple, cast import attrs import numpy as np @@ -2135,6 +2135,46 @@ def fxt_uploaded_video_task_with_segments( ) -> Generator[tuple[_TaskSpec, int], None, None]: yield from self._uploaded_video_task_fxt_base(request=request, segment_size=4) + def _compute_segment_params(self, task_spec: _TaskSpec) -> list[tuple[int, int]]: + segment_params = [] + segment_size = getattr(task_spec, "segment_size", 0) or task_spec.size + start_frame = getattr(task_spec, "start_frame", 0) + end_frame = (getattr(task_spec, "stop_frame", None) or (task_spec.size - 1)) + 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 * task_spec.frame_step + + segment_end = segment_start + task_spec.frame_step * segment_size + + segment_params.append((segment_start, min(segment_end, end_frame) - 1)) + 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) + _default_task_cases = [ fixture_ref("fxt_uploaded_images_task"), fixture_ref("fxt_uploaded_images_task_with_segments"), @@ -2148,12 +2188,8 @@ def test_can_get_task_meta(self, task_spec: _TaskSpec, task_id: int): (task_meta, _) = api_client.tasks_api.retrieve_data_meta(task_id) assert task_meta.size == task_spec.size - assert task_meta.start_frame == getattr(task_spec, "start_frame", 0) - - if getattr(task_spec, "stop_frame", None): - assert task_meta.stop_frame == task_spec.stop_frame - + assert task_meta.stop_frame == getattr(task_spec, "stop_frame", None) or task_spec.size assert task_meta.frame_filter == getattr(task_spec, "frame_filter", "") task_frame_set = set( @@ -2174,35 +2210,40 @@ def test_can_get_task_frames(self, task_spec: _TaskSpec, task_id: int): with make_api_client(self._USERNAME) as api_client: (task_meta, _) = api_client.tasks_api.retrieve_data_meta(task_id) - for quality, frame_id in product(["original", "compressed"], range(task_meta.size)): + for quality, abs_frame_id in product( + ["original", "compressed"], + range(task_meta.start_frame, task_meta.stop_frame + 1, task_spec.frame_step), + ): + rel_frame_id = ( + abs_frame_id - getattr(task_spec, "start_frame", 0) // task_spec.frame_step + ) (_, response) = api_client.tasks_api.retrieve_data( task_id, type="frame", quality=quality, - number=frame_id, + number=rel_frame_id, _parse_response=False, ) if task_spec.source_data_type == _SourceDataType.video: - frame_size = (task_meta.frames[0].height, task_meta.frames[0].width) + frame_size = (task_meta.frames[0].width, task_meta.frames[0].height) else: frame_size = ( - task_meta.frames[frame_id].height, - task_meta.frames[frame_id].width, + task_meta.frames[rel_frame_id].width, + task_meta.frames[rel_frame_id].height, ) - expected_pixels = np.array(task_spec.read_frame(frame_id)) - assert frame_size == expected_pixels.shape[:2] - - frame_pixels = np.array(Image.open(io.BytesIO(response.data))) - assert frame_size == frame_pixels.shape[:2] + frame = Image.open(io.BytesIO(response.data)) + assert frame_size == frame.size - if ( - quality == "original" - and task_spec.source_data_type == _SourceDataType.images - # video chunks can have slightly changed colors, due to codec specifics - ): - assert np.array_equal(frame_pixels, expected_pixels) + self._compare_images( + task_spec.read_frame(abs_frame_id), + frame, + must_be_identical=( + task_spec.source_data_type == _SourceDataType.images + and quality == "original" + ), + ) @parametrize("task_spec, task_id", _default_task_cases) def test_can_get_task_chunks(self, task_spec: _TaskSpec, task_id: int): @@ -2243,25 +2284,166 @@ def test_can_get_task_chunks(self, task_spec: _TaskSpec, task_id: int): ) for name in chunk_archive.namelist() } - + chunk_images = dict(sorted(chunk_images.items(), key=lambda e: e[0])) else: chunk_images = dict(enumerate(read_video_file(chunk_file))) - assert sorted(chunk_images.keys()) == list( - v - chunk_id * task_meta.chunk_size for v in expected_chunk_frame_ids + assert sorted(chunk_images.keys()) == list(range(len(expected_chunk_frame_ids))) + + for chunk_frame, abs_frame_id in zip(chunk_images, expected_chunk_frame_ids): + self._compare_images( + task_spec.read_frame(abs_frame_id), + chunk_images[chunk_frame], + must_be_identical=( + task_spec.source_data_type == _SourceDataType.images + and quality == "original" + ), + ) + + @parametrize("task_spec, task_id", _default_task_cases) + def test_can_get_job_meta(self, task_spec: _TaskSpec, task_id: int): + segment_params = self._compute_segment_params(task_spec) + with make_api_client(self._USERNAME) as api_client: + jobs = sorted( + get_paginated_collection(api_client.jobs_api.list_endpoint, task_id=task_id), + key=lambda j: j.start_frame, + ) + assert len(jobs) == len(segment_params) + + for (segment_start, segment_end), job in zip(segment_params, jobs): + (job_meta, _) = api_client.jobs_api.retrieve_data_meta(job.id) + + assert (job_meta.start_frame, job_meta.stop_frame) == (segment_start, segment_end) + assert job_meta.frame_filter == getattr(task_spec, "frame_filter", "") + + segment_size = segment_end - segment_start + 1 + assert job_meta.size == segment_size + + task_frame_set = set( + range(job_meta.start_frame, job_meta.stop_frame + 1, task_spec.frame_step) ) + assert len(task_frame_set) == job_meta.size + + if getattr(task_spec, "chunk_size", None): + assert job_meta.chunk_size == task_spec.chunk_size + + if task_spec.source_data_type == _SourceDataType.video: + assert len(job_meta.frames) == 1 + else: + assert len(job_meta.frames) == job_meta.size + + @parametrize("task_spec, task_id", _default_task_cases) + def test_can_get_job_frames(self, task_spec: _TaskSpec, task_id: int): + with make_api_client(self._USERNAME) as api_client: + jobs = sorted( + get_paginated_collection(api_client.jobs_api.list_endpoint, task_id=task_id), + key=lambda j: j.start_frame, + ) + for job in jobs: + (job_meta, _) = api_client.jobs_api.retrieve_data_meta(job.id) - for chunk_frame, frame_id in zip(chunk_images, expected_chunk_frame_ids): - expected_pixels = np.array(task_spec.read_frame(frame_id)) - chunk_frame_pixels = np.array(chunk_images[chunk_frame]) - assert expected_pixels.shape == chunk_frame_pixels.shape + for quality, (frame_pos, abs_frame_id) in product( + ["original", "compressed"], + enumerate(range(job_meta.start_frame, job_meta.stop_frame)), + ): + rel_frame_id = ( + abs_frame_id - getattr(task_spec, "start_frame", 0) // task_spec.frame_step + ) + (_, response) = api_client.jobs_api.retrieve_data( + job.id, + type="frame", + quality=quality, + number=rel_frame_id, + _parse_response=False, + ) - if ( - quality == "original" - and task_spec.source_data_type == _SourceDataType.images - # video chunks can have slightly changed colors, due to codec specifics + if task_spec.source_data_type == _SourceDataType.video: + frame_size = (job_meta.frames[0].width, job_meta.frames[0].height) + else: + frame_size = ( + job_meta.frames[frame_pos].width, + job_meta.frames[frame_pos].height, + ) + + frame = Image.open(io.BytesIO(response.data)) + assert frame_size == frame.size + + self._compare_images( + task_spec.read_frame(abs_frame_id), + frame, + must_be_identical=( + task_spec.source_data_type == _SourceDataType.images + and quality == "original" + ), + ) + + @parametrize("task_spec, task_id", _default_task_cases) + def test_can_get_job_chunks(self, task_spec: _TaskSpec, task_id: int): + with make_api_client(self._USERNAME) as api_client: + jobs = sorted( + get_paginated_collection(api_client.jobs_api.list_endpoint, task_id=task_id), + key=lambda j: j.start_frame, + ) + for job in jobs: + (job_meta, _) = api_client.jobs_api.retrieve_data_meta(job.id) + + if task_spec.source_data_type == _SourceDataType.images: + assert job.data_original_chunk_type == "imageset" + assert job.data_compressed_chunk_type == "imageset" + elif task_spec.source_data_type == _SourceDataType.video: + assert job.data_original_chunk_type == "video" + + if getattr(task_spec, "use_zip_chunks", False): + assert job.data_compressed_chunk_type == "imageset" + else: + assert job.data_compressed_chunk_type == "video" + else: + assert False + + chunk_count = math.ceil(job_meta.size / job_meta.chunk_size) + for quality, chunk_id in product(["original", "compressed"], range(chunk_count)): + expected_chunk_abs_frame_ids = range( + job_meta.start_frame + + chunk_id * job_meta.chunk_size * task_spec.frame_step, + job_meta.start_frame + + min((chunk_id + 1) * job_meta.chunk_size, job_meta.size) + * task_spec.frame_step, + ) + + (_, response) = api_client.jobs_api.retrieve_data( + job.id, + type="chunk", + quality=quality, + number=chunk_id, + _parse_response=False, + ) + + chunk_file = io.BytesIO(response.data) + if zipfile.is_zipfile(chunk_file): + with zipfile.ZipFile(chunk_file, "r") as chunk_archive: + chunk_images = { + int(os.path.splitext(name)[0]): np.array( + Image.open(io.BytesIO(chunk_archive.read(name))) + ) + for name in chunk_archive.namelist() + } + chunk_images = dict(sorted(chunk_images.items(), key=lambda e: e[0])) + else: + chunk_images = dict(enumerate(read_video_file(chunk_file))) + + assert sorted(chunk_images.keys()) == list(range(job_meta.size)) + + for chunk_frame, abs_frame_id in zip( + chunk_images, expected_chunk_abs_frame_ids ): - assert np.array_equal(chunk_frame_pixels, expected_pixels) + self._compare_images( + task_spec.read_frame(abs_frame_id), + chunk_images[chunk_frame], + must_be_identical=( + task_spec.source_data_type == _SourceDataType.images + and quality == "original" + ), + ) @pytest.mark.usefixtures("restore_db_per_function") From f5661e4bfd4e95634237a4d596ad2d02235a5563 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Sat, 17 Aug 2024 01:57:54 +0300 Subject: [PATCH 073/115] Update test assets --- tests/python/shared/assets/jobs.json | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/python/shared/assets/jobs.json b/tests/python/shared/assets/jobs.json index 82dd91cc0d9..9d26222e7ac 100644 --- a/tests/python/shared/assets/jobs.json +++ b/tests/python/shared/assets/jobs.json @@ -10,6 +10,7 @@ "created_date": "2024-07-15T15:34:53.594000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 1, "guide_id": null, @@ -51,6 +52,7 @@ "created_date": "2024-07-15T15:33:10.549000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 1, "guide_id": null, @@ -92,6 +94,7 @@ "created_date": "2024-03-21T20:50:05.838000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 3, "guide_id": null, @@ -125,6 +128,7 @@ "created_date": "2024-03-21T20:50:05.815000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 1, "guide_id": null, @@ -158,6 +162,7 @@ "created_date": "2024-03-21T20:50:05.811000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 5, "guide_id": null, @@ -191,6 +196,7 @@ "created_date": "2024-03-21T20:50:05.805000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 5, "guide_id": null, @@ -224,6 +230,7 @@ "created_date": "2023-05-26T16:11:23.946000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 3, "guide_id": null, @@ -257,6 +264,7 @@ "created_date": "2023-05-26T16:11:23.880000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 11, "guide_id": null, @@ -290,6 +298,7 @@ "created_date": "2023-03-27T19:08:07.649000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 4, "guide_id": null, @@ -331,6 +340,7 @@ "created_date": "2023-03-27T19:08:07.649000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 6, "guide_id": null, @@ -372,6 +382,7 @@ "created_date": "2023-03-10T11:57:31.614000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 2, "guide_id": null, @@ -413,6 +424,7 @@ "created_date": "2023-03-10T11:56:33.757000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 2, "guide_id": null, @@ -454,6 +466,7 @@ "created_date": "2023-03-01T15:36:26.668000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 2, "guide_id": null, @@ -495,6 +508,7 @@ "created_date": "2023-02-10T14:05:25.947000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 5, "guide_id": null, @@ -528,6 +542,7 @@ "created_date": "2022-12-01T12:53:10.425000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "video", "dimension": "2d", "frame_count": 25, "guide_id": null, @@ -569,6 +584,7 @@ "created_date": "2022-09-22T14:22:25.820000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 8, "guide_id": null, @@ -610,6 +626,7 @@ "created_date": "2022-06-08T08:33:06.505000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 5, "guide_id": null, @@ -649,6 +666,7 @@ "created_date": "2022-03-05T10:32:19.149000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 11, "guide_id": null, @@ -690,6 +708,7 @@ "created_date": "2022-03-05T09:33:10.420000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 5, "guide_id": null, @@ -723,6 +742,7 @@ "created_date": "2022-03-05T09:33:10.420000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 5, "guide_id": null, @@ -756,6 +776,7 @@ "created_date": "2022-03-05T09:33:10.420000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 5, "guide_id": null, @@ -795,6 +816,7 @@ "created_date": "2022-03-05T09:33:10.420000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 5, "guide_id": null, @@ -834,6 +856,7 @@ "created_date": "2022-03-05T08:30:48.612000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 14, "guide_id": null, @@ -867,6 +890,7 @@ "created_date": "2022-02-21T10:31:52.429000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 11, "guide_id": null, @@ -900,6 +924,7 @@ "created_date": "2022-02-16T06:26:54.631000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "3d", "frame_count": 1, "guide_id": null, @@ -939,6 +964,7 @@ "created_date": "2022-02-16T06:25:48.168000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "video", "dimension": "2d", "frame_count": 25, "guide_id": null, @@ -978,6 +1004,7 @@ "created_date": "2021-12-14T18:50:29.458000Z", "data_chunk_size": 72, "data_compressed_chunk_type": "imageset", + "data_original_chunk_type": "imageset", "dimension": "2d", "frame_count": 23, "guide_id": null, From 754757fd3cde40d389f0466b6fcb236b900d547d Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Sat, 17 Aug 2024 02:00:06 +0300 Subject: [PATCH 074/115] Clean imports --- tests/python/rest_api/test_tasks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py index a4c247ba93b..82a0c1c9bb5 100644 --- a/tests/python/rest_api/test_tasks.py +++ b/tests/python/rest_api/test_tasks.py @@ -22,7 +22,7 @@ from pathlib import Path from tempfile import NamedTemporaryFile, TemporaryDirectory from time import sleep, time -from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Tuple, cast +from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Tuple import attrs import numpy as np From 0c001a52a0e1bf3663d6c7f8f796df5b7c47baa3 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Sat, 17 Aug 2024 02:30:13 +0300 Subject: [PATCH 075/115] Python 3.8 compatibility --- tests/python/shared/utils/helpers.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/python/shared/utils/helpers.py b/tests/python/shared/utils/helpers.py index 2200fd7f4b2..ac5948182d7 100644 --- a/tests/python/shared/utils/helpers.py +++ b/tests/python/shared/utils/helpers.py @@ -67,13 +67,11 @@ def read_video_file(file: BytesIO) -> Generator[Image.Image, None, None]: with av.open(file) as container: video_stream = container.streams.video[0] - with ( - closing(video_stream.codec_context), # pyav has a memory leak in stream.close() - closing(container.demux(video_stream)) as demux_iter, - ): - for packet in demux_iter: - for frame in packet.decode(): - yield frame.to_image() + with closing(video_stream.codec_context): # pyav has a memory leak in stream.close() + with closing(container.demux(video_stream)) as demux_iter: + for packet in demux_iter: + for frame in packet.decode(): + yield frame.to_image() def generate_manifest(path: str) -> None: From a9390eb67425b7610e00b0898d7558847e6a4bda Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Sat, 17 Aug 2024 12:24:33 +0300 Subject: [PATCH 076/115] Python 3.8 compatibility --- tests/python/rest_api/test_tasks.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py index 82a0c1c9bb5..9e73b5aaacc 100644 --- a/tests/python/rest_api/test_tasks.py +++ b/tests/python/rest_api/test_tasks.py @@ -22,7 +22,7 @@ from pathlib import Path from tempfile import NamedTemporaryFile, TemporaryDirectory from time import sleep, time -from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Tuple +from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Tuple, Union import attrs import numpy as np @@ -1992,8 +1992,8 @@ def read_frame(self, i: int) -> Image.Image: ... @attrs.define class _TaskSpecBase(_TaskSpec): - _params: dict | models.TaskWriteRequest - _data_params: dict | models.DataRequest + _params: Union[dict, models.TaskWriteRequest] + _data_params: Union[dict, models.DataRequest] size: int = attrs.field(kw_only=True) @property From d2b138564f080371d62a96b27e401a9c8cc7fb69 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Sat, 17 Aug 2024 13:08:33 +0300 Subject: [PATCH 077/115] Python 3.8 compatibility --- tests/python/rest_api/test_tasks.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py index 9e73b5aaacc..c66cdd2ed60 100644 --- a/tests/python/rest_api/test_tasks.py +++ b/tests/python/rest_api/test_tasks.py @@ -1992,8 +1992,8 @@ def read_frame(self, i: int) -> Image.Image: ... @attrs.define class _TaskSpecBase(_TaskSpec): - _params: Union[dict, models.TaskWriteRequest] - _data_params: Union[dict, models.DataRequest] + _params: Union[Dict, models.TaskWriteRequest] + _data_params: Union[Dict, models.DataRequest] size: int = attrs.field(kw_only=True) @property @@ -2052,7 +2052,7 @@ def _uploaded_images_task_fxt_base( *, frame_count: int = 10, segment_size: Optional[int] = None, - ) -> Generator[tuple[_TaskSpec, int], None, None]: + ) -> Generator[Tuple[_TaskSpec, int], None, None]: task_params = { "name": request.node.name, "labels": [{"name": "a"}], @@ -2081,13 +2081,13 @@ def get_frame(i: int) -> bytes: @pytest.fixture(scope="class") def fxt_uploaded_images_task( self, request: pytest.FixtureRequest - ) -> Generator[tuple[_TaskSpec, int], None, None]: + ) -> Generator[Tuple[_TaskSpec, int], None, None]: yield from self._uploaded_images_task_fxt_base(request=request) @pytest.fixture(scope="class") def fxt_uploaded_images_task_with_segments( self, request: pytest.FixtureRequest - ) -> Generator[tuple[_TaskSpec, int], None, None]: + ) -> Generator[Tuple[_TaskSpec, int], None, None]: yield from self._uploaded_images_task_fxt_base(request=request, segment_size=4) def _uploaded_video_task_fxt_base( @@ -2096,7 +2096,7 @@ def _uploaded_video_task_fxt_base( *, frame_count: int = 10, segment_size: Optional[int] = None, - ) -> Generator[tuple[_TaskSpec, int], None, None]: + ) -> Generator[Tuple[_TaskSpec, int], None, None]: task_params = { "name": request.node.name, "labels": [{"name": "a"}], @@ -2126,16 +2126,16 @@ def get_video_file() -> io.BytesIO: def fxt_uploaded_video_task( self, request: pytest.FixtureRequest, - ) -> Generator[tuple[_TaskSpec, int], None, None]: + ) -> Generator[Tuple[_TaskSpec, 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[_TaskSpec, int], None, None]: + ) -> Generator[Tuple[_TaskSpec, int], None, None]: yield from self._uploaded_video_task_fxt_base(request=request, segment_size=4) - def _compute_segment_params(self, task_spec: _TaskSpec) -> list[tuple[int, int]]: + def _compute_segment_params(self, task_spec: _TaskSpec) -> List[Tuple[int, int]]: segment_params = [] segment_size = getattr(task_spec, "segment_size", 0) or task_spec.size start_frame = getattr(task_spec, "start_frame", 0) From 621afa7f4737842dccc2bf4ff19891c380bc7187 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 19 Aug 2024 17:10:30 +0300 Subject: [PATCH 078/115] Add logging into shell command runs, fix invalid redis-cli invocation for k8s deployment, add failures on redis-cli command failures --- tests/python/shared/fixtures/init.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/python/shared/fixtures/init.py b/tests/python/shared/fixtures/init.py index 0b4a13f5ec9..45fa543eb92 100644 --- a/tests/python/shared/fixtures/init.py +++ b/tests/python/shared/fixtures/init.py @@ -96,12 +96,20 @@ def pytest_addoption(parser): def _run(command, capture_output=True): _command = command.split() if isinstance(command, str) else command try: + logger.debug(f"Executing a command: {_command}") + stdout, stderr = "", "" if capture_output: proc = run(_command, check=True, stdout=PIPE, stderr=PIPE) # nosec stdout, stderr = proc.stdout.decode(), proc.stderr.decode() else: proc = run(_command) # nosec + + if stdout: + logger.debug(f"Output (stdout): {stdout}") + if stderr: + logger.debug(f"Output (stderr): {stderr}") + return stdout, stderr except CalledProcessError as exc: message = f"Command failed: {' '.join(map(shlex.quote, _command))}." @@ -232,20 +240,20 @@ def kube_restore_clickhouse_db(): def docker_restore_redis_inmem(): - docker_exec_redis_inmem(["redis-cli", "flushall"]) + docker_exec_redis_inmem(["redis-cli", "-e", "flushall"]) def kube_restore_redis_inmem(): - kube_exec_redis_inmem(["redis-cli", "flushall"]) + kube_exec_redis_inmem(["redis-cli", "-e", "flushall"]) def docker_restore_redis_ondisk(): - docker_exec_redis_ondisk(["redis-cli", "-p", "6666", "flushall"]) + docker_exec_redis_ondisk(["redis-cli", "-e", "-p", "6666", "flushall"]) def kube_restore_redis_ondisk(): kube_exec_redis_ondisk( - ["redis-cli", "-p", "6666", "-a", "${CVAT_REDIS_ONDISK_PASSWORD}", "flushall"] + ["sh", "-c", 'redis-cli -e -p 6666 -a "${CVAT_REDIS_ONDISK_PASSWORD}" flushall'] ) From 441d0e75ca2a640d8dedba440a5accdfc5c4680e Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 21 Aug 2024 16:32:48 +0300 Subject: [PATCH 079/115] Allow calling flushall in redis in helm tests --- helm-chart/test.values.yaml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/helm-chart/test.values.yaml b/helm-chart/test.values.yaml index 5a5fa8fe6ba..24802a0cc54 100644 --- a/helm-chart/test.values.yaml +++ b/helm-chart/test.values.yaml @@ -27,6 +27,12 @@ cvat: frontend: imagePullPolicy: Never +redis: + master: + # The "flushall" command, which we use in tests, is disabled in helm by default + # https://github.com/helm/charts/tree/master/stable/redis#parameters + disableCommands: [] + keydb: resources: requests: From 0963f9490451f7949ff27a63ecc7a006e705f23f Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 21 Aug 2024 17:19:50 +0300 Subject: [PATCH 080/115] Update comment --- helm-chart/test.values.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/helm-chart/test.values.yaml b/helm-chart/test.values.yaml index 24802a0cc54..73edaa815d7 100644 --- a/helm-chart/test.values.yaml +++ b/helm-chart/test.values.yaml @@ -30,7 +30,7 @@ cvat: redis: master: # The "flushall" command, which we use in tests, is disabled in helm by default - # https://github.com/helm/charts/tree/master/stable/redis#parameters + # https://artifacthub.io/packages/helm/bitnami/redis#redis-master-configuration-parameters disableCommands: [] keydb: From 0d78e6380b688707f23c40cc14b0d68a0ce46acc Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 21 Aug 2024 17:20:08 +0300 Subject: [PATCH 081/115] Update redis cleanup command --- tests/python/shared/fixtures/init.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/shared/fixtures/init.py b/tests/python/shared/fixtures/init.py index 45fa543eb92..cf5aeabbbf3 100644 --- a/tests/python/shared/fixtures/init.py +++ b/tests/python/shared/fixtures/init.py @@ -244,7 +244,7 @@ def docker_restore_redis_inmem(): def kube_restore_redis_inmem(): - kube_exec_redis_inmem(["redis-cli", "-e", "flushall"]) + kube_exec_redis_inmem(["sh", "-c", 'redis-cli -e -a "${REDIS_PASSWORD}" flushall']) def docker_restore_redis_ondisk(): From e4db8ad8af0b6432664557e18eee9f68314a8fb9 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 28 Aug 2024 17:59:02 +0300 Subject: [PATCH 082/115] Reuse _get --- cvat/apps/engine/cache.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 4995c046664..2b4fa5de63d 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -77,14 +77,7 @@ def create_item() -> _CacheItem: return item - slogger.glob.info(f"Starting to get chunk from cache: key {key}") - try: - item = self._cache.get(key) - except pickle.UnpicklingError: - slogger.glob.error(f"Unable to get item from cache: key {key}", exc_info=True) - item = None - slogger.glob.info(f"Ending to get chunk from cache: key {key}, is_cached {bool(item)}") - + item = self._get(key) if not item: item = create_item() else: From b1c54f951856aba3bb624983cef9dd20c3af3f5e Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 28 Aug 2024 17:59:31 +0300 Subject: [PATCH 083/115] Make get_checksum private --- cvat/apps/engine/cache.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 2b4fa5de63d..b856eebcda3 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -58,7 +58,7 @@ class MediaCache: def __init__(self) -> None: self._cache = caches["media"] - def get_checksum(self, value: bytes) -> int: + def _get_checksum(self, value: bytes) -> int: return zlib.crc32(value) def _get_or_set_cache_item( @@ -70,7 +70,7 @@ def create_item() -> _CacheItem: slogger.glob.info(f"Ending to prepare chunk: key {key}") if item_data[0]: - item = (item_data[0], item_data[1], self.get_checksum(item_data[0].getbuffer())) + item = (item_data[0], item_data[1], self._get_checksum(item_data[0].getbuffer())) self._cache.set(key, item) else: item = (item_data[0], item_data[1], None) @@ -84,7 +84,7 @@ def create_item() -> _CacheItem: # compare checksum item_data = item[0].getbuffer() if isinstance(item[0], io.BytesIO) else item[0] item_checksum = item[2] if len(item) == 3 else None - if item_checksum != self.get_checksum(item_data): + if item_checksum != self._get_checksum(item_data): slogger.glob.info(f"Recreating cache item {key} due to checksum mismatch") item = create_item() From 5312b0005a197f5858cc035eb35fd64448462d9c Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 28 Aug 2024 18:00:06 +0300 Subject: [PATCH 084/115] Add get_raw_data_dirname to the Data model --- cvat/apps/engine/cache.py | 17 ++++------------- cvat/apps/engine/models.py | 7 +++++++ 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index b856eebcda3..62301b6253c 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -170,11 +170,9 @@ def _read_raw_images( db_task: models.Task, frame_ids: Sequence[int], *, - raw_data_dir: str, manifest_path: str, ): db_data = db_task.data - dimension = db_task.dimension if os.path.isfile(manifest_path) and db_data.storage == models.StorageChoice.CLOUD_STORAGE: reader = ImageReaderWithManifest(manifest_path) @@ -235,6 +233,7 @@ def _read_raw_images( .all() ) + raw_data_dir = db_data.get_raw_data_dirname() media = [] for frame_id, frame_path in db_images: if frame_id == next_requested_frame_id: @@ -248,7 +247,7 @@ def _read_raw_images( assert next_requested_frame_id is None - if dimension == models.DimensionType.DIM_2D: + if db_task.dimension == models.DimensionType.DIM_2D: media = preload_images(media) yield from media @@ -263,16 +262,10 @@ def _read_raw_frames( db_data = db_task.data - raw_data_dir = { - models.StorageChoice.LOCAL: db_data.get_upload_dirname(), - models.StorageChoice.SHARE: settings.SHARE_ROOT, - models.StorageChoice.CLOUD_STORAGE: db_data.get_upload_dirname(), - }[db_data.storage] - manifest_path = db_data.get_manifest_path() if hasattr(db_data, "video"): - source_path = os.path.join(raw_data_dir, db_data.video.path) + source_path = os.path.join(db_data.get_raw_data_dirname(), db_data.video.path) reader = VideoReaderWithManifest( manifest_path=manifest_path, @@ -298,9 +291,7 @@ def _read_raw_frames( for frame_tuple in reader.iterate_frames(frame_filter=frame_ids): yield frame_tuple else: - yield from self._read_raw_images( - db_task, frame_ids, raw_data_dir=raw_data_dir, manifest_path=manifest_path - ) + yield from self._read_raw_images(db_task, frame_ids, manifest_path=manifest_path) def prepare_segment_chunk( self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality diff --git a/cvat/apps/engine/models.py b/cvat/apps/engine/models.py index b1f01ab64a7..c88c9360bdf 100644 --- a/cvat/apps/engine/models.py +++ b/cvat/apps/engine/models.py @@ -252,6 +252,13 @@ def get_data_dirname(self): def get_upload_dirname(self): return os.path.join(self.get_data_dirname(), "raw") + def get_raw_data_dirname(self) -> str: + return { + StorageChoice.LOCAL: self.get_upload_dirname(), + StorageChoice.SHARE: settings.SHARE_ROOT, + StorageChoice.CLOUD_STORAGE: self.get_upload_dirname(), + }[self.storage] + def get_compressed_cache_dirname(self): return os.path.join(self.get_data_dirname(), "compressed") From 3c117fe1fa7153eebb85691727a29e37ab62a0aa Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 28 Aug 2024 18:00:40 +0300 Subject: [PATCH 085/115] Make SegmentFrameProvider available in make_frame_provider --- cvat/apps/engine/cache.py | 9 ++++----- cvat/apps/engine/frame_provider.py | 20 ++++++++++++++++++-- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 62301b6253c..e5b09c72c25 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -391,14 +391,13 @@ def _prepare_segment_preview(self, db_segment: models.Segment) -> DataWithMime: os.path.join(os.path.dirname(__file__), "assets/3d_preview.jpeg") ) else: - from cvat.apps.engine.frame_provider import ( + from cvat.apps.engine.frame_provider import ( # avoid circular import FrameOutputType, - SegmentFrameProvider, - TaskFrameProvider, + make_frame_provider, ) - task_frame_provider = TaskFrameProvider(db_segment.task) - segment_frame_provider = SegmentFrameProvider(db_segment) + task_frame_provider = make_frame_provider(db_segment.task) + segment_frame_provider = make_frame_provider(db_segment) preview = segment_frame_provider.get_frame( task_frame_provider.get_rel_frame_number(min(db_segment.frame_set)), quality=FrameQuality.COMPRESSED, diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 11ea5539e29..0821271fd1a 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -12,7 +12,7 @@ from dataclasses import dataclass from enum import Enum, auto from io import BytesIO -from typing import Any, Callable, Generic, Iterator, Optional, Tuple, Type, TypeVar, Union +from typing import Any, Callable, Generic, Iterator, Optional, Tuple, Type, TypeVar, Union, overload import av import cv2 @@ -561,9 +561,25 @@ def __init__(self, db_job: models.Job) -> None: super().__init__(db_job.segment) -def make_frame_provider(data_source: Union[models.Job, models.Task, Any]) -> IFrameProvider: +@overload +def make_frame_provider(data_source: models.Job) -> JobFrameProvider: ... + + +@overload +def make_frame_provider(data_source: models.Segment) -> SegmentFrameProvider: ... + + +@overload +def make_frame_provider(data_source: models.Task) -> TaskFrameProvider: ... + + +def make_frame_provider( + data_source: Union[models.Job, models.Segment, models.Task, Any] +) -> IFrameProvider: if isinstance(data_source, models.Task): frame_provider = TaskFrameProvider(data_source) + elif isinstance(data_source, models.Segment): + frame_provider = SegmentFrameProvider(data_source) elif isinstance(data_source, models.Job): frame_provider = JobFrameProvider(data_source) else: From 98eff81384549a56b37f893b6ab1f916dde08a79 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 28 Aug 2024 18:00:58 +0300 Subject: [PATCH 086/115] Remove extra variable --- cvat/apps/engine/task.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index 96c78d9bbc6..f24cd686a58 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -923,11 +923,9 @@ def _update_status(msg: str) -> None: db_data.compressed_chunk_type = models.DataChoice.VIDEO if task_mode == 'interpolation' and not data['use_zip_chunks'] else models.DataChoice.IMAGESET db_data.original_chunk_type = models.DataChoice.VIDEO if task_mode == 'interpolation' else models.DataChoice.IMAGESET - compressed_chunk_writer_class = Mpeg4CompressedChunkWriter if db_data.compressed_chunk_type == models.DataChoice.VIDEO else ZipCompressedChunkWriter - # calculate chunk size if it isn't specified if db_data.chunk_size is None: - if issubclass(compressed_chunk_writer_class, ZipCompressedChunkWriter): + if db_data.compressed_chunk_type == models.DataChoice.IMAGESET: first_image_idx = db_data.start_frame if not is_data_in_cloud: w, h = extractor.get_image_size(first_image_idx) From 316ec785c933f966213c9f800b93f06f0a9ab1ce Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 28 Aug 2024 18:41:51 +0300 Subject: [PATCH 087/115] Include both cases of CVAT_ALLOW_STATIC_CACHE in CI checks --- .github/workflows/full.yml | 2 ++ .github/workflows/main.yml | 3 ++- .github/workflows/schedule.yml | 6 ++++++ .../docker-compose.configurable_static_cache.yml | 16 ++++++++++++++++ tests/python/rest_api/test_jobs.py | 2 +- tests/python/rest_api/test_queues.py | 2 +- .../rest_api/test_resource_import_export.py | 2 +- tests/python/rest_api/test_tasks.py | 14 ++++++++++---- tests/python/shared/fixtures/init.py | 3 ++- 9 files changed, 41 insertions(+), 9 deletions(-) create mode 100644 tests/docker-compose.configurable_static_cache.yml diff --git a/.github/workflows/full.yml b/.github/workflows/full.yml index 4042828ccfb..f6e0d56a83a 100644 --- a/.github/workflows/full.yml +++ b/.github/workflows/full.yml @@ -172,6 +172,8 @@ jobs: id: run_tests run: | pytest tests/python/ + ONE_RUNNING_JOB_IN_QUEUE_PER_USER="true" pytest tests/python/rest_api/test_queues.py + CVAT_ALLOW_STATIC_CACHE="true" pytest -k "TestTaskData" tests/python - name: Creating a log file from cvat containers if: failure() && steps.run_tests.conclusion == 'failure' diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 822934ee9ba..3ab279d1371 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -184,8 +184,9 @@ jobs: COVERAGE_PROCESS_START: ".coveragerc" run: | pytest tests/python/ --cov --cov-report=json - for COVERAGE_FILE in `find -name "coverage*.json" -type f -printf "%f\n"`; do mv ${COVERAGE_FILE} "${COVERAGE_FILE%%.*}_0.json"; done ONE_RUNNING_JOB_IN_QUEUE_PER_USER="true" pytest tests/python/rest_api/test_queues.py --cov --cov-report=json + CVAT_ALLOW_STATIC_CACHE="true" pytest -k "TestTaskData" tests/python --cov --cov-report=json + for COVERAGE_FILE in `find -name "coverage*.json" -type f -printf "%f\n"`; do mv ${COVERAGE_FILE} "${COVERAGE_FILE%%.*}_0.json"; done - name: Uploading code coverage results as an artifact uses: actions/upload-artifact@v3.1.1 diff --git a/.github/workflows/schedule.yml b/.github/workflows/schedule.yml index 7fd9e6f6304..faf1e8a15eb 100644 --- a/.github/workflows/schedule.yml +++ b/.github/workflows/schedule.yml @@ -170,6 +170,12 @@ jobs: pytest tests/python/ pytest tests/python/ --stop-services + ONE_RUNNING_JOB_IN_QUEUE_PER_USER="true" pytest tests/python/rest_api/test_queues.py + pytest tests/python/ --stop-services + + CVAT_ALLOW_STATIC_CACHE="true" pytest tests/python + pytest tests/python/ --stop-services + - name: Unit tests env: HOST_COVERAGE_DATA_DIR: ${{ github.workspace }} diff --git a/tests/docker-compose.configurable_static_cache.yml b/tests/docker-compose.configurable_static_cache.yml new file mode 100644 index 00000000000..5afa4347080 --- /dev/null +++ b/tests/docker-compose.configurable_static_cache.yml @@ -0,0 +1,16 @@ +services: + cvat_server: + environment: + CVAT_ALLOW_STATIC_CACHE: '${CVAT_ALLOW_STATIC_CACHE:-no}' + + cvat_worker_import: + environment: + CVAT_ALLOW_STATIC_CACHE: '${CVAT_ALLOW_STATIC_CACHE:-no}' + + cvat_worker_export: + environment: + CVAT_ALLOW_STATIC_CACHE: '${CVAT_ALLOW_STATIC_CACHE:-no}' + + cvat_worker_annotation: + environment: + CVAT_ALLOW_STATIC_CACHE: '${CVAT_ALLOW_STATIC_CACHE:-no}' diff --git a/tests/python/rest_api/test_jobs.py b/tests/python/rest_api/test_jobs.py index 1cf761270f3..167b0d63c3b 100644 --- a/tests/python/rest_api/test_jobs.py +++ b/tests/python/rest_api/test_jobs.py @@ -361,7 +361,7 @@ def _test_destroy_job_fails(self, user, job_id, *, expected_status: int, **kwarg assert response.status == expected_status return response - @pytest.mark.usefixtures("restore_cvat_data") + @pytest.mark.usefixtures("restore_cvat_data_per_function") @pytest.mark.parametrize("job_type, allow", (("ground_truth", True), ("annotation", False))) def test_destroy_job(self, admin_user, jobs, job_type, allow): job = next(j for j in jobs if j["type"] == job_type) diff --git a/tests/python/rest_api/test_queues.py b/tests/python/rest_api/test_queues.py index a1729cf6f25..5d0a3190f16 100644 --- a/tests/python/rest_api/test_queues.py +++ b/tests/python/rest_api/test_queues.py @@ -18,7 +18,7 @@ @pytest.mark.usefixtures("restore_db_per_function") -@pytest.mark.usefixtures("restore_cvat_data") +@pytest.mark.usefixtures("restore_cvat_data_per_function") @pytest.mark.usefixtures("restore_redis_inmem_per_function") class TestRQQueueWorking: _USER_1 = "admin1" diff --git a/tests/python/rest_api/test_resource_import_export.py b/tests/python/rest_api/test_resource_import_export.py index 833661fcfab..39f4be22a01 100644 --- a/tests/python/rest_api/test_resource_import_export.py +++ b/tests/python/rest_api/test_resource_import_export.py @@ -177,7 +177,7 @@ def test_user_cannot_export_to_cloud_storage_with_specific_location_without_acce @pytest.mark.usefixtures("restore_db_per_function") -@pytest.mark.usefixtures("restore_cvat_data") +@pytest.mark.usefixtures("restore_cvat_data_per_function") class TestImportResourceFromS3(_S3ResourceTest): @pytest.mark.usefixtures("restore_redis_inmem_per_function") @pytest.mark.parametrize("cloud_storage_id", [3]) diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py index 4d01a93f244..236270efea5 100644 --- a/tests/python/rest_api/test_tasks.py +++ b/tests/python/rest_api/test_tasks.py @@ -910,7 +910,7 @@ def test_uses_subset_name( @pytest.mark.usefixtures("restore_db_per_function") -@pytest.mark.usefixtures("restore_cvat_data") +@pytest.mark.usefixtures("restore_cvat_data_per_function") @pytest.mark.usefixtures("restore_redis_ondisk_per_function") class TestPostTaskData: _USERNAME = "admin1" @@ -2107,7 +2107,7 @@ def read_frame(self, i: int) -> Image.Image: @pytest.mark.usefixtures("restore_db_per_class") @pytest.mark.usefixtures("restore_redis_ondisk_per_class") -@pytest.mark.usefixtures("restore_cvat_data") +@pytest.mark.usefixtures("restore_cvat_data_per_function") class TestTaskData: _USERNAME = "admin1" @@ -2712,7 +2712,7 @@ def test_admin_can_add_skeleton(self, tasks, admin_user): @pytest.mark.usefixtures("restore_db_per_function") -@pytest.mark.usefixtures("restore_cvat_data") +@pytest.mark.usefixtures("restore_cvat_data_per_function") @pytest.mark.usefixtures("restore_redis_ondisk_per_function") class TestWorkWithTask: _USERNAME = "admin1" @@ -2772,7 +2772,13 @@ def _make_client(self) -> Client: return Client(BASE_URL, config=Config(status_check_period=0.01)) @pytest.fixture(autouse=True) - def setup(self, restore_db_per_function, restore_cvat_data, tmp_path: Path, admin_user: str): + def setup( + self, + restore_db_per_function, + restore_cvat_data_per_function, + tmp_path: Path, + admin_user: str, + ): self.tmp_dir = tmp_path self.client = self._make_client() diff --git a/tests/python/shared/fixtures/init.py b/tests/python/shared/fixtures/init.py index cf5aeabbbf3..99f1f02f8e0 100644 --- a/tests/python/shared/fixtures/init.py +++ b/tests/python/shared/fixtures/init.py @@ -31,6 +31,7 @@ "tests/docker-compose.file_share.yml", "tests/docker-compose.minio.yml", "tests/docker-compose.test_servers.yml", + "tests/docker-compose.configurable_static_cache.yml", ] @@ -559,7 +560,7 @@ def restore_db_per_class(request): @pytest.fixture(scope="function") -def restore_cvat_data(request): +def restore_cvat_data_per_function(request): platform = request.config.getoption("--platform") if platform == "local": docker_restore_data_volumes() From 2b6e98761d92dcb111ce27f81fc21998819df07b Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 28 Aug 2024 18:48:17 +0300 Subject: [PATCH 088/115] Remove extra import --- cvat/apps/engine/cache.py | 1 - 1 file changed, 1 deletion(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index e5b09c72c25..3c1b54e2cd4 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -21,7 +21,6 @@ import cv2 import PIL.Image import PIL.ImageOps -from django.conf import settings from django.core.cache import caches from rest_framework.exceptions import NotFound, ValidationError From f67a1a2cd64ea8f0ce9efba3a97f7cf6044afac4 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 5 Sep 2024 16:47:22 +0300 Subject: [PATCH 089/115] Update changelog --- changelog.d/20240812_161617_mzhiltso_job_chunks.md | 12 ++++++++++++ changelog.d/20240812_161734_mzhiltso_job_chunks.md | 4 ---- changelog.d/20240812_161912_mzhiltso_job_chunks.md | 6 ------ 3 files changed, 12 insertions(+), 10 deletions(-) delete mode 100644 changelog.d/20240812_161734_mzhiltso_job_chunks.md delete mode 100644 changelog.d/20240812_161912_mzhiltso_job_chunks.md diff --git a/changelog.d/20240812_161617_mzhiltso_job_chunks.md b/changelog.d/20240812_161617_mzhiltso_job_chunks.md index f78376d9443..b54c5fe9e22 100644 --- a/changelog.d/20240812_161617_mzhiltso_job_chunks.md +++ b/changelog.d/20240812_161617_mzhiltso_job_chunks.md @@ -2,3 +2,15 @@ - A server setting to disable media chunks on the local filesystem () + +### Changed + +- \[Server API\] Chunk ids in each job now start from 0, instead of using ones from the task + () + +### Fixed + +- Various memory leaks in video reading on the server + () +- Job assignees will not receive frames from adjacent jobs in the boundary chunks + () diff --git a/changelog.d/20240812_161734_mzhiltso_job_chunks.md b/changelog.d/20240812_161734_mzhiltso_job_chunks.md deleted file mode 100644 index 2a587593b4f..00000000000 --- a/changelog.d/20240812_161734_mzhiltso_job_chunks.md +++ /dev/null @@ -1,4 +0,0 @@ -### Changed - -- Jobs now have separate chunk ids starting from 0, instead of using ones from the task - () diff --git a/changelog.d/20240812_161912_mzhiltso_job_chunks.md b/changelog.d/20240812_161912_mzhiltso_job_chunks.md deleted file mode 100644 index 2a0198dd8f7..00000000000 --- a/changelog.d/20240812_161912_mzhiltso_job_chunks.md +++ /dev/null @@ -1,6 +0,0 @@ -### Fixed - -- Various memory leaks in video reading on the server - () -- Job assignees will not receive frames from adjacent jobs in the boundary chunks - () From d72fe85d6f9bf722628c530bc2b21c6dd76c20e0 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 5 Sep 2024 17:22:18 +0300 Subject: [PATCH 090/115] Refactor cache keys in media cache --- cvat/apps/engine/cache.py | 50 +++++++++++++++++++++++++++++---------- 1 file changed, 37 insertions(+), 13 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 3c1b54e2cd4..ab2c1aacc11 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -100,25 +100,49 @@ def _get(self, key: str) -> Optional[DataWithMime]: return item + def _make_cache_key_prefix( + self, obj: Union[models.Task, models.Segment, models.Job, models.CloudStorage] + ) -> str: + if isinstance(obj, models.Task): + return f"task_{obj.id}" + elif isinstance(obj, models.Segment): + return f"segment_{obj.id}" + elif isinstance(obj, models.Job): + return f"job_{obj.id}" + elif isinstance(obj, models.CloudStorage): + return f"cloudstorage_{obj.id}" + else: + assert False, f"Unexpected object type {type(obj)}" + + def _make_chunk_key( + self, + db_obj: Union[models.Task, models.Segment, models.Job], + chunk_number: int, + *, + quality: FrameQuality, + ) -> str: + return f"{self._make_cache_key_prefix(db_obj)}_{chunk_number}_{quality}" + + def _make_preview_key(self, db_obj: Union[models.Segment, models.CloudStorage]) -> str: + return f"{self._make_cache_key_prefix(db_obj)}_preview" + + def _make_context_image_preview_key(self, db_data: models.Data, frame_number: int) -> str: + return f"context_image_{db_data.id}_{frame_number}_preview" + def get_segment_chunk( self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality ) -> DataWithMime: return self._get_or_set_cache_item( - key=f"segment_{db_segment.id}_{chunk_number}_{quality}", + key=self._make_chunk_key(db_segment, chunk_number, quality=quality), create_callback=lambda: self.prepare_segment_chunk( db_segment, chunk_number, quality=quality ), ) - def _make_task_chunk_key( - self, db_task: models.Task, chunk_number: int, *, quality: FrameQuality - ) -> str: - return f"task_{db_task.id}_{chunk_number}_{quality}" - def get_task_chunk( self, db_task: models.Task, chunk_number: int, *, quality: FrameQuality ) -> Optional[DataWithMime]: - return self._get(key=self._make_task_chunk_key(db_task, chunk_number, quality=quality)) + return self._get(key=self._make_chunk_key(db_task, chunk_number, quality=quality)) def get_or_set_task_chunk( self, @@ -129,7 +153,7 @@ def get_or_set_task_chunk( set_callback: Callable[[], DataWithMime], ) -> DataWithMime: return self._get_or_set_cache_item( - key=self._make_task_chunk_key(db_task, chunk_number, quality=quality), + key=self._make_chunk_key(db_task, chunk_number, quality=quality), create_callback=set_callback, ) @@ -137,7 +161,7 @@ def get_selective_job_chunk( self, db_job: models.Job, chunk_number: int, *, quality: FrameQuality ) -> DataWithMime: return self._get_or_set_cache_item( - key=f"job_{db_job.id}_{chunk_number}_{quality}", + key=self._make_chunk_key(db_job, chunk_number, quality=quality), create_callback=lambda: self.prepare_masked_range_segment_chunk( db_job.segment, chunk_number, quality=quality ), @@ -145,22 +169,22 @@ def get_selective_job_chunk( def get_or_set_segment_preview(self, db_segment: models.Segment) -> DataWithMime: return self._get_or_set_cache_item( - f"segment_preview_{db_segment.id}", + self._make_preview_key(db_segment), create_callback=lambda: self._prepare_segment_preview(db_segment), ) def get_cloud_preview(self, db_storage: models.CloudStorage) -> Optional[DataWithMime]: - return self._get(f"cloudstorage_{db_storage.id}_preview") + return self._get(self._make_preview_key(db_storage)) def get_or_set_cloud_preview(self, db_storage: models.CloudStorage) -> DataWithMime: return self._get_or_set_cache_item( - f"cloudstorage_{db_storage.id}_preview", + self._make_preview_key(db_storage), create_callback=lambda: self._prepare_cloud_preview(db_storage), ) def get_frame_context_images(self, db_data: models.Data, frame_number: int) -> DataWithMime: return self._get_or_set_cache_item( - key=f"context_image_{db_data.id}_{frame_number}", + key=self._make_context_image_preview_key(db_data, frame_number), create_callback=lambda: self.prepare_context_images(db_data, frame_number), ) From d5bfb888c44b0e871b08ca7d2a4925327674726a Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 5 Sep 2024 17:23:20 +0300 Subject: [PATCH 091/115] Refactor selective segment chunk creation --- cvat/apps/engine/cache.py | 80 +++++++++++++++++++-------------------- 1 file changed, 40 insertions(+), 40 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index ab2c1aacc11..c1116d8d88a 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -335,7 +335,7 @@ def prepare_range_segment_chunk( db_data = db_task.data chunk_size = db_data.chunk_size - chunk_frame_ids = db_segment.frame_set[ + chunk_frame_ids = list(db_segment.frame_set)[ chunk_size * chunk_number : chunk_size * (chunk_number + 1) ] @@ -348,15 +348,14 @@ def prepare_masked_range_segment_chunk( db_task = db_segment.task db_data = db_task.data - from cvat.apps.engine.frame_provider import TaskFrameProvider - - frame_provider = TaskFrameProvider(db_task) - - frame_set = db_segment.frame_set + chunk_size = db_data.chunk_size + chunk_frame_ids = list(db_segment.frame_set)[ + chunk_size * chunk_number : chunk_size * (chunk_number + 1) + ] frame_step = db_data.get_frame_step() - chunk_frames = [] writer = ZipCompressedChunkWriter(db_data.image_quality, dimension=db_task.dimension) + dummy_frame = io.BytesIO() PIL.Image.new("RGB", (1, 1)).save(dummy_frame, writer.IMAGE_EXT) @@ -365,46 +364,47 @@ def prepare_masked_range_segment_chunk( else: frame_size = None - for frame_idx in range(db_data.chunk_size): - frame_idx = ( - db_data.start_frame + chunk_number * db_data.chunk_size + frame_idx * frame_step - ) - if db_data.stop_frame < frame_idx: - break - - frame_bytes = None - - if frame_idx in frame_set: - frame_bytes = frame_provider.get_frame(frame_idx, quality=quality).data - - if frame_size is not None: - # Decoded video frames can have different size, restore the original one + def get_frames(): + with closing( + self._read_raw_frames(db_task, frame_ids=chunk_frame_ids) + ) as read_frame_iter: + for frame_idx in range(db_data.chunk_size): + frame_idx = ( + db_data.start_frame + + chunk_number * db_data.chunk_size + + frame_idx * frame_step + ) + if db_data.stop_frame < frame_idx: + break - frame = PIL.Image.open(frame_bytes) - if frame.size != frame_size: - frame = frame.resize(frame_size) + if frame_idx in chunk_frame_ids: + frame = next(read_frame_iter)[0] - frame_bytes = io.BytesIO() - frame.save(frame_bytes, writer.IMAGE_EXT) - frame_bytes.seek(0) + if hasattr(db_data, "video"): + # Decoded video frames can have different size, restore the original one - else: - # Populate skipped frames with placeholder data, - # this is required for video chunk decoding implementation in UI - frame_bytes = io.BytesIO(dummy_frame.getvalue()) + frame = frame.to_image() + if frame.size != frame_size: + frame = frame.resize(frame_size) + else: + # Populate skipped frames with placeholder data, + # this is required for video chunk decoding implementation in UI + # TODO: try to fix decoding in UI + frame = io.BytesIO(dummy_frame.getvalue()) - if frame_bytes is not None: - chunk_frames.append((frame_bytes, None, None)) + yield (frame, None, None) buff = io.BytesIO() - writer.save_as_chunk( - chunk_frames, - buff, - compress_frames=False, - zip_compress_level=1, # there are likely to be many skips in SPECIFIC_FRAMES segments - ) - buff.seek(0) + with closing(get_frames()) as frame_iter: + writer.save_as_chunk( + frame_iter, + buff, + zip_compress_level=1, + # there are likely to be many skips with repeated placeholder frames + # in SPECIFIC_FRAMES segments, it makes sense to compress the archive + ) + buff.seek(0) return buff, get_chunk_mime_type_for_writer(writer) def _prepare_segment_preview(self, db_segment: models.Segment) -> DataWithMime: From c5a1197f41b6664b098b1d8b68a57dc55f5495ef Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 6 Sep 2024 19:58:43 +0300 Subject: [PATCH 092/115] Remove the breaking change in the chunk retrieval API, add a new index parameter --- .../20240812_161617_mzhiltso_job_chunks.md | 14 ++- cvat/apps/engine/cache.py | 24 +++- cvat/apps/engine/frame_provider.py | 117 ++++++++++++++++-- cvat/apps/engine/views.py | 60 +++++++-- cvat/schema.yml | 9 +- tests/python/rest_api/test_tasks.py | 67 ++++++++-- 6 files changed, 257 insertions(+), 34 deletions(-) diff --git a/changelog.d/20240812_161617_mzhiltso_job_chunks.md b/changelog.d/20240812_161617_mzhiltso_job_chunks.md index b54c5fe9e22..6a3f609c02e 100644 --- a/changelog.d/20240812_161617_mzhiltso_job_chunks.md +++ b/changelog.d/20240812_161617_mzhiltso_job_chunks.md @@ -2,15 +2,23 @@ - A server setting to disable media chunks on the local filesystem () +- \[Server API\] `GET /api/jobs/{id}/data/?type=chunk&index=x` parameter combination. + The new `index` parameter allows to retrieve job chunks using 0-based index in each job, + instead of the `number` parameter, which used task chunk ids. + () ### Changed -- \[Server API\] Chunk ids in each job now start from 0, instead of using ones from the task +- Job assignees will not receive frames from adjacent jobs in chunks + () + +### Deprecated + +- \[Server API\] `GET /api/jobs/{id}/data/?type=chunk&number=x` parameter combination () + ### Fixed - Various memory leaks in video reading on the server () -- Job assignees will not receive frames from adjacent jobs in the boundary chunks - () diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index c1116d8d88a..6ace102040b 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -121,11 +121,20 @@ def _make_chunk_key( *, quality: FrameQuality, ) -> str: - return f"{self._make_cache_key_prefix(db_obj)}_{chunk_number}_{quality}" + return f"{self._make_cache_key_prefix(db_obj)}_chunk_{chunk_number}_{quality}" def _make_preview_key(self, db_obj: Union[models.Segment, models.CloudStorage]) -> str: return f"{self._make_cache_key_prefix(db_obj)}_preview" + def _make_segment_task_chunk_key( + self, + db_obj: models.Segment, + chunk_number: int, + *, + quality: FrameQuality, + ) -> str: + return f"{self._make_cache_key_prefix(db_obj)}_task_chunk_{chunk_number}_{quality}" + def _make_context_image_preview_key(self, db_data: models.Data, frame_number: int) -> str: return f"context_image_{db_data.id}_{frame_number}_preview" @@ -157,6 +166,19 @@ def get_or_set_task_chunk( create_callback=set_callback, ) + def get_or_set_segment_task_chunk( + self, + db_segment: models.Segment, + chunk_number: int, + *, + quality: FrameQuality, + set_callback: Callable[[], DataWithMime], + ) -> DataWithMime: + return self._get_or_set_cache_item( + key=self._make_segment_task_chunk_key(db_segment, chunk_number, quality=quality), + create_callback=set_callback, + ) + def get_selective_job_chunk( self, db_job: models.Job, chunk_number: int, *, quality: FrameQuality ) -> DataWithMime: diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 0821271fd1a..71414490c8a 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -12,11 +12,24 @@ from dataclasses import dataclass from enum import Enum, auto from io import BytesIO -from typing import Any, Callable, Generic, Iterator, Optional, Tuple, Type, TypeVar, Union, overload +from typing import ( + Any, + Callable, + Generic, + Iterator, + Optional, + Sequence, + Tuple, + Type, + TypeVar, + Union, + overload, +) import av import cv2 import numpy as np +from datumaro.util import take_by from django.conf import settings from PIL import Image from rest_framework.exceptions import ValidationError @@ -255,13 +268,12 @@ def get_chunk( return_type = DataWithMeta[BytesIO] chunk_number = self.validate_chunk_number(chunk_number) - db_data = self._db_task.data - cache = MediaCache() cached_chunk = cache.get_task_chunk(self._db_task, chunk_number, quality=quality) if cached_chunk: return return_type(cached_chunk[0], cached_chunk[1]) + db_data = self._db_task.data step = db_data.get_frame_step() task_chunk_start_frame = chunk_number * db_data.chunk_size task_chunk_stop_frame = (chunk_number + 1) * db_data.chunk_size - 1 @@ -273,7 +285,7 @@ def get_chunk( ) ) - matching_segments = sorted( + matching_segments: list[models.Segment] = sorted( [ s for s in self._db_task.segment_set.all() @@ -285,13 +297,15 @@ def get_chunk( assert matching_segments # Don't put this into set_callback to avoid data duplication in the cache - if len(matching_segments) == 1 and task_chunk_frame_set == set( - matching_segments[0].frame_set - ): + + if len(matching_segments) == 1: segment_frame_provider = SegmentFrameProvider(matching_segments[0]) - return segment_frame_provider.get_chunk( - segment_frame_provider.get_chunk_number(task_chunk_start_frame), quality=quality + matching_chunk_index = segment_frame_provider.find_matching_chunk( + sorted(task_chunk_frame_set) ) + if matching_chunk_index is not None: + # The requested frames match one of the job chunks, we can use it directly + return segment_frame_provider.get_chunk(matching_chunk_index, quality=quality) def _set_callback() -> DataWithMime: # Create and return a joined / cleaned chunk @@ -469,6 +483,20 @@ def validate_frame_number(self, frame_number: int) -> Tuple[int, int, int]: def get_chunk_number(self, frame_number: int) -> int: return int(frame_number) // self._db_segment.task.data.chunk_size + def find_matching_chunk(self, frames: Sequence[int]) -> Optional[int]: + return next( + ( + i + for i, chunk_frames in enumerate( + take_by( + sorted(self._db_segment.frame_set), self._db_segment.task.data.chunk_size + ) + ) + if frames == set(chunk_frames) + ), + None, + ) + def validate_chunk_number(self, chunk_number: int) -> int: segment_size = self._db_segment.frame_count last_chunk = math.ceil(segment_size / self._db_segment.task.data.chunk_size) - 1 @@ -560,6 +588,77 @@ class JobFrameProvider(SegmentFrameProvider): def __init__(self, db_job: models.Job) -> None: super().__init__(db_job.segment) + def get_chunk( + self, + chunk_number: int, + *, + quality: FrameQuality = FrameQuality.ORIGINAL, + is_task_chunk: bool = False, + ) -> DataWithMeta[BytesIO]: + if not is_task_chunk: + return super().get_chunk(chunk_number, quality=quality) + + task_frame_provider = TaskFrameProvider(self._db_segment.task) + segment_start_chunk = task_frame_provider.get_chunk_number(self._db_segment.start_frame) + segment_stop_chunk = task_frame_provider.get_chunk_number(self._db_segment.stop_frame) + if not segment_start_chunk <= chunk_number <= segment_stop_chunk: + raise ValidationError( + f"Invalid chunk number '{chunk_number}'. " + "The chunk number should be in the " + f"[{segment_start_chunk}, {segment_stop_chunk}] range" + ) + + # Reproduce the task chunks, limited by this job + return_type = DataWithMeta[BytesIO] + + cache = MediaCache() + cached_chunk = cache.get_task_chunk(self._db_segment.task, chunk_number, quality=quality) + if cached_chunk: + return return_type(cached_chunk[0], cached_chunk[1]) + + db_data = self._db_segment.task.data + step = db_data.get_frame_step() + task_chunk_start_frame = chunk_number * db_data.chunk_size + task_chunk_stop_frame = (chunk_number + 1) * db_data.chunk_size - 1 + task_chunk_frame_set = set( + range( + db_data.start_frame + task_chunk_start_frame * step, + min(db_data.start_frame + task_chunk_stop_frame * step, db_data.stop_frame) + step, + step, + ) + ) + + # Don't put this into set_callback to avoid data duplication in the cache + matching_chunk = self.find_matching_chunk(sorted(task_chunk_frame_set)) + if matching_chunk is not None: + return self.get_chunk(matching_chunk, quality=quality) + + def _set_callback() -> DataWithMime: + # Create and return a joined / cleaned chunk + segment_frame_set = set(self._db_segment.frame_set) + task_chunk_frames = {} + for task_chunk_frame_id in sorted(task_chunk_frame_set): + if task_chunk_frame_id not in segment_frame_set: + continue + + frame, frame_name, _ = self._get_raw_frame( + task_frame_provider.get_rel_frame_number(task_chunk_frame_id), quality=quality + ) + task_chunk_frames[task_chunk_frame_id] = (frame, frame_name, None) + + return prepare_chunk( + task_chunk_frames.values(), + quality=quality, + db_task=self._db_segment.task, + dump_unchanged=True, + ) + + buffer, mime_type = cache.get_or_set_segment_task_chunk( + self._db_segment, chunk_number, quality=quality, set_callback=_set_callback + ) + + return return_type(data=buffer, mime=mime_type) + @overload def make_frame_provider(data_source: models.Job) -> JobFrameProvider: ... diff --git a/cvat/apps/engine/views.py b/cvat/apps/engine/views.py index 1bc2921605a..ebda8c7fae3 100644 --- a/cvat/apps/engine/views.py +++ b/cvat/apps/engine/views.py @@ -685,8 +685,7 @@ def __init__( if data_quality == 'compressed' else FrameQuality.ORIGINAL @abstractmethod - def _get_frame_provider(self) -> IFrameProvider: - ... + def _get_frame_provider(self) -> IFrameProvider: ... def __call__(self): frame_provider = self._get_frame_provider() @@ -694,7 +693,7 @@ def __call__(self): try: if self.type == 'chunk': data = frame_provider.get_chunk(self.number, quality=self.quality) - return HttpResponse(data.data.getvalue(), content_type=data.mime) # TODO: add new headers + return HttpResponse(data.data.getvalue(), content_type=data.mime) elif self.type == 'frame' or self.type == 'preview': if self.type == 'preview': data = frame_provider.get_preview() @@ -729,7 +728,7 @@ def __init__( super().__init__(data_type=data_type, data_num=data_num, data_quality=data_quality) self._db_task = db_task - def _get_frame_provider(self) -> IFrameProvider: + def _get_frame_provider(self) -> TaskFrameProvider: return TaskFrameProvider(self._db_task) @@ -741,13 +740,47 @@ def __init__( data_type: str, data_quality: str, data_num: Optional[Union[str, int]] = None, + data_index: Optional[Union[str, int]] = None, ) -> None: - super().__init__(data_type=data_type, data_num=data_num, data_quality=data_quality) + possible_data_type_values = ('chunk', 'frame', 'preview', 'context_image') + possible_quality_values = ('compressed', 'original') + + if not data_type or data_type not in possible_data_type_values: + raise ValidationError('Data type not specified or has wrong value') + elif data_type == 'chunk' or data_type == 'frame' or data_type == 'preview': + if data_type == 'chunk': + if data_num is None and data_index is None: + raise ValidationError('Number or Index is not specified') + if data_num is not None and data_index is not None: + raise ValidationError('Number and Index cannot be used together') + elif data_num is None and data_type != 'preview': + raise ValidationError('Number is not specified') + elif data_quality not in possible_quality_values: + raise ValidationError('Wrong quality value') + + self.type = data_type + + self.number = int(data_index) if data_index is not None else None + self.task_chunk_number = int(data_num) if data_num is not None else None + + self.quality = FrameQuality.COMPRESSED \ + if data_quality == 'compressed' else FrameQuality.ORIGINAL + self._db_job = db_job - def _get_frame_provider(self) -> IFrameProvider: + def _get_frame_provider(self) -> JobFrameProvider: return JobFrameProvider(self._db_job) + def __call__(self): + if self.type == 'chunk' and self.task_chunk_number is not None: + # Reproduce the task chunk indexing + frame_provider = self._get_frame_provider() + data = frame_provider.get_chunk( + self.task_chunk_number, quality=self.quality, is_task_chunk=True + ) + return HttpResponse(data.data.getvalue(), content_type=data.mime) + else: + return super().__call__() @extend_schema(tags=['tasks']) @extend_schema_view( @@ -1987,8 +2020,14 @@ def get_export_callback(self, save_images: bool) -> Callable: OpenApiParameter('quality', location=OpenApiParameter.QUERY, required=False, type=OpenApiTypes.STR, enum=['compressed', 'original'], description="Specifies the quality level of the requested data"), - OpenApiParameter('number', location=OpenApiParameter.QUERY, required=False, type=OpenApiTypes.INT, - description="A unique number value identifying chunk or frame"), + OpenApiParameter('number', + location=OpenApiParameter.QUERY, required=False, type=OpenApiTypes.INT, + description="A unique number value identifying chunk or frame. " + "The numbers are the same as for the task. " + "Deprecated for chunks in favor of 'index'"), + OpenApiParameter('index', + location=OpenApiParameter.QUERY, required=False, type=OpenApiTypes.INT, + description="A unique number value identifying chunk, starts from 0 for each job"), ], responses={ '200': OpenApiResponse(OpenApiTypes.BINARY, description='Data of a specific type'), @@ -2000,10 +2039,13 @@ def data(self, request, pk): db_job = self.get_object() # call check_object_permissions as well data_type = request.query_params.get('type', None) data_num = request.query_params.get('number', None) + data_index = request.query_params.get('index', None) data_quality = request.query_params.get('quality', 'compressed') data_getter = _JobDataGetter( - db_job, data_type=data_type, data_num=data_num, data_quality=data_quality + db_job, + data_type=data_type, data_quality=data_quality, + data_index=data_index, data_num=data_num ) return data_getter() diff --git a/cvat/schema.yml b/cvat/schema.yml index 96ece3d07ca..0290ce2a5dd 100644 --- a/cvat/schema.yml +++ b/cvat/schema.yml @@ -2322,11 +2322,18 @@ paths: type: integer description: A unique integer value identifying this job. required: true + - in: query + name: index + schema: + type: integer + description: A unique number value identifying chunk, starts from 0 for each + job - in: query name: number schema: type: integer - description: A unique number value identifying chunk or frame + description: A unique number value identifying chunk or frame. The numbers + are the same as for the task. Deprecated for chunks in favor of 'index' - in: query name: quality schema: diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py index 5a3d28c131c..4c3c3eee5f5 100644 --- a/tests/python/rest_api/test_tasks.py +++ b/tests/python/rest_api/test_tasks.py @@ -22,7 +22,7 @@ from pathlib import Path from tempfile import NamedTemporaryFile, TemporaryDirectory from time import sleep, time -from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Tuple, Union +from typing import Any, Callable, ClassVar, Dict, Generator, List, Optional, Sequence, Tuple, Union import attrs import numpy as np @@ -2444,12 +2444,16 @@ def test_can_get_job_frames(self, task_spec: _TaskSpec, task_id: int): ) @parametrize("task_spec, task_id", _default_task_cases) - def test_can_get_job_chunks(self, task_spec: _TaskSpec, task_id: int): + @parametrize("indexing", ["absolute", "relative"]) + def test_can_get_job_chunks(self, task_spec: _TaskSpec, task_id: int, indexing: str): with make_api_client(self._USERNAME) as api_client: jobs = sorted( get_paginated_collection(api_client.jobs_api.list_endpoint, task_id=task_id), key=lambda j: j.start_frame, ) + + (task_meta, _) = api_client.tasks_api.retrieve_data_meta(task_id) + for job in jobs: (job_meta, _) = api_client.jobs_api.retrieve_data_meta(job.id) @@ -2466,21 +2470,62 @@ def test_can_get_job_chunks(self, task_spec: _TaskSpec, task_id: int): else: assert False - chunk_count = math.ceil(job_meta.size / job_meta.chunk_size) - for quality, chunk_id in product(["original", "compressed"], range(chunk_count)): - expected_chunk_abs_frame_ids = range( - job_meta.start_frame - + chunk_id * job_meta.chunk_size * task_spec.frame_step, - job_meta.start_frame - + min((chunk_id + 1) * job_meta.chunk_size, job_meta.size) - * task_spec.frame_step, + if indexing == "absolute": + chunk_count = math.ceil(task_meta.size / job_meta.chunk_size) + + def get_task_chunk_abs_frame_ids(chunk_id: int) -> Sequence[int]: + return range( + task_meta.start_frame + + chunk_id * task_meta.chunk_size * task_spec.frame_step, + task_meta.start_frame + + min((chunk_id + 1) * task_meta.chunk_size, task_meta.size) + * task_spec.frame_step, + ) + + def get_job_frame_ids() -> Sequence[int]: + return range( + job_meta.start_frame, job_meta.stop_frame + 1, task_spec.frame_step + ) + + def get_expected_chunk_abs_frame_ids(chunk_id: int): + return sorted( + set(get_task_chunk_abs_frame_ids(chunk_id)) & set(get_job_frame_ids()) + ) + + job_chunk_ids = ( + task_chunk_id + for task_chunk_id in range(chunk_count) + if get_expected_chunk_abs_frame_ids(task_chunk_id) ) + else: + chunk_count = math.ceil(job_meta.size / job_meta.chunk_size) + job_chunk_ids = range(chunk_count) + + def get_expected_chunk_abs_frame_ids(chunk_id: int): + return range( + job_meta.start_frame + + chunk_id * job_meta.chunk_size * task_spec.frame_step, + job_meta.start_frame + + min((chunk_id + 1) * job_meta.chunk_size, job_meta.size) + * task_spec.frame_step, + ) + + for quality, chunk_id in product(["original", "compressed"], job_chunk_ids): + expected_chunk_abs_frame_ids = get_expected_chunk_abs_frame_ids(chunk_id) + + kwargs = {} + if indexing == "absolute": + kwargs["number"] = chunk_id + elif indexing == "relative": + kwargs["index"] = chunk_id + else: + assert False (_, response) = api_client.jobs_api.retrieve_data( job.id, type="chunk", quality=quality, - number=chunk_id, + **kwargs, _parse_response=False, ) From a5cf3b77f489ba99d75c121ca2d8adb935a4731b Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Sat, 7 Sep 2024 12:42:12 +0300 Subject: [PATCH 093/115] Update UI to use the new chunk index parameter --- cvat-core/src/frames.ts | 2 +- cvat-core/src/server-proxy.ts | 2 +- cvat-core/src/session-implementation.ts | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index ff6200ce91b..d2547f0010c 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -587,7 +587,7 @@ export async function getFrame( isPlaying: boolean, step: number, dimension: DimensionType, - getChunk: (chunkNumber: number, quality: ChunkQuality) => Promise, + getChunk: (chunkIndex: number, quality: ChunkQuality) => Promise, ): Promise { if (!(jobID in frameDataCache)) { const blockType = chunkType === 'video' ? BlockType.MP4VIDEO : BlockType.ARCHIVE; diff --git a/cvat-core/src/server-proxy.ts b/cvat-core/src/server-proxy.ts index 91dc52a7182..51309426198 100644 --- a/cvat-core/src/server-proxy.ts +++ b/cvat-core/src/server-proxy.ts @@ -1438,7 +1438,7 @@ async function getData(jid: number, chunk: number, quality: ChunkQuality, retry ...enableOrganization(), quality, type: 'chunk', - number: chunk, + index: chunk, }, responseType: 'arraybuffer', }); diff --git a/cvat-core/src/session-implementation.ts b/cvat-core/src/session-implementation.ts index 4483c111393..11430241caa 100644 --- a/cvat-core/src/session-implementation.ts +++ b/cvat-core/src/session-implementation.ts @@ -189,7 +189,7 @@ export function implementJob(Job: typeof JobClass): typeof JobClass { isPlaying, step, this.dimension, - (chunkNumber, quality) => this.frames.chunk(chunkNumber, quality), + (chunkIndex, quality) => this.frames.chunk(chunkIndex, quality), ); }, }); @@ -273,10 +273,10 @@ export function implementJob(Job: typeof JobClass): typeof JobClass { Object.defineProperty(Job.prototype.frames.chunk, 'implementation', { value: function chunkImplementation( this: JobClass, - chunkNumber: Parameters[0], + chunkIndex: Parameters[0], quality: Parameters[1], ): ReturnType { - return serverProxy.frames.getData(this.id, chunkNumber, quality); + return serverProxy.frames.getData(this.id, chunkIndex, quality); }, }); @@ -824,7 +824,7 @@ export function implementTask(Task: typeof TaskClass): typeof TaskClass { isPlaying, step, this.dimension, - (chunkNumber, quality) => job.frames.chunk(chunkNumber, quality), + (chunkIndex, quality) => job.frames.chunk(chunkIndex, quality), ); return result; }, From cfdde3f86320ae498a4eef7b09de969bb9ae8d41 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Sat, 7 Sep 2024 12:50:30 +0300 Subject: [PATCH 094/115] Update test initialization --- cvat/apps/engine/tests/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cvat/apps/engine/tests/utils.py b/cvat/apps/engine/tests/utils.py index 87c5911e12b..3d2a533d1e9 100644 --- a/cvat/apps/engine/tests/utils.py +++ b/cvat/apps/engine/tests/utils.py @@ -13,7 +13,7 @@ from django.core.cache import caches from django.http.response import HttpResponse from PIL import Image -from rest_framework.test import APIClient, APITestCase +from rest_framework.test import APITestCase import av import django_rq import numpy as np @@ -112,7 +112,7 @@ def setUp(self): self._clear_temp_data() super().setUp() - self.client = APIClient() + self.client = self.client_class() def generate_image_file(filename, size=(100, 100)): From 843b9571696edb9ff234935da6bdd50de7dd17f0 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Sat, 7 Sep 2024 12:58:24 +0300 Subject: [PATCH 095/115] Update changelog --- changelog.d/20240812_161617_mzhiltso_job_chunks.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog.d/20240812_161617_mzhiltso_job_chunks.md b/changelog.d/20240812_161617_mzhiltso_job_chunks.md index 6a3f609c02e..af931641d6d 100644 --- a/changelog.d/20240812_161617_mzhiltso_job_chunks.md +++ b/changelog.d/20240812_161617_mzhiltso_job_chunks.md @@ -1,6 +1,6 @@ ### Added -- A server setting to disable media chunks on the local filesystem +- A server setting to enable or disable storage of permanent media chunks on the server filesystem () - \[Server API\] `GET /api/jobs/{id}/data/?type=chunk&index=x` parameter combination. The new `index` parameter allows to retrieve job chunks using 0-based index in each job, From feb92cd8ba99ef46873ece894bb63651a4a329f2 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 9 Sep 2024 19:33:42 +0300 Subject: [PATCH 096/115] Add backward compatibility for chunk "number" in GT jobs, remove placeholder frames from "index" response --- cvat/apps/engine/cache.py | 120 ++++++++++++++++++++++++---- cvat/apps/engine/frame_provider.py | 42 +++++----- tests/python/rest_api/test_jobs.py | 82 ++++++++++--------- tests/python/rest_api/test_tasks.py | 16 ++-- 4 files changed, 181 insertions(+), 79 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 6ace102040b..8dd3f8f02b5 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -14,8 +14,19 @@ import zlib from contextlib import ExitStack, closing from datetime import datetime, timezone -from itertools import pairwise -from typing import Any, Callable, Generator, Iterator, Optional, Sequence, Tuple, Type, Union +from itertools import groupby, pairwise +from typing import ( + Any, + Callable, + Collection, + Generator, + Iterator, + Optional, + Sequence, + Tuple, + Type, + Union, +) import av import cv2 @@ -100,6 +111,9 @@ def _get(self, key: str) -> Optional[DataWithMime]: return item + def _has_key(self, key: str) -> bool: + return self._cache.has_key(key) + def _make_cache_key_prefix( self, obj: Union[models.Task, models.Segment, models.Job, models.CloudStorage] ) -> str: @@ -166,6 +180,13 @@ def get_or_set_task_chunk( create_callback=set_callback, ) + def get_segment_task_chunk( + self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality + ) -> Optional[DataWithMime]: + return self._get( + key=self._make_segment_task_chunk_key(db_segment, chunk_number, quality=quality) + ) + def get_or_set_segment_task_chunk( self, db_segment: models.Segment, @@ -361,7 +382,12 @@ def prepare_range_segment_chunk( chunk_size * chunk_number : chunk_size * (chunk_number + 1) ] - with closing(self._read_raw_frames(db_task, frame_ids=chunk_frame_ids)) as frame_iter: + return self.prepare_custom_range_segment_chunk(db_task, chunk_frame_ids, quality=quality) + + def prepare_custom_range_segment_chunk( + self, db_task: models.Task, frame_ids: Sequence[int], *, quality: FrameQuality + ) -> DataWithMime: + with closing(self._read_raw_frames(db_task, frame_ids=frame_ids)) as frame_iter: return prepare_chunk(frame_iter, quality=quality, db_task=db_task) def prepare_masked_range_segment_chunk( @@ -371,41 +397,101 @@ def prepare_masked_range_segment_chunk( db_data = db_task.data chunk_size = db_data.chunk_size - chunk_frame_ids = list(db_segment.frame_set)[ + chunk_frame_ids = sorted(db_segment.frame_set)[ chunk_size * chunk_number : chunk_size * (chunk_number + 1) ] + + return self.prepare_custom_masked_range_segment_chunk( + db_task, chunk_frame_ids, chunk_number, quality=quality + ) + + def prepare_custom_masked_range_segment_chunk( + self, + db_task: models.Task, + frame_ids: Collection[int], + chunk_number: int, + *, + quality: FrameQuality, + insert_placeholders: bool = False, + ) -> DataWithMime: + db_data = db_task.data + frame_step = db_data.get_frame_step() - writer = ZipCompressedChunkWriter(db_data.image_quality, dimension=db_task.dimension) + image_quality = 100 if quality == FrameQuality.ORIGINAL else db_data.image_quality + writer = ZipCompressedChunkWriter(image_quality, dimension=db_task.dimension) dummy_frame = io.BytesIO() PIL.Image.new("RGB", (1, 1)).save(dummy_frame, writer.IMAGE_EXT) + # Optimize frame access if all the required frames are already cached + # Otherwise we might need to download files. + # This is not needed for video tasks, as it will reduce performance + from cvat.apps.engine.frame_provider import FrameOutputType, TaskFrameProvider + + task_frame_provider = TaskFrameProvider(db_task) + + use_cached_data = False + if db_task.mode != "interpolation": + required_frame_set = set(frame_ids) + all_chunks_available = all( + self._has_key(self._make_chunk_key(db_segment, chunk_number, quality=quality)) + for db_segment in db_task.segment_set.filter(type=models.SegmentType.RANGE).all() + for chunk_number, _ in groupby( + required_frame_set.intersection(db_segment.frame_set), + key=lambda frame: frame // db_data.chunk_size, + ) + ) + use_cached_data = all_chunks_available + if hasattr(db_data, "video"): frame_size = (db_data.video.width, db_data.video.height) else: frame_size = None def get_frames(): - with closing( - self._read_raw_frames(db_task, frame_ids=chunk_frame_ids) - ) as read_frame_iter: - for frame_idx in range(db_data.chunk_size): - frame_idx = ( - db_data.start_frame - + chunk_number * db_data.chunk_size - + frame_idx * frame_step + with ExitStack() as es: + es.callback(task_frame_provider.unload) + + if insert_placeholders: + frame_range = ( + ( + db_data.start_frame + + chunk_number * db_data.chunk_size + + chunk_frame_idx * frame_step + ) + for chunk_frame_idx in range(db_data.chunk_size) ) - if db_data.stop_frame < frame_idx: + else: + frame_range = frame_ids + + if not use_cached_data: + frames_gen = self._read_raw_frames(db_task, frame_ids) + frames_iter = iter(es.enter_context(closing(frames_gen))) + + for abs_frame_idx in frame_range: + if db_data.stop_frame < abs_frame_idx: break - if frame_idx in chunk_frame_ids: - frame = next(read_frame_iter)[0] + if abs_frame_idx in frame_ids: + if use_cached_data: + frame_data = task_frame_provider.get_frame( + task_frame_provider.get_rel_frame_number(abs_frame_idx), + quality=quality, + out_type=FrameOutputType.BUFFER, + ) + frame = frame_data.data + else: + frame, _, _ = next(frames_iter) if hasattr(db_data, "video"): # Decoded video frames can have different size, restore the original one - frame = frame.to_image() + if isinstance(frame, av.VideoFrame): + frame = frame.to_image() + else: + frame = PIL.Image.open(frame) + if frame.size != frame_size: frame = frame.resize(frame_size) else: diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 71414490c8a..314c8cf5b23 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -598,6 +598,10 @@ def get_chunk( if not is_task_chunk: return super().get_chunk(chunk_number, quality=quality) + # Backward compatibility for the "number" parameter + # Reproduce the task chunks, limited by this job + return_type = DataWithMeta[BytesIO] + task_frame_provider = TaskFrameProvider(self._db_segment.task) segment_start_chunk = task_frame_provider.get_chunk_number(self._db_segment.start_frame) segment_stop_chunk = task_frame_provider.get_chunk_number(self._db_segment.stop_frame) @@ -608,11 +612,8 @@ def get_chunk( f"[{segment_start_chunk}, {segment_stop_chunk}] range" ) - # Reproduce the task chunks, limited by this job - return_type = DataWithMeta[BytesIO] - cache = MediaCache() - cached_chunk = cache.get_task_chunk(self._db_segment.task, chunk_number, quality=quality) + cached_chunk = cache.get_segment_task_chunk(self._db_segment, chunk_number, quality=quality) if cached_chunk: return return_type(cached_chunk[0], cached_chunk[1]) @@ -635,23 +636,26 @@ def get_chunk( def _set_callback() -> DataWithMime: # Create and return a joined / cleaned chunk - segment_frame_set = set(self._db_segment.frame_set) - task_chunk_frames = {} - for task_chunk_frame_id in sorted(task_chunk_frame_set): - if task_chunk_frame_id not in segment_frame_set: - continue + segment_chunk_frame_ids = sorted( + task_chunk_frame_set.intersection(self._db_segment.frame_set) + ) - frame, frame_name, _ = self._get_raw_frame( - task_frame_provider.get_rel_frame_number(task_chunk_frame_id), quality=quality + if self._db_segment.type == models.SegmentType.RANGE: + return cache.prepare_custom_range_segment_chunk( + db_task=self._db_segment.task, + frame_ids=segment_chunk_frame_ids, + quality=quality, ) - task_chunk_frames[task_chunk_frame_id] = (frame, frame_name, None) - - return prepare_chunk( - task_chunk_frames.values(), - quality=quality, - db_task=self._db_segment.task, - dump_unchanged=True, - ) + elif self._db_segment.type == models.SegmentType.SPECIFIC_FRAMES: + return cache.prepare_custom_masked_range_segment_chunk( + db_task=self._db_segment.task, + frame_ids=segment_chunk_frame_ids, + chunk_number=chunk_number, + quality=quality, + insert_placeholders=True, + ) + else: + assert False buffer, mime_type = cache.get_or_set_segment_task_chunk( self._db_segment, chunk_number, quality=quality, set_callback=_set_callback diff --git a/tests/python/rest_api/test_jobs.py b/tests/python/rest_api/test_jobs.py index 167b0d63c3b..a6cd225a5d5 100644 --- a/tests/python/rest_api/test_jobs.py +++ b/tests/python/rest_api/test_jobs.py @@ -11,7 +11,7 @@ from copy import deepcopy from http import HTTPStatus from io import BytesIO -from itertools import product +from itertools import groupby, product from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np @@ -603,12 +603,8 @@ def test_get_gt_job_in_org_task( self._test_get_job_403(user["username"], job["id"]) -@pytest.mark.usefixtures( - # if the db is restored per test, there are conflicts with the server data cache - # if we don't clean the db, the gt jobs created will be reused, and their - # ids won't conflict - "restore_db_per_class" -) +@pytest.mark.usefixtures("restore_db_per_class") +@pytest.mark.usefixtures("restore_redis_ondisk_per_class") class TestGetGtJobData: def _delete_gt_job(self, user, gt_job_id): with make_api_client(user) as api_client: @@ -715,7 +711,10 @@ def test_can_get_gt_job_meta_with_complex_frame_setup(self, admin_user, request) @pytest.mark.parametrize("task_mode", ["annotation", "interpolation"]) @pytest.mark.parametrize("quality", ["compressed", "original"]) - def test_can_get_gt_job_chunk(self, admin_user, tasks, jobs, task_mode, quality, request): + @pytest.mark.parametrize("indexing", ["absolute", "relative"]) + def test_can_get_gt_job_chunk( + self, admin_user, tasks, jobs, task_mode, quality, request, indexing + ): user = admin_user job_frame_count = 4 task = next( @@ -732,40 +731,49 @@ def test_can_get_gt_job_chunk(self, admin_user, tasks, jobs, task_mode, quality, (task_meta, _) = api_client.tasks_api.retrieve_data_meta(task_id) frame_step = int(task_meta.frame_filter.split("=")[-1]) if task_meta.frame_filter else 1 - job_frame_ids = list(range(task_meta.start_frame, task_meta.stop_frame, frame_step))[ - :job_frame_count - ] + task_frame_ids = range(task_meta.start_frame, task_meta.stop_frame + 1, frame_step) + rng = np.random.Generator(np.random.MT19937(42)) + job_frame_ids = sorted(rng.choice(task_frame_ids, job_frame_count, replace=False).tolist()) + gt_job = self._create_gt_job(admin_user, task_id, job_frame_ids) request.addfinalizer(lambda: self._delete_gt_job(admin_user, gt_job.id)) - with make_api_client(admin_user) as api_client: - (chunk_file, response) = api_client.jobs_api.retrieve_data( - gt_job.id, number=0, quality=quality, type="chunk" - ) - assert response.status == HTTPStatus.OK + if indexing == "absolute": + chunk_iter = groupby(task_frame_ids, key=lambda f: f // task_meta.chunk_size) + else: + chunk_iter = groupby(job_frame_ids, key=lambda f: f // task_meta.chunk_size) - frame_range = range( - task_meta.start_frame, min(task_meta.stop_frame + 1, task_meta.chunk_size), frame_step - ) - included_frames = job_frame_ids + for chunk_id, chunk_frames in chunk_iter: + chunk_frames = list(chunk_frames) - # The frame count is the same as in the whole range - # with placeholders in the frames outside the job. - # This is required by the UI implementation - with zipfile.ZipFile(chunk_file) as chunk: - assert set(chunk.namelist()) == set("{:06d}.jpeg".format(i) for i in frame_range) - - for file_info in chunk.filelist: - with chunk.open(file_info) as image_file: - image = Image.open(image_file) - image_data = np.array(image) - - if int(os.path.splitext(file_info.filename)[0]) not in included_frames: - assert image.size == (1, 1) - assert np.all(image_data == 0), image_data - else: - assert image.size > (1, 1) - assert np.any(image_data != 0) + if indexing == "absolute": + kwargs = {"number": chunk_id} + else: + kwargs = {"index": chunk_id} + + with make_api_client(admin_user) as api_client: + (chunk_file, response) = api_client.jobs_api.retrieve_data( + gt_job.id, **kwargs, quality=quality, type="chunk" + ) + assert response.status == HTTPStatus.OK + + # The frame count is the same as in the whole range + # with placeholders in the frames outside the job. + # This is required by the UI implementation + with zipfile.ZipFile(chunk_file) as chunk: + assert set(chunk.namelist()) == set( + f"{i:06d}.jpeg" for i in range(len(chunk_frames)) + ) + + for file_info in chunk.filelist: + with chunk.open(file_info) as image_file: + image = Image.open(image_file) + + chunk_frame_id = int(os.path.splitext(file_info.filename)[0]) + if chunk_frames[chunk_frame_id] not in job_frame_ids: + assert image.size == (1, 1) + else: + assert image.size > (1, 1) def _create_gt_job(self, user, task_id, frames): with make_api_client(user) as api_client: diff --git a/tests/python/rest_api/test_tasks.py b/tests/python/rest_api/test_tasks.py index 3baa4f4228f..eda54b8ddd0 100644 --- a/tests/python/rest_api/test_tasks.py +++ b/tests/python/rest_api/test_tasks.py @@ -2496,12 +2496,16 @@ def get_expected_chunk_abs_frame_ids(chunk_id: int): job_chunk_ids = range(chunk_count) def get_expected_chunk_abs_frame_ids(chunk_id: int): - return range( - job_meta.start_frame - + chunk_id * job_meta.chunk_size * task_spec.frame_step, - job_meta.start_frame - + min((chunk_id + 1) * job_meta.chunk_size, job_meta.size) - * task_spec.frame_step, + return sorted( + frame + for frame in range( + job_meta.start_frame + + chunk_id * job_meta.chunk_size * task_spec.frame_step, + job_meta.start_frame + + min((chunk_id + 1) * job_meta.chunk_size, job_meta.size) + * task_spec.frame_step, + ) + if not job_meta.included_frames or frame in job_meta.included_frames ) for quality, chunk_id in product(["original", "compressed"], job_chunk_ids): From 2424f2b7821a0706915b5af57e7550bda3ae59ed Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 9 Sep 2024 19:36:35 +0300 Subject: [PATCH 097/115] Update UI to support job chunks with non-sequential frame ids --- cvat-core/src/frames.ts | 83 ++++++++++++++++++++++--------- cvat-data/src/ts/cvat-data.ts | 94 +++++++++++++++++++++++------------ cvat/apps/engine/cache.py | 1 - 3 files changed, 123 insertions(+), 55 deletions(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index d2547f0010c..9b11fbfec34 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -40,6 +40,10 @@ const frameDataCache: Record> = {}; +function rangeArray(start: number, end: number): number[] { + return Array.from({ length: end - start }, (v, k) => k + start); +} + export class FramesMetaData { public chunkSize: number; public deletedFrames: Record; @@ -144,6 +148,35 @@ export class FramesMetaData { resetUpdated(): void { this.#updateTrigger.reset(); } + + getFrameIndex(frameNumber: number): number { + if (frameNumber < this.startFrame || frameNumber > this.stopFrame) { + throw new ArgumentError(`Frame number ${frameNumber} doesn't belong to the job`); + } + + let frameIndex = null; + if (this.includedFrames) { + frameIndex = this.includedFrames.indexOf(frameNumber); // TODO: use binary search + if (frameIndex === -1) { + throw new ArgumentError(`Frame number ${frameNumber} doesn't belong to the job`); + } + } else { + frameIndex = frameNumber - this.startFrame; + } + return frameIndex; + } + + getFrameChunkIndex(frame_number: number): number { + return Math.floor(this.getFrameIndex(frame_number) / this.chunkSize); + } + + getFrameSequence(): number[] { + if (this.includedFrames) { + return this.includedFrames; + } + + return rangeArray(this.startFrame, this.stopFrame + 1); + } } export class FrameData { @@ -206,14 +239,12 @@ export class FrameData { } class PrefetchAnalyzer { - #chunkSize: number; #requestedFrames: number[]; - #startFrame: number; + #meta: FramesMetaData; - constructor(chunkSize, startFrame) { - this.#chunkSize = chunkSize; + constructor(meta: FramesMetaData) { this.#requestedFrames = []; - this.#startFrame = startFrame; + this.#meta = meta; } shouldPrefetchNext(current: number, isPlaying: boolean, isChunkCached: (chunk) => boolean): boolean { @@ -221,13 +252,13 @@ class PrefetchAnalyzer { return true; } - const currentChunk = Math.floor((current - this.#startFrame) / this.#chunkSize); + const currentChunk = this.#meta.getFrameChunkIndex(current); const { length } = this.#requestedFrames; const isIncreasingOrder = this.#requestedFrames .every((val, index) => index === 0 || val > this.#requestedFrames[index - 1]); if ( length && (isIncreasingOrder && current > this.#requestedFrames[length - 1]) && - ((current - this.#startFrame) % this.#chunkSize) >= Math.ceil(this.#chunkSize / 2) && + (this.#meta.getFrameIndex(current) % this.#meta.chunkSize) >= Math.ceil(this.#meta.chunkSize / 2) && !isChunkCached(currentChunk + 1) ) { // is increasing order including the current frame @@ -249,7 +280,7 @@ class PrefetchAnalyzer { this.#requestedFrames.push(frame); // only half of chunk size is considered in this logic - const limit = Math.ceil(this.#chunkSize / 2); + const limit = Math.ceil(this.#meta.chunkSize / 2); if (this.#requestedFrames.length > limit) { this.#requestedFrames.shift(); } @@ -264,24 +295,27 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { imageData: ImageBitmap | Blob; } | Blob>((resolve, reject) => { const { - provider, prefetchAnalyzer, chunkSize, startFrame, stopFrame, + meta, provider, prefetchAnalyzer, chunkSize, startFrame, stopFrame, decodeForward, forwardStep, decodedBlocksCacheSize, } = frameDataCache[this.jobID]; const requestId = +_.uniqueId(); - const chunkNumber = Math.floor((this.number - startFrame) / chunkSize); + const chunkNumber = meta.getFrameChunkIndex(this.number); const frame = provider.frame(this.number); - function findTheNextNotDecodedChunk(searchFrom: number): number { - let firstFrameInNextChunk = searchFrom + forwardStep; - let nextChunkNumber = Math.floor((firstFrameInNextChunk - startFrame) / chunkSize); + function findTheNextNotDecodedChunk( + searchFrom: number, isIndex: boolean = false, + ): number { + const currentFrameIndex = isIndex ? searchFrom : meta.getFrameIndex(searchFrom); + let nextFrameIndex = currentFrameIndex + forwardStep; + let nextChunkNumber = Math.floor(nextFrameIndex / chunkSize); while (nextChunkNumber === chunkNumber) { - firstFrameInNextChunk += forwardStep; - nextChunkNumber = Math.floor((firstFrameInNextChunk - startFrame) / chunkSize); + nextFrameIndex += forwardStep; + nextChunkNumber = Math.floor(nextFrameIndex / chunkSize); } if (provider.isChunkCached(nextChunkNumber)) { - return findTheNextNotDecodedChunk(firstFrameInNextChunk); + return findTheNextNotDecodedChunk(nextFrameIndex, true); } return nextChunkNumber; @@ -319,8 +353,10 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { provider.cleanup(1); provider.requestDecodeBlock( chunk, - startFrame + nextChunkNumber * chunkSize, - Math.min(stopFrame, startFrame + (nextChunkNumber + 1) * chunkSize - 1), + meta.getFrameSequence().slice( + nextChunkNumber * chunkSize, + (nextChunkNumber + 1) * chunkSize, + ), () => {}, releasePromise, releasePromise, @@ -381,8 +417,10 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { provider .requestDecodeBlock( chunk, - startFrame + chunkNumber * chunkSize, - Math.min(stopFrame, startFrame + (chunkNumber + 1) * chunkSize - 1), + meta.getFrameSequence().slice( + chunkNumber * chunkSize, + (chunkNumber + 1) * chunkSize, + ), (_frame: number, bitmap: ImageBitmap | Blob) => { if (decodeForward) { // resolve immediately only if is not playing @@ -613,12 +651,11 @@ export async function getFrame( forwardStep: step, provider: new FrameDecoder( blockType, - chunkSize, decodedBlocksCacheSize, - startFrame, + meta.getFrameChunkIndex.bind(meta), dimension, ), - prefetchAnalyzer: new PrefetchAnalyzer(chunkSize, startFrame), + prefetchAnalyzer: new PrefetchAnalyzer(meta), decodedBlocksCacheSize, activeChunkRequest: null, activeContextRequest: null, diff --git a/cvat-data/src/ts/cvat-data.ts b/cvat-data/src/ts/cvat-data.ts index ec9fc5cccc5..05aa359dc45 100644 --- a/cvat-data/src/ts/cvat-data.ts +++ b/cvat-data/src/ts/cvat-data.ts @@ -72,8 +72,8 @@ export function decodeContextImages( decodeContextImages.mutex = new Mutex(); interface BlockToDecode { - start: number; - end: number; + frameNumbers: number[]; + chunkNumber: number; block: ArrayBuffer; onDecodeAll(): void; onDecode(frame: number, bitmap: ImageBitmap | Blob): void; @@ -82,7 +82,6 @@ interface BlockToDecode { export class FrameDecoder { private blockType: BlockType; - private chunkSize: number; /* ImageBitmap when decode zip or video chunks Blob when 3D dimension @@ -100,13 +99,12 @@ export class FrameDecoder { private renderHeight: number; private zipWorker: Worker | null; private videoWorker: Worker | null; - private startFrame: number; + private getChunkNumber: (frame: number) => number; constructor( blockType: BlockType, - chunkSize: number, cachedBlockCount: number, - startFrame: number, + getChunkNumber: (frame: number) => number, dimension: DimensionType = DimensionType.DIMENSION_2D, ) { this.mutex = new Mutex(); @@ -119,8 +117,7 @@ export class FrameDecoder { this.renderWidth = 1920; this.renderHeight = 1080; - this.chunkSize = chunkSize; - this.startFrame = startFrame; + this.getChunkNumber = getChunkNumber; this.blockType = blockType; this.decodedChunks = {}; @@ -158,17 +155,43 @@ export class FrameDecoder { } } + private validateFrameNumbers(frameNumbers: number[]): void { + if (!frameNumbers || !frameNumbers.length) { + throw new Error('frameNumbers must not be empty'); + } + + // ensure is ordered + for (let i = 1; i < frameNumbers.length; ++i) { + const prev = frameNumbers[i - 1]; + const current = frameNumbers[i]; + if (current <= prev) { + throw new Error( + 'frameNumbers must be sorted in ascending order, ' + + `got a (${prev}, ${current}) pair instead`, + ); + } + } + } + + private arraysEqual(a: number[], b: number[]): boolean { + return ( + a.length === b.length && + a.every((element, index) => element === b[index]) + ); + } + requestDecodeBlock( block: ArrayBuffer, - start: number, - end: number, + frameNumbers: number[], onDecode: (frame: number, bitmap: ImageBitmap | Blob) => void, onDecodeAll: () => void, onReject: (e: Error) => void, ): void { + this.validateFrameNumbers(frameNumbers); + if (this.requestedChunkToDecode !== null) { // a chunk was already requested to be decoded, but decoding didn't start yet - if (start === this.requestedChunkToDecode.start && end === this.requestedChunkToDecode.end) { + if (this.arraysEqual(frameNumbers, this.requestedChunkToDecode.frameNumbers)) { // it was the same chunk this.requestedChunkToDecode.onReject(new RequestOutdatedError()); @@ -178,12 +201,14 @@ export class FrameDecoder { // it was other chunk this.requestedChunkToDecode.onReject(new RequestOutdatedError()); } - } else if (this.chunkIsBeingDecoded === null || this.chunkIsBeingDecoded.start !== start) { + } else if (this.chunkIsBeingDecoded === null || + !this.arraysEqual(frameNumbers, this.requestedChunkToDecode.frameNumbers) + ) { // everything was decoded or decoding other chunk is in process this.requestedChunkToDecode = { + frameNumbers, + chunkNumber: this.getChunkNumber(frameNumbers[0]), block, - start, - end, onDecode, onDecodeAll, onReject, @@ -206,7 +231,7 @@ export class FrameDecoder { } frame(frameNumber: number): ImageBitmap | Blob | null { - const chunkNumber = Math.floor((frameNumber - this.startFrame) / this.chunkSize); + const chunkNumber = this.getChunkNumber(frameNumber); if (chunkNumber in this.decodedChunks) { return this.decodedChunks[chunkNumber][frameNumber]; } @@ -256,8 +281,8 @@ export class FrameDecoder { releaseMutex(); }; try { - const { start, end, block } = this.requestedChunkToDecode; - if (start !== blockToDecode.start) { + const { frameNumbers, chunkNumber, block } = this.requestedChunkToDecode; + if (!this.arraysEqual(frameNumbers, blockToDecode.frameNumbers)) { // request is not relevant, another block was already requested // it happens when A is being decoded, B comes and wait for mutex, C comes and wait for mutex // B is not necessary anymore, because C already was requested @@ -265,7 +290,8 @@ export class FrameDecoder { throw new RequestOutdatedError(); } - const chunkNumber = Math.floor((start - this.startFrame) / this.chunkSize); + const getFrameNumber = (chunkFrameIndex: number): number => frameNumbers[chunkFrameIndex]; + this.orderedStack = [chunkNumber, ...this.orderedStack]; this.cleanup(); const decodedFrames: Record = {}; @@ -276,7 +302,7 @@ export class FrameDecoder { this.videoWorker = new Worker( new URL('./3rdparty/Decoder.worker', import.meta.url), ); - let index = start; + let index = 0; this.videoWorker.onmessage = (e) => { if (e.data.consoleLog) { @@ -284,6 +310,7 @@ export class FrameDecoder { return; } const keptIndex = index; + const frameNumber = getFrameNumber(keptIndex); // do not use e.data.height and e.data.width because they might be not correct // instead, try to understand real height and width of decoded image via scale factor @@ -298,10 +325,10 @@ export class FrameDecoder { width, height, )).then((bitmap) => { - decodedFrames[keptIndex] = bitmap; - this.chunkIsBeingDecoded.onDecode(keptIndex, decodedFrames[keptIndex]); + decodedFrames[frameNumber] = bitmap; + this.chunkIsBeingDecoded.onDecode(frameNumber, decodedFrames[frameNumber]); - if (keptIndex === end) { + if (keptIndex === frameNumbers.length - 1) { this.decodedChunks[chunkNumber] = decodedFrames; this.chunkIsBeingDecoded.onDecodeAll(); this.chunkIsBeingDecoded = null; @@ -346,7 +373,7 @@ export class FrameDecoder { this.zipWorker = this.zipWorker || new Worker( new URL('./unzip_imgs.worker', import.meta.url), ); - let index = start; + let decodedCount = 0; this.zipWorker.onmessage = async (event) => { if (event.data.error) { @@ -356,16 +383,18 @@ export class FrameDecoder { return; } - decodedFrames[event.data.index] = event.data.data as ImageBitmap | Blob; - this.chunkIsBeingDecoded.onDecode(event.data.index, decodedFrames[event.data.index]); + const frameNumber = getFrameNumber(event.data.index); + decodedFrames[frameNumber] = event.data.data as ImageBitmap | Blob; + this.chunkIsBeingDecoded.onDecode(frameNumber, decodedFrames[frameNumber]); - if (index === end) { + if (decodedCount === frameNumbers.length - 1) { this.decodedChunks[chunkNumber] = decodedFrames; this.chunkIsBeingDecoded.onDecodeAll(); this.chunkIsBeingDecoded = null; release(); } - index++; + + decodedCount++; }; this.zipWorker.onerror = (event: ErrorEvent) => { @@ -376,8 +405,8 @@ export class FrameDecoder { this.zipWorker.postMessage({ block, - start, - end, + start: 0, + end: frameNumbers.length - 1, dimension: this.dimension, dimension2D: DimensionType.DIMENSION_2D, }); @@ -403,8 +432,11 @@ export class FrameDecoder { } public cachedChunks(includeInProgress = false): number[] { - const chunkIsBeingDecoded = includeInProgress && this.chunkIsBeingDecoded ? - Math.floor(this.chunkIsBeingDecoded.start / this.chunkSize) : null; + const chunkIsBeingDecoded = ( + includeInProgress && this.chunkIsBeingDecoded ? + this.chunkIsBeingDecoded.chunkNumber : + null + ); return Object.keys(this.decodedChunks).map((chunkNumber: string) => +chunkNumber).concat( ...(chunkIsBeingDecoded !== null ? [chunkIsBeingDecoded] : []), ).sort((a, b) => a - b); diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 8dd3f8f02b5..0d1699b4fdd 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -497,7 +497,6 @@ def get_frames(): else: # Populate skipped frames with placeholder data, # this is required for video chunk decoding implementation in UI - # TODO: try to fix decoding in UI frame = io.BytesIO(dummy_frame.getvalue()) yield (frame, None, None) From fe60bdf7cfd3c76fd148f7c79b67d2c1b2e63dfd Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 9 Sep 2024 21:14:11 +0300 Subject: [PATCH 098/115] Fix job frame retrieval --- cvat/apps/engine/views.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/cvat/apps/engine/views.py b/cvat/apps/engine/views.py index ebda8c7fae3..e30a857d301 100644 --- a/cvat/apps/engine/views.py +++ b/cvat/apps/engine/views.py @@ -760,8 +760,8 @@ def __init__( self.type = data_type - self.number = int(data_index) if data_index is not None else None - self.task_chunk_number = int(data_num) if data_num is not None else None + self.index = int(data_index) if data_index is not None else None + self.number = int(data_num) if data_num is not None else None self.quality = FrameQuality.COMPRESSED \ if data_quality == 'compressed' else FrameQuality.ORIGINAL @@ -772,12 +772,19 @@ def _get_frame_provider(self) -> JobFrameProvider: return JobFrameProvider(self._db_job) def __call__(self): - if self.type == 'chunk' and self.task_chunk_number is not None: + if self.type == 'chunk': # Reproduce the task chunk indexing frame_provider = self._get_frame_provider() - data = frame_provider.get_chunk( - self.task_chunk_number, quality=self.quality, is_task_chunk=True - ) + + if self.index is not None: + data = frame_provider.get_chunk( + self.index, quality=self.quality, is_task_chunk=False + ) + else: + data = frame_provider.get_chunk( + self.number, quality=self.quality, is_task_chunk=True + ) + return HttpResponse(data.data.getvalue(), content_type=data.mime) else: return super().__call__() From 6ddb6bf8fef3276e3667646b7f08c8da01d768f2 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 10 Sep 2024 00:40:09 +0300 Subject: [PATCH 099/115] Fix 3d task chunk writing --- cvat/apps/engine/media_extractors.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 1d5c6fde76e..9ddbad10e3a 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -914,7 +914,10 @@ def save_as_chunk(self, images: Iterator[tuple[Image.Image|io.IOBase|str, str, s else: output = path else: - output, ext = self._write_pcd_file(path)[0:2] + if isinstance(image, io.BytesIO): + output, ext = self._write_pcd_file(image)[0:2] + else: + output, ext = self._write_pcd_file(path)[0:2] arcname = '{:06d}.{}'.format(idx, ext) if isinstance(output, io.BytesIO): @@ -945,7 +948,11 @@ def save_as_chunk( w, h = img.size extension = self.IMAGE_EXT else: - image_buf, extension, w, h = self._write_pcd_file(path) + if isinstance(image, io.BytesIO): + image_buf, extension, w, h = self._write_pcd_file(image) + else: + image_buf, extension, w, h = self._write_pcd_file(path) + image_sizes.append((w, h)) arcname = '{:06d}.{}'.format(idx, extension) zip_chunk.writestr(arcname, image_buf.getvalue()) From 4fa7b9762fccf1be981d8266c3722d97b8340168 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 10 Sep 2024 20:48:45 +0300 Subject: [PATCH 100/115] Fix frame retrieval in UI --- cvat-core/src/frames.ts | 109 ++++++++++++++++------ cvat-ui/src/actions/annotation-actions.ts | 3 +- 2 files changed, 82 insertions(+), 30 deletions(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index 9b11fbfec34..bcba1899acb 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -40,8 +40,11 @@ const frameDataCache: Record> = {}; -function rangeArray(start: number, end: number): number[] { - return Array.from({ length: end - start }, (v, k) => k + start); +function rangeArray(start: number, end: number, step: number = 1): number[] { + return Array.from( + { length: +(start < end) * Math.ceil((end - start) / step) }, + (v, k) => k * step + start, + ); } export class FramesMetaData { @@ -149,33 +152,46 @@ export class FramesMetaData { this.#updateTrigger.reset(); } - getFrameIndex(frameNumber: number): number { - if (frameNumber < this.startFrame || frameNumber > this.stopFrame) { - throw new ArgumentError(`Frame number ${frameNumber} doesn't belong to the job`); + getFrameIndex(dataFrameNumber: number): number { + // TODO: migrate to local frame numbers to simplify code + + if (dataFrameNumber < this.startFrame || dataFrameNumber > this.stopFrame) { + throw new ArgumentError(`Frame number ${dataFrameNumber} doesn't belong to the job`); } let frameIndex = null; if (this.includedFrames) { - frameIndex = this.includedFrames.indexOf(frameNumber); // TODO: use binary search + frameIndex = this.includedFrames.indexOf(dataFrameNumber); // TODO: use binary search if (frameIndex === -1) { - throw new ArgumentError(`Frame number ${frameNumber} doesn't belong to the job`); + throw new ArgumentError(`Frame number ${dataFrameNumber} doesn't belong to the job`); } } else { - frameIndex = frameNumber - this.startFrame; + frameIndex = Math.floor((dataFrameNumber - this.startFrame) / this.getFrameStep()); } return frameIndex; } - getFrameChunkIndex(frame_number: number): number { - return Math.floor(this.getFrameIndex(frame_number) / this.chunkSize); + getFrameChunkIndex(dataFrameNumber: number): number { + return Math.floor(this.getFrameIndex(dataFrameNumber) / this.chunkSize); } - getFrameSequence(): number[] { + getFrameStep(): number { + if (this.frameFilter) { + const frameStepParts = this.frameFilter.split('=', 2); + if (frameStepParts.length !== 2) { + throw new Error(`Invalid frame filter '${this.frameFilter}'`); + } + return parseInt(frameStepParts[1], 10); + } + return 1; + } + + getDataFrameNumbers(): number[] { if (this.includedFrames) { return this.includedFrames; } - return rangeArray(this.startFrame, this.stopFrame + 1); + return rangeArray(this.startFrame, this.stopFrame + 1, this.getFrameStep()); } } @@ -241,10 +257,12 @@ export class FrameData { class PrefetchAnalyzer { #requestedFrames: number[]; #meta: FramesMetaData; + #getDataFrameNumber: (frameNumber: number) => number; - constructor(meta: FramesMetaData) { + constructor(meta: FramesMetaData, dataFrameNumberGetter: (frameNumber: number) => number) { this.#requestedFrames = []; this.#meta = meta; + this.#getDataFrameNumber = dataFrameNumberGetter; } shouldPrefetchNext(current: number, isPlaying: boolean, isChunkCached: (chunk) => boolean): boolean { @@ -252,13 +270,16 @@ class PrefetchAnalyzer { return true; } - const currentChunk = this.#meta.getFrameChunkIndex(current); + const currentDataFrameNumber = this.#getDataFrameNumber(current); + const currentChunk = this.#meta.getFrameChunkIndex(currentDataFrameNumber); const { length } = this.#requestedFrames; const isIncreasingOrder = this.#requestedFrames .every((val, index) => index === 0 || val > this.#requestedFrames[index - 1]); if ( length && (isIncreasingOrder && current > this.#requestedFrames[length - 1]) && - (this.#meta.getFrameIndex(current) % this.#meta.chunkSize) >= Math.ceil(this.#meta.chunkSize / 2) && + ( + this.#meta.getFrameIndex(currentDataFrameNumber) % this.#meta.chunkSize + ) >= Math.ceil(this.#meta.chunkSize / 2) && !isChunkCached(currentChunk + 1) ) { // is increasing order including the current frame @@ -287,6 +308,18 @@ class PrefetchAnalyzer { } } +function getDataStartFrame(meta: FramesMetaData, localStartFrame: number): number { + return meta.startFrame - localStartFrame * meta.getFrameStep(); +} + +function getDataFrameNumber(frameNumber: number, dataStartFrame: number, step: number): number { + return frameNumber * step + dataStartFrame; +} + +function getFrameNumber(dataFrameNumber: number, dataStartFrame: number, step: number): number { + return (dataFrameNumber - dataStartFrame) / step; +} + Object.defineProperty(FrameData.prototype.data, 'implementation', { value(this: FrameData, onServerRequest) { return new Promise<{ @@ -300,14 +333,18 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { } = frameDataCache[this.jobID]; const requestId = +_.uniqueId(); - const chunkNumber = meta.getFrameChunkIndex(this.number); + const dataStartFrame = getDataStartFrame(meta, startFrame); + const requestedDataFrameNumber = getDataFrameNumber( + this.number, dataStartFrame, meta.getFrameStep(), + ); + const chunkNumber = meta.getFrameChunkIndex(requestedDataFrameNumber); + const segmentFrameNumbers = meta.getDataFrameNumbers().map( + (dataFrameNumber: number) => getFrameNumber(dataFrameNumber, dataStartFrame, meta.getFrameStep()), + ); const frame = provider.frame(this.number); - function findTheNextNotDecodedChunk( - searchFrom: number, isIndex: boolean = false, - ): number { - const currentFrameIndex = isIndex ? searchFrom : meta.getFrameIndex(searchFrom); - let nextFrameIndex = currentFrameIndex + forwardStep; + function findTheNextNotDecodedChunk(searchFrom: number): number { + let nextFrameIndex = searchFrom + forwardStep; let nextChunkNumber = Math.floor(nextFrameIndex / chunkSize); while (nextChunkNumber === chunkNumber) { nextFrameIndex += forwardStep; @@ -315,7 +352,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { } if (provider.isChunkCached(nextChunkNumber)) { - return findTheNextNotDecodedChunk(nextFrameIndex, true); + return findTheNextNotDecodedChunk(nextFrameIndex); } return nextChunkNumber; @@ -329,7 +366,9 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { (chunk) => provider.isChunkCached(chunk), ) && decodedBlocksCacheSize > 1 && !frameDataCache[this.jobID].activeChunkRequest ) { - const nextChunkNumber = findTheNextNotDecodedChunk(this.number); + const nextChunkNumber = findTheNextNotDecodedChunk( + meta.getFrameIndex(requestedDataFrameNumber), + ); const predecodeChunksMax = Math.floor(decodedBlocksCacheSize / 2); if (startFrame + nextChunkNumber * chunkSize <= stopFrame && nextChunkNumber <= chunkNumber + predecodeChunksMax @@ -353,7 +392,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { provider.cleanup(1); provider.requestDecodeBlock( chunk, - meta.getFrameSequence().slice( + segmentFrameNumbers.slice( nextChunkNumber * chunkSize, (nextChunkNumber + 1) * chunkSize, ), @@ -417,7 +456,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { provider .requestDecodeBlock( chunk, - meta.getFrameSequence().slice( + segmentFrameNumbers.slice( chunkNumber * chunkSize, (chunkNumber + 1) * chunkSize, ), @@ -641,6 +680,13 @@ export async function getFrame( const decodedBlocksCacheSize = Math.min( Math.floor((2048 * 1024 * 1024) / ((mean + stdDev) * 4 * chunkSize)) || 1, 10, ); + + // TODO: migrate to local frame numbers + const dataStartFrame = getDataStartFrame(meta, startFrame); + const dataFrameNumberGetter = (frameNumber: number): number => ( + getDataFrameNumber(frameNumber, dataStartFrame, meta.getFrameStep()) + ); + frameDataCache[jobID] = { meta, chunkSize, @@ -652,10 +698,12 @@ export async function getFrame( provider: new FrameDecoder( blockType, decodedBlocksCacheSize, - meta.getFrameChunkIndex.bind(meta), + (frameNumber: number): number => ( + meta.getFrameChunkIndex(dataFrameNumberGetter(frameNumber)) + ), dimension, ), - prefetchAnalyzer: new PrefetchAnalyzer(meta), + prefetchAnalyzer: new PrefetchAnalyzer(meta, dataFrameNumberGetter), decodedBlocksCacheSize, activeChunkRequest: null, activeContextRequest: null, @@ -729,8 +777,11 @@ export async function findFrame( let lastUndeletedFrame = null; const check = (frame): boolean => { if (meta.includedFrames) { - return (meta.includedFrames.includes(frame)) && - (!filters.notDeleted || !(frame in meta.deletedFrames)); + // meta.includedFrames contains input frame numbers now + const dataStartFrame = meta.startFrame; // this is only true when includedFrames is set + return (meta.includedFrames.includes( + getDataFrameNumber(frame, dataStartFrame, meta.getFrameStep())) + ) && (!filters.notDeleted || !(frame in meta.deletedFrames)); } if (filters.notDeleted) { return !(frame in meta.deletedFrames); diff --git a/cvat-ui/src/actions/annotation-actions.ts b/cvat-ui/src/actions/annotation-actions.ts index a5594093db5..c54d548b78d 100644 --- a/cvat-ui/src/actions/annotation-actions.ts +++ b/cvat-ui/src/actions/annotation-actions.ts @@ -905,7 +905,8 @@ export function getJobAsync({ // frame query parameter does not work for GT job const frameNumber = Number.isInteger(initialFrame) && gtJob?.id !== job.id ? - initialFrame as number : (await job.frames.search( + initialFrame as number : + (await job.frames.search( { notDeleted: !showDeletedFrames }, job.startFrame, job.stopFrame, )) || job.startFrame; From 0e95b4016ffedfbec93ab7f4170b2a5b2b2e52b4 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Wed, 11 Sep 2024 11:25:50 +0300 Subject: [PATCH 101/115] Fix chunk availability check --- cvat/apps/engine/cache.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 0d1699b4fdd..b32dbeef00e 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -434,15 +434,15 @@ def prepare_custom_masked_range_segment_chunk( use_cached_data = False if db_task.mode != "interpolation": required_frame_set = set(frame_ids) - all_chunks_available = all( + available_chunks = [ self._has_key(self._make_chunk_key(db_segment, chunk_number, quality=quality)) for db_segment in db_task.segment_set.filter(type=models.SegmentType.RANGE).all() for chunk_number, _ in groupby( - required_frame_set.intersection(db_segment.frame_set), + sorted(required_frame_set.intersection(db_segment.frame_set)), key=lambda frame: frame // db_data.chunk_size, ) - ) - use_cached_data = all_chunks_available + ] + use_cached_data = bool(available_chunks) and all(available_chunks) if hasattr(db_data, "video"): frame_size = (db_data.video.width, db_data.video.height) From 79bb1f74be75bbb3485c8f772b0e1cf3cae540af Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 12 Sep 2024 16:14:36 +0300 Subject: [PATCH 102/115] Remove array comparisons --- cvat-core/src/frames.ts | 2 ++ cvat-data/src/ts/cvat-data.ts | 36 ++++++++++++++++------------------- 2 files changed, 18 insertions(+), 20 deletions(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index bcba1899acb..47739b7e5ca 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -392,6 +392,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { provider.cleanup(1); provider.requestDecodeBlock( chunk, + nextChunkNumber, segmentFrameNumbers.slice( nextChunkNumber * chunkSize, (nextChunkNumber + 1) * chunkSize, @@ -456,6 +457,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { provider .requestDecodeBlock( chunk, + chunkNumber, segmentFrameNumbers.slice( chunkNumber * chunkSize, (chunkNumber + 1) * chunkSize, diff --git a/cvat-data/src/ts/cvat-data.ts b/cvat-data/src/ts/cvat-data.ts index 05aa359dc45..5b9e9286a2f 100644 --- a/cvat-data/src/ts/cvat-data.ts +++ b/cvat-data/src/ts/cvat-data.ts @@ -72,7 +72,7 @@ export function decodeContextImages( decodeContextImages.mutex = new Mutex(); interface BlockToDecode { - frameNumbers: number[]; + chunkFrameNumbers: number[]; chunkNumber: number; block: ArrayBuffer; onDecodeAll(): void; @@ -173,25 +173,19 @@ export class FrameDecoder { } } - private arraysEqual(a: number[], b: number[]): boolean { - return ( - a.length === b.length && - a.every((element, index) => element === b[index]) - ); - } - requestDecodeBlock( block: ArrayBuffer, - frameNumbers: number[], + chunkNumber: number, + chunkFrameNumbers: number[], onDecode: (frame: number, bitmap: ImageBitmap | Blob) => void, onDecodeAll: () => void, onReject: (e: Error) => void, ): void { - this.validateFrameNumbers(frameNumbers); + this.validateFrameNumbers(chunkFrameNumbers); if (this.requestedChunkToDecode !== null) { // a chunk was already requested to be decoded, but decoding didn't start yet - if (this.arraysEqual(frameNumbers, this.requestedChunkToDecode.frameNumbers)) { + if (chunkNumber === this.requestedChunkToDecode.chunkNumber) { // it was the same chunk this.requestedChunkToDecode.onReject(new RequestOutdatedError()); @@ -202,12 +196,12 @@ export class FrameDecoder { this.requestedChunkToDecode.onReject(new RequestOutdatedError()); } } else if (this.chunkIsBeingDecoded === null || - !this.arraysEqual(frameNumbers, this.requestedChunkToDecode.frameNumbers) + chunkNumber !== this.chunkIsBeingDecoded.chunkNumber ) { // everything was decoded or decoding other chunk is in process this.requestedChunkToDecode = { - frameNumbers, - chunkNumber: this.getChunkNumber(frameNumbers[0]), + chunkFrameNumbers, + chunkNumber, block, onDecode, onDecodeAll, @@ -281,8 +275,8 @@ export class FrameDecoder { releaseMutex(); }; try { - const { frameNumbers, chunkNumber, block } = this.requestedChunkToDecode; - if (!this.arraysEqual(frameNumbers, blockToDecode.frameNumbers)) { + const { chunkFrameNumbers, chunkNumber, block } = this.requestedChunkToDecode; + if (chunkNumber !== blockToDecode.chunkNumber) { // request is not relevant, another block was already requested // it happens when A is being decoded, B comes and wait for mutex, C comes and wait for mutex // B is not necessary anymore, because C already was requested @@ -290,7 +284,9 @@ export class FrameDecoder { throw new RequestOutdatedError(); } - const getFrameNumber = (chunkFrameIndex: number): number => frameNumbers[chunkFrameIndex]; + const getFrameNumber = (chunkFrameIndex: number): number => ( + chunkFrameNumbers[chunkFrameIndex] + ); this.orderedStack = [chunkNumber, ...this.orderedStack]; this.cleanup(); @@ -328,7 +324,7 @@ export class FrameDecoder { decodedFrames[frameNumber] = bitmap; this.chunkIsBeingDecoded.onDecode(frameNumber, decodedFrames[frameNumber]); - if (keptIndex === frameNumbers.length - 1) { + if (keptIndex === chunkFrameNumbers.length - 1) { this.decodedChunks[chunkNumber] = decodedFrames; this.chunkIsBeingDecoded.onDecodeAll(); this.chunkIsBeingDecoded = null; @@ -387,7 +383,7 @@ export class FrameDecoder { decodedFrames[frameNumber] = event.data.data as ImageBitmap | Blob; this.chunkIsBeingDecoded.onDecode(frameNumber, decodedFrames[frameNumber]); - if (decodedCount === frameNumbers.length - 1) { + if (decodedCount === chunkFrameNumbers.length - 1) { this.decodedChunks[chunkNumber] = decodedFrames; this.chunkIsBeingDecoded.onDecodeAll(); this.chunkIsBeingDecoded = null; @@ -406,7 +402,7 @@ export class FrameDecoder { this.zipWorker.postMessage({ block, start: 0, - end: frameNumbers.length - 1, + end: chunkFrameNumbers.length - 1, dimension: this.dimension, dimension2D: DimensionType.DIMENSION_2D, }); From 55a8424ccb1db668308e536670c6b845a3b701a2 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 12 Sep 2024 16:17:26 +0300 Subject: [PATCH 103/115] Update validateFrameNumbers --- cvat-data/src/ts/cvat-data.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cvat-data/src/ts/cvat-data.ts b/cvat-data/src/ts/cvat-data.ts index 5b9e9286a2f..97e796af9df 100644 --- a/cvat-data/src/ts/cvat-data.ts +++ b/cvat-data/src/ts/cvat-data.ts @@ -156,8 +156,8 @@ export class FrameDecoder { } private validateFrameNumbers(frameNumbers: number[]): void { - if (!frameNumbers || !frameNumbers.length) { - throw new Error('frameNumbers must not be empty'); + if (!Array.isArray(frameNumbers) || !frameNumbers.length) { + throw new Error('chunkFrameNumbers must not be empty'); } // ensure is ordered @@ -166,7 +166,7 @@ export class FrameDecoder { const current = frameNumbers[i]; if (current <= prev) { throw new Error( - 'frameNumbers must be sorted in ascending order, ' + + 'chunkFrameNumbers must be sorted in the ascending order, ' + `got a (${prev}, ${current}) pair instead`, ); } From add5ae6e5cfe9375666d04d284d26611f214288b Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 12 Sep 2024 16:51:36 +0300 Subject: [PATCH 104/115] Use builtins for range and binary search, convert frame step into a cached field --- cvat-core/src/frames.ts | 58 +++++++++++++++++++++-------------------- 1 file changed, 30 insertions(+), 28 deletions(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index 47739b7e5ca..02cc63bee54 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -3,7 +3,7 @@ // // SPDX-License-Identifier: MIT -import _ from 'lodash'; +import _, { range, sortedIndexOf } from 'lodash'; import { FrameDecoder, BlockType, DimensionType, ChunkQuality, decodeContextImages, RequestOutdatedError, } from 'cvat-data'; @@ -40,13 +40,6 @@ const frameDataCache: Record> = {}; -function rangeArray(start: number, end: number, step: number = 1): number[] { - return Array.from( - { length: +(start < end) * Math.ceil((end - start) / step) }, - (v, k) => k * step + start, - ); -} - export class FramesMetaData { public chunkSize: number; public deletedFrames: Record; @@ -62,6 +55,7 @@ export class FramesMetaData { public size: number; public startFrame: number; public stopFrame: number; + public frameStep: number; #updateTrigger: FieldUpdateTrigger; @@ -110,6 +104,17 @@ export class FramesMetaData { } } + const frameStep: number = (() => { + if (data.frame_filter) { + const frameStepParts = data.frame_filter.split('=', 2); + if (frameStepParts.length !== 2) { + throw new ArgumentError(`Invalid frame filter '${data.frame_filter}'`); + } + return +frameStepParts[1]; + } + return 1; + })(); + Object.defineProperties( this, Object.freeze({ @@ -140,6 +145,9 @@ export class FramesMetaData { stopFrame: { get: () => data.stop_frame, }, + frameStep: { + get: () => frameStep, + }, }), ); } @@ -153,7 +161,10 @@ export class FramesMetaData { } getFrameIndex(dataFrameNumber: number): number { - // TODO: migrate to local frame numbers to simplify code + // Here we use absolute (task source data) frame numbers. + // TODO: migrate from data frame numbers to local frame numbers to simplify code. + // Requires server changes in api/jobs/{id}/data/meta/ + // for included_frames, start_frame, stop_frame fields if (dataFrameNumber < this.startFrame || dataFrameNumber > this.stopFrame) { throw new ArgumentError(`Frame number ${dataFrameNumber} doesn't belong to the job`); @@ -161,12 +172,12 @@ export class FramesMetaData { let frameIndex = null; if (this.includedFrames) { - frameIndex = this.includedFrames.indexOf(dataFrameNumber); // TODO: use binary search + frameIndex = sortedIndexOf(this.includedFrames, dataFrameNumber); if (frameIndex === -1) { throw new ArgumentError(`Frame number ${dataFrameNumber} doesn't belong to the job`); } } else { - frameIndex = Math.floor((dataFrameNumber - this.startFrame) / this.getFrameStep()); + frameIndex = Math.floor((dataFrameNumber - this.startFrame) / this.frameStep); } return frameIndex; } @@ -175,23 +186,12 @@ export class FramesMetaData { return Math.floor(this.getFrameIndex(dataFrameNumber) / this.chunkSize); } - getFrameStep(): number { - if (this.frameFilter) { - const frameStepParts = this.frameFilter.split('=', 2); - if (frameStepParts.length !== 2) { - throw new Error(`Invalid frame filter '${this.frameFilter}'`); - } - return parseInt(frameStepParts[1], 10); - } - return 1; - } - getDataFrameNumbers(): number[] { if (this.includedFrames) { return this.includedFrames; } - return rangeArray(this.startFrame, this.stopFrame + 1, this.getFrameStep()); + return range(this.startFrame, this.stopFrame + 1, this.frameStep); } } @@ -309,7 +309,7 @@ class PrefetchAnalyzer { } function getDataStartFrame(meta: FramesMetaData, localStartFrame: number): number { - return meta.startFrame - localStartFrame * meta.getFrameStep(); + return meta.startFrame - localStartFrame * meta.frameStep; } function getDataFrameNumber(frameNumber: number, dataStartFrame: number, step: number): number { @@ -335,11 +335,13 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { const requestId = +_.uniqueId(); const dataStartFrame = getDataStartFrame(meta, startFrame); const requestedDataFrameNumber = getDataFrameNumber( - this.number, dataStartFrame, meta.getFrameStep(), + this.number, dataStartFrame, meta.frameStep, ); const chunkNumber = meta.getFrameChunkIndex(requestedDataFrameNumber); const segmentFrameNumbers = meta.getDataFrameNumbers().map( - (dataFrameNumber: number) => getFrameNumber(dataFrameNumber, dataStartFrame, meta.getFrameStep()), + (dataFrameNumber: number) => getFrameNumber( + dataFrameNumber, dataStartFrame, meta.frameStep, + ), ); const frame = provider.frame(this.number); @@ -686,7 +688,7 @@ export async function getFrame( // TODO: migrate to local frame numbers const dataStartFrame = getDataStartFrame(meta, startFrame); const dataFrameNumberGetter = (frameNumber: number): number => ( - getDataFrameNumber(frameNumber, dataStartFrame, meta.getFrameStep()) + getDataFrameNumber(frameNumber, dataStartFrame, meta.frameStep) ); frameDataCache[jobID] = { @@ -782,7 +784,7 @@ export async function findFrame( // meta.includedFrames contains input frame numbers now const dataStartFrame = meta.startFrame; // this is only true when includedFrames is set return (meta.includedFrames.includes( - getDataFrameNumber(frame, dataStartFrame, meta.getFrameStep())) + getDataFrameNumber(frame, dataStartFrame, meta.frameStep)) ) && (!filters.notDeleted || !(frame in meta.deletedFrames)); } if (filters.notDeleted) { From df90b338e4354ee68de0e92c4a45dc33b235b572 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 12 Sep 2024 19:45:32 +0300 Subject: [PATCH 105/115] Fix cached chunk indicators in frame player --- cvat-core/src/frames.ts | 12 ++++++++++++ cvat-core/src/session-implementation.ts | 9 +++++++++ cvat-core/src/session.ts | 6 ++++++ cvat-ui/src/actions/annotation-actions.ts | 7 ++++--- 4 files changed, 31 insertions(+), 3 deletions(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index 8d2c0117d48..37fa51242e3 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -822,6 +822,18 @@ export function getCachedChunks(jobID): number[] { return frameDataCache[jobID].provider.cachedChunks(true); } +export function getJobFrameNumbers(jobID): number[] { + if (!(jobID in frameDataCache)) { + return []; + } + + const { meta, startFrame } = frameDataCache[jobID]; + const dataStartFrame = getDataStartFrame(meta, startFrame); + return meta.getDataFrameNumbers().map((dataFrameNumber: number): number => ( + getFrameNumber(dataFrameNumber, dataStartFrame, meta.frameStep) + )); +} + export function clear(jobID: number): void { if (jobID in frameDataCache) { frameDataCache[jobID].provider.close(); diff --git a/cvat-core/src/session-implementation.ts b/cvat-core/src/session-implementation.ts index 4afcce184e8..96177170872 100644 --- a/cvat-core/src/session-implementation.ts +++ b/cvat-core/src/session-implementation.ts @@ -18,6 +18,7 @@ import { deleteFrame, restoreFrame, getCachedChunks, + getJobFrameNumbers, clear as clearFrames, findFrame, getContextImage, @@ -244,6 +245,14 @@ export function implementJob(Job: typeof JobClass): typeof JobClass { }, }); + Object.defineProperty(Job.prototype.frames.frameNumbers, 'implementation', { + value: function includedFramesImplementation( + this: JobClass, + ): ReturnType { + return Promise.resolve(getJobFrameNumbers(this.id)); + }, + }); + Object.defineProperty(Job.prototype.frames.preview, 'implementation', { value: function previewImplementation( this: JobClass, diff --git a/cvat-core/src/session.ts b/cvat-core/src/session.ts index 1985a72b268..d0d4e5e0218 100644 --- a/cvat-core/src/session.ts +++ b/cvat-core/src/session.ts @@ -233,6 +233,10 @@ function buildDuplicatedAPI(prototype) { const result = await PluginRegistry.apiWrapper.call(this, prototype.frames.cachedChunks); return result; }, + async frameNumbers() { + const result = await PluginRegistry.apiWrapper.call(this, prototype.frames.frameNumbers); + return result; + }, async preview() { const result = await PluginRegistry.apiWrapper.call(this, prototype.frames.preview); return result; @@ -380,6 +384,7 @@ export class Session { restore: (frame: number) => Promise; save: () => Promise; cachedChunks: () => Promise; + frameNumbers: () => Promise; preview: () => Promise; contextImage: (frame: number) => Promise>; search: ( @@ -443,6 +448,7 @@ export class Session { restore: Object.getPrototypeOf(this).frames.restore.bind(this), save: Object.getPrototypeOf(this).frames.save.bind(this), cachedChunks: Object.getPrototypeOf(this).frames.cachedChunks.bind(this), + frameNumbers: Object.getPrototypeOf(this).frames.frameNumbers.bind(this), preview: Object.getPrototypeOf(this).frames.preview.bind(this), search: Object.getPrototypeOf(this).frames.search.bind(this), contextImage: Object.getPrototypeOf(this).frames.contextImage.bind(this), diff --git a/cvat-ui/src/actions/annotation-actions.ts b/cvat-ui/src/actions/annotation-actions.ts index 739acac410a..b3fa8b503aa 100644 --- a/cvat-ui/src/actions/annotation-actions.ts +++ b/cvat-ui/src/actions/annotation-actions.ts @@ -587,12 +587,13 @@ export function confirmCanvasReadyAsync(): ThunkAction { const { instance: job } = state.annotation.job; const { changeFrameEvent } = state.annotation.player.frame; const chunks = await job.frames.cachedChunks() as number[]; - const { startFrame, stopFrame, dataChunkSize } = job; + const includedFrames = await job.frames.frameNumbers() as number[]; + const { frameCount, dataChunkSize } = job; const ranges = chunks.map((chunk) => ( [ - Math.max(startFrame, chunk * dataChunkSize), - Math.min(stopFrame, (chunk + 1) * dataChunkSize - 1), + includedFrames[chunk * dataChunkSize], + includedFrames[Math.min(frameCount - 1, (chunk + 1) * dataChunkSize - 1)], ] )).reduce>((acc, val) => { if (acc.length && acc[acc.length - 1][1] + 1 === val[0]) { From 6ccb7dba08d245cccb33e29b599f4027b941e209 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 13 Sep 2024 11:50:19 +0300 Subject: [PATCH 106/115] Fix chunk predecode logic --- cvat-core/src/frames.ts | 57 ++++++++++++++++++++++++++--------------- 1 file changed, 37 insertions(+), 20 deletions(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index 37fa51242e3..ed948f7a263 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -56,6 +56,7 @@ export class FramesMetaData { public startFrame: number; public stopFrame: number; public frameStep: number; + public chunkCount: number; #updateTrigger: FieldUpdateTrigger; @@ -150,6 +151,17 @@ export class FramesMetaData { }, }), ); + + const chunkCount: number = Math.ceil(this.getDataFrameNumbers().length / this.chunkSize); + + Object.defineProperties( + this, + Object.freeze({ + chunkCount: { + get: () => chunkCount, + }, + }), + ); } getUpdated(): Record { @@ -328,7 +340,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { imageData: ImageBitmap | Blob; } | Blob>((resolve, reject) => { const { - meta, provider, prefetchAnalyzer, chunkSize, startFrame, stopFrame, + meta, provider, prefetchAnalyzer, chunkSize, startFrame, decodeForward, forwardStep, decodedBlocksCacheSize, } = frameDataCache[this.jobID]; @@ -337,7 +349,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { const requestedDataFrameNumber = getDataFrameNumber( this.number, dataStartFrame, meta.frameStep, ); - const chunkNumber = meta.getFrameChunkIndex(requestedDataFrameNumber); + const chunkIndex = meta.getFrameChunkIndex(requestedDataFrameNumber); const segmentFrameNumbers = meta.getDataFrameNumbers().map( (dataFrameNumber: number) => getFrameNumber( dataFrameNumber, dataStartFrame, meta.frameStep, @@ -345,19 +357,24 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { ); const frame = provider.frame(this.number); - function findTheNextNotDecodedChunk(searchFrom: number): number { - let nextFrameIndex = searchFrom + forwardStep; - let nextChunkNumber = Math.floor(nextFrameIndex / chunkSize); - while (nextChunkNumber === chunkNumber) { + function findTheNextNotDecodedChunk(currentFrameIndex: number): number | null { + const { chunkCount } = meta; + let nextFrameIndex = currentFrameIndex + forwardStep; + let nextChunkIndex = Math.floor(nextFrameIndex / chunkSize); + while (nextChunkIndex === chunkIndex) { nextFrameIndex += forwardStep; - nextChunkNumber = Math.floor(nextFrameIndex / chunkSize); + nextChunkIndex = Math.floor(nextFrameIndex / chunkSize); + } + + if (chunkCount <= nextChunkIndex) { + return null; } - if (provider.isChunkCached(nextChunkNumber)) { + if (provider.isChunkCached(nextChunkIndex)) { return findTheNextNotDecodedChunk(nextFrameIndex); } - return nextChunkNumber; + return nextChunkIndex; } if (frame) { @@ -368,12 +385,12 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { (chunk) => provider.isChunkCached(chunk), ) && decodedBlocksCacheSize > 1 && !frameDataCache[this.jobID].activeChunkRequest ) { - const nextChunkNumber = findTheNextNotDecodedChunk( + const nextChunkIndex = findTheNextNotDecodedChunk( meta.getFrameIndex(requestedDataFrameNumber), ); const predecodeChunksMax = Math.floor(decodedBlocksCacheSize / 2); - if (startFrame + nextChunkNumber * chunkSize <= stopFrame && - nextChunkNumber <= chunkNumber + predecodeChunksMax + if (nextChunkIndex !== null && + nextChunkIndex <= chunkIndex + predecodeChunksMax ) { frameDataCache[this.jobID].activeChunkRequest = new Promise((resolveForward) => { const releasePromise = (): void => { @@ -382,7 +399,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { }; frameDataCache[this.jobID].getChunk( - nextChunkNumber, ChunkQuality.COMPRESSED, + nextChunkIndex, ChunkQuality.COMPRESSED, ).then((chunk: ArrayBuffer) => { if (!(this.jobID in frameDataCache)) { // check if frameDataCache still exist @@ -394,10 +411,10 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { provider.cleanup(1); provider.requestDecodeBlock( chunk, - nextChunkNumber, + nextChunkIndex, segmentFrameNumbers.slice( - nextChunkNumber * chunkSize, - (nextChunkNumber + 1) * chunkSize, + nextChunkIndex * chunkSize, + (nextChunkIndex + 1) * chunkSize, ), () => {}, releasePromise, @@ -445,7 +462,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { ) => { let wasResolved = false; frameDataCache[this.jobID].getChunk( - chunkNumber, ChunkQuality.COMPRESSED, + chunkIndex, ChunkQuality.COMPRESSED, ).then((chunk: ArrayBuffer) => { try { if (!(this.jobID in frameDataCache)) { @@ -459,10 +476,10 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { provider .requestDecodeBlock( chunk, - chunkNumber, + chunkIndex, segmentFrameNumbers.slice( - chunkNumber * chunkSize, - (chunkNumber + 1) * chunkSize, + chunkIndex * chunkSize, + (chunkIndex + 1) * chunkSize, ), (_frame: number, bitmap: ImageBitmap | Blob) => { if (decodeForward) { From 1fb68bcb2860616ff56d5bb240d581cb4a2b60b3 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 13 Sep 2024 11:58:50 +0300 Subject: [PATCH 107/115] Rename chunkNumber to chunkIndex where necessary --- cvat-core/src/frames.ts | 2 +- cvat-core/src/session.ts | 4 ++-- cvat-data/src/ts/cvat-data.ts | 40 +++++++++++++++++------------------ 3 files changed, 23 insertions(+), 23 deletions(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index ed948f7a263..a5d3c43be4a 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -34,7 +34,7 @@ const frameDataCache: Record; - getChunk: (chunkNumber: number, quality: ChunkQuality) => Promise; + getChunk: (chunkIndex: number, quality: ChunkQuality) => Promise; }> = {}; // frame meta data storage by job id diff --git a/cvat-core/src/session.ts b/cvat-core/src/session.ts index d0d4e5e0218..54133ff6b66 100644 --- a/cvat-core/src/session.ts +++ b/cvat-core/src/session.ts @@ -259,11 +259,11 @@ function buildDuplicatedAPI(prototype) { ); return result; }, - async chunk(chunkNumber, quality) { + async chunk(chunkIndex, quality) { const result = await PluginRegistry.apiWrapper.call( this, prototype.frames.chunk, - chunkNumber, + chunkIndex, quality, ); return result; diff --git a/cvat-data/src/ts/cvat-data.ts b/cvat-data/src/ts/cvat-data.ts index 97e796af9df..baf00ac443c 100644 --- a/cvat-data/src/ts/cvat-data.ts +++ b/cvat-data/src/ts/cvat-data.ts @@ -73,7 +73,7 @@ decodeContextImages.mutex = new Mutex(); interface BlockToDecode { chunkFrameNumbers: number[]; - chunkNumber: number; + chunkIndex: number; block: ArrayBuffer; onDecodeAll(): void; onDecode(frame: number, bitmap: ImageBitmap | Blob): void; @@ -99,12 +99,12 @@ export class FrameDecoder { private renderHeight: number; private zipWorker: Worker | null; private videoWorker: Worker | null; - private getChunkNumber: (frame: number) => number; + private getChunkIndex: (frame: number) => number; constructor( blockType: BlockType, cachedBlockCount: number, - getChunkNumber: (frame: number) => number, + getChunkIndex: (frame: number) => number, dimension: DimensionType = DimensionType.DIMENSION_2D, ) { this.mutex = new Mutex(); @@ -117,7 +117,7 @@ export class FrameDecoder { this.renderWidth = 1920; this.renderHeight = 1080; - this.getChunkNumber = getChunkNumber; + this.getChunkIndex = getChunkIndex; this.blockType = blockType; this.decodedChunks = {}; @@ -125,8 +125,8 @@ export class FrameDecoder { this.chunkIsBeingDecoded = null; } - isChunkCached(chunkNumber: number): boolean { - return chunkNumber in this.decodedChunks; + isChunkCached(chunkIndex: number): boolean { + return chunkIndex in this.decodedChunks; } hasFreeSpace(): boolean { @@ -175,7 +175,7 @@ export class FrameDecoder { requestDecodeBlock( block: ArrayBuffer, - chunkNumber: number, + chunkIndex: number, chunkFrameNumbers: number[], onDecode: (frame: number, bitmap: ImageBitmap | Blob) => void, onDecodeAll: () => void, @@ -185,7 +185,7 @@ export class FrameDecoder { if (this.requestedChunkToDecode !== null) { // a chunk was already requested to be decoded, but decoding didn't start yet - if (chunkNumber === this.requestedChunkToDecode.chunkNumber) { + if (chunkIndex === this.requestedChunkToDecode.chunkIndex) { // it was the same chunk this.requestedChunkToDecode.onReject(new RequestOutdatedError()); @@ -196,12 +196,12 @@ export class FrameDecoder { this.requestedChunkToDecode.onReject(new RequestOutdatedError()); } } else if (this.chunkIsBeingDecoded === null || - chunkNumber !== this.chunkIsBeingDecoded.chunkNumber + chunkIndex !== this.chunkIsBeingDecoded.chunkIndex ) { // everything was decoded or decoding other chunk is in process this.requestedChunkToDecode = { chunkFrameNumbers, - chunkNumber, + chunkIndex, block, onDecode, onDecodeAll, @@ -225,9 +225,9 @@ export class FrameDecoder { } frame(frameNumber: number): ImageBitmap | Blob | null { - const chunkNumber = this.getChunkNumber(frameNumber); - if (chunkNumber in this.decodedChunks) { - return this.decodedChunks[chunkNumber][frameNumber]; + const chunkIndex = this.getChunkIndex(frameNumber); + if (chunkIndex in this.decodedChunks) { + return this.decodedChunks[chunkIndex][frameNumber]; } return null; @@ -275,8 +275,8 @@ export class FrameDecoder { releaseMutex(); }; try { - const { chunkFrameNumbers, chunkNumber, block } = this.requestedChunkToDecode; - if (chunkNumber !== blockToDecode.chunkNumber) { + const { chunkFrameNumbers, chunkIndex, block } = this.requestedChunkToDecode; + if (chunkIndex !== blockToDecode.chunkIndex) { // request is not relevant, another block was already requested // it happens when A is being decoded, B comes and wait for mutex, C comes and wait for mutex // B is not necessary anymore, because C already was requested @@ -288,7 +288,7 @@ export class FrameDecoder { chunkFrameNumbers[chunkFrameIndex] ); - this.orderedStack = [chunkNumber, ...this.orderedStack]; + this.orderedStack = [chunkIndex, ...this.orderedStack]; this.cleanup(); const decodedFrames: Record = {}; this.chunkIsBeingDecoded = this.requestedChunkToDecode; @@ -325,7 +325,7 @@ export class FrameDecoder { this.chunkIsBeingDecoded.onDecode(frameNumber, decodedFrames[frameNumber]); if (keptIndex === chunkFrameNumbers.length - 1) { - this.decodedChunks[chunkNumber] = decodedFrames; + this.decodedChunks[chunkIndex] = decodedFrames; this.chunkIsBeingDecoded.onDecodeAll(); this.chunkIsBeingDecoded = null; release(); @@ -384,7 +384,7 @@ export class FrameDecoder { this.chunkIsBeingDecoded.onDecode(frameNumber, decodedFrames[frameNumber]); if (decodedCount === chunkFrameNumbers.length - 1) { - this.decodedChunks[chunkNumber] = decodedFrames; + this.decodedChunks[chunkIndex] = decodedFrames; this.chunkIsBeingDecoded.onDecodeAll(); this.chunkIsBeingDecoded = null; release(); @@ -430,10 +430,10 @@ export class FrameDecoder { public cachedChunks(includeInProgress = false): number[] { const chunkIsBeingDecoded = ( includeInProgress && this.chunkIsBeingDecoded ? - this.chunkIsBeingDecoded.chunkNumber : + this.chunkIsBeingDecoded.chunkIndex : null ); - return Object.keys(this.decodedChunks).map((chunkNumber: string) => +chunkNumber).concat( + return Object.keys(this.decodedChunks).map((chunkIndex: string) => +chunkIndex).concat( ...(chunkIsBeingDecoded !== null ? [chunkIsBeingDecoded] : []), ).sort((a, b) => a - b); } From 92d0c7a519e05632b0e82698e4717c508940366f Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Fri, 13 Sep 2024 12:39:48 +0300 Subject: [PATCH 108/115] Fix potential prefetch problem with reverse playback --- cvat-core/src/frames.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat-core/src/frames.ts b/cvat-core/src/frames.ts index a5d3c43be4a..dda847cf7a7 100644 --- a/cvat-core/src/frames.ts +++ b/cvat-core/src/frames.ts @@ -366,7 +366,7 @@ Object.defineProperty(FrameData.prototype.data, 'implementation', { nextChunkIndex = Math.floor(nextFrameIndex / chunkSize); } - if (chunkCount <= nextChunkIndex) { + if (nextChunkIndex < 0 || chunkCount <= nextChunkIndex) { return null; } From 3cdc4dc1dac23a7424e9bc5d51661fcd0898fd74 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 16 Sep 2024 13:41:41 +0300 Subject: [PATCH 109/115] Move env variable into docker-compose.yml --- docker-compose.yml | 1 + .../docker-compose.configurable_static_cache.yml | 16 ---------------- tests/python/shared/fixtures/init.py | 1 - 3 files changed, 1 insertion(+), 17 deletions(-) delete mode 100644 tests/docker-compose.configurable_static_cache.yml diff --git a/docker-compose.yml b/docker-compose.yml index 051bd0bfd8c..569e163e9fe 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -10,6 +10,7 @@ x-backend-env: &backend-env CVAT_REDIS_ONDISK_HOST: cvat_redis_ondisk CVAT_REDIS_ONDISK_PORT: 6666 CVAT_LOG_IMPORT_ERRORS: 'true' + CVAT_ALLOW_STATIC_CACHE: '${CVAT_ALLOW_STATIC_CACHE:-no}' DJANGO_LOG_SERVER_HOST: vector DJANGO_LOG_SERVER_PORT: 80 no_proxy: clickhouse,grafana,vector,nuclio,opa,${no_proxy:-} diff --git a/tests/docker-compose.configurable_static_cache.yml b/tests/docker-compose.configurable_static_cache.yml deleted file mode 100644 index 5afa4347080..00000000000 --- a/tests/docker-compose.configurable_static_cache.yml +++ /dev/null @@ -1,16 +0,0 @@ -services: - cvat_server: - environment: - CVAT_ALLOW_STATIC_CACHE: '${CVAT_ALLOW_STATIC_CACHE:-no}' - - cvat_worker_import: - environment: - CVAT_ALLOW_STATIC_CACHE: '${CVAT_ALLOW_STATIC_CACHE:-no}' - - cvat_worker_export: - environment: - CVAT_ALLOW_STATIC_CACHE: '${CVAT_ALLOW_STATIC_CACHE:-no}' - - cvat_worker_annotation: - environment: - CVAT_ALLOW_STATIC_CACHE: '${CVAT_ALLOW_STATIC_CACHE:-no}' diff --git a/tests/python/shared/fixtures/init.py b/tests/python/shared/fixtures/init.py index 99f1f02f8e0..4a17454617d 100644 --- a/tests/python/shared/fixtures/init.py +++ b/tests/python/shared/fixtures/init.py @@ -31,7 +31,6 @@ "tests/docker-compose.file_share.yml", "tests/docker-compose.minio.yml", "tests/docker-compose.test_servers.yml", - "tests/docker-compose.configurable_static_cache.yml", ] From bc5ed39a1960a6e3699d8e026e0b8252dd65ce93 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 17 Sep 2024 12:40:43 +0300 Subject: [PATCH 110/115] Fix invalid cached chunk display in GT jobs --- .../top-bar/player-navigation.tsx | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/cvat-ui/src/components/annotation-page/top-bar/player-navigation.tsx b/cvat-ui/src/components/annotation-page/top-bar/player-navigation.tsx index 2088d14d7cc..f1a2e9cf289 100644 --- a/cvat-ui/src/components/annotation-page/top-bar/player-navigation.tsx +++ b/cvat-ui/src/components/annotation-page/top-bar/player-navigation.tsx @@ -169,17 +169,14 @@ function PlayerNavigation(props: Props): JSX.Element { {!!ranges && ( {ranges.split(';').map((range) => { - const [start, end] = range.split(':').map((num) => +num); - const adjustedStart = Math.max(0, start - 1); - let totalSegments = stopFrame - startFrame; - if (totalSegments === 0) { - // corner case for jobs with one image - totalSegments = 1; - } + const [rangeStart, rangeStop] = range.split(':').map((num) => +num); + const totalSegments = stopFrame - startFrame + 1; const segmentWidth = 1000 / totalSegments; - const width = Math.max((end - adjustedStart), 1) * segmentWidth; - const offset = (Math.max((adjustedStart - startFrame), 0) / totalSegments) * 1000; - return (); + const width = (rangeStop - rangeStart + 1) * segmentWidth; + const offset = (Math.max((rangeStart - startFrame), 0) / totalSegments) * 1000; + return ( + + ); })} )} From 08ddd288f7bd2d16ce07a9ddb754f58802cc53f6 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 17 Sep 2024 12:41:24 +0300 Subject: [PATCH 111/115] Fix invalid task preview generation --- cvat/apps/engine/frame_provider.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 314c8cf5b23..fac158b3639 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -260,7 +260,7 @@ def get_rel_frame_number(self, abs_frame_number: int) -> int: return super()._get_rel_frame_number(self._db_task.data, abs_frame_number) def get_preview(self) -> DataWithMeta[BytesIO]: - return self._get_segment_frame_provider(self._db_task.data.start_frame).get_preview() + return self._get_segment_frame_provider(0).get_preview() def get_chunk( self, chunk_number: int, *, quality: FrameQuality = FrameQuality.ORIGINAL From 1d969bd7f41c217d813c55d3cfdc9c8b20b737bb Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 17 Sep 2024 12:43:21 +0300 Subject: [PATCH 112/115] Refactor CS previews, context image chunk generation, media cache creation and retrieval --- cvat/apps/engine/cache.py | 155 +++++++++++++++++------------ cvat/apps/engine/frame_provider.py | 28 +++--- cvat/apps/engine/views.py | 2 +- 3 files changed, 109 insertions(+), 76 deletions(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index b32dbeef00e..0bb921fd7b1 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -26,6 +26,7 @@ Tuple, Type, Union, + overload, ) import av @@ -73,21 +74,20 @@ def _get_checksum(self, value: bytes) -> int: def _get_or_set_cache_item( self, key: str, create_callback: Callable[[], DataWithMime] - ) -> DataWithMime: + ) -> _CacheItem: def create_item() -> _CacheItem: slogger.glob.info(f"Starting to prepare chunk: key {key}") item_data = create_callback() slogger.glob.info(f"Ending to prepare chunk: key {key}") - if item_data[0]: - item = (item_data[0], item_data[1], self._get_checksum(item_data[0].getbuffer())) + item_data_bytes = item_data[0].getvalue() + item = (item_data[0], item_data[1], self._get_checksum(item_data_bytes)) + if item_data_bytes: self._cache.set(key, item) - else: - item = (item_data[0], item_data[1], None) return item - item = self._get(key) + item = self._get_cache_item(key) if not item: item = create_item() else: @@ -98,9 +98,9 @@ def create_item() -> _CacheItem: slogger.glob.info(f"Recreating cache item {key} due to checksum mismatch") item = create_item() - return item[0], item[1] + return item - def _get(self, key: str) -> Optional[DataWithMime]: + def _get_cache_item(self, key: str) -> Optional[_CacheItem]: slogger.glob.info(f"Starting to get chunk from cache: key {key}") try: item = self._cache.get(key) @@ -152,20 +152,36 @@ def _make_segment_task_chunk_key( def _make_context_image_preview_key(self, db_data: models.Data, frame_number: int) -> str: return f"context_image_{db_data.id}_{frame_number}_preview" - def get_segment_chunk( + @overload + def _to_data_with_mime(self, cache_item: _CacheItem) -> DataWithMime: ... + + @overload + def _to_data_with_mime(self, cache_item: Optional[_CacheItem]) -> Optional[DataWithMime]: ... + + def _to_data_with_mime(self, cache_item: Optional[_CacheItem]) -> Optional[DataWithMime]: + if not cache_item: + return None + + return cache_item[:2] + + def get_or_set_segment_chunk( self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality ) -> DataWithMime: - return self._get_or_set_cache_item( - key=self._make_chunk_key(db_segment, chunk_number, quality=quality), - create_callback=lambda: self.prepare_segment_chunk( - db_segment, chunk_number, quality=quality - ), + return self._to_data_with_mime( + self._get_or_set_cache_item( + key=self._make_chunk_key(db_segment, chunk_number, quality=quality), + create_callback=lambda: self.prepare_segment_chunk( + db_segment, chunk_number, quality=quality + ), + ) ) def get_task_chunk( self, db_task: models.Task, chunk_number: int, *, quality: FrameQuality ) -> Optional[DataWithMime]: - return self._get(key=self._make_chunk_key(db_task, chunk_number, quality=quality)) + return self._to_data_with_mime( + self._get_cache_item(key=self._make_chunk_key(db_task, chunk_number, quality=quality)) + ) def get_or_set_task_chunk( self, @@ -175,16 +191,20 @@ def get_or_set_task_chunk( quality: FrameQuality, set_callback: Callable[[], DataWithMime], ) -> DataWithMime: - return self._get_or_set_cache_item( - key=self._make_chunk_key(db_task, chunk_number, quality=quality), - create_callback=set_callback, + return self._to_data_with_mime( + self._get_or_set_cache_item( + key=self._make_chunk_key(db_task, chunk_number, quality=quality), + create_callback=set_callback, + ) ) def get_segment_task_chunk( self, db_segment: models.Segment, chunk_number: int, *, quality: FrameQuality ) -> Optional[DataWithMime]: - return self._get( - key=self._make_segment_task_chunk_key(db_segment, chunk_number, quality=quality) + return self._to_data_with_mime( + self._get_cache_item( + key=self._make_segment_task_chunk_key(db_segment, chunk_number, quality=quality) + ) ) def get_or_set_segment_task_chunk( @@ -195,40 +215,52 @@ def get_or_set_segment_task_chunk( quality: FrameQuality, set_callback: Callable[[], DataWithMime], ) -> DataWithMime: - return self._get_or_set_cache_item( - key=self._make_segment_task_chunk_key(db_segment, chunk_number, quality=quality), - create_callback=set_callback, + return self._to_data_with_mime( + self._get_or_set_cache_item( + key=self._make_segment_task_chunk_key(db_segment, chunk_number, quality=quality), + create_callback=set_callback, + ) ) - def get_selective_job_chunk( + def get_or_set_selective_job_chunk( self, db_job: models.Job, chunk_number: int, *, quality: FrameQuality ) -> DataWithMime: - return self._get_or_set_cache_item( - key=self._make_chunk_key(db_job, chunk_number, quality=quality), - create_callback=lambda: self.prepare_masked_range_segment_chunk( - db_job.segment, chunk_number, quality=quality - ), + return self._to_data_with_mime( + self._get_or_set_cache_item( + key=self._make_chunk_key(db_job, chunk_number, quality=quality), + create_callback=lambda: self.prepare_masked_range_segment_chunk( + db_job.segment, chunk_number, quality=quality + ), + ) ) def get_or_set_segment_preview(self, db_segment: models.Segment) -> DataWithMime: - return self._get_or_set_cache_item( - self._make_preview_key(db_segment), - create_callback=lambda: self._prepare_segment_preview(db_segment), + return self._to_data_with_mime( + self._get_or_set_cache_item( + self._make_preview_key(db_segment), + create_callback=lambda: self._prepare_segment_preview(db_segment), + ) ) def get_cloud_preview(self, db_storage: models.CloudStorage) -> Optional[DataWithMime]: - return self._get(self._make_preview_key(db_storage)) + return self._to_data_with_mime(self._get_cache_item(self._make_preview_key(db_storage))) def get_or_set_cloud_preview(self, db_storage: models.CloudStorage) -> DataWithMime: - return self._get_or_set_cache_item( - self._make_preview_key(db_storage), - create_callback=lambda: self._prepare_cloud_preview(db_storage), + return self._to_data_with_mime( + self._get_or_set_cache_item( + self._make_preview_key(db_storage), + create_callback=lambda: self._prepare_cloud_preview(db_storage), + ) ) - def get_frame_context_images(self, db_data: models.Data, frame_number: int) -> DataWithMime: - return self._get_or_set_cache_item( - key=self._make_context_image_preview_key(db_data, frame_number), - create_callback=lambda: self.prepare_context_images(db_data, frame_number), + def get_or_set_frame_context_images_chunk( + self, db_data: models.Data, frame_number: int + ) -> DataWithMime: + return self._to_data_with_mime( + self._get_or_set_cache_item( + key=self._make_context_image_preview_key(db_data, frame_number), + create_callback=lambda: self.prepare_context_images_chunk(db_data, frame_number), + ) ) def _read_raw_images( @@ -536,59 +568,56 @@ def _prepare_segment_preview(self, db_segment: models.Segment) -> DataWithMime: return prepare_preview_image(preview) - def _prepare_cloud_preview(self, db_storage): + def _prepare_cloud_preview(self, db_storage: models.CloudStorage) -> DataWithMime: storage = db_storage_to_storage_instance(db_storage) if not db_storage.manifests.count(): raise ValidationError("Cannot get the cloud storage preview. There is no manifest file") + preview_path = None - for manifest_model in db_storage.manifests.all(): - manifest_prefix = os.path.dirname(manifest_model.filename) + for db_manifest in db_storage.manifests.all(): + manifest_prefix = os.path.dirname(db_manifest.filename) + full_manifest_path = os.path.join( - db_storage.get_storage_dirname(), manifest_model.filename + db_storage.get_storage_dirname(), db_manifest.filename ) if not os.path.exists(full_manifest_path) or datetime.fromtimestamp( os.path.getmtime(full_manifest_path), tz=timezone.utc - ) < storage.get_file_last_modified(manifest_model.filename): - storage.download_file(manifest_model.filename, full_manifest_path) + ) < storage.get_file_last_modified(db_manifest.filename): + storage.download_file(db_manifest.filename, full_manifest_path) + manifest = ImageManifestManager( - os.path.join(db_storage.get_storage_dirname(), manifest_model.filename), + os.path.join(db_storage.get_storage_dirname(), db_manifest.filename), db_storage.get_storage_dirname(), ) # need to update index manifest.set_index() if not len(manifest): continue + preview_info = manifest[0] preview_filename = "".join([preview_info["name"], preview_info["extension"]]) preview_path = os.path.join(manifest_prefix, preview_filename) break + if not preview_path: msg = "Cloud storage {} does not contain any images".format(db_storage.pk) slogger.cloud_storage[db_storage.pk].info(msg) raise NotFound(msg) buff = storage.download_fileobj(preview_path) - mime_type = mimetypes.guess_type(preview_path)[0] - - return buff, mime_type + image = PIL.Image.open(buff) + return prepare_preview_image(image) - def prepare_context_images( - self, db_data: models.Data, frame_number: int - ) -> Optional[DataWithMime]: + def prepare_context_images_chunk(self, db_data: models.Data, frame_number: int) -> DataWithMime: zip_buffer = io.BytesIO() - try: - image = models.Image.objects.get(data_id=db_data.id, frame=frame_number) - except models.Image.DoesNotExist: - return None - if not image.related_files.count(): - return None, None + related_images = db_data.related_files.filter(primary_image__frame=frame_number).all() + if not related_images: + return zip_buffer, "" with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED, False) as zip_file: - common_path = os.path.commonpath( - list(map(lambda x: str(x.path), image.related_files.all())) - ) - for i in image.related_files.all(): + common_path = os.path.commonpath(list(map(lambda x: str(x.path), related_images))) + for i in related_images: path = os.path.realpath(str(i.path)) name = os.path.relpath(str(i.path), common_path) image = cv2.imread(path) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index fac158b3639..ea14b40a75a 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -200,7 +200,7 @@ def get_frame( ) -> DataWithMeta[AnyFrame]: ... @abstractmethod - def get_frame_context_images( + def get_frame_context_images_chunk( self, frame_number: int, ) -> Optional[DataWithMeta[BytesIO]]: ... @@ -350,11 +350,13 @@ def get_frame( frame_number, quality=quality, out_type=out_type ) - def get_frame_context_images( + def get_frame_context_images_chunk( self, frame_number: int, ) -> Optional[DataWithMeta[BytesIO]]: - return self._get_segment_frame_provider(frame_number).get_frame_context_images(frame_number) + return self._get_segment_frame_provider(frame_number).get_frame_context_images_chunk( + frame_number + ) def iterate_frames( self, @@ -432,7 +434,7 @@ def __init__(self, db_segment: models.Segment) -> None: self._loaders[FrameQuality.COMPRESSED] = _BufferChunkLoader( reader_class=reader_class[db_data.compressed_chunk_type][0], reader_params=reader_class[db_data.compressed_chunk_type][1], - get_chunk_callback=lambda chunk_idx: cache.get_segment_chunk( + get_chunk_callback=lambda chunk_idx: cache.get_or_set_segment_chunk( db_segment, chunk_idx, quality=FrameQuality.COMPRESSED ), ) @@ -440,7 +442,7 @@ def __init__(self, db_segment: models.Segment) -> None: self._loaders[FrameQuality.ORIGINAL] = _BufferChunkLoader( reader_class=reader_class[db_data.original_chunk_type][0], reader_params=reader_class[db_data.original_chunk_type][1], - get_chunk_callback=lambda chunk_idx: cache.get_segment_chunk( + get_chunk_callback=lambda chunk_idx: cache.get_or_set_segment_chunk( db_segment, chunk_idx, quality=FrameQuality.ORIGINAL ), ) @@ -549,19 +551,21 @@ def get_frame( return return_type(frame, mime=mimetypes.guess_type(frame_name)[0]) - def get_frame_context_images( + def get_frame_context_images_chunk( self, frame_number: int, ) -> Optional[DataWithMeta[BytesIO]]: - # TODO: refactor, optimize - cache = MediaCache() + self.validate_frame_number(frame_number) - if self._db_segment.task.data.storage_method == models.StorageMethodChoice.CACHE: - data, mime = cache.get_frame_context_images(self._db_segment.task.data, frame_number) + db_data = self._db_segment.task.data + + cache = MediaCache() + if db_data.storage_method == models.StorageMethodChoice.CACHE: + data, mime = cache.get_or_set_frame_context_images_chunk(db_data, frame_number) else: - data, mime = cache.prepare_context_images(self._db_segment.task.data, frame_number) + data, mime = cache.prepare_context_images_chunk(db_data, frame_number) - if not data: + if not data.getvalue(): return None return DataWithMeta[BytesIO](data, mime=mime) diff --git a/cvat/apps/engine/views.py b/cvat/apps/engine/views.py index e30a857d301..e6b51c40880 100644 --- a/cvat/apps/engine/views.py +++ b/cvat/apps/engine/views.py @@ -703,7 +703,7 @@ def __call__(self): return HttpResponse(data.data.getvalue(), content_type=data.mime) elif self.type == 'context_image': - data = frame_provider.get_frame_context_images(self.number) + data = frame_provider.get_frame_context_images_chunk(self.number) if not data: return HttpResponseNotFound() From d135475e4792e0ffb745d9a53cc2777a32ae1626 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 17 Sep 2024 12:46:19 +0300 Subject: [PATCH 113/115] Remove extra import --- cvat/apps/engine/cache.py | 1 - 1 file changed, 1 deletion(-) diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 0bb921fd7b1..bc4c8616bd7 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -54,7 +54,6 @@ ZipChunkWriter, ZipCompressedChunkWriter, ) -from cvat.apps.engine.mime_types import mimetypes from cvat.apps.engine.utils import md5_hash, preload_images from utils.dataset_manifest import ImageManifestManager From a1638c94dade78a6e2835f51057c6ca60c6d1c01 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 17 Sep 2024 15:24:50 +0300 Subject: [PATCH 114/115] Fix CS preview in response --- cvat/apps/engine/views.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cvat/apps/engine/views.py b/cvat/apps/engine/views.py index e6b51c40880..8ca23ab9488 100644 --- a/cvat/apps/engine/views.py +++ b/cvat/apps/engine/views.py @@ -2726,10 +2726,10 @@ def preview(self, request, pk): result = cache.get_cloud_preview(db_storage) if not result: return HttpResponseNotFound('Cloud storage preview not found') - return HttpResponse(result[0], result[1]) + return HttpResponse(result[0].getvalue(), result[1]) preview, mime = cache.get_or_set_cloud_preview(db_storage) - return HttpResponse(preview, mime) + return HttpResponse(preview.getvalue(), mime) except CloudStorageModel.DoesNotExist: message = f"Storage {pk} does not exist" slogger.glob.error(message) From fc89c015ed5aa7fdd1d85d492715cdff295bc9f7 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Tue, 17 Sep 2024 17:34:04 +0300 Subject: [PATCH 115/115] Add reverse migration --- .../migrations/0083_move_to_segment_chunks.py | 95 +++++++++++++++---- 1 file changed, 79 insertions(+), 16 deletions(-) diff --git a/cvat/apps/engine/migrations/0083_move_to_segment_chunks.py b/cvat/apps/engine/migrations/0083_move_to_segment_chunks.py index c9f59593d23..8ef887d4c54 100644 --- a/cvat/apps/engine/migrations/0083_move_to_segment_chunks.py +++ b/cvat/apps/engine/migrations/0083_move_to_segment_chunks.py @@ -1,39 +1,67 @@ # Generated by Django 4.2.13 on 2024-08-12 09:49 import os +from itertools import islice +from typing import Iterable, TypeVar + from django.db import migrations -from cvat.apps.engine.log import get_migration_logger, get_migration_log_dir + +from cvat.apps.engine.log import get_migration_log_dir, get_migration_logger + +T = TypeVar("T") + + +def take_by(iterable: Iterable[T], count: int) -> Iterable[T]: + """ + Returns elements from the input iterable by batches of N items. + ('abcdefg', 3) -> ['a', 'b', 'c'], ['d', 'e', 'f'], ['g'] + """ + + it = iter(iterable) + while True: + batch = list(islice(it, count)) + if len(batch) == 0: + break + + yield batch + + +def get_migration_name() -> str: + return os.path.splitext(os.path.basename(__file__))[0] + + +def get_updated_ids_filename(log_dir: str, migration_name: str) -> str: + return os.path.join(log_dir, migration_name + "-data_ids.log") + + +MIGRATION_LOG_HEADER = ( + 'The following Data ids have been switched from using "filesystem" chunk storage ' 'to "cache":' +) + def switch_tasks_with_static_chunks_to_dynamic_chunks(apps, schema_editor): - migration_name = os.path.splitext(os.path.basename(__file__))[0] + migration_name = get_migration_name() migration_log_dir = get_migration_log_dir() with get_migration_logger(migration_name) as common_logger: Data = apps.get_model("engine", "Data") - data_with_static_cache_query = ( - Data.objects - .filter(storage_method="file_system") - ) + data_with_static_cache_query = Data.objects.filter(storage_method="file_system") data_with_static_cache_ids = list( v[0] for v in ( - data_with_static_cache_query - .order_by('id') - .values_list('id') + data_with_static_cache_query.order_by("id") + .values_list("id") .iterator(chunk_size=100000) ) ) data_with_static_cache_query.update(storage_method="cache") - updated_ids_filename = os.path.join(migration_log_dir, migration_name + "-data_ids.log") + updated_ids_filename = get_updated_ids_filename(migration_log_dir, migration_name) with open(updated_ids_filename, "w") as data_ids_file: - print( - "The following Data ids have been switched from using \"filesystem\" chunk storage " - "to \"cache\":", - file=data_ids_file - ) + print(MIGRATION_LOG_HEADER, file=data_ids_file) + for data_id in data_with_static_cache_ids: print(data_id, file=data_ids_file) @@ -44,6 +72,38 @@ def switch_tasks_with_static_chunks_to_dynamic_chunks(apps, schema_editor): ) ) + +def revert_switch_tasks_with_static_chunks_to_dynamic_chunks(apps, schema_editor): + migration_name = get_migration_name() + migration_log_dir = get_migration_log_dir() + + updated_ids_filename = get_updated_ids_filename(migration_log_dir, migration_name) + if not os.path.isfile(updated_ids_filename): + raise FileNotFoundError( + "Can't revert the migration: can't file forward migration logfile at " + f"'{updated_ids_filename}'." + ) + + with open(updated_ids_filename, "r") as data_ids_file: + header = data_ids_file.readline().strip() + if header != MIGRATION_LOG_HEADER: + raise ValueError( + "Can't revert the migration: the migration log file has unexpected header" + ) + + forward_updated_ids = tuple(map(int, data_ids_file)) + + if not forward_updated_ids: + return + + Data = apps.get_model("engine", "Data") + + for id_batch in take_by(forward_updated_ids, 1000): + Data.objects.filter(storage_method="cache", id__in=id_batch).update( + storage_method="file_system" + ) + + class Migration(migrations.Migration): dependencies = [ @@ -51,5 +111,8 @@ class Migration(migrations.Migration): ] operations = [ - migrations.RunPython(switch_tasks_with_static_chunks_to_dynamic_chunks) + migrations.RunPython( + switch_tasks_with_static_chunks_to_dynamic_chunks, + reverse_code=revert_switch_tasks_with_static_chunks_to_dynamic_chunks, + ) ]