Update media extractors

CHANGELOG.md

@@ -31,6 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 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>)
 
 ### Security

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)