Support more jpeg SOF markers for size reading the dimensions correctly (#49)

* Add more markers for reading jpg dimensions efficiently

* Add more tests for getting dimensions from various jpeg files

* Address feedback

Author: IgorSusmelj

src/labelformat/utils.py

@@ -20,6 +20,25 @@
     ".webp",
 }
 
+# JPEG SOF (Start of Frame) markers that contain image dimensions.
+# Excludes 0xC4 (DHT - Define Huffman Table) and 0xC8 (JPG reserved).
+# List from https://www.disktuna.com/list-of-jpeg-markers/
+JPEG_SOF_MARKERS = {
+    0xC0,  # SOF0 - Baseline DCT
+    0xC1,  # SOF1 - Extended Sequential DCT
+    0xC2,  # SOF2 - Progressive DCT
+    0xC3,  # SOF3 - Lossless (sequential)
+    0xC5,  # SOF5 - Differential sequential DCT
+    0xC6,  # SOF6 - Differential progressive DCT
+    0xC7,  # SOF7 - Differential lossless (sequential)
+    0xC9,  # SOF9 - Extended sequential DCT, Arithmetic coding
+    0xCA,  # SOF10 - Progressive DCT, Arithmetic coding
+    0xCB,  # SOF11 - Lossless (sequential), Arithmetic coding
+    0xCD,  # SOF13 - Differential sequential DCT, Arithmetic coding
+    0xCE,  # SOF14 - Differential progressive DCT, Arithmetic coding
+    0xCF,  # SOF15 - Differential lossless (sequential), Arithmetic coding
+}
+
 
 class ImageDimensionError(Exception):
     """Raised when unable to extract image dimensions using fast methods."""
@@ -59,7 +78,7 @@ def get_jpeg_dimensions(file_path: Path) -> Tuple[int, int]:
                 if len(marker) < 2:
                     raise ImageDimensionError("Invalid JPEG format")
                 # Find SOFn marker
-                if 0xFF == marker[0] and marker[1] in range(0xC0, 0xCF):
+                if marker[0] == 0xFF and marker[1] in JPEG_SOF_MARKERS:
                     # Skip marker length
                     img_file.seek(3, 1)
                     h = int.from_bytes(img_file.read(2), "big")

tests/unit/test_utils.py

@@ -1,5 +1,7 @@
 from pathlib import Path
+from typing import Tuple
 
+import PIL.Image
 import pytest
 
 from labelformat.utils import (
@@ -12,62 +14,116 @@
 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)
+def test_get_jpeg_dimensions() -> 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_get_jpeg_dimensions__baseline(tmp_path: Path) -> None:
+    # Tests SOF0 (0xC0) - Baseline DCT.
+    jpeg_path = tmp_path / "baseline.jpg"
+    _create_test_jpeg(path=jpeg_path, size=(800, 600))
+
+    width, height = get_jpeg_dimensions(jpeg_path)
+    assert width == 800
+    assert height == 600
+
+
+def test_get_jpeg_dimensions__progressive(tmp_path: Path) -> None:
+    # Tests SOF2 (0xC2) - Progressive DCT with DHT markers before SOF.
+    jpeg_path = tmp_path / "progressive.jpg"
+    _create_test_jpeg(path=jpeg_path, size=(1920, 1440), progressive=True)
+
+    width, height = get_jpeg_dimensions(jpeg_path)
+    assert width == 1920
+    assert height == 1440
+
+
+def test_get_jpeg_dimensions__optimized(tmp_path: Path) -> None:
+    # Tests SOF0 (0xC0) with custom Huffman tables (more DHT markers before SOF).
+    jpeg_path = tmp_path / "optimized.jpg"
+    _create_test_jpeg(path=jpeg_path, size=(1024, 768), optimize=True)
+
+    width, height = get_jpeg_dimensions(jpeg_path)
+    assert width == 1024
+    assert height == 768
+
+
+def test_get_jpeg_dimensions__progressive_optimized(tmp_path: Path) -> None:
+    # Tests SOF2 (0xC2) with custom Huffman tables.
+    jpeg_path = tmp_path / "progressive_optimized.jpg"
+    _create_test_jpeg(
+        path=jpeg_path, size=(2048, 1536), progressive=True, optimize=True
+    )
+
+    width, height = get_jpeg_dimensions(jpeg_path)
+    assert width == 2048
+    assert height == 1536
+
+
+def test_get_jpeg_dimensions__nonexistent_file() -> None:
+    with pytest.raises(ImageDimensionError):
+        get_jpeg_dimensions(Path("nonexistent.jpg"))
+
+
+def test_get_jpeg_dimensions__invalid_format() -> None:
+    yaml_file = FIXTURES_DIR / "object_detection/YOLOv8/example.yaml"
+    with pytest.raises(ImageDimensionError):
+        get_jpeg_dimensions(yaml_file)
+
+
+def test_get_png_dimensions() -> 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_get_png_dimensions__nonexistent_file() -> None:
+    with pytest.raises(ImageDimensionError):
+        get_png_dimensions(Path("nonexistent.png"))
+
+
+def test_get_png_dimensions__invalid_format() -> 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() -> 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() -> 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() -> 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() -> None:
+    yaml_file = FIXTURES_DIR / "object_detection/YOLOv8/example.yaml"
+    with pytest.raises(Exception):
+        get_image_dimensions(yaml_file)
+
+
+def _create_test_jpeg(
+    path: Path,
+    size: Tuple[int, int],
+    progressive: bool = False,
+    optimize: bool = False,
+) -> None:
+    img = PIL.Image.new("RGB", size, color="red")
+    img.save(path, "JPEG", progressive=progressive, optimize=optimize)