Merge branch 'cache' of https://github.com/Marishka17/cvat into Marishka17-cache
Nikita Manovich
6 years ago
commit
28225f9dc6
@ -0,0 +1,57 @@
|
||||
# Copyright (C) 2020 Intel Corporation
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from diskcache import Cache
|
||||
from django.conf import settings
|
||||
from cvat.apps.engine.media_extractors import (Mpeg4ChunkWriter, ZipChunkWriter,
|
||||
Mpeg4CompressedChunkWriter, ZipCompressedChunkWriter)
|
||||
from cvat.apps.engine.models import DataChoice
|
||||
from .prepare import PrepareInfo
|
||||
import os
|
||||
from io import BytesIO
|
||||
|
||||
class CacheInteraction:
|
||||
def __init__(self):
|
||||
self._cache = Cache(settings.CACHE_ROOT)
|
||||
|
||||
def __del__(self):
|
||||
self._cache.close()
|
||||
|
||||
def get_buff_mime(self, chunk_number, quality, db_data):
|
||||
chunk, tag = self._cache.get('{}_{}_{}'.format(db_data.id, chunk_number, quality), tag=True)
|
||||
|
||||
if not chunk:
|
||||
chunk, tag = self.prepare_chunk_buff(db_data, quality, chunk_number)
|
||||
self.save_chunk(db_data.id, chunk_number, quality, chunk, tag)
|
||||
return chunk, tag
|
||||
|
||||
def prepare_chunk_buff(self, db_data, quality, chunk_number):
|
||||
from cvat.apps.engine.frame_provider import FrameProvider
|
||||
extractor_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,
|
||||
}
|
||||
|
||||
image_quality = 100 if extractor_classes[quality] in [Mpeg4ChunkWriter, ZipChunkWriter] else db_data.image_quality
|
||||
mime_type = 'video/mp4' if extractor_classes[quality] in [Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter] else 'application/zip'
|
||||
|
||||
extractor = extractor_classes[quality](image_quality)
|
||||
|
||||
images = []
|
||||
buff = BytesIO()
|
||||
if os.path.exists(db_data.get_meta_path()):
|
||||
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):
|
||||
images.append(frame)
|
||||
extractor.save_as_chunk([(image, source_path, None) for image in images], buff)
|
||||
else:
|
||||
with open(db_data.get_dummy_chunk_path(chunk_number), 'r') as dummy_file:
|
||||
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):
|
||||
self._cache.set('{}_{}_{}'.format(db_data_id, chunk_number, quality), buff, tag=mime_type)
|
||||
@ -0,0 +1,36 @@
|
||||
# Generated by Django 2.2.13 on 2020-08-13 05:49
|
||||
|
||||
from cvat.apps.engine.media_extractors import _is_archive, _is_zip
|
||||
import cvat.apps.engine.models
|
||||
from django.conf import settings
|
||||
from django.db import migrations, models
|
||||
import os
|
||||
from pyunpack import Archive
|
||||
|
||||
def unzip(apps, schema_editor):
|
||||
Data = apps.get_model("engine", "Data")
|
||||
data_q_set = Data.objects.all()
|
||||
archive_paths = []
|
||||
|
||||
for data_instance in data_q_set:
|
||||
for root, _, files in os.walk(os.path.join(settings.MEDIA_DATA_ROOT, '{}/raw/'.format(data_instance.id))):
|
||||
archive_paths.extend([os.path.join(root, file) for file in files if _is_archive(file) or _is_zip(file)])
|
||||
|
||||
for path in archive_paths:
|
||||
Archive(path).extractall(os.path.dirname(path))
|
||||
os.remove(path)
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
|
||||
dependencies = [
|
||||
('engine', '0028_labelcolor'),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.AddField(
|
||||
model_name='data',
|
||||
name='storage_method',
|
||||
field=models.CharField(choices=[('cache', 'CACHE'), ('file_system', 'FILE_SYSTEM')], default=cvat.apps.engine.models.StorageMethodChoice('file_system'), max_length=15),
|
||||
),
|
||||
migrations.RunPython(unzip),
|
||||
]
|
||||
@ -0,0 +1,155 @@
|
||||
# Copyright (C) 2020 Intel Corporation
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import av
|
||||
import hashlib
|
||||
|
||||
class WorkWithVideo:
|
||||
def __init__(self, **kwargs):
|
||||
if not kwargs.get('source_path'):
|
||||
raise Exception('No sourse path')
|
||||
self.source_path = kwargs.get('source_path')
|
||||
|
||||
def _open_video_container(self, sourse_path, mode, options=None):
|
||||
return av.open(sourse_path, mode=mode, options=options)
|
||||
|
||||
def _close_video_container(self, container):
|
||||
container.close()
|
||||
|
||||
def _get_video_stream(self, container):
|
||||
video_stream = next(stream for stream in container.streams if stream.type == 'video')
|
||||
video_stream.thread_type = 'AUTO'
|
||||
return video_stream
|
||||
|
||||
|
||||
class AnalyzeVideo(WorkWithVideo):
|
||||
def check_type_first_frame(self):
|
||||
container = self._open_video_container(self.source_path, mode='r')
|
||||
video_stream = self._get_video_stream(container)
|
||||
|
||||
for packet in container.demux(video_stream):
|
||||
for frame in packet.decode():
|
||||
self._close_video_container(container)
|
||||
assert frame.pict_type.name == 'I', 'First frame is not key frame'
|
||||
return
|
||||
|
||||
def check_video_timestamps_sequences(self):
|
||||
container = self._open_video_container(self.source_path, mode='r')
|
||||
video_stream = self._get_video_stream(container)
|
||||
|
||||
frame_pts = -1
|
||||
frame_dts = -1
|
||||
for packet in container.demux(video_stream):
|
||||
for frame in packet.decode():
|
||||
|
||||
if None not in [frame.pts, frame_pts] and frame.pts <= frame_pts:
|
||||
self._close_video_container(container)
|
||||
raise Exception('Invalid pts sequences')
|
||||
|
||||
if None not in [frame.dts, frame_dts] and frame.dts <= frame_dts:
|
||||
self._close_video_container(container)
|
||||
raise Exception('Invalid dts sequences')
|
||||
|
||||
frame_pts, frame_dts = frame.pts, frame.dts
|
||||
self._close_video_container(container)
|
||||
|
||||
def md5_hash(frame):
|
||||
return hashlib.md5(frame.to_image().tobytes()).hexdigest()
|
||||
|
||||
class PrepareInfo(WorkWithVideo):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
if not kwargs.get('meta_path'):
|
||||
raise Exception('No meta path')
|
||||
|
||||
self.meta_path = kwargs.get('meta_path')
|
||||
self.key_frames = {}
|
||||
self.frames = 0
|
||||
|
||||
def get_task_size(self):
|
||||
return self.frames
|
||||
|
||||
def check_seek_key_frames(self):
|
||||
container = self._open_video_container(self.source_path, mode='r')
|
||||
video_stream = self._get_video_stream(container)
|
||||
|
||||
key_frames_copy = self.key_frames.copy()
|
||||
|
||||
for index, key_frame in key_frames_copy.items():
|
||||
container.seek(offset=key_frame.pts, stream=video_stream)
|
||||
flag = True
|
||||
for packet in container.demux(video_stream):
|
||||
for frame in packet.decode():
|
||||
if md5_hash(frame) != md5_hash(key_frame) or frame.pts != key_frame.pts:
|
||||
self.key_frames.pop(index)
|
||||
flag = False
|
||||
break
|
||||
if not flag:
|
||||
break
|
||||
|
||||
#TODO: correct ratio of number of frames to keyframes
|
||||
if len(self.key_frames) == 0:
|
||||
raise Exception('Too few keyframes')
|
||||
|
||||
def save_key_frames(self):
|
||||
container = self._open_video_container(self.source_path, mode='r')
|
||||
video_stream = self._get_video_stream(container)
|
||||
frame_number = 0
|
||||
|
||||
for packet in container.demux(video_stream):
|
||||
for frame in packet.decode():
|
||||
if frame.key_frame:
|
||||
self.key_frames[frame_number] = frame
|
||||
frame_number += 1
|
||||
|
||||
self.frames = frame_number
|
||||
self._close_video_container(container)
|
||||
|
||||
def save_meta_info(self):
|
||||
with open(self.meta_path, 'w') as meta_file:
|
||||
for index, frame in self.key_frames.items():
|
||||
meta_file.write('{} {}\n'.format(index, frame.pts))
|
||||
|
||||
def get_nearest_left_key_frame(self, start_chunk_frame_number):
|
||||
start_decode_frame_number = 0
|
||||
start_decode_timestamp = 0
|
||||
|
||||
with open(self.meta_path, 'r') as file:
|
||||
for line in file:
|
||||
frame_number, timestamp = line.strip().split(' ')
|
||||
|
||||
if int(frame_number) <= start_chunk_frame_number:
|
||||
start_decode_frame_number = frame_number
|
||||
start_decode_timestamp = timestamp
|
||||
else:
|
||||
break
|
||||
|
||||
return int(start_decode_frame_number), int(start_decode_timestamp)
|
||||
|
||||
def decode_needed_frames(self, chunk_number, db_data):
|
||||
step = db_data.get_frame_step()
|
||||
start_chunk_frame_number = db_data.start_frame + chunk_number * db_data.chunk_size * step
|
||||
end_chunk_frame_number = min(start_chunk_frame_number + (db_data.chunk_size - 1) * step + 1, db_data.stop_frame + 1)
|
||||
start_decode_frame_number, start_decode_timestamp = self.get_nearest_left_key_frame(start_chunk_frame_number)
|
||||
container = self._open_video_container(self.source_path, mode='r')
|
||||
video_stream = self._get_video_stream(container)
|
||||
container.seek(offset=start_decode_timestamp, stream=video_stream)
|
||||
|
||||
frame_number = start_decode_frame_number - 1
|
||||
for packet in container.demux(video_stream):
|
||||
for frame in packet.decode():
|
||||
frame_number += 1
|
||||
if frame_number < start_chunk_frame_number:
|
||||
continue
|
||||
elif frame_number < end_chunk_frame_number and not ((frame_number - start_chunk_frame_number) % step):
|
||||
yield frame
|
||||
elif (frame_number - start_chunk_frame_number) % step:
|
||||
continue
|
||||
else:
|
||||
self._close_video_container(container)
|
||||
return
|
||||
|
||||
self._close_video_container(container)
|
||||
Loading…
Reference in New Issue