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/extraction_workflows/pdf/pdfium_helper.py

@@ -3,7 +3,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-import io
 import logging
 from math import ceil
 from math import floor
@@ -14,15 +13,16 @@
 import numpy as np
 import pypdfium2 as libpdfium
 import tritonclient.grpc as grpcclient
-from PIL import Image
 
 from nv_ingest.extraction_workflows.pdf import yolox_utils
 from nv_ingest.schemas.metadata_schema import AccessLevelEnum
 from nv_ingest.schemas.metadata_schema import TextTypeEnum
 from nv_ingest.schemas.pdf_extractor_schema import PDFiumConfigSchema
-from nv_ingest.util.converters import bytetools
 from nv_ingest.util.image_processing.table_and_chart import join_cached_and_deplot_output
 from nv_ingest.util.image_processing.transforms import numpy_to_base64
+from nv_ingest.util.nim.helpers import call_image_inference_model
+from nv_ingest.util.nim.helpers import create_inference_client
+from nv_ingest.util.nim.helpers import perform_model_inference
 from nv_ingest.util.pdf.metadata_aggregators import Base64Image
 from nv_ingest.util.pdf.metadata_aggregators import ImageChart
 from nv_ingest.util.pdf.metadata_aggregators import ImageTable
@@ -33,9 +33,6 @@
 from nv_ingest.util.pdf.pdfium import PDFIUM_PAGEOBJ_MAPPING
 from nv_ingest.util.pdf.pdfium import pdfium_pages_to_numpy
 from nv_ingest.util.pdf.pdfium import pdfium_try_get_bitmap_as_numpy
-from nv_ingest.util.nim.helpers import call_image_inference_model
-from nv_ingest.util.nim.helpers import create_inference_client
-from nv_ingest.util.nim.helpers import perform_model_inference
 
 # Copyright (c) 2024, NVIDIA CORPORATION.
 #
@@ -309,10 +306,7 @@ def handle_table_chart_extraction(
             h1, w1, h2, w2 = bbox * np.array([height, width, height, width])
             cropped = original_image[floor(w1) : ceil(w2), floor(h1) : ceil(h2)]  # noqa: E203
 
-            img = Image.fromarray(cropped.astype(np.uint8))
-            with io.BytesIO() as buffer:
-                img.save(buffer, format="PNG")
-                base64_img = bytetools.base64frombytes(buffer.getvalue())
+            base64_img = numpy_to_base64(cropped)
 
             if label == "table":
                 table_content = call_image_inference_model(paddle_client, "paddle", cropped)

src/nv_ingest/util/image_processing/transforms.py

@@ -138,6 +138,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