Pil image support. Configurable default wait. (#31)

* Adding PIL image support.

* Fixup import error.  Version bump.

* Fixing PIL image support.  Making DEFAULT_WAIT configurable in groundlight client.

* Adding test on invalid image format, and fixing a PIL import bug there.

* Adding Image.Image to supported input types in type hint.

* Fixing import.  Documenting RGB expectation.

* Automatically reformatting code with black

* Fixing incorrect documentation on our expected tensor dimensions for numpy image inputs.

---------

Co-authored-by: Auto-format Bot <runner@fv-az365-486.rrd34hpwizlenh5trlchfhotwb.cx.internal.cloudapp.net>

Author: robotrapta

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "groundlight"
-version = "0.7.1"
+version = "0.7.2"
 license = "MIT"
 readme = "README.md"
 homepage = "https://www.groundlight.ai"

src/groundlight/client.py

@@ -4,6 +4,7 @@
 from io import BufferedReader, BytesIO
 from typing import Optional, Union
 
+from groundlight.optional_imports import Image
 from model import Detector, ImageQuery, PaginatedDetectorList, PaginatedImageQueryList
 from openapi_client import ApiClient, Configuration
 from openapi_client.api.detectors_api import DetectorsApi
@@ -12,7 +13,7 @@
 
 from groundlight.binary_labels import convert_display_label_to_internal, convert_internal_label_to_display
 from groundlight.config import API_TOKEN_VARIABLE_NAME, API_TOKEN_WEB_URL, DEFAULT_ENDPOINT
-from groundlight.images import buffer_from_jpeg_file, jpeg_from_numpy
+from groundlight.images import buffer_from_jpeg_file, jpeg_from_numpy, parse_supported_image_types
 from groundlight.internalapi import GroundlightApiClient, sanitize_endpoint_url
 from groundlight.optional_imports import np
 
@@ -38,6 +39,8 @@ class Groundlight:
     ```
     """
 
+    DEFAULT_WAIT = 30
+
     BEFORE_POLLING_DELAY = 3.0  # Expected minimum time for a label to post
     POLLING_INITIAL_DELAY = 0.5
     POLLING_EXPONENTIAL_BACKOFF = 1.3  # This still has the nice backoff property that the max number of requests
@@ -122,38 +125,28 @@ def list_image_queries(self, page: int = 1, page_size: int = 10) -> PaginatedIma
     def submit_image_query(
         self,
         detector: Union[Detector, str],
-        image: Union[str, bytes, BytesIO, BufferedReader, np.ndarray],
-        wait: float = 30,
+        image: Union[str, bytes, Image.Image, BytesIO, BufferedReader, np.ndarray],
+        wait: Optional[float] = None,
     ) -> ImageQuery:
         """Evaluates an image with Groundlight.
         :param detector: the Detector object, or string id of a detector like `det_12345`
         :param image: The image, in several possible formats:
-            - a filename (string) of a jpeg file
-            - a byte array or BytesIO with jpeg bytes
-            - a numpy array in the 0-255 range (gets converted to jpeg)
+          - filename (string) of a jpeg file
+          - byte array or BytesIO or BufferedReader with jpeg bytes
+          - numpy array with values 0-255 and dimensions (H,W,3) in RGB order
+            (Note OpenCV uses BGR not RGB. `img[:, :, ::-1]` will reverse the channels)
+          - PIL Image
+          Any binary format must be JPEG-encoded already.  Any pixel format will get
+          converted to JPEG at high quality before sending to service.
         :param wait: How long to wait (in seconds) for a confident answer
         """
+        if wait is None:
+            wait = self.DEFAULT_WAIT
         if isinstance(detector, Detector):
             detector_id = detector.id
         else:
             detector_id = detector
-        image_bytesio: Union[BytesIO, BufferedReader]
-        # TODO: support PIL Images
-        if isinstance(image, str):
-            # Assume it is a filename
-            image_bytesio = buffer_from_jpeg_file(image)
-        elif isinstance(image, bytes):
-            # Create a BytesIO object
-            image_bytesio = BytesIO(image)
-        elif isinstance(image, BytesIO) or isinstance(image, BufferedReader):
-            # Already in the right format
-            image_bytesio = image
-        elif isinstance(image, np.ndarray):
-            image_bytesio = BytesIO(jpeg_from_numpy(image))
-        else:
-            raise TypeError(
-                "Unsupported type for image. We only support numpy arrays (3,W,H) or JPEG images specified through a filename, bytes, BytesIO, or BufferedReader object."
-            )
+        image_bytesio: Union[BytesIO, BufferedReader] = parse_supported_image_types(image)
 
         raw_image_query = self.image_queries_api.submit_image_query(detector_id=detector_id, body=image_bytesio)
         image_query = ImageQuery.parse_obj(raw_image_query.to_dict())

src/groundlight/images.py

@@ -1,10 +1,11 @@
 import imghdr
-import io
+from io import BufferedReader, BytesIO
+from typing import Union
 
-from groundlight.optional_imports import np, Image
+from groundlight.optional_imports import Image, np
 
 
-def buffer_from_jpeg_file(image_filename: str) -> io.BufferedReader:
+def buffer_from_jpeg_file(image_filename: str) -> BufferedReader:
     """
     Get a buffer from an jpeg image file.
 
@@ -21,8 +22,36 @@ def buffer_from_jpeg_file(image_filename: str) -> io.BufferedReader:
 def jpeg_from_numpy(img: np.ndarray, jpeg_quality: int = 95) -> bytes:
     """Converts a numpy array to BytesIO"""
     pilim = Image.fromarray(img.astype("uint8"), "RGB")
-    with io.BytesIO() as buf:
-        buf = io.BytesIO()
+    with BytesIO() as buf:
+        buf = BytesIO()
         pilim.save(buf, "jpeg", quality=jpeg_quality)
         out = buf.getvalue()
         return out
+
+
+def parse_supported_image_types(
+    image: Union[str, bytes, Image.Image, BytesIO, BufferedReader, np.ndarray], jpeg_quality: int = 95
+) -> Union[BytesIO, BufferedReader]:
+    """Parse the many supported image types into a bytes-stream objects.
+    In some cases we have to JPEG compress."""
+    if isinstance(image, str):
+        # Assume it is a filename
+        return buffer_from_jpeg_file(image)
+    elif isinstance(image, bytes):
+        # Create a BytesIO object
+        return BytesIO(image)
+    elif isinstance(image, Image.Image):
+        # Save PIL image as jpeg in BytesIO
+        bytesio = BytesIO()
+        image.save(bytesio, "jpeg", quality=jpeg_quality)
+        bytesio.seek(0)
+        return bytesio
+    elif isinstance(image, BytesIO) or isinstance(image, BufferedReader):
+        # Already in the right format
+        return image
+    elif isinstance(image, np.ndarray):
+        return BytesIO(jpeg_from_numpy(image, jpeg_quality=jpeg_quality))
+    else:
+        raise TypeError(
+            "Unsupported type for image. Must be PIL, numpy (H,W,3) RGB, or a JPEG as a filename (str), bytes, BytesIO, or BufferedReader."
+        )

src/groundlight/optional_imports.py

@@ -45,6 +45,7 @@ def __getattr__(self, key):
 except ImportError as e:
     PIL = UnavailableModule(e)
     Image = PIL
+    Image.Image = PIL  # for type-hinting
     MISSING_PIL = True
 
 

test/integration/test_groundlight.py

@@ -12,7 +12,9 @@
 @pytest.fixture
 def gl() -> Groundlight:
     """Creates a Groundlight client object for testing."""
-    return Groundlight()
+    gl = Groundlight()
+    gl.DEFAULT_WAIT = 0.1
+    return gl
 
 
 @pytest.fixture
@@ -102,6 +104,18 @@ def test_submit_image_query_bad_jpeg_file(gl: Groundlight, detector: Detector):
     assert "jpeg" in str(exc_info).lower()
 
 
+@pytest.mark.skipif(MISSING_PIL, reason="Needs pillow")
+def test_submit_image_query_pil(gl: Groundlight, detector: Detector):
+    # generates a pil image and submits it
+    from PIL import Image
+
+    dog = Image.open("test/assets/dog.jpeg")
+    _image_query = gl.submit_image_query(detector=detector.id, image=dog)
+
+    black = Image.new("RGB", (640, 480))
+    _image_query = gl.submit_image_query(detector=detector.id, image=black)
+
+
 def test_list_image_queries(gl: Groundlight):
     image_queries = gl.list_image_queries()
     assert str(image_queries)

test/unit/test_imagefuncs.py

@@ -1,3 +1,5 @@
+import tempfile
+
 import pytest
 
 from groundlight.images import *
@@ -17,3 +19,49 @@ def test_jpeg_from_numpy():
     np_img = np.random.uniform(0, 255, (768, 1024, 3))
     jpeg3 = jpeg_from_numpy(np_img, jpeg_quality=50)
     assert len(jpeg2) > len(jpeg3)
+
+
+def test_unsupported_imagetype():
+    with pytest.raises(TypeError):
+        parse_supported_image_types(1)
+
+    with pytest.raises(TypeError):
+        parse_supported_image_types(None)
+
+    with pytest.raises(TypeError):
+        parse_supported_image_types(pytest)
+
+
+@pytest.mark.skipif(MISSING_PIL, reason="Needs pillow")
+def test_pil_support():
+    from PIL import Image
+
+    img = Image.new("RGB", (640, 480))
+    jpeg = parse_supported_image_types(img)
+    assert isinstance(jpeg, BytesIO)
+
+    # Now try to parse the BytesIO object as an image
+    jpeg_bytes = jpeg.getvalue()
+    # save the image to a tempfile
+    with tempfile.TemporaryFile() as f:
+        f.write(jpeg_bytes)
+        f.seek(0)
+        img2 = Image.open(f)
+        assert img2.size == (640, 480)
+
+
+@pytest.mark.skipif(MISSING_PIL, reason="Needs pillow")
+def test_pil_support_ref():
+    # Similar to test_pil_support, but uses a known-good file
+    from PIL import Image
+
+    fn = "test/assets/dog.jpeg"
+    parsed = parse_supported_image_types(fn)
+    # Now try to parse the BytesIO object as an image
+    jpeg_bytes = parsed.read()
+    # save the image to a tempfile
+    with tempfile.TemporaryFile() as f:
+        f.write(jpeg_bytes)
+        f.seek(0)
+        img2 = Image.open(f)
+        assert img2.size == (509, 339)