Merge pull request #1811 from wkentaro/perf/large-image-loading

perf: fix large image loading performance

Author: wkentaro

labelme/_label_file.py

@@ -4,6 +4,7 @@
 import io
 import json
 import os.path as osp
+import time
 from pathlib import PureWindowsPath
 from typing import TypedDict
 
@@ -139,20 +140,26 @@ def __init__(self, filename=None):
 
     @staticmethod
     def load_image_file(filename):
+        t0 = time.time()
         image_pil = PIL.Image.open(filename)
 
-        # apply orientation to image according to exif
-        image_pil = utils.apply_exif_orientation(image_pil)
-
-        with io.BytesIO() as f:
-            ext = osp.splitext(filename)[1].lower()
-            if ext in [".jpg", ".jpeg"]:
-                format = "JPEG"
-            else:
-                format = "PNG"
-            image_pil.save(f, format=format)
-            f.seek(0)
-            return f.read()
+        oriented: PIL.Image.Image = utils.apply_exif_orientation(image_pil)
+        ext = osp.splitext(filename)[1].lower()
+        if oriented is image_pil and ext in (".jpg", ".jpeg", ".png"):
+            # no encoding needed
+            with builtins.open(filename, "rb") as f:
+                imageData = f.read()
+        else:
+            with io.BytesIO() as f:
+                format = "PNG" if "A" in oriented.mode else "JPEG"
+                oriented.save(f, format=format, quality=95)
+                f.seek(0)
+                imageData = f.read()
+
+        logger.debug(
+            "Loaded image file: {!r} in {:.0f}ms", filename, (time.time() - t0) * 1000
+        )
+        return imageData
 
     def load(self, filename):
         keys = [
@@ -180,7 +187,7 @@ def load(self, filename):
                 )
             flags = data.get("flags") or {}
             self._check_image_height_and_width(
-                base64.b64encode(imageData).decode("utf-8"),
+                imageData,
                 data.get("imageHeight"),
                 data.get("imageWidth"),
             )
@@ -205,19 +212,20 @@ def load(self, filename):
 
     @staticmethod
     def _check_image_height_and_width(imageData, imageHeight, imageWidth):
-        img_arr = utils.img_b64_to_arr(imageData)
-        if imageHeight is not None and img_arr.shape[0] != imageHeight:
+        img_pil = utils.img_data_to_pil(imageData)
+        actual_w, actual_h = img_pil.size
+        if imageHeight is not None and actual_h != imageHeight:
             logger.error(
                 "imageHeight does not match with imageData or imagePath, "
                 "so getting imageHeight from actual image."
             )
-            imageHeight = img_arr.shape[0]
-        if imageWidth is not None and img_arr.shape[1] != imageWidth:
+            imageHeight = actual_h
+        if imageWidth is not None and actual_w != imageWidth:
             logger.error(
                 "imageWidth does not match with imageData or imagePath, "
                 "so getting imageWidth from actual image."
             )
