From e54eb68a6e9b31292bbbd2f7095e67626fb2b825 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 6 Apr 2020 15:45:06 +0300 Subject: [PATCH 1/2] Add chunk iterator cache --- cvat/apps/engine/frame_provider.py | 51 +++++++++++++++++++++++------- 1 file changed, 39 insertions(+), 12 deletions(-) diff --git a/cvat/apps/engine/frame_provider.py b/cvat/apps/engine/frame_provider.py index 7bd60a8a9fbb..25575ea51d36 100644 --- a/cvat/apps/engine/frame_provider.py +++ b/cvat/apps/engine/frame_provider.py @@ -2,7 +2,6 @@ # # SPDX-License-Identifier: MIT -import itertools import math from enum import Enum from io import BytesIO @@ -15,6 +14,33 @@ from cvat.apps.engine.models import DataChoice +class RandomAccessIterator: + def __init__(self, iterable): + self.iterable = iterable + self.iterator = None + self.pos = -1 + + def __iter__(self): + return self + + def __next__(self): + return self[self.pos + 1] + + def __getitem__(self, idx): + 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.iterator = iter(self.iterable) + self.pos = -1 + class FrameProvider: class Quality(Enum): COMPRESSED = 0 @@ -35,7 +61,8 @@ def __init__(self, reader_class, path_getter): def load(self, chunk_id): if self.chunk_id != chunk_id: self.chunk_id = chunk_id - self.chunk_reader = self.reader_class([self.get_chunk_path(chunk_id)]) + self.chunk_reader = RandomAccessIterator( + self.reader_class([self.get_chunk_path(chunk_id)])) return self.chunk_reader def __init__(self, db_data): @@ -104,18 +131,18 @@ def get_chunk(self, chunk_number, quality=Quality.ORIGINAL): chunk_number = self._validate_chunk_number(chunk_number) return self._loaders[quality].get_chunk_path(chunk_number) - def get_frame(self, frame_number, quality=Quality.ORIGINAL): + def get_frame(self, frame_number, quality=Quality.ORIGINAL, + out_type=Type.BUFFER): _, 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] - chunk_reader = self._loaders[quality].load(chunk_number) - - frame, frame_name, _ = next(itertools.islice(chunk_reader, frame_offset, None)) - if self._loaders[quality].reader_class is VideoReader: - return (self._av_frame_to_png_bytes(frame), 'image/png') + frame = self._convert_frame(frame, loader.reader_class, out_type) + if loader.reader_class is VideoReader: + return (frame, 'image/png') return (frame, mimetypes.guess_type(frame_name)) def get_frames(self, quality=Quality.ORIGINAL, out_type=Type.BUFFER): - loader = self._loaders[quality] - for chunk_idx in range(math.ceil(self._db_data.size / self._db_data.chunk_size)): - for frame, _, _ in loader.load(chunk_idx): - yield self._convert_frame(frame, loader.reader_class, out_type) + for idx in range(self._db_data.size): + yield self.get_frame(idx, quality=quality, out_type=out_type) From a242f777b72d71600034a2832d21f7746c6d8781 Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Mon, 6 Apr 2020 16:04:46 +0300 Subject: [PATCH 2/2] fix --- 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 7a7b80b27859..53a103f02a68 100644 --- a/cvat/apps/dataset_manager/bindings.py +++ b/cvat/apps/dataset_manager/bindings.py @@ -26,7 +26,7 @@ def __iter__(self): frames = self._frame_provider.get_frames( self._frame_provider.Quality.ORIGINAL, self._frame_provider.Type.NUMPY_ARRAY) - for item_id, image in enumerate(frames): + for item_id, (image, _) in enumerate(frames): yield datumaro.DatasetItem( id=item_id, image=Image(image),