nv-ingest-api 2025.4.16.dev20250416__py3-none-any.whl → 2025.4.18.dev20250418__py3-none-any.whl

nv_ingest_api/internal/extract/image/image_helpers/common.py

@@ -0,0 +1,403 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Copyright (c) 2024, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import io
+import logging
+from datetime import datetime
+from typing import Dict, IO, Any
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+import numpy as np
+from PIL import Image
+
+# from wand.image import Image as WandImage
+
+from nv_ingest_api.internal.enums.common import AccessLevelEnum
+from nv_ingest_api.internal.primitives.nim.model_interface.yolox import (
+    YoloxPageElementsModelInterface,
+    get_yolox_model_name,
+)
+from nv_ingest_api.internal.schemas.extract.extract_image_schema import ImageConfigSchema
+from nv_ingest_api.util.image_processing.transforms import crop_image, numpy_to_base64
+from nv_ingest_api.util.metadata.aggregators import (
+    CroppedImageWithContent,
+    construct_page_element_metadata,
+    construct_image_metadata_from_base64,
+)
+from nv_ingest_api.util.nim import create_inference_client
+
+logger = logging.getLogger(__name__)
+
+YOLOX_MAX_BATCH_SIZE = 8
+
+RAW_FILE_FORMATS = ["jpeg", "jpg", "png", "tiff", "bmp"]
+PREPROC_FILE_FORMATS = ["svg"]
+
+SUPPORTED_FILE_TYPES = RAW_FILE_FORMATS + ["svg"]
+
+
+def load_and_preprocess_image(image_stream: io.BytesIO) -> np.ndarray:
+    """
+    Loads and preprocesses a JPEG, JPG, or PNG image from a bytestream.
+
+    Parameters
+    ----------
+    image_stream : io.BytesIO
+        A bytestream of the image file.
+
+    Returns
+    -------
+    np.ndarray
+        Preprocessed image as a numpy array.
+    """
+    # Load image from the byte stream
+    processed_image = Image.open(image_stream).convert("RGB")
+
+    # Convert image to numpy uint8 array
+    image_array = np.asarray(processed_image)
+
+    return image_array
+
+
+def convert_svg_to_bitmap(image_stream: io.BytesIO) -> np.ndarray:
+    """
+    Converts an SVG image from a bytestream to a bitmap format.
+
+    Parameters
+    ----------
+    image_stream : io.BytesIO
+        A bytestream of the SVG file.
+
+    Returns
+    -------
+    np.ndarray
+        Preprocessed image as a numpy array in bitmap format.
+    """
+
+    pass
+    # Convert SVG to PNG using Wand (ImageMagick)
+    # with WandImage(blob=image_stream.read(), format="svg") as img:
+    #    img.format = "png"
+    #    png_data = img.make_blob()
+
+    ## Reload the PNG as a PIL Image
+    # processed_image = Image.open(io.BytesIO(png_data)).convert("RGB")
+
+    ## Convert image to numpy array and normalize pixel values
+    # image_array = np.asarray(processed_image, dtype=np.float32)
+
+    # return image_array
+
+
+def extract_page_element_images(
+    annotation_dict: Dict[str, List[List[float]]],
+    original_image: np.ndarray,
+    page_idx: int,
+    page_elements: List[Tuple[int, "CroppedImageWithContent"]],
+) -> None:
+    """
+    Handle the extraction of tables and charts from the inference results and run additional model inference.
+
+    Parameters
+    ----------
+    annotation_dict : dict of {str : list of list of float}
+        A dictionary containing detected objects and their bounding boxes. Keys should include "table" and "chart",
+        and each key's value should be a list of bounding boxes, with each bounding box represented as a list of floats.
+    original_image : np.ndarray
+        The original image from which objects were detected, expected to be in RGB format with shape (H, W, 3).
+    page_idx : int
+        The index of the current page being processed.
+    page_elements : list of tuple of (int, CroppedImageWithContent)
+        A list to which extracted tables and charts will be appended. Each item in the list is a tuple where the first
+        element is the page index, and the second is an instance of CroppedImageWithContent representing a cropped image
+        and associated metadata.
+
+    Returns
+    -------
+    None
+
+    Notes
+    -----
+    This function iterates over detected objects labeled as "table" or "chart". For each object, it crops the original
+    image according to the bounding box coordinates, then creates an instance of `CroppedImageWithContent` containing
+    the cropped image and metadata, and appends it to `page_elements`.
+
+    Examples
+    --------
+    >>> annotation_dict = {"table": [[0.1, 0.1, 0.5, 0.5, 0.8]], "chart": [[0.6, 0.6, 0.9, 0.9, 0.9]]}
+    >>> original_image = np.random.rand(1536, 1536, 3)
+    >>> page_elements = []
+    >>> extract_page_element_images(annotation_dict, original_image, 0, page_elements)
+    >>> len(page_elements)
+    2
+    """
+
+    width, height, *_ = original_image.shape
+    for label in ["table", "chart"]:
+        if not annotation_dict or label not in annotation_dict:
+            continue
+
+        objects = annotation_dict[label]
+        for idx, bboxes in enumerate(objects):
+            *bbox, _ = bboxes
+            h1, w1, h2, w2 = bbox
+
+            cropped_img = crop_image(original_image, (int(h1), int(w1), int(h2), int(w2)))
+            base64_img = numpy_to_base64(cropped_img) if cropped_img is not None else None
+
+            table_data = CroppedImageWithContent(
+                content="",
+                image=base64_img,
+                bbox=(int(w1), int(h1), int(w2), int(h2)),
+                max_width=width,
+                max_height=height,
+                type_string=label,
+            )
+            page_elements.append((page_idx, table_data))
+
+
+def extract_page_elements_from_images(
+    images: List[np.ndarray],
+    config: ImageConfigSchema,
+    trace_info: Optional[List] = None,
+) -> List[Tuple[int, object]]:
+    """
+    Detect and extract tables/charts from a list of NumPy images using YOLOX.
+
+    Parameters
+    ----------
+    images : List[np.ndarray]
+        List of images in NumPy array format.
+    config : ImageConfigSchema
+        Configuration object containing YOLOX endpoints, auth token, etc.
+    trace_info : Optional[List], optional
+        Optional tracing data for debugging/performance profiling.
+
+    Returns
+    -------
+    List[Tuple[int, object]]
+        A list of (image_index, CroppedImageWithContent) representing extracted
+        table/chart data from each image.
+    """
+    page_elements = []
+    yolox_client = None
+
+    # Obtain yolox_version
+    # Assuming that the http endpoint is at index 1
+    yolox_http_endpoint = config.yolox_endpoints[1]
+    yolox_model_name = get_yolox_model_name(yolox_http_endpoint)
+
+    try:
+        model_interface = YoloxPageElementsModelInterface(yolox_model_name=yolox_model_name)
+        yolox_client = create_inference_client(
+            config.yolox_endpoints,
+            model_interface,
+            config.auth_token,
+            config.yolox_infer_protocol,
+        )
+
+        # Prepare the payload with all images.
+        data = {"images": images}
+
+        # Perform inference in a single call. The NimClient handles batching internally.
+        inference_results = yolox_client.infer(
+            data,
+            model_name="yolox",
+            max_batch_size=YOLOX_MAX_BATCH_SIZE,
+            trace_info=trace_info,
+            stage_name="pdf_content_extractor",
+        )
+
+        # Process each result along with its corresponding image.
+        for i, (annotation_dict, original_image) in enumerate(zip(inference_results, images)):
+            extract_page_element_images(
+                annotation_dict,
+                original_image,
+                i,
+                page_elements,
+            )
+
+    except TimeoutError:
+        logger.error("Timeout error during table/chart extraction.")
+        raise
+
+    except Exception as e:
+        logger.exception(f"Unhandled error during table/chart extraction: {str(e)}")
+        raise
+
+    finally:
+        if yolox_client:
+            yolox_client.close()
+
+    logger.debug(f"Extracted {len(page_elements)} tables and charts from image.")
+    return page_elements
+
+
+def unstructured_image_extractor(
+    *,
+    image_stream: IO[bytes],
+    extract_text: bool,
+    extract_images: bool,
+    extract_infographics: bool,
+    extract_tables: bool,
+    extract_charts: bool,
+    extraction_config: Dict[str, Any],
+    extraction_trace_log: Optional[Dict[str, Any]] = None,
+) -> List[Any]:
+    """
+    Extract primitives from an unstructured image bytestream.
+
+    This helper function processes an image bytestream according to the provided extraction
+    configuration. It supports extraction of tables, charts, and infographics from the image.
+    (Note: text and additional image extraction are not supported yet for raw images.)
+
+    Parameters
+    ----------
+    image_stream : IO[bytes]
+        A bytestream (e.g. io.BytesIO) containing the image file data.
+    image_stream : io.BytesIO
+        A bytestream for the image file.
+    document_type : str
+        Specifies the type of the image document ('png', 'jpeg', 'jpg', 'svg', 'tiff', 'bmp').
+    extract_text : bool
+        Flag specifying whether to extract text (currently not supported for raw images).
+    extract_images : bool
+        Flag specifying whether to extract images (currently not supported for raw images).
+    extract_infographics : bool
+        Flag specifying whether to extract infographics.
+    extract_tables : bool
+        Flag specifying whether to extract tables.
+    extract_charts : bool
+        Flag specifying whether to extract charts.
+    extraction_config : Dict[str, Any]
+        A dictionary containing additional extraction parameters and configurations.
+        Expected keys include "document_type", "row_data", "metadata_column", and
+        "image_extraction_config".
+    extraction_trace_log : Optional[Dict[str, Any]], optional
+        An optional dictionary containing trace information for logging or debugging,
+        by default None.
+
+    Returns
+    -------
+    List[Any]
+        A list of extracted data items (e.g., metadata dictionaries) from the image.
+
+    Raises
+    ------
+    ValueError
+        If the document type is unsupported.
+    Exception
+        If an error occurs during extraction.
+    """
+    # Note: extract_infographics flag is not currently used in extraction
+    _ = extract_infographics
+
+    # Determine the type of the document from the extraction config.
+    document_type: str = extraction_config.get("document_type", "unknown")
+    logger.debug(f"Extracting {document_type.upper()} image with image extractor.")
+
+    # Ensure the document type is supported.
+    if document_type not in SUPPORTED_FILE_TYPES:
+        raise ValueError(f"Unsupported document type: {document_type}")
+
+    # Retrieve additional row-specific data and source identifier.
+    row_data: Dict[str, Any] = extraction_config.get("row_data", {})
+    source_id: str = row_data.get("source_id", "unknown_source")
+
+    # Build source metadata based on row data.
+    base_unified_metadata: Dict[str, Any] = row_data.get(extraction_config.get("metadata_column", "metadata"), {})
+    current_iso_datetime: str = datetime.now().isoformat()
+    source_metadata: Dict[str, Any] = {
+        "source_name": source_id if os.path.splitext(source_id)[1] else f"{source_id}.{document_type}",
+        "source_id": source_id,
+        "source_location": row_data.get("source_location", ""),
+        "source_type": document_type,
+        "collection_id": row_data.get("collection_id", ""),
+        "date_created": row_data.get("date_created", current_iso_datetime),
+        "last_modified": row_data.get("last_modified", current_iso_datetime),
+        "summary": f"Raw {document_type} image extracted from source {source_id}",
+        "partition_id": row_data.get("partition_id", -1),
+        "access_level": row_data.get("access_level", AccessLevelEnum.UNKNOWN),
+    }
+
+    # Optionally update the extract_infographics flag based on extraction_config.
+    extract_infographics = extraction_config.get("extract_infographics", False)
+
+    # Log which primitives are requested for extraction.
+    logger.debug(f"Extract text: {extract_text} (not supported yet for raw images)")
+    logger.debug(f"Extract images: {extract_images} (not supported yet for raw images)")
+    logger.debug(f"Extract tables: {extract_tables}")
+    logger.debug(f"Extract charts: {extract_charts}")
+    logger.debug(f"Extract infographics: {extract_infographics}")
+
+    # Preprocess the image based on the document type.
+    if document_type in RAW_FILE_FORMATS:
+        logger.debug(f"Loading and preprocessing {document_type} image.")
+        image_array = load_and_preprocess_image(image_stream)
+    elif document_type in PREPROC_FILE_FORMATS:
+        logger.debug(f"Converting {document_type} to bitmap.")
+        image_array = convert_svg_to_bitmap(image_stream)
+    else:
+        raise ValueError(f"Unsupported document type: {document_type}")
+
+    extracted_data: List[Any] = []
+
+    # Text extraction stub (not supported for raw images)
+    if extract_text:
+        logger.warning("Text extraction is not supported for raw images.")
+
+    # Extract tables, charts, or infographics if requested.
+    if extract_tables or extract_charts or extract_infographics:
+        try:
+            page_elements = extract_page_elements_from_images(
+                [image_array],
+                config=extraction_config.get("image_extraction_config"),
+                trace_info=extraction_trace_log,
+            )
+            for item in page_elements:
+                table_chart_data = item[1]
+                extracted_data.append(
+                    construct_page_element_metadata(
+                        table_chart_data,
+                        page_idx=0,  # Treat single image as one page.
+                        page_count=1,
+                        source_metadata=source_metadata,
+                        base_unified_metadata=base_unified_metadata,
+                    )
+                )
+        except Exception as e:
+            logger.error(f"Error extracting tables/charts from image: {e}")
+            raise
+
+    # Image extraction stub: if no structured elements were extracted and image extraction is requested.
+    if extract_images and not extracted_data:
+        extracted_data.append(
+            construct_image_metadata_from_base64(
+                numpy_to_base64(image_array),
+                page_idx=0,  # Treat single image as one page.
+                page_count=1,
+                source_metadata=source_metadata,
+                base_unified_metadata=base_unified_metadata,
+            )
+        )
+
+    logger.debug(f"Extracted {len(extracted_data)} items from the image.")
+    return extracted_data

