Refactor code structure for improved readability and maintainability

Author: patches-kroll

athena-protobufs

@@ -244,7 +244,7 @@ async def classify_single(
             correlation_id=correlation_id,
             encoding=request_encoding,
             data=processed_image.data,
-            format=ImageFormat.IMAGE_FORMAT_RAW_UINT8,
+            format=ImageFormat.IMAGE_FORMAT_RAW_UINT8_BGR,
             hashes=[
                 ImageHash(
                     value=hash_value,

src/resolver_athena_client/client/athena_client.py

@@ -56,7 +56,7 @@ def _create_classification_input(
             ),
             data=image_data.data,
             encoding=self.request_encoding,
-            format=ImageFormat.IMAGE_FORMAT_RAW_UINT8,
+            format=ImageFormat.IMAGE_FORMAT_RAW_UINT8_BGR,
         )
 
     @override

src/resolver_athena_client/client/transformers/classification_input.py

@@ -19,8 +19,8 @@
 _expected_raw_size = EXPECTED_WIDTH * EXPECTED_HEIGHT * 3
 
 
-def _is_raw_brg_expected_size(data: bytes) -> bool:
-    """Detect if data is already a raw BRG array of expected size."""
+def _is_raw_bgr_expected_size(data: bytes) -> bool:
+    """Detect if data is already a raw BGR array of expected size."""
     return len(data) == _expected_raw_size
 
 
@@ -37,7 +37,7 @@ async def resize_image(image_data: ImageData) -> ImageData:
 
     def process_image() -> tuple[bytes, bool]:
         # Fast path for raw RGB arrays of correct size
-        if _is_raw_brg_expected_size(image_data.data):
+        if _is_raw_bgr_expected_size(image_data.data):
             return image_data.data, False  # No transformation needed
 
         # Try to load the image data directly
@@ -57,15 +57,15 @@ def process_image() -> tuple[bytes, bool]:
 
             rgb_bytes = resized_image.tobytes()
 
-            # Convert RGB to BRG by swapping channels
-            brg_bytes = bytearray(len(rgb_bytes))
+            # Convert RGB to BGR by swapping channels
+            bgr_bytes = bytearray(len(rgb_bytes))
 
             for i in range(0, len(rgb_bytes), 3):
-                brg_bytes[i] = rgb_bytes[i + 2]
-                brg_bytes[i + 1] = rgb_bytes[i]
-                brg_bytes[i + 2] = rgb_bytes[i + 1]
+                bgr_bytes[i] = rgb_bytes[i + 2]
+                bgr_bytes[i + 1] = rgb_bytes[i + 1]
+                bgr_bytes[i + 2] = rgb_bytes[i]
 
-            return bytes(brg_bytes), True  # Data was transformed
+            return bytes(bgr_bytes), True  # Data was transformed
 
     # Use thread pool for CPU-intensive processing
     resized_bytes, was_transformed = await asyncio.to_thread(process_image)

src/resolver_athena_client/client/transformers/core.py

@@ -49,7 +49,7 @@ async def test_classification_input_transform(
     )  # Should be a non-empty string
     assert result.data == test_data.data
     assert result.encoding == RequestEncoding.REQUEST_ENCODING_BROTLI
-    assert result.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8
+    assert result.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8_BGR
 
 
 @pytest.mark.asyncio
@@ -69,7 +69,7 @@ async def test_classification_input_iteration(
     assert result.data == b"test1"
     assert result.affiliate == transformer_config.affiliate
     assert result.encoding == RequestEncoding.REQUEST_ENCODING_BROTLI
-    assert result.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8
+    assert result.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8_BGR
 
     # Test second item
     result = await anext(transformer)
@@ -106,7 +106,7 @@ async def test_classification_input_empty(
     assert result.data == b""
     assert result.affiliate == transformer_config.affiliate
     assert result.encoding == RequestEncoding.REQUEST_ENCODING_BROTLI
-    assert result.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8
+    assert result.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8_BGR
 
 
 @pytest.mark.parametrize(
@@ -133,4 +133,4 @@ async def test_classification_input_encodings(
     test_data = ImageData(b"test")
     result = await transformer.transform(test_data)
     assert result.encoding == encoding
-    assert result.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8
+    assert result.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8_BGR

tests/client/transformers/test_classification_input.py

@@ -100,7 +100,7 @@ async def test_classify_single_success(
     assert call_args.affiliate == "test-affiliate"
     assert call_args.data == b"fake_image_bytes"
     assert call_args.encoding == RequestEncoding.REQUEST_ENCODING_UNCOMPRESSED
-    assert call_args.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8
+    assert call_args.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8_BGR
     assert len(call_args.hashes) == 1
     assert call_args.hashes[0].type == HashType.HASH_TYPE_MD5
     assert len(call_args.hashes[0].value) > 0  # MD5 hash should be generated

tests/test_classify_single.py

@@ -10,6 +10,14 @@
 
 from resolver_athena_client.client.models import ImageData
 
+
+def _rgb_to_bgr(raw_rgb: bytes) -> bytes:
+    data = bytearray(raw_rgb)
+    for i in range(0, len(data), 3):
+        data[i], data[i + 2] = data[i + 2], data[i]
+    return bytes(data)
+
+
 # Global cache for reusable objects and constants
 _image_cache: dict[
     tuple[int, int], tuple[Image.Image, ImageDraw.ImageDraw]
@@ -61,8 +69,8 @@ def create_random_image(
     x2, y2 = (width * 3) // 4, (height * 3) // 4
     draw.rectangle([x1, y1, x2, y2], fill=accent_color)
 
-    if img_format.upper() == "RAW_UINT8":
-        return image.tobytes()
+    if img_format.upper() in {"RAW_UINT8", "RAW_UINT8_BGR"}:
+        return _rgb_to_bgr(image.tobytes())
 
     # Convert to PNG bytes
     buffer = io.BytesIO()