From 61c6a016b91208cc2d8c3ee80fcebab8f017bf66 Mon Sep 17 00:00:00 2001 From: Andrey Zhavoronkov Date: Mon, 9 Dec 2024 09:33:41 +0300 Subject: [PATCH] Memory optimization for image chunk preparation (#8778) ### Motivation and context ### How has this been tested? ### Checklist - [x] I submit my changes into the `develop` branch - [x] I have created a changelog fragment - [ ] I have updated the documentation accordingly - [ ] I have added tests to cover my changes - [ ] I have linked related issues (see [GitHub docs]( https://help.github.com/en/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)) - [ ] I have increased versions of npm packages if it is necessary ([cvat-canvas](https://github.com/cvat-ai/cvat/tree/develop/cvat-canvas#versioning), [cvat-core](https://github.com/cvat-ai/cvat/tree/develop/cvat-core#versioning), [cvat-data](https://github.com/cvat-ai/cvat/tree/develop/cvat-data#versioning) and [cvat-ui](https://github.com/cvat-ai/cvat/tree/develop/cvat-ui#versioning)) ### License - [x] I submit _my code changes_ under the same [MIT License]( https://github.com/cvat-ai/cvat/blob/develop/LICENSE) that covers the project. Feel free to contact the maintainers if that's a concern. ## Summary by CodeRabbit ## Release Notes - **New Features** - Enhanced image loading functionality with a new `load_image` feature, improving handling of various image formats and EXIF rotation. - Improved segment preview generation, specifically for 3D segments. - **Bug Fixes** - Resolved issues with TIFF image handling, ensuring correct image rotation and processing. - **Improvements** - Refined task creation process with better validation checks and progress tracking for long-running tasks. - Updated utility functions for better type safety and error handling. --- changelog.d/20241205_165408_andrey_memory.md | 4 +++ cvat/apps/engine/cache.py | 8 ++---- cvat/apps/engine/media_extractors.py | 30 ++++++++++++-------- cvat/apps/engine/task.py | 5 ++-- cvat/apps/engine/utils.py | 4 --- 5 files changed, 27 insertions(+), 24 deletions(-) create mode 100644 changelog.d/20241205_165408_andrey_memory.md diff --git a/changelog.d/20241205_165408_andrey_memory.md b/changelog.d/20241205_165408_andrey_memory.md new file mode 100644 index 000000000000..fa7d959977f8 --- /dev/null +++ b/changelog.d/20241205_165408_andrey_memory.md @@ -0,0 +1,4 @@ +### Fixed + +- Memory consumption during preparation of image chunks + () diff --git a/cvat/apps/engine/cache.py b/cvat/apps/engine/cache.py index 0ecd7fcc010c..e27b697e41c4 100644 --- a/cvat/apps/engine/cache.py +++ b/cvat/apps/engine/cache.py @@ -62,14 +62,10 @@ VideoReaderWithManifest, ZipChunkWriter, ZipCompressedChunkWriter, -) -from cvat.apps.engine.rq_job_handler import RQJobMetaField -from cvat.apps.engine.utils import ( - CvatChunkTimestampMismatchError, - get_rq_lock_for_job, load_image, - md5_hash, ) +from cvat.apps.engine.rq_job_handler import RQJobMetaField +from cvat.apps.engine.utils import CvatChunkTimestampMismatchError, get_rq_lock_for_job, md5_hash from utils.dataset_manifest import ImageManifestManager slogger = ServerLogManager(__name__) diff --git a/cvat/apps/engine/media_extractors.py b/cvat/apps/engine/media_extractors.py index c923083b18b3..3e7b8e17a31b 100644 --- a/cvat/apps/engine/media_extractors.py +++ b/cvat/apps/engine/media_extractors.py @@ -101,6 +101,12 @@ def image_size_within_orientation(img: Image.Image): def has_exif_rotation(img: Image.Image): return img.getexif().get(ORIENTATION_EXIF_TAG, ORIENTATION.NORMAL_HORIZONTAL) != ORIENTATION.NORMAL_HORIZONTAL + +def load_image(image: tuple[str, str, str])-> tuple[Image.Image, str, str]: + with Image.open(image[0]) as pil_img: + pil_img.load() + return pil_img, image[1], image[2] + _T = TypeVar("_T") @@ -837,13 +843,15 @@ def _compress_image(source_image: av.VideoFrame | io.IOBase | Image.Image, quali if isinstance(source_image, av.VideoFrame): image = source_image.to_image() elif isinstance(source_image, io.IOBase): - with Image.open(source_image) as _img: - image = ImageOps.exif_transpose(_img) + image, _, _ = load_image((source_image, None, None)) elif isinstance(source_image, Image.Image): - image = ImageOps.exif_transpose(source_image) + image = source_image assert image is not None + if has_exif_rotation(image): + image = ImageOps.exif_transpose(image) + # Ensure image data fits into 8bit per pixel before RGB conversion as PIL clips values on conversion if image.mode == "I": # Image mode is 32bit integer pixels. @@ -868,16 +876,14 @@ def _compress_image(source_image: av.VideoFrame | io.IOBase | Image.Image, quali image = Image.fromarray(image, mode="L") # 'L' := Unsigned Integer 8, Grayscale image = ImageOps.equalize(image) # The Images need equalization. High resolution with 16-bit but only small range that actually contains information - converted_image = image.convert('RGB') + if image.mode != 'RGB' and image.mode != 'L': + image = image.convert('RGB') - try: - buf = io.BytesIO() - converted_image.save(buf, format='JPEG', quality=quality, optimize=True) - buf.seek(0) - width, height = converted_image.size - return width, height, buf - finally: - converted_image.close() + buf = io.BytesIO() + image.save(buf, format='JPEG', quality=quality, optimize=True) + buf.seek(0) + + return image.width, image.height, buf @abstractmethod def save_as_chunk(self, images, chunk_path): diff --git a/cvat/apps/engine/task.py b/cvat/apps/engine/task.py index c1a27c38f392..4e10ea9e5dab 100644 --- a/cvat/apps/engine/task.py +++ b/cvat/apps/engine/task.py @@ -34,12 +34,13 @@ from cvat.apps.engine.media_extractors import ( MEDIA_TYPES, CachingMediaIterator, IMediaReader, ImageListReader, Mpeg4ChunkWriter, Mpeg4CompressedChunkWriter, RandomAccessIterator, - ValidateDimension, ZipChunkWriter, ZipCompressedChunkWriter, get_mime, sort + ValidateDimension, ZipChunkWriter, ZipCompressedChunkWriter, get_mime, sort, + load_image, ) from cvat.apps.engine.models import RequestAction, RequestTarget from cvat.apps.engine.utils import ( av_scan_paths, format_list,get_rq_job_meta, - define_dependent_job, get_rq_lock_by_user, load_image + define_dependent_job, get_rq_lock_by_user ) from cvat.apps.engine.rq_job_handler import RQId from cvat.utils.http import make_requests_session, PROXIES_FOR_UNTRUSTED_URLS diff --git a/cvat/apps/engine/utils.py b/cvat/apps/engine/utils.py index 59409ceb69dd..5d383df3465e 100644 --- a/cvat/apps/engine/utils.py +++ b/cvat/apps/engine/utils.py @@ -377,10 +377,6 @@ def sendfile( return _sendfile(request, filename, attachment, attachment_filename, mimetype, encoding) -def load_image(image: tuple[str, str, str])-> tuple[Image.Image, str, str]: - pil_img = Image.open(image[0]) - pil_img.load() - return pil_img, image[1], image[2] def build_backup_file_name( *,