Repository: NVIDIA/nv-ingest

Issue:
[BUG]: Error extracting image: Failed to convert NumPy array to image: Cannot handle this data type: (1, 1, 1), |u1
### Version

24.08

### Which installation method(s) does this occur on?

_No response_

### Describe the bug.

The log shows error:
```
ERROR - Error extracting image: Failed to convert NumPy array to image: Cannot handle this data type: (1, 1, 1), |u1
```
with CLI command:
```
nv-ingest-cli --client_host redis --client_port 6379 --task 'extract:{"document_type":"pdf","extract_text":true,"extract_tables":false,"extract_images":true}' --doc "data/bo_20/*.pdf" --output_directory ./processed
```

### Minimum reproducible example

_No response_

### Relevant log output

_No response_

### Other/Misc.

_No response_

Discussion:

Patch:
diff --git a/src/nv_ingest/extraction_workflows/pdf/doughnut_utils.py b/src/nv_ingest/extraction_workflows/pdf/doughnut_utils.py
index 8ba3d2c..d1f7ac2 100644
--- a/src/nv_ingest/extraction_workflows/pdf/doughnut_utils.py
+++ b/src/nv_ingest/extraction_workflows/pdf/doughnut_utils.py
@@ -3,7 +3,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import re
-from io import BytesIO
 from math import ceil
 from math import floor
 from typing import List
@@ -11,10 +10,8 @@ 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
 
diff --git a/src/nv_ingest/util/image_processing/transforms.py b/src/nv_ingest/util/image_processing/transforms.py
index 6e1e6cc..6b0acb0 100644
--- a/src/nv_ingest/util/image_processing/transforms.py
+++ b/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
diff --git a/src/nv_ingest/util/nim/helpers.py b/src/nv_ingest/util/nim/helpers.py
index 9376085..7f964a6 100644
--- a/src/nv_ingest/util/nim/helpers.py
+++ b/src/nv_ingest/util/nim/helpers.py
@@ -2,7 +2,6 @@
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import io
 import logging
 from typing import Optional
 from typing import Tuple
@@ -10,9 +9,8 @@ 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"]


Test patch:
diff --git a/tests/nv_ingest/util/image_processing/test_transforms.py b/tests/nv_ingest/util/image_processing/test_transforms.py
new file mode 100644
index 0000000..b9b7e27
--- /dev/null
+++ b/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