nv_ingest_api/internal/extract/image/infographic_extractor.py

@@ -0,0 +1,253 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+import pandas as pd
+
+from nv_ingest_api.internal.primitives.nim import NimClient
+from nv_ingest_api.internal.primitives.nim.model_interface.paddle import PaddleOCRModelInterface
+from nv_ingest_api.internal.schemas.extract.extract_infographic_schema import (
+    InfographicExtractorSchema,
+)
+from nv_ingest_api.util.image_processing.transforms import base64_to_numpy
+from nv_ingest_api.util.nim import create_inference_client
+
+logger = logging.getLogger(__name__)
+
+PADDLE_MIN_WIDTH = 32
+PADDLE_MIN_HEIGHT = 32
+
+
+def _filter_infographic_images(
+    base64_images: List[str],
+) -> Tuple[List[str], List[int], List[Tuple[str, Optional[Any], Optional[Any]]]]:
+    """
+    Filters base64-encoded images based on minimum size requirements.
+
+    Parameters
+    ----------
+    base64_images : List[str]
+        List of base64-encoded image strings.
+
+    Returns
+    -------
+    Tuple[List[str], List[int], List[Tuple[str, Optional[Any], Optional[Any]]]]
+        - valid_images: List of images that meet the size requirements.
+        - valid_indices: Original indices of valid images.
+        - results: Initialized results list, with invalid images marked as (img, None, None).
+    """
+    results: List[Tuple[str, Optional[Any], Optional[Any]]] = [("", None, None)] * len(base64_images)
+    valid_images: List[str] = []
+    valid_indices: List[int] = []
+
+    for i, img in enumerate(base64_images):
+        array = base64_to_numpy(img)
+        height, width = array.shape[0], array.shape[1]
+        if width >= PADDLE_MIN_WIDTH and height >= PADDLE_MIN_HEIGHT:
+            valid_images.append(img)
+            valid_indices.append(i)
+        else:
+            # Mark image as skipped if it does not meet size requirements.
+            results[i] = (img, None, None)
+    return valid_images, valid_indices, results
+
+
+def _update_infographic_metadata(
+    base64_images: List[str],
+    paddle_client: NimClient,
+    worker_pool_size: int = 8,  # Not currently used
+    trace_info: Optional[Dict] = None,
+) -> List[Tuple[str, Optional[Any], Optional[Any]]]:
+    """
+    Filters base64-encoded images and uses PaddleOCR to extract infographic data.
+
+    For each image that meets the minimum size, calls paddle_client.infer to obtain
+    (text_predictions, bounding_boxes). Invalid images are marked as skipped.
+
+    Parameters
+    ----------
+    base64_images : List[str]
+        List of base64-encoded images.
+    paddle_client : NimClient
+        Client instance for PaddleOCR inference.
+    worker_pool_size : int, optional
+        Worker pool size (currently not used), by default 8.
+    trace_info : Optional[Dict], optional
+        Optional trace information for debugging.
+
+    Returns
+    -------
+    List[Tuple[str, Optional[Any], Optional[Any]]]
+        List of tuples in the same order as base64_images, where each tuple contains:
+        (base64_image, text_predictions, bounding_boxes).
+    """
+    logger.debug(f"Running infographic extraction using protocol {paddle_client.protocol}")
+
+    valid_images, valid_indices, results = _filter_infographic_images(base64_images)
+    data_paddle = {"base64_images": valid_images}
+
+    # worker_pool_size is not used in current implementation.
+    _ = worker_pool_size
+
+    try:
+        paddle_results = paddle_client.infer(
+            data=data_paddle,
+            model_name="paddle",
+            stage_name="infographic_data_extraction",
+            max_batch_size=1 if paddle_client.protocol == "grpc" else 2,
+            trace_info=trace_info,
+        )
+    except Exception as e:
+        logger.error(f"Error calling paddle_client.infer: {e}", exc_info=True)
+        raise
+
+    if len(paddle_results) != len(valid_images):
+        raise ValueError(f"Expected {len(valid_images)} paddle results, got {len(paddle_results)}")
+
+    for idx, paddle_res in enumerate(paddle_results):
+        original_index = valid_indices[idx]
+        # Each paddle_res is expected to be a tuple (text_predictions, bounding_boxes)
+        results[original_index] = (base64_images[original_index], paddle_res[0], paddle_res[1])
+
+    return results
+
+
+def _create_clients(
+    paddle_endpoints: Tuple[str, str],
+    paddle_protocol: str,
+    auth_token: str,
+) -> NimClient:
+    paddle_model_interface = PaddleOCRModelInterface()
+
+    logger.debug(f"Inference protocols: paddle={paddle_protocol}")
+
+    paddle_client = create_inference_client(
+        endpoints=paddle_endpoints,
+        model_interface=paddle_model_interface,
+        auth_token=auth_token,
+        infer_protocol=paddle_protocol,
+    )
+
+    return paddle_client
+
+
+def _meets_infographic_criteria(row: pd.Series) -> bool:
+    """
+    Determines if a DataFrame row meets the criteria for infographic extraction.
+
+    A row qualifies if:
+      - It contains a 'metadata' dictionary.
+      - The 'content_metadata' in metadata has type "structured" and subtype "infographic".
+      - The 'table_metadata' is not None.
+      - The 'content' is not None or an empty string.
+
+    Parameters
+    ----------
+    row : pd.Series
+        A row from the DataFrame.
+
+    Returns
+    -------
+    bool
+        True if the row meets all criteria; False otherwise.
+    """
+    metadata = row.get("metadata", {})
+    if not metadata:
+        return False
+
+    content_md = metadata.get("content_metadata", {})
+    if (
+        content_md.get("type") == "structured"
+        and content_md.get("subtype") == "infographic"
+        and metadata.get("table_metadata") is not None
+        and metadata.get("content") not in [None, ""]
+    ):
+        return True
+
+    return False
+
+
+def extract_infographic_data_from_image_internal(
+    df_extraction_ledger: pd.DataFrame,
+    task_config: Dict[str, Any],
+    extraction_config: InfographicExtractorSchema,
+    execution_trace_log: Optional[Dict] = None,
+) -> Tuple[pd.DataFrame, Dict]:
+    """
+    Extracts infographic data from a DataFrame in bulk, following the chart extraction pattern.
+
+    Parameters
+    ----------
+    df_extraction_ledger : pd.DataFrame
+        DataFrame containing the content from which infographic data is to be extracted.
+    task_config : Dict[str, Any]
+        Dictionary containing task properties and configurations.
+    extraction_config : Any
+        The validated configuration object for infographic extraction.
+    execution_trace_log : Optional[Dict], optional
+        Optional trace information for debugging or logging. Defaults to None.
+
+    Returns
+    -------
+    Tuple[pd.DataFrame, Dict]
+        A tuple containing the updated DataFrame and the trace information.
+    """
+    _ = task_config  # Unused
+
+    if execution_trace_log is None:
+        execution_trace_log = {}
+        logger.debug("No trace_info provided. Initialized empty trace_info dictionary.")
+
+    if df_extraction_ledger.empty:
+        return df_extraction_ledger, execution_trace_log
+
+    endpoint_config = extraction_config.endpoint_config
+    paddle_client = _create_clients(
+        endpoint_config.paddle_endpoints,
+        endpoint_config.paddle_infer_protocol,
+        endpoint_config.auth_token,
+    )
+
+    try:
+        # Identify rows that meet the infographic criteria.
+        mask = df_extraction_ledger.apply(_meets_infographic_criteria, axis=1)
+        valid_indices = df_extraction_ledger[mask].index.tolist()
+
+        # If no rows meet the criteria, return early.
+        if not valid_indices:
+            return df_extraction_ledger, {"trace_info": execution_trace_log}
+
+        # Extract base64 images from valid rows.
+        base64_images = [df_extraction_ledger.at[idx, "metadata"]["content"] for idx in valid_indices]
+
+        # Call bulk update to extract infographic data.
+        bulk_results = _update_infographic_metadata(
+            base64_images=base64_images,
+            paddle_client=paddle_client,
+            worker_pool_size=endpoint_config.workers_per_progress_engine,
+            trace_info=execution_trace_log,
+        )
+
+        # Write the extracted results back into the DataFrame.
+        for result_idx, df_idx in enumerate(valid_indices):
+            # Unpack result: (base64_image, paddle_bounding_boxes, paddle_text_predictions)
+            _, _, text_predictions = bulk_results[result_idx]
+            table_content = " ".join(text_predictions) if text_predictions else None
+            df_extraction_ledger.at[df_idx, "metadata"]["table_metadata"]["table_content"] = table_content
+
+        return df_extraction_ledger, {"trace_info": execution_trace_log}
+
+    except Exception:
+        err_msg = "Error occurred while extracting infographic data."
+        logger.exception(err_msg)
+        raise
+
+    finally:
+        paddle_client.close()