-            imageWidth = img_arr.shape[1]
+            imageWidth = actual_w
         return imageHeight, imageWidth
 
     def save(
@@ -232,10 +240,10 @@ def save(
         flags=None,
     ):
         if imageData is not None:
-            imageData = base64.b64encode(imageData).decode("utf-8")
             imageHeight, imageWidth = self._check_image_height_and_width(
                 imageData, imageHeight, imageWidth
             )
+            imageData = base64.b64encode(imageData).decode("utf-8")
         if otherData is None:
             otherData = {}
         if flags is None:

labelme/app.py

@@ -9,6 +9,7 @@
 import platform
 import re
 import subprocess
+import time
 import types
 import webbrowser
 from pathlib import Path
@@ -1718,12 +1719,6 @@ def brightnessContrast(self, value: bool, is_initial_load: bool = False):
             logger.warning("filename is None, cannot set brightness/contrast")
             return
 
-        dialog = BrightnessContrastDialog(
-            utils.img_data_to_pil(self.imageData).convert("RGB"),
-            self.onNewBrightnessContrast,
-            parent=self,
-        )
-
         brightness: int | None
         contrast: int | None
         brightness, contrast = self._brightness_contrast_values.get(
@@ -1735,6 +1730,20 @@ def brightnessContrast(self, value: bool, is_initial_load: bool = False):
                 brightness, contrast = self._brightness_contrast_values.get(
                     prev_filename, (None, None)
                 )
+            if brightness is None and contrast is None:
+                return
+
+        logger.debug(
+            "Opening brightness/contrast dialog with brightness={}, contrast={}",
+            brightness,
+            contrast,
+        )
+        dialog = BrightnessContrastDialog(
+            utils.img_data_to_pil(self.imageData).convert("RGB"),
+            self.onNewBrightnessContrast,
+            parent=self,
+        )
+
         if brightness is not None:
             dialog.slider_brightness.setValue(brightness)
         if contrast is not None:
@@ -1748,6 +1757,12 @@ def brightnessContrast(self, value: bool, is_initial_load: bool = False):
             contrast = dialog.slider_contrast.value()
 
         self._brightness_contrast_values[self.filename] = (brightness, contrast)
+        logger.debug(
+            "Updated states for {}: brightness={}, contrast={}",
+            self.filename,
+            brightness,
+            contrast,
+        )
 
     def togglePolygons(self, value):
         flag = value
@@ -1786,6 +1801,7 @@ def _load_file(self, filename=None):
             return False
         # assumes same name, but json extension
         self.show_status_message(self.tr("Loading %s...") % osp.basename(str(filename)))
+        t0_load_file = time.time()
         label_file = f"{osp.splitext(filename)[0]}.json"
         if self.output_dir:
             label_file_without_path = osp.basename(label_file)
@@ -1830,7 +1846,9 @@ def _load_file(self, filename=None):
                 self.imagePath = filename
             self.labelFile = None
         assert self.imageData is not None
+        t0 = time.time()
         image = QtGui.QImage.fromData(self.imageData)
+        logger.debug("Created QImage in {:.0f}ms", (time.time() - t0) * 1000)
 
         if image.isNull():
             formats = [
@@ -1848,7 +1866,9 @@ def _load_file(self, filename=None):
             return False
         self.image = image
         self.filename = filename
+        t0 = time.time()
         self.canvas.loadPixmap(QtGui.QPixmap.fromImage(image))
+        logger.debug("Loaded pixmap in {:.0f}ms", (time.time() - t0) * 1000)
         flags = {k: False for k in self._config["flags"] or []}
         if self.labelFile:
             self._load_shape_dicts(shape_dicts=self.labelFile.shapes)
@@ -1881,7 +1901,11 @@ def _load_file(self, filename=None):
         self.toggleActions(True)
         self.canvas.setFocus()
         self.show_status_message(self.tr("Loaded %s") % osp.basename(filename))
-        logger.debug("loaded file: {!r}", filename)
+        logger.info(
+            "Loaded file: {!r} in {:.0f}ms",
+            filename,
+            (time.time() - t0_load_file) * 1000,
+        )
         return True
 
     def resizeEvent(self, a0: QtGui.QResizeEvent) -> None:

tests/unit/load_image_file_test.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import PIL.Image
+
+from labelme._label_file import LabelFile
+
+
+def _make_image(tmp_path: Path, filename: str, mode: str = "RGB", size=(100, 100)):
+    channels = 4 if mode == "RGBA" else 3
+    arr = np.random.randint(0, 255, (size[1], size[0], channels), dtype=np.uint8)
+    path = tmp_path / filename
+    PIL.Image.fromarray(arr, mode=mode).save(str(path))
+    return path
+
+
+def test_tiff_without_alpha_encoded_as_jpeg(tmp_path):
+    path = _make_image(tmp_path, "test.tiff")
+    data = LabelFile.load_image_file(str(path))
+    assert data[:2] == b"\xff\xd8"
+
+
+def test_tiff_with_alpha_encoded_as_png(tmp_path):
+    path = _make_image(tmp_path, "test.tiff", mode="RGBA")
+    data = LabelFile.load_image_file(str(path))
+    assert data[:4] == b"\x89PNG"
+
+
+def test_jpeg_returns_raw_bytes(tmp_path):
+    path = _make_image(tmp_path, "test.jpg")
+    data = LabelFile.load_image_file(str(path))
+    assert data == path.read_bytes()
+
+
+def test_png_returns_raw_bytes(tmp_path):
+    path = _make_image(tmp_path, "test.png")
+    data = LabelFile.load_image_file(str(path))
+    assert data == path.read_bytes()