Merge branch 'cache' of https://github.com/Marishka17/cvat into origin/cache

cvat/apps/engine/prepare.py

@@ -0,0 +1,76 @@
+import av
+
+class PrepareInfo:
+
+    def __init__(self, source_path, meta_path):
+        self.source_path = source_path
+        self.meta_path = meta_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
+
+    #@get_execution_time
+    def save_meta_info(self):
+        container = self._open_video_container(self.source_path, mode='r')
+        video_stream = self._get_video_stream(container)
+
+        with open(self.meta_path, 'w') as file:
+            frame_number = 0
+
+            for packet in  container.demux(video_stream):
+                for frame in packet.decode():
+                    frame_number += 1
+                    if frame.key_frame:
+                        file.write('{} {}\n'.format(frame_number, frame.pts))
+
+        self._close_video_container(container)
+        return frame_number# == task_size
+
+    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(' ')
+
+                #TODO: исправить если вдруг ключевой кадр окажется не первым
+                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, chunk_size):
+        start_chunk_frame_number = (chunk_number - 1) * chunk_size + 1
+        end_chunk_frame_number = start_chunk_frame_number + chunk_size #- 1
+        start_decode_frame_number, start_decode_timestamp = self.get_nearest_left_key_frame(start_chunk_frame_number)
+        extra_frames = start_chunk_frame_number - start_decode_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:
+                    yield frame
+                else:
+                    self._close_video_container(container)
+                    return
+
+        self._close_video_container(container)

cvat/apps/engine/task.py

@@ -24,6 +24,8 @@ from distutils.dir_util import copy_tree
 
 from . import models
 from .log import slogger
+from .prepare import PrepareInfo
+from diskcache import Cache
 
 ############################# Low Level server API
 
@@ -230,6 +232,22 @@ def _create_thread(tid, data):
     job.meta['status'] = 'Media files are being extracted...'
     job.save_meta()
 
+    if settings.USE_CACHE:
+        for media_type, media_files in media.items():
+            if media_files:
+                if task_mode == MEDIA_TYPES['video']['mode']:
+                    meta_info = PrepareInfo(source_path=os.path.join(upload_dir, media_files[0]),
+                                        meta_path=os.path.join(upload_dir, 'meta_info.txt'))
+                    meta_info.save_meta_info()
+                # else:
+                #     with Cache(settings.CACHE_ROOT) as cache:
+                #         counter_ = itertools.count(start=1)
+
+                          #TODO: chunk size
+                #         for chunk_number, media_paths in itertools.groupby(media_files, lambda x: next(counter_) // db_data.chunk_size):
+                #             cache.set('{}_{}'.format(tid, chunk_number), media_paths, tag='dummy')
+    #else:
+
     db_images = []
     extractor = None
 

cvat/settings/base.py

@@ -333,6 +333,9 @@ os.makedirs(DATA_ROOT, exist_ok=True)
 MEDIA_DATA_ROOT = os.path.join(DATA_ROOT, 'data')
 os.makedirs(MEDIA_DATA_ROOT, exist_ok=True)
 
+CACHE_ROOT = os.path.join(DATA_ROOT, 'cache')
+os.makedirs(CACHE_ROOT, exist_ok=True)
+
 TASKS_ROOT = os.path.join(DATA_ROOT, 'tasks')
 os.makedirs(TASKS_ROOT, exist_ok=True)
 
@@ -431,3 +434,17 @@ RESTRICTIONS = {
         'engine.role.admin',
         ),
 }
+
+CACHES = {
+   'default' : {
+       'BACKEND' : 'diskcache.DjangoCache',
+       'LOCATION' : CACHE_ROOT,
+       'TIMEOUT' : None,
+       'OPTIONS' : {
+            #'statistics' :True,
+            'size_limit' : 2 ** 40, # 1 тб
+       }
+   }
+}
+
+USE_CACHE = True