Update media extractors (#2154)

* Fix PDF reading

* Fixes

* update changelog

* t

* fix len

Co-authored-by: Nikita Manovich <nikita.manovich@intel.com>

CHANGELOG.md

@@ -39,6 +39,7 @@ filters and searching the nearest frame without any annotations (<https://github
 - Fixed multiple errors which arises when polygon is of length 5 or less (<https://github.com/opencv/cvat/pull/2100>)
 - Fixed task creation from PDF (<https://github.com/opencv/cvat/pull/2141>)
 - Fixed CVAT format import for frame stepped tasks (<https://github.com/openvinotoolkit/cvat/pull/2151>)
+- Fixed the reading problem with large PDFs (<https://github.com/openvinotoolkit/cvat/pull/2154>)
 - Fixed unnecessary pyhash dependency (<https://github.com/openvinotoolkit/cvat/pull/2170>)
 - Fixed Data is not getting cleared, even after deleting the Task from Django Admin App(<https://github.com/openvinotoolkit/cvat/issues/1925>)
 - Fixed blinking message: "Some tasks have not been showed because they do not have any data" (<https://github.com/openvinotoolkit/cvat/pull/2200>)

cvat/apps/engine/cache.py

@@ -2,14 +2,17 @@
 #
 # SPDX-License-Identifier: MIT
 
+import os
+from io import BytesIO
+
 from diskcache import Cache
 from django.conf import settings
-from cvat.apps.engine.media_extractors import (Mpeg4ChunkWriter, ZipChunkWriter,
-    Mpeg4CompressedChunkWriter, ZipCompressedChunkWriter)
+
+from cvat.apps.engine.media_extractors import (Mpeg4ChunkWriter,
+    Mpeg4CompressedChunkWriter, ZipChunkWriter, ZipCompressedChunkWriter)
 from cvat.apps.engine.models import DataChoice
-from .prepare import PrepareInfo
-import os
-from io import BytesIO
+from cvat.apps.engine.prepare import PrepareInfo
+
 
 class CacheInteraction:
     def __init__(self):
@@ -27,7 +30,7 @@ def get_buff_mime(self, chunk_number, quality, db_data):
         return chunk, tag
 
     def prepare_chunk_buff(self, db_data, quality, chunk_number):
-        from cvat.apps.engine.frame_provider import FrameProvider
+        from cvat.apps.engine.frame_provider import FrameProvider # TODO: remove circular dependency
         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,
@@ -54,4 +57,4 @@ def prepare_chunk_buff(self, db_data, quality, chunk_number):
         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)
+        self._cache.set('{}_{}_{}'.format(db_data_id, chunk_number, quality), buff, tag=mime_type)

cvat/apps/engine/frame_provider.py

@@ -9,10 +9,11 @@
 import numpy as np
 from PIL import Image
 
+from cvat.apps.engine.cache import CacheInteraction
 from cvat.apps.engine.media_extractors import VideoReader, ZipReader
 from cvat.apps.engine.mime_types import mimetypes
 from cvat.apps.engine.models import DataChoice, StorageMethodChoice
-from .cache import CacheInteraction
+
 
 class RandomAccessIterator:
     def __init__(self, iterable):

cvat/apps/engine/media_extractors.py

@@ -7,6 +7,7 @@
 import shutil
 import zipfile
 import io
+import itertools
 from abc import ABC, abstractmethod
 
 import av
@@ -65,9 +66,16 @@ def _get_preview(obj):
         return preview.convert('RGB')
 
     @abstractmethod
-    def get_image_size(self):
+    def get_image_size(self, i):
         pass
 
+    def __len__(self):
+        return len(self.frame_range)
+
+    @property
+    def frame_range(self):
+        return range(self._start, self._stop, self._step)
+
 class ImageListReader(IMediaReader):
     def __init__(self, source_path, step=1, start=0, stop=None):
         if not source_path:
@@ -104,8 +112,8 @@ def get_preview(self):
         fp = open(self._source_path[0], "rb")
         return self._get_preview(fp)
 
-    def get_image_size(self):
-        img = Image.open(self._source_path[0])
+    def get_image_size(self, i):
+        img = Image.open(self._source_path[i])
         return img.width, img.height
 
 class DirectoryReader(ImageListReader):
@@ -127,44 +135,45 @@ class ArchiveReader(DirectoryReader):
     def __init__(self, source_path, step=1, start=0, stop=None):
         self._archive_source = source_path[0]
         Archive(self._archive_source).extractall(os.path.dirname(source_path[0]))
+        os.remove(self._archive_source)
         super().__init__(
             source_path=[os.path.dirname(source_path[0])],
             step=step,
             start=start,
             stop=stop,
         )
 
-    def __del__(self):
-        os.remove(self._archive_source)
-
-class PdfReader(DirectoryReader):
+class PdfReader(ImageListReader):
     def __init__(self, source_path, step=1, start=0, stop=None):
         if not source_path:
             raise Exception('No PDF found')
 
-        from pdf2image import convert_from_path
         self._pdf_source = source_path[0]
-        self._tmp_dir = create_tmp_dir()
-        file_ = convert_from_path(self._pdf_source)
-        basename = os.path.splitext(os.path.basename(self._pdf_source))[0]
-        for page_num, page in enumerate(file_):
-            output = os.path.join(self._tmp_dir, '{}{:09d}.jpeg'.format(basename, page_num))
-            page.save(output, 'JPEG')
+
+        _basename = os.path.splitext(os.path.basename(self._pdf_source))[0]
+        _counter = itertools.count()
+        def _make_name():
+            for page_num in _counter:
+                yield '{}{:09d}.jpeg'.format(_basename, page_num)
+
+        from pdf2image import convert_from_path
+        self._tmp_dir = os.path.dirname(source_path[0])
+        os.makedirs(self._tmp_dir, exist_ok=True)
+
+        # Avoid OOM: https://github.com/openvinotoolkit/cvat/issues/940
+        paths = convert_from_path(self._pdf_source,
+            last_page=stop, paths_only=True,
+            output_folder=self._tmp_dir, fmt="jpeg", output_file=_make_name())
+
+        os.remove(source_path[0])
 
         super().__init__(
-            source_path=[self._tmp_dir],
+            source_path=paths,
             step=step,
             start=start,
             stop=stop,
         )
 
-    def __del__(self):
-        delete_tmp_dir(self._tmp_dir)
-
-    def get_path(self, i):
-        base_dir = os.path.dirname(self._pdf_source)
-        return os.path.join(base_dir, os.path.relpath(self._source_path[i], self._tmp_dir))
-
 class ZipReader(ImageListReader):
     def __init__(self, source_path, step=1, start=0, stop=None):
         self._zip_source = zipfile.ZipFile(source_path[0], mode='r')
@@ -178,8 +187,8 @@ def get_preview(self):
         io_image = io.BytesIO(self._zip_source.read(self._source_path[0]))
         return self._get_preview(io_image)
 
-    def get_image_size(self):
-        img = Image.open(io.BytesIO(self._zip_source.read(self._source_path[0])))
+    def get_image_size(self, i):
+        img = Image.open(io.BytesIO(self._zip_source.read(self._source_path[i])))
         return img.width, img.height
 
     def get_image(self, i):
@@ -243,7 +252,7 @@ def get_preview(self):
         preview = next(container.decode(stream))
         return self._get_preview(preview.to_image())
 
-    def get_image_size(self):
+    def get_image_size(self, i):
         image = (next(iter(self)))[0]
         return image.width, image.height
 

cvat/apps/engine/task.py

@@ -273,7 +273,7 @@ def update_progress(progress):
     # calculate chunk size if it isn't specified
     if db_data.chunk_size is None:
         if isinstance(compressed_chunk_writer, ZipCompressedChunkWriter):
-            w, h = extractor.get_image_size()
+            w, h = extractor.get_image_size(0)
             area = h * w
             db_data.chunk_size = max(2, min(72, 36 * 1920 * 1080 // area))
         else:
@@ -285,58 +285,48 @@ def update_progress(progress):
 
     if settings.USE_CACHE and db_data.storage_method == StorageMethodChoice.CACHE:
        for media_type, media_files in media.items():
-            if media_files:
-                if task_mode == MEDIA_TYPES['video']['mode']:
-                    try:
-                        analyzer = AnalyzeVideo(source_path=os.path.join(upload_dir, media_files[0]))
-                        analyzer.check_type_first_frame()
-                        analyzer.check_video_timestamps_sequences()
-
-                        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_key_frames()
-                        meta_info.check_seek_key_frames()
-                        meta_info.save_meta_info()
-
-                        all_frames = meta_info.get_task_size()
-                        db_data.size = len(range(db_data.start_frame, min(data['stop_frame'] + 1 if data['stop_frame'] else all_frames, all_frames), db_data.get_frame_step()))
-                        video_path = os.path.join(upload_dir, media_files[0])
-                        frame = meta_info.key_frames.get(next(iter(meta_info.key_frames)))
-                        video_size = (frame.width, frame.height)
-
-                    except Exception:
-                        db_data.storage_method = StorageMethodChoice.FILE_SYSTEM
-
-                else:#images,archive
-                    counter_ = itertools.count()
-                    if isinstance(extractor, MEDIA_TYPES['archive']['extractor']):
-                        media_files = [os.path.relpath(path, upload_dir) for path in extractor._source_path]
-                    elif isinstance(extractor, (MEDIA_TYPES['zip']['extractor'], MEDIA_TYPES['pdf']['extractor'])):
-                        media_files = extractor._source_path
-
-                    numbers_sequence = range(db_data.start_frame, min(data['stop_frame'] if data['stop_frame'] else len(media_files), len(media_files)), db_data.get_frame_step())
-                    m_paths = []
-                    m_paths = [(path, numb) for numb, path in enumerate(sorted(media_files)) if numb in numbers_sequence]
-
-                    for chunk_number, media_paths in itertools.groupby(m_paths, lambda x: next(counter_) // db_data.chunk_size):
-                        media_paths = list(media_paths)
-                        img_sizes = []
-                        from PIL import Image
-                        with open(db_data.get_dummy_chunk_path(chunk_number), 'w') as dummy_chunk:
-                            for path, _ in media_paths:
-                                dummy_chunk.write(path+'\n')
-                                img_sizes += [Image.open(os.path.join(upload_dir, path)).size]
-
-                        db_data.size += len(media_paths)
-                        db_images.extend([
-                            models.Image(
-                                data=db_data,
-                                path=data[0],
-                                frame=data[1],
-                                width=size[0],
-                                height=size[1])
-                            for data, size in zip(media_paths, img_sizes)
-                        ])
+            if not media_files:
+                continue
+
+            if task_mode == MEDIA_TYPES['video']['mode']:
+                try:
+                    analyzer = AnalyzeVideo(source_path=os.path.join(upload_dir, media_files[0]))
+                    analyzer.check_type_first_frame()
+                    analyzer.check_video_timestamps_sequences()
+
+                    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_key_frames()
+                    meta_info.check_seek_key_frames()
+                    meta_info.save_meta_info()
+
+                    all_frames = meta_info.get_task_size()
+                    db_data.size = len(range(db_data.start_frame, min(data['stop_frame'] + 1 if data['stop_frame'] else all_frames, all_frames), db_data.get_frame_step()))
+                    video_path = os.path.join(upload_dir, media_files[0])
+                    frame = meta_info.key_frames.get(next(iter(meta_info.key_frames)))
+                    video_size = (frame.width, frame.height)
+
+                except Exception:
+                    db_data.storage_method = StorageMethodChoice.FILE_SYSTEM
+
+            else:#images,archive
+                db_data.size = len(extractor)
+
+                counter = itertools.count()
+                for chunk_number, chunk_frames in itertools.groupby(extractor.frame_range, lambda x: next(counter) // db_data.chunk_size):
+                    chunk_paths = [(extractor.get_path(i), i) for i in chunk_frames]
+                    img_sizes = []
+                    with open(db_data.get_dummy_chunk_path(chunk_number), 'w') as dummy_chunk:
+                        for path, frame_id in chunk_paths:
+                            dummy_chunk.write(path + '\n')
+                            img_sizes.append(extractor.get_image_size(frame_id))
+
+                    db_images.extend([
+                        models.Image(data=db_data,
+                            path=os.path.relpath(path, upload_dir),
+                            frame=frame, width=w, height=h)
+                        for (path, frame), (w, h) in zip(chunk_paths, img_sizes)
+                    ])
 
     if db_data.storage_method == StorageMethodChoice.FILE_SYSTEM or not settings.USE_CACHE:
         counter = itertools.count()
@@ -383,5 +373,5 @@ def update_progress(progress):
     preview = extractor.get_preview()
     preview.save(db_data.get_preview_path())
 
-    slogger.glob.info("Founded frames {} for Data #{}".format(db_data.size, db_data.id))
+    slogger.glob.info("Found frames {} for Data #{}".format(db_data.size, db_data.id))
     _save_task_to_db(db_task)