Faster image size extraction is

.github/workflows/run-tests.yml

@@ -11,29 +11,29 @@ on:
 jobs:
   test:
     name: Tests
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         python: ["3.7", "3.10"]
 
     steps:
-    - name: Checkout code
-      uses: actions/checkout@v3
+      - name: Checkout code
+        uses: actions/checkout@v3
 
-    - name: Set up Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: ${{ matrix.python }}
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python }}
 
-    - name: Install Poetry
-      uses: snok/install-poetry@v1
-      with:
-        version: 1.4.2
+      - name: Install Poetry
+        uses: snok/install-poetry@v1
+        with:
+          version: 1.4.2
 
-    - name: Install the package and dependencies
-      run: |
-        poetry install
+      - name: Install the package and dependencies
+        run: |
+          poetry install
 
-    - name: Run tests
-      run: |
-        poetry run make all-checks
+      - name: Run tests
+        run: |
+          poetry run make all-checks

src/labelformat/utils.py

@@ -1,6 +1,6 @@
 import logging
 from pathlib import Path
-from typing import Iterable
+from typing import Iterable, Tuple
 
 import PIL.Image
 
@@ -21,6 +21,124 @@
 }
 
 
+class ImageDimensionError(Exception):
+    """Raised when unable to extract image dimensions using fast methods."""
+
+    pass
+
+
+def get_jpeg_dimensions(file_path: Path) -> Tuple[int, int]:
+    """Try to efficiently get JPEG dimensions from file headers without decoding the image.
+
+    This method reads only the JPEG file headers looking for the Start Of Frame (SOFn)
+    marker which contains the dimensions. This is much faster than decoding the entire
+    image as it:
+    - Only reads the file headers (typically a few KB) instead of the entire file
+    - Doesn't perform any image decompression
+    - Doesn't load the pixel data into memory
+
+    This works for most standard JPEG files (including progressive JPEGs) but may fail
+    for some unusual formats or corrupted files. In those cases, an ImageDimensionError
+    is raised and a full image decode may be needed as fallback.
+
+    Args:
+        file_path: Path to the JPEG file
+
+    Returns:
+        Tuple of (width, height)
+
+    Raises:
+        ImageDimensionError: If dimensions cannot be extracted from headers
+    """
+    try:
+        with open(file_path, "rb") as img_file:
+            # Skip SOI marker
+            img_file.seek(2)
+            while True:
+                marker = img_file.read(2)
+                if len(marker) < 2:
+                    raise ImageDimensionError("Invalid JPEG format")
+                # Find SOFn marker
+                if 0xFF == marker[0] and marker[1] in range(0xC0, 0xCF):
+                    # Skip marker length
+                    img_file.seek(3, 1)
+                    h = int.from_bytes(img_file.read(2), "big")
+                    w = int.from_bytes(img_file.read(2), "big")
+                    return w, h
+                # Skip to next marker
+                length = int.from_bytes(img_file.read(2), "big")
+                img_file.seek(length - 2, 1)
+    except Exception as e:
+        raise ImageDimensionError(f"Failed to read JPEG dimensions: {str(e)}")
+
+
+def get_png_dimensions(file_path: Path) -> Tuple[int, int]:
+    """Try to efficiently get PNG dimensions from file headers without decoding the image.
+
+    This method reads only the PNG IHDR (Image Header) chunk which is always the first
+    chunk after the PNG signature. This is much faster than decoding the entire image as it:
+    - Only reads the first ~30 bytes of the file
+    - Doesn't perform any image decompression
+    - Doesn't load the pixel data into memory
+
+    This works for all valid PNG files since the IHDR chunk is mandatory and must appear
+    first according to the PNG specification. However, it may fail for corrupted files
+    or files that don't follow the PNG spec. In those cases, an ImageDimensionError is
+    raised and a full image decode may be needed as fallback.
+
+    Args:
+        file_path: Path to the PNG file
+
+    Returns:
+        Tuple of (width, height)
+
+    Raises:
+        ImageDimensionError: If dimensions cannot be extracted from headers
+    """
+    try:
+        with open(file_path, "rb") as img_file:
+            # Skip PNG signature
+            img_file.seek(8)
+            # Read IHDR chunk
+            chunk_length = int.from_bytes(img_file.read(4), "big")
+            chunk_type = img_file.read(4)
+            if chunk_type == b"IHDR":
+                w = int.from_bytes(img_file.read(4), "big")
+                h = int.from_bytes(img_file.read(4), "big")
+                return w, h
+            raise ImageDimensionError("Invalid PNG format")
+    except Exception as e:
+        raise ImageDimensionError(f"Failed to read PNG dimensions: {str(e)}")
+
+
+def get_image_dimensions(image_path: Path) -> Tuple[int, int]:
+    """Get image dimensions using the most efficient method available.
+
+    Args:
+        image_path: Path to the image file
+
+    Returns:
+        Tuple of (width, height)
+
+    Raises:
+        Exception: If image dimensions cannot be extracted using any method
+    """
+    suffix = image_path.suffix.lower()
+    if suffix in {".jpg", ".jpeg"}:
+        try:
+            return get_jpeg_dimensions(image_path)
+        except ImageDimensionError:
+            pass
+    elif suffix == ".png":
+        try:
+            return get_png_dimensions(image_path)
+        except ImageDimensionError:
+            pass
+
+    with PIL.Image.open(image_path) as img:
+        return img.size
+
+
 def get_images_from_folder(folder: Path) -> Iterable[Image]:
     """Yields an Image structure for all images in the given folder.
 
@@ -36,8 +154,7 @@ def get_images_from_folder(folder: Path) -> Iterable[Image]:
             logger.debug(f"Skipping non-image file '{image_path}'")
             continue
         image_filename = str(image_path.relative_to(folder))
-        with PIL.Image.open(image_path) as img:
-            image_width, image_height = img.size
+        image_width, image_height = get_image_dimensions(image_path)
         yield Image(
             id=image_id,
             filename=image_filename,

tests/unit/test_utils.py

@@ -0,0 +1,73 @@
+from pathlib import Path
+
+import pytest
+
+from labelformat.utils import (
+    ImageDimensionError,
+    get_image_dimensions,
+    get_jpeg_dimensions,
+    get_png_dimensions,
+)
+
+FIXTURES_DIR = Path(__file__).parent.parent / "fixtures"
+
+
+class TestImageDimensions:
+    def test_jpeg_dimensions_valid_file(self) -> None:
+        image_path = (
+            FIXTURES_DIR / "instance_segmentation/YOLOv8/images/000000109005.jpg"
+        )
+        width, height = get_jpeg_dimensions(image_path)
+        assert width == 640
+        assert height == 428
+
+    def test_jpeg_dimensions_nonexistent_file(self) -> None:
+        with pytest.raises(ImageDimensionError):
+            get_jpeg_dimensions(Path("nonexistent.jpg"))
+
+    def test_jpeg_dimensions_invalid_format(self) -> None:
+        yaml_file = FIXTURES_DIR / "object_detection/YOLOv8/example.yaml"
+        with pytest.raises(ImageDimensionError):
+            get_jpeg_dimensions(yaml_file)
+
+    def test_png_dimensions_valid_file(self) -> None:
+        png_path = FIXTURES_DIR / "image_file_loading/0001.png"
+        width, height = get_png_dimensions(png_path)
+        assert width == 278
+        assert height == 181
+
+    def test_png_dimensions_nonexistent_file(self) -> None:
+        with pytest.raises(ImageDimensionError):
+            get_png_dimensions(Path("nonexistent.png"))
+
+    def test_png_dimensions_invalid_format(self) -> None:
+        yaml_file = FIXTURES_DIR / "object_detection/YOLOv8/example.yaml"
+        with pytest.raises(ImageDimensionError):
+            get_png_dimensions(yaml_file)
+
+    def test_get_image_dimensions_jpeg_first_file(self) -> None:
+        jpeg_path = (
+            FIXTURES_DIR / "instance_segmentation/YOLOv8/images/000000109005.jpg"
+        )
+        width, height = get_image_dimensions(jpeg_path)
+        assert width == 640
+        assert height == 428
+
+    def test_get_image_dimensions_jpeg_second_file(self) -> None:
+        jpeg_path = (
+            FIXTURES_DIR / "instance_segmentation/YOLOv8/images/000000036086.jpg"
+        )
+        width, height = get_image_dimensions(jpeg_path)
+        assert width == 482
+        assert height == 640
+
+    def test_get_image_dimensions_png(self) -> None:
+        png_path = FIXTURES_DIR / "image_file_loading/0001.png"
+        width, height = get_image_dimensions(png_path)
+        assert width == 278
+        assert height == 181
+
+    def test_get_image_dimensions_unsupported_format(self) -> None:
+        yaml_file = FIXTURES_DIR / "object_detection/YOLOv8/example.yaml"
+        with pytest.raises(Exception):
+            get_image_dimensions(yaml_file)