Process grayscale images

src/nv_ingest/extraction_workflows/pdf/doughnut_utils.py

@@ -3,18 +3,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import re
-from io import BytesIO
 from math import ceil
 from math import floor
 from typing import List
 from typing import Optional
 from typing import Tuple
 
 import numpy as np
-from PIL import Image
 
-from nv_ingest.util.converters import bytetools
-from nv_ingest.util.pdf.metadata_aggregators import LatexTable
+from nv_ingest.util.image_processing.transforms import numpy_to_base64
 
 DEFAULT_DPI = 300
 DEFAULT_MAX_WIDTH = 1024
@@ -112,10 +109,7 @@ def crop_image(array: np.array, bbox: Tuple[int, int, int, int], format="PNG") -
     if (w2 - w1 <= 0) or (h2 - h1 <= 0):
         return None
     cropped = array[h1:h2, w1:w2]
-    pil_image = Image.fromarray(cropped.astype(np.uint8))
-    with BytesIO() as buffer:
-        pil_image.save(buffer, format="PNG")
-        base64_img = bytetools.base64frombytes(buffer.getvalue())
+    base64_img = numpy_to_base64(cropped)
 
     return base64_img
 

src/nv_ingest/util/image_processing/transforms.py

@@ -146,6 +146,12 @@ def numpy_to_base64(array: np.ndarray) -> str:
     >>> isinstance(encoded_str, str)
     True
     """
+    # If the array represents a grayscale image, drop the redundant axis in
+    # (h, w, 1). PIL.Image.fromarray() expects an array of form (h, w) if it's
+    # a grayscale image.
+    if array.ndim == 3 and array.shape[2] == 1:
+        array = np.squeeze(array, axis=2)
+
     # Check if the array is valid and can be converted to an image
     try:
         # Convert the NumPy array to a PIL image

src/nv_ingest/util/nim/helpers.py

@@ -2,17 +2,15 @@
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import io
 import logging
 from typing import Optional
 from typing import Tuple
 
 import numpy as np
 import requests
 import tritonclient.grpc as grpcclient
-from PIL import Image
 
-from nv_ingest.util.converters import bytetools
+from nv_ingest.util.image_processing.transforms import numpy_to_base64
 
 logger = logging.getLogger(__name__)
 
@@ -90,10 +88,7 @@ def call_image_inference_model(client, model_name: str, image_data):
             logger.error(err_msg)
             raise RuntimeError(err_msg)
     else:
-        image = Image.fromarray(image_data)
-        with io.BytesIO() as buffer:
-            image.save(buffer, format="PNG")
-            base64_img = bytetools.base64frombytes(buffer.getvalue())
+        base64_img = numpy_to_base64(image_data)
 
         try:
             url = client["endpoint_url"]

tests/nv_ingest/util/image_processing/test_transforms.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+from nv_ingest.util.image_processing.transforms import numpy_to_base64
+
+
+def test_numpy_to_base64_valid_rgba_image():
+    array = np.random.randint(0, 255, (100, 100, 4), dtype=np.uint8)
+    result = numpy_to_base64(array)
+
+    assert isinstance(result, str)
+    assert len(result) > 0
+
+
+def test_numpy_to_base64_valid_rgb_image():
+    array = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
+    result = numpy_to_base64(array)
+
+    assert isinstance(result, str)
+    assert len(result) > 0
+
+
+def test_numpy_to_base64_grayscale_redundant_axis():
+    array = np.random.randint(0, 255, (100, 100, 1), dtype=np.uint8)
+    result = numpy_to_base64(array)
+
+    assert isinstance(result, str)
+    assert len(result) > 0