diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 34753557..2f8b0a62 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -9,6 +9,7 @@ from cvat.apps.engine.media_extractors import (Mpeg4ChunkWriter, ZipChunkWriter, from cvat.apps.engine.models import DataChoice from .prepare import PrepareInfo import os +from io import BytesIO class CacheInteraction: def __init__(self): @@ -37,18 +38,19 @@ class CacheInteraction: extractor = extractor_classes[quality](image_quality) + images = [] + buff = BytesIO() if os.path.exists(db_data.get_meta_path()): - meta = PrepareInfo(source_path=os.path.join(db_data.get_upload_dirname(), db_data.video.path), - meta_path=db_data.get_meta_path()) - frames = [] + source_path = os.path.join(db_data.get_upload_dirname(), db_data.video.path) + meta = PrepareInfo(source_path=source_path, meta_path=db_data.get_meta_path()) for frame in meta.decode_needed_frames(chunk_number, db_data): - frames.append(frame) - buff = extractor.save_as_chunk_to_buff(frames) + images.append(frame) + extractor.save_as_chunk([(image, source_path, None) for image in images], buff) else: - img_paths = None with open(db_data.get_dummy_chunk_path(chunk_number), 'r') as dummy_file: - img_paths = [os.path.join(db_data.get_upload_dirname(), line.strip()) for line in dummy_file] - buff = extractor.save_as_chunk_to_buff(img_paths) + images = [os.path.join(db_data.get_upload_dirname(), line.strip()) for line in dummy_file] + extractor.save_as_chunk([(image, image, None) for image in images], buff) + buff.seek(0) return buff, mime_type def save_chunk(self, db_data_id, chunk_number, quality, buff, mime_type): diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index 49ee9c9f..08a660c5 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -287,19 +287,6 @@ class ZipChunkWriter(IChunkWriter): # and does not decode it to know img size. return [] - def save_as_chunk_to_buff(self, images): - buff = io.BytesIO() - - with zipfile.ZipFile(buff, 'w') as zip_file: - for idx, image in enumerate(images): - arcname = '{:06d}.{}'.format(idx, os.path.splitext(image)[1]) - if isinstance(image, av.VideoFrame): - zip_file.writestr(arcname, image.to_image().tobytes().getvalue()) - else: - zip_file.write(filename=image, arcname=arcname) - buff.seek(0) - return buff - class ZipCompressedChunkWriter(IChunkWriter): def save_as_chunk(self, images, chunk_path): image_sizes = [] @@ -312,23 +299,13 @@ class ZipCompressedChunkWriter(IChunkWriter): return image_sizes - def save_as_chunk_to_buff(self, images): - buff = io.BytesIO() - with zipfile.ZipFile(buff, 'x') as zip_file: - for idx, image in enumerate(images): - (_, _, image_buf) = self._compress_image(image, self._image_quality) - arcname = '{:06d}.jpeg'.format(idx) - zip_file.writestr(arcname, image_buf.getvalue()) - buff.seek(0) - return buff - class Mpeg4ChunkWriter(IChunkWriter): def __init__(self, _): super().__init__(17) self._output_fps = 25 @staticmethod - def _create_av_container(path, w, h, rate, options, f=None): + def _create_av_container(path, w, h, rate, options, f='mp4'): # x264 requires width and height must be divisible by 2 for yuv420p if h % 2: h += 1 @@ -366,41 +343,6 @@ class Mpeg4ChunkWriter(IChunkWriter): output_container.close() return [(input_w, input_h)] - def save_as_chunk_to_buff(self, frames): - if not frames: - raise Exception('no images to save') - - buff = io.BytesIO() - input_w = frames[0].width - input_h = frames[0].height - - output_container, output_v_stream = self._create_av_container( - path=buff, - w=input_w, - h=input_h, - rate=self._output_fps, - options={ - "crf": str(self._image_quality), - "preset": "ultrafast", - }, - f='mp4', - ) - - for frame in frames: - # let libav set the correct pts and time_base - frame.pts = None - frame.time_base = None - - for packet in output_v_stream.encode(frame): - output_container.mux(packet) - - # Flush streams - for packet in output_v_stream.encode(): - output_container.mux(packet) - output_container.close() - buff.seek(0) - return buff - @staticmethod def _encode_images(images, container, stream): for frame, _, _ in images: @@ -454,52 +396,6 @@ class Mpeg4CompressedChunkWriter(Mpeg4ChunkWriter): output_container.close() return [(input_w, input_h)] - def save_as_chunk_to_buff(self, frames): - if not frames: - raise Exception('no images to save') - - buff = io.BytesIO() - input_w = frames[0].width - input_h = frames[0].height - - downscale_factor = 1 - while input_h / downscale_factor >= 1080: - downscale_factor *= 2 - - output_h = input_h // downscale_factor - output_w = input_w // downscale_factor - - - output_container, output_v_stream = self._create_av_container( - path=buff, - w=output_w, - h=output_h, - rate=self._output_fps, - options={ - 'profile': 'baseline', - 'coder': '0', - 'crf': str(self._image_quality), - 'wpredp': '0', - 'flags': '-loop' - }, - f='mp4', - ) - - for frame in frames: - # let libav set the correct pts and time_base - frame.pts = None - frame.time_base = None - - for packet in output_v_stream.encode(frame): - output_container.mux(packet) - - # Flush streams - for packet in output_v_stream.encode(): - output_container.mux(packet) - output_container.close() - buff.seek(0) - return buff - def _is_archive(path): mime = mimetypes.guess_type(path) mime_type = mime[0]