refactor: preprocess_image + tests (#702)

Author: maciejmajek

src/rai_core/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rai_core"
-version = "2.5.2"
+version = "2.5.3"
 description = "Core functionality for RAI framework"
 authors = ["Maciej Majek <maciej.majek@robotec.ai>", "Bartłomiej Boczek <bartlomiej.boczek@robotec.ai>", "Kajetan Rachwał <kajetan.rachwal@robotec.ai>"]
 readme = "README.md"

src/rai_core/rai/messages/conversion.py

@@ -14,40 +14,104 @@
 
 
 import base64
-from typing import Any, Callable, Union
+import os
+from io import BytesIO
+from typing import Callable, Union
 
-import cv2
 import numpy as np
 import requests
+from PIL import Image as PILImage
 from PIL.Image import Image
 
 
 def preprocess_image(
-    image: Union[Image, str, bytes, np.ndarray[Any, np.dtype[np.uint8 | np.float_]]],
-    encoding_function: Callable[[Any], str] = lambda x: base64.b64encode(x).decode(
+    image: Union[Image, str, bytes, np.ndarray],
+    encoding_function: Callable[[bytes], str] = lambda b: base64.b64encode(b).decode(
         "utf-8"
     ),
 ) -> str:
-    if isinstance(image, Image):
-        image = np.array(image)
-        _, image_data = cv2.imencode(".png", image)
-        encoding_function = lambda x: base64.b64encode(x).decode("utf-8")
-    elif isinstance(image, str) and image.startswith(("http://", "https://")):
-        response = requests.get(image)
-        response.raise_for_status()
-        image_data = response.content
-    elif isinstance(image, str):
-        with open(image, "rb") as image_file:
-            image_data = image_file.read()
-    elif isinstance(image, bytes):
-        image_data = image
-        encoding_function = lambda x: x.decode("utf-8")
-    elif isinstance(image, np.ndarray):  # type: ignore
-        if image.dtype == np.float32 or image.dtype == np.float64:
-            image = (image * 255).astype(np.uint8)
-        _, image_data = cv2.imencode(".png", image)
-        encoding_function = lambda x: base64.b64encode(x).decode("utf-8")
-    else:
-        image_data = image
-
-    return encoding_function(image_data)
+    """Convert various image inputs into a base64-encoded PNG string.
+
+    Parameters
+    ----------
+    image : PIL.Image.Image or str or bytes or numpy.ndarray
+        Supported inputs:
+        - PIL Image
+        - Path to a file, ``file://`` URL, or HTTP(S) URL
+        - Raw bytes containing image data
+        - ``numpy.ndarray`` with dtype ``uint8`` or ``float32``/``float64``;
+          grayscale or 3/4-channel arrays are supported.
+    encoding_function : callable, optional
+        Function that converts PNG bytes to the final string representation.
+        By default, returns base64-encoded UTF-8 string.
+
+    Returns
+    -------
+    str
+        Base64-encoded PNG string.
+
+    Raises
+    ------
+    FileNotFoundError
+        If a file path (or ``file://`` URL) does not exist.
+    TypeError
+        If the input type is not supported.
+    requests.HTTPError
+        If fetching an HTTP(S) URL fails with a non-2xx response.
+    requests.RequestException
+        If a network error occurs while fetching an HTTP(S) URL.
+    OSError
+        If the input cannot be decoded as an image by Pillow.
+
+    Notes
+    -----
+    - All inputs are decoded and re-encoded to PNG to guarantee consistent output.
+    - Float arrays are assumed to be in [0, 1] and are scaled to ``uint8``.
+    - Network requests use a timeout of ``(5, 15)`` seconds (connect, read).
+
+    Examples
+    --------
+    >>> b64_png = preprocess_image("path/to/image.jpg")
+    >>> import numpy as np
+    >>> arr = np.random.rand(64, 64, 3).astype(np.float32)
+    >>> b64_png = preprocess_image(arr)
+    """
+
+    def _to_pil_from_ndarray(arr: np.ndarray) -> Image:
+        a = arr
+        if a.dtype in (np.float32, np.float64):
+            a = np.clip(a, 0.0, 1.0)
+            a = (a * 255.0).round().astype(np.uint8)
+        a = np.ascontiguousarray(a)
+        return PILImage.fromarray(a)
+
+    def _ensure_pil(img: Union[Image, str, bytes, np.ndarray]) -> Image:
+        if isinstance(img, Image):
+            return img
+        if isinstance(img, np.ndarray):  # type: ignore
+            return _to_pil_from_ndarray(img)
+        if isinstance(img, str):
+            if img.startswith(("http://", "https://")):
+                response = requests.get(img, timeout=(5, 15))
+                response.raise_for_status()
+                return PILImage.open(BytesIO(response.content))
+            if img.startswith("file://"):
+                file_path = img[len("file://") :]
+            else:
+                # fallback to file path if not marked with file://
+                file_path = img
+            if not os.path.exists(file_path):
+                raise FileNotFoundError(f"File not found: {file_path}")
+            return PILImage.open(file_path)
+        if isinstance(img, bytes):
+            return PILImage.open(BytesIO(img))
+        raise TypeError(f"Unsupported image type: {type(img).__name__}")
+
+    pil_image = _ensure_pil(image)
+
+    # Normalize to PNG bytes
+    with BytesIO() as buffer:
+        pil_image.save(buffer, format="PNG")
+        png_bytes = buffer.getvalue()
+
+    return encoding_function(png_bytes)

tests/messages/test_utils.py

@@ -13,7 +13,11 @@
 # limitations under the License.
 
 import base64
+import threading
+from functools import partial
+from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
 from io import BytesIO
+from pathlib import Path
 
 import numpy as np
 import pytest
@@ -30,10 +34,74 @@ def decode_image(base64_string: str) -> Image.Image:
 @pytest.mark.parametrize(
     "test_image",
     [
-        np.zeros((300, 300, 3)),
+        np.zeros((300, 300), dtype=np.uint8),
+        np.zeros((300, 300, 3), dtype=np.uint8),
+        np.zeros((300, 300, 4), dtype=np.uint8),
+        np.random.rand(300, 300).astype(np.float32),
+        np.random.rand(300, 300, 3).astype(np.float32),
+        np.random.rand(300, 300, 4).astype(np.float32),
+        np.random.rand(300, 300).astype(np.float64),
+        np.random.rand(300, 300, 3).astype(np.float64),
+        np.random.rand(300, 300, 4).astype(np.float64),
         "tests/resources/image.png",
     ],
 )
-def test_preprocess_image(test_image):
+def test_preprocess_image_always_png(test_image):
     base64_image = preprocess_image(test_image)
-    _ = decode_image(base64_image)  # noqa: F841
+    img = decode_image(base64_image)
+    assert img.format == "PNG"
+    assert img.size[0] > 0 and img.size[1] > 0
+
+
+def test_preprocess_image_from_bytes_and_file_url(tmp_path: Path):
+    # Create a temporary PNG file
+    arr = (np.random.rand(32, 32, 3) * 255).astype(np.uint8)
+    pil_img = Image.fromarray(arr)
+    file_path = tmp_path / "tmp.png"
+    pil_img.save(file_path, format="PNG")
+
+    # bytes input
+    with open(file_path, "rb") as f:
+        raw_bytes = f.read()
+    b64_bytes = preprocess_image(raw_bytes)
+    img_from_bytes = decode_image(b64_bytes)
+    assert img_from_bytes.format == "PNG"
+    assert img_from_bytes.size == (32, 32)
+
+    # file:// URL input
+    file_url = f"file://{file_path}"
+    b64_file_url = preprocess_image(file_url)
+    img_from_url = decode_image(b64_file_url)
+    assert img_from_url.format == "PNG"
+    assert img_from_url.size == (32, 32)
+
+
+def test_preprocess_image_unsupported_type():
+    class NotAnImage:
+        pass
+
+    with pytest.raises(TypeError):
+        _ = preprocess_image(NotAnImage())
+
+
+def test_preprocess_image_http_url():
+    resources_dir = Path(__file__).resolve().parents[1] / "resources"
+    image_path = resources_dir / "image.png"
+    assert image_path.exists(), "tests/resources/image.png must exist for this test"
+
+    handler = partial(SimpleHTTPRequestHandler, directory=str(resources_dir))
+    httpd = ThreadingHTTPServer(("127.0.0.1", 0), handler)
+
+    try:
+        port = httpd.server_address[1]
+        thread = threading.Thread(target=httpd.serve_forever, daemon=True)
+        thread.start()
+
+        url = f"http://127.0.0.1:{port}/image.png"
+        base64_image = preprocess_image(url)
+        img = decode_image(base64_image)
+        assert img.format == "PNG"
+        assert img.size[0] > 0 and img.size[1] > 0
+    finally:
+        httpd.shutdown()
+        httpd.server_close()