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