refactor: refactor image blob to uri (#3523)

jina/types/document/__init__.py

@@ -36,7 +36,6 @@
     random_identity,
     download_mermaid_url,
     dunder_get,
-    cached_property,
 )
 from ...importer import ImportExtensions
 from ...logging.predefined import default_logger
@@ -1007,14 +1006,23 @@ def convert_image_buffer_to_blob(self, color_axis: int = -1):
         self.blob = to_image_blob(io.BytesIO(self.buffer), color_axis)
 
     def convert_image_blob_to_uri(
-        self, width: int, height: int, resize_method: str = 'BILINEAR'
+        self,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        resize_method: str = 'BILINEAR',
+        color_axis: int = -1,
     ):
         """Assuming :attr:`blob` is a _valid_ image, set :attr:`uri` accordingly
-        :param width: the width of the blob
-        :param height: the height of the blob
+        :param width: the width of the blob, if None, interpret from :attr:`blob` shape.
+        :param height: the height of the blob, if None, interpret from :attr:`blob` shape.
         :param resize_method: the resize method name
+        :param color_axis: the axis id of the color channel, ``-1`` indicates the color channel info at the last axis
+
+        ..note::
+            if both :attr:`width` and :attr:`height` were provided, will not resize. Otherwise, will get image size
+            by :attr:`self.blob` shape and apply resize method :attr:`resize_method`.
         """
-        png_bytes = png_to_buffer(self.blob, width, height, resize_method)
+        png_bytes = png_to_buffer(self.blob, width, height, resize_method, color_axis)
         self.uri = 'data:image/png;base64,' + base64.b64encode(png_bytes).decode()
 
     def convert_image_uri_to_blob(

jina/types/document/converters.py

@@ -1,5 +1,6 @@
 import struct
 import zlib
+from typing import Optional
 
 import numpy as np
 
@@ -46,31 +47,63 @@ def _pillow_image_to_buffer(image, image_format: str) -> bytes:
     return img_byte_arr
 
 
-def png_to_buffer(arr: 'np.ndarray', width: int, height: int, resize_method: str):
+def png_to_buffer(
+    arr: 'np.ndarray',
+    width: Optional[int] = None,
+    height: Optional[int] = None,
+    resize_method: str = 'BILINEAR',
+    color_axis: int = -1,
+):
     """
     Convert png to buffer bytes.
 
     :param arr: Data representations of the png.
-    :param width: Width of the png.
-    :param height: Height of the png.
+    :param width: the width of the :attr:`arr`, if None, interpret from :attr:`arr` shape.
+    :param height: the height of the :attr:`arr`, if None, interpret from :attr:`arr` shape.
     :param resize_method: Resize methods (e.g. `NEAREST`, `BILINEAR`, `BICUBIC`, and `LANCZOS`).
+    :param color_axis: the axis id of the color channel, ``-1`` indicates the color channel info at the last axis
     :return: Png in buffer bytes.
+
+    ..note::
+        if both :attr:`width` and :attr:`height` were provided, will not resize. Otherwise, will get image size
+        by :attr:`arr` shape and apply resize method :attr:`resize_method`.
     """
     arr = arr.astype(np.uint8)
+    is_height_width_set = bool(height and width)
 
     if arr.ndim == 1:
+        if not is_height_width_set:
+            height, width = arr.shape[0], 1
         png_bytes = _png_to_buffer_1d(arr, width, height)
     elif arr.ndim == 2:
         from PIL import Image
 
-        im = Image.fromarray(arr).convert('L')
-        im = im.resize((width, height), getattr(Image, resize_method))
+        if not is_height_width_set:
+            height, width = arr.shape
+            im = Image.fromarray(arr).convert('L')
+            im = im.resize((width, height), getattr(Image, resize_method))
+        else:
+            im = Image.fromarray(arr).convert('L')
         png_bytes = _pillow_image_to_buffer(im, image_format='PNG')
     elif arr.ndim == 3:
         from PIL import Image
 
-        im = Image.fromarray(arr).convert('RGB')
-        im = im.resize((width, height), getattr(Image, resize_method))
+        if color_axis != -1:
+            arr = np.moveaxis(arr, color_axis, -1)
+
+        if not is_height_width_set:
+            height, width, num_channels = arr.shape
+        else:
+            _, _, num_channels = arr.shape
+
+        if num_channels == 1:  # greyscale image
+            im = Image.fromarray((arr[0] * 255).astype(np.uint8))
+        else:
+            im = Image.fromarray(arr).convert('RGB')
+
+        if not is_height_width_set:
+            im = im.resize((width, height), getattr(Image, resize_method))
+
         png_bytes = _pillow_image_to_buffer(im, image_format='PNG')
     else:
         raise ValueError(f'ndim={len(arr.shape)} array is not supported')

tests/unit/types/document/test_converters.py

@@ -48,14 +48,31 @@ def test_convert_buffer_to_blob():
     np.testing.assert_almost_equal(doc.content.reshape([10, 10]), array)
 
 
+@pytest.mark.parametrize('resize_method', ['BILINEAR', 'NEAREST', 'BICUBIC', 'LANCZOS'])
 @pytest.mark.parametrize(
-    'arr_size,mode', [(32 * 28, 'L'), ([32, 28], 'L'), ([32, 28, 3], 'RGB')]
+    'arr_size, color_axis, height, width',
+    [
+        ((32 * 28), -1, None, None),  # single line
+        ([32, 28], -1, None, None),  # without channel info
+        ([32, 28, 3], -1, None, None),  # h, w, c (rgb)
+        ([3, 32, 28], 0, None, None),  # c, h, w  (rgb)
+        ([1, 32, 28], 0, None, None),  # c, h, w, (greyscale)
+        ([32, 28, 1], -1, None, None),  # h, w, c, (greyscale)
+        ((32 * 28), -1, 896, 1),  # single line
+        ([32, 28], -1, 32, 28),  # without channel info
+        ([32, 28, 3], -1, 32, 28),  # h, w, c (rgb)
+        ([3, 32, 28], 0, 32, 28),  # c, h, w  (rgb)
+        ([1, 32, 28], 0, 32, 28),  # c, h, w, (greyscale)
+        ([32, 28, 1], -1, 32, 28),  # h, w, c, (greyscale)
+    ],
 )
-def test_convert_blob_to_uri(arr_size, mode):
+def test_convert_image_blob_to_uri(arr_size, color_axis, width, height, resize_method):
     doc = Document(content=np.random.randint(0, 255, arr_size))
     assert doc.blob.any()
     assert not doc.uri
-    doc.convert_image_blob_to_uri(32, 28)
+    doc.convert_image_blob_to_uri(
+        color_axis=color_axis, width=width, height=height, resize_method=resize_method
+    )
     assert doc.uri.startswith('data:image/png;base64,')
     assert doc.mime_type == 'image/png'