unstructured-io--unstructured/unstructured/partition/pdf_image/analysis/bbox_visualisation.py

import logging
import math
import tempfile
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
from io import BytesIO
from pathlib import Path
from typing import Any, Generator, List, Optional, TypeVar, Union

import numpy as np
from matplotlib import colors, font_manager
from numba import njit
from PIL import Image, ImageDraw, ImageFont
from unstructured_inference.constants import ElementType

from unstructured.partition.pdf_image.analysis.processor import AnalysisProcessor
from unstructured.partition.pdf_image.pdf_image_utils import convert_pdf_to_image

PageImage = TypeVar("PageImage", Image.Image, np.ndarray)


def get_font():
    preferred_fonts = ["Arial.ttf"]
    available_fonts = font_manager.findSystemFonts()
    if not available_fonts:
        raise ValueError("No fonts available")
    for font in preferred_fonts:
        for available_font in available_fonts:
            if font in available_font:
                return available_font
    return available_fonts[0]


COLOR_WHITE = ("white", (255, 255, 255))
COLOR_BLACK = ("black", (0, 0, 0))


class TextAlignment(Enum):
    TOP_LEFT = "top_left"
    TOP_RIGHT = "top_right"
    BOTTOM_LEFT = "bottom_left"
    BOTTOM_RIGHT = "bottom_right"
    CENTER = "center"


@dataclass
class BboxLabels:
    top_left: Optional[str] = None
    top_right: Optional[str] = None
    bottom_left: Optional[str] = None
    bottom_right: Optional[str] = None
    center: Optional[str] = None


@dataclass
class BBox:
    points: tuple[int, int, int, int]
    labels: Optional[BboxLabels] = None


def get_rgb_color(color: str) -> tuple[int, int, int]:
    """Convert a color name to RGB values.

    Args:
        color: A color name supported by matplotlib.

    Returns:
        A tuple of three integers representing the RGB values of the color.
    """
    try:
        rgb_colors = colors.to_rgb(color)
    except ValueError:
        print("Error")
        raise
    return int(rgb_colors[0] * 255), int(rgb_colors[1] * 255), int(rgb_colors[2] * 255)


@njit(cache=True, fastmath=True)
def _get_bbox_to_page_ratio(bbox: tuple[int, int, int, int], page_size: tuple[int, int]) -> float:
    """Compute the ratio of the bounding box to the page size.

    Args:
        bbox: Tuple containing coordinates of the bbox: (x1, y1, x2, y2).
        page_size: Tuple containing page size: (width, height).

    Returns:
        The ratio of the bounding box to the page size.
    """
    x1, y1, x2, y2 = bbox
    page_width, page_height = page_size
    page_diagonal = math.sqrt(page_height**2 + page_width**2)
    bbox_width = x2 - x1
    bbox_height = y2 - y1
    bbox_diagonal = math.sqrt(bbox_height**2 + bbox_width**2)
    return bbox_diagonal / page_diagonal


def _get_optimal_value_for_bbox(
    bbox: tuple[int, int, int, int],
    page_size: tuple[int, int],
    min_value: int,
    max_value: int,
    ratio_for_min_value: float = 0.01,
    ratio_for_max_value: float = 0.5,
) -> int:
    """Compute the optimal value for a given bounding box using a linear function
    generated for given min and max values and ratios

    Args:
        bbox: Tuple containing coordinates of the bbox: (x1, y1, x2, y2).
        page_size: Tuple containing page size: (width, height).
        min_value: The minimum value returned by the function.
        max_value: The maximum value returned by the function.
        ratio_for_min_value: The ratio of the bbox to page size for the min value.
        ratio_for_max_value: The ratio of the bbox to page size for the max value.

    Returns:
        The optimal value for the given bounding box and parameters given.
    """
    bbox_to_page_ratio = _get_bbox_to_page_ratio(bbox, page_size)
    slope, intercept = _linear_polyfit_2point(
        ratio_for_min_value, ratio_for_max_value, min_value, max_value
    )
    value = int(bbox_to_page_ratio * slope + intercept)
    return max(min_value, min(max_value, value))


def get_bbox_text_size(
    bbox: tuple[int, int, int, int],
    page_size: tuple[int, int],
    min_font_size: int = 16,
    max_font_size: int = 32,
) -> int:
    """Compute the optimal font size for a given bounding box.

    Args:
        bbox: Tuple containing coordinates of the bbox: (x1, y1, x2, y2).
        page_size: Tuple containing page size: (width, height).
        min_font_size: The minimum font size returned by the function.
        max_font_size: The maximum font size returned by the function.

    Returns:
        The optimal font size for the given bounding box.
    """
    return _get_optimal_value_for_bbox(
        bbox=bbox,
        page_size=page_size,
        min_value=min_font_size,
        max_value=max_font_size,
    )


def get_bbox_thickness(
    bbox: tuple[int, int, int, int],
    page_size: tuple[int, int],
    min_thickness: int = 1,
    max_thickness: int = 4,
) -> float:
    """Compute the optimal thickness for a given bounding box.

    Args:
        bbox: Tuple containing coordinates of the bbox: (x1, y1, x2, y2).
        page_size: Tuple containing page size: (width, height).
        min_thickness: The minimum font size returned by the function.
        max_thickness: The maximum font size returned by the function.

    Returns:

    """
    return _get_optimal_value_for_bbox(
        bbox=bbox,
        page_size=page_size,
        min_value=min_thickness,
        max_value=max_thickness,
    )


def get_text_color(
    background_color: Union[str, tuple[int, int, int]], brightness_threshold: float = 0.5
) -> tuple[str, tuple[int, int, int]]:
    """Returns the contrastive text color (black or white) for a given background color.

    Args:
        background_color: Tuple containing RGB values of the background color.

    Returns:
        Tuple containing RGB values of the text color.
    """
    if isinstance(background_color, str):
        background_color = get_rgb_color(background_color)
    background_brightness = (
        0.299 * background_color[0] + 0.587 * background_color[1] + 0.114 * background_color[0]
    ) / 255
    if background_brightness > brightness_threshold:
        return COLOR_BLACK
    else:
        return COLOR_WHITE


def get_label_rect_and_coords(
    alignment: TextAlignment,
    bbox_points: tuple[int, int, int, int],
    text_width: int,
    text_height: int,
):
    indent = max(int(text_width * 0.2), 10)
    vertical_correction = max(int(text_height * 0.3), 10)

    # with this the text should be centered in the rectangle
    rect_width = text_width + indent * 2
    # we apply a correction to the height to make the text not overlap over the rectangle
    # (the text height getter looks to return too small value)
    rect_height = text_height + vertical_correction
    x1, y1, x2, y2 = bbox_points
    if alignment is TextAlignment.CENTER:
        # center:
        horizontal_half = int(rect_width / 2 * 1.05)
        vertical_half = int(rect_height / 2 * 1.05)
        center_point = x1 + (x2 - x1) // 2, y1 + (y2 - y1) // 2
        # resize rectangle to make it look better in the center
        label_rectangle = (
            (
                center_point[0] - horizontal_half,
                center_point[1] - vertical_half,
            ),
            (
                center_point[0] + horizontal_half,
                center_point[1] + vertical_half,
            ),
        )
        label_coords = (
            center_point[0] - horizontal_half + int(indent * 1.05),
            center_point[1] - vertical_half * 1.05,
        )
    elif alignment is TextAlignment.TOP_LEFT:
        label_rectangle = (
            (x1, y1 - rect_height),
            (x1 + rect_width, y1),
        )
        label_coords = (x1 + indent, y1 - rect_height)
    elif alignment is TextAlignment.TOP_RIGHT:
        label_rectangle = (
            (x2 - rect_width, y1),
            (x2, y1 + rect_height),
        )
        label_coords = (x2 - text_width - indent, y1)
    elif alignment is TextAlignment.BOTTOM_LEFT:
        label_rectangle = (
            (x1, y2 - rect_height),
            (x1 + rect_width, y2),
        )
        label_coords = (x1 + indent, y2 - rect_height)
    elif alignment is TextAlignment.BOTTOM_RIGHT:
        label_rectangle = (
            (x2 - rect_width, y2 - rect_height),
            (x2, y2),
        )
        label_coords = (x2 - text_width - indent, y2 - rect_height)
    else:
        raise ValueError(f"Unknown alignment {alignment}")
    return label_rectangle, label_coords


def draw_bbox_label(
    image_draw: ImageDraw.ImageDraw,
    text: str,
    bbox_points: tuple[int, int, int, int],
    alignment: TextAlignment,
    font_size: int,
    background_color: str,
):
    """Draw a label stick to a bounding box.
    The alignment parameter specifies where the label should be placed.

    Args:
        image_draw:         ImageDraw object to draw on the image.
        text:               Text to draw.
        bbox_points:        Bounding box points.
        alignment:          Text alignment.
        font_size:          Font size of the text.
        background_color:   RGB values of the background color.
    """
    font = ImageFont.truetype(get_font(), font_size)
    text_x1, text_y1, text_x2, text_y2 = image_draw.textbbox(
        (0, 0), text, font=font, align="center"
    )
    text_width = text_x2 - text_x1
    text_height = text_y2 - text_y1

    label_rectangle, label_coords = get_label_rect_and_coords(
        alignment, bbox_points, text_width, text_height
    )

    rgb_background_color = get_rgb_color(background_color)
    try:
        image_draw.rectangle(
            label_rectangle,
            fill=background_color,
            outline=background_color,
        )
    except TypeError:
        image_draw.rectangle(
            label_rectangle,
            fill=rgb_background_color,
            outline=rgb_background_color,
        )
    text_color, text_color_rgb = get_text_color(background_color)
    try:
        image_draw.text(label_coords, text, fill=text_color, font=font, align="center")
    except TypeError:
        image_draw.text(label_coords, text, fill=text_color_rgb, font=font, align="center")


def draw_bbox_on_image(
    image_draw: ImageDraw.ImageDraw,
    bbox: BBox,
    color: str,
):
    """Draw bbox with additional labels on the image..

    Args:
        image_draw:     ImageDraw object to draw on the image.
        bbox:           Bounding box to draw.
        color:          RGB values of the color of the bounding box (edges + label backgrounds).
    """
    x1, y1, x2, y2 = bbox.points
    if x1 >= x2 or y1 >= y2:
        print(f"Invalid bbox coordinates: {bbox.points}")
        return
    top_left = x1, y1  # the main
    bottom_right = x2, y2
    box_thickness = get_bbox_thickness(bbox=bbox.points, page_size=image_draw.im.size)
    font_size = get_bbox_text_size(bbox=bbox.points, page_size=image_draw.im.size)

    try:
        image_draw.rectangle((top_left, bottom_right), outline=color, width=box_thickness)
    except TypeError:
        rgb_color = get_rgb_color(color)
        image_draw.rectangle((top_left, bottom_right), outline=rgb_color, width=box_thickness)

    if bbox.labels is not None:
        if top_left_label := bbox.labels.top_left:
            draw_bbox_label(
                image_draw,
                top_left_label,
                bbox_points=bbox.points,
                alignment=TextAlignment.TOP_LEFT,
                font_size=font_size,
                background_color=color,
            )
        if top_right_label := bbox.labels.top_right:
            draw_bbox_label(
                image_draw,
                top_right_label,
                bbox_points=bbox.points,
                alignment=TextAlignment.TOP_RIGHT,
                font_size=font_size,
                background_color=color,
            )
        if bottom_left_label := bbox.labels.bottom_left:
            draw_bbox_label(
                image_draw,
                bottom_left_label,
                bbox_points=bbox.points,
                alignment=TextAlignment.BOTTOM_LEFT,
                font_size=font_size,
                background_color=color,
            )
        if bottom_right_label := bbox.labels.bottom_right:
            draw_bbox_label(
                image_draw,
                bottom_right_label,
                bbox_points=bbox.points,
                alignment=TextAlignment.BOTTOM_RIGHT,
                font_size=font_size,
                background_color=color,
            )
        if center_label := bbox.labels.center:
            draw_bbox_label(
                image_draw,
                center_label,
                bbox_points=bbox.points,
                alignment=TextAlignment.CENTER,
                font_size=font_size * 2,
                background_color=color,
            )


@njit(cache=True, fastmath=True)
def _linear_polyfit_2point(x0: float, x1: float, y0: float, y1: float):
    """Compute slope and intercept for a line passing through (x0, y0), (x1, y1)."""
    if x1 == x0:
        slope = 0.0
        intercept = (y0 + y1) / 2.0
    else:
        slope = (y1 - y0) / (x1 - x0)
        intercept = y0 - slope * x0
    return slope, intercept


class LayoutDrawer(ABC):
    layout_source: str = "unknown"
    laytout_dump: dict

    def __init__(self, layout_dump: dict):
        self.layout_dump = layout_dump

    def draw_layout_on_page(self, page_image: Image.Image, page_num: int) -> Image.Image:
        """Draw the layout bboxes with additional metadata on the image."""
        layout_pages = self.layout_dump.get("pages")
        if not layout_pages:
            print(f"Warning: layout in drawer {self.__class__.__name__} is empty - skipping")
            return page_image
        if len(layout_pages) < page_num:
            print(f"Error! Page {page_num} not found in layout (pages: {len(layout_pages)})")
            return page_image
        image_draw = ImageDraw.ImageDraw(page_image)
        page_layout_dump = layout_pages[page_num - 1]
        if page_num != page_layout_dump.get("number"):
            dump_page_num = page_layout_dump.get("number")
            print(f"Warning: Requested page num {page_num} differs from dump {dump_page_num}")
        for idx, elements in enumerate(page_layout_dump["elements"], 1):
            self.render_element_on_page(idx, image_draw, elements)
        return page_image

    @abstractmethod
    def render_element_on_page(self, idx: int, image_draw: ImageDraw, elements: dict[str, Any]):
        """Draw a single element on the image."""


class SimpleLayoutDrawer(LayoutDrawer, ABC):
    color: str
    show_order: bool = False
    show_text_length: bool = False

    def render_element_on_page(self, idx: int, image_draw: ImageDraw, elements: dict[str, Any]):
        text_len = len(elements["text"]) if elements.get("text") else 0
        element_prob = elements.get("prob")
        element_order = f"{idx}" if self.show_order else None
        text_len = f"len: {text_len}" if self.show_text_length else None
        bbox = BBox(
            points=elements["bbox"],
            labels=BboxLabels(
                top_right=f"prob: {element_prob:.2f}" if element_prob else None,
                bottom_left=text_len,
                center=element_order,
            ),
        )
        draw_bbox_on_image(image_draw, bbox, color=self.color)


class PdfminerLayoutDrawer(SimpleLayoutDrawer):
    layout_source = "pdfminer"

    def __init__(self, layout_dump: dict, color: str = "red"):
        self.layout_dump = layout_dump
        self.color = color
        self.show_order = True
        super().__init__(layout_dump)


class OCRLayoutDrawer(SimpleLayoutDrawer):
    layout_source = "ocr"

    def __init__(self, layout_dump: dict, color: str = "red"):
        self.color = color
        self.show_order = False
        self.show_text_length = False
        super().__init__(layout_dump)


class ODModelLayoutDrawer(LayoutDrawer):
    layout_source = "od_model"

    color_map = {
        ElementType.CAPTION: "salmon",
        ElementType.FOOTNOTE: "orange",
        ElementType.FORMULA: "mediumpurple",
        ElementType.LIST_ITEM: "navy",
        ElementType.PAGE_FOOTER: "deeppink",
        ElementType.PAGE_HEADER: "green",
        ElementType.PICTURE: "sienna",
        ElementType.SECTION_HEADER: "darkorange",
        ElementType.TABLE: "blue",
        ElementType.TEXT: "turquoise",
        ElementType.TITLE: "greenyellow",
    }

    def render_element_on_page(self, idx: int, image_draw: ImageDraw, elements: dict[str, Any]):
        element_type = elements["type"]
        element_prob = elements.get("prob")
        bbox_points = elements["bbox"]
        color = self.get_element_type_color(element_type)
        bbox = BBox(
            points=bbox_points,
            labels=BboxLabels(
                top_left=f"{element_type}",
                top_right=f"prob: {element_prob:.2f}" if element_prob else None,
            ),
        )
        draw_bbox_on_image(image_draw, bbox, color=color)

    def get_element_type_color(self, element_type: str) -> str:
        return self.color_map.get(element_type, "cyan")


class FinalLayoutDrawer(LayoutDrawer):
    layout_source = "final"

    color_map = {
        "CheckBox": "brown",
        "ListItem": "red",
        "Title": "greenyellow",
        "NarrativeText": "turquoise",
        "Header": "green",
        "Footer": "orange",
        "FigureCaption": "sienna",
        "Image": "sienna",
        "Table": "blue",
        "Address": "gold",
        "EmailAddress": "lightskyblue",
        "Formula": "mediumpurple",
        "CodeSnippet": "magenta",
        "PageNumber": "crimson",
    }

    def __init__(self, layout_dump: dict):
        self.layout_dump = layout_dump

    def render_element_on_page(self, idx: int, image_draw: ImageDraw, elements: dict[str, Any]):
        element_type = elements["type"]
        element_prob = elements.get("prob")
        text_len = len(elements["text"]) if elements.get("text") else 0
        bbox_points = elements["bbox"]
        color = self.get_element_type_color(element_type)
        cluster = elements.get("cluster")
        bbox = BBox(
            points=bbox_points,
            labels=BboxLabels(
                top_left=f"{element_type}",
                top_right=f"prob: {element_prob:.2f}" if element_prob else None,
                bottom_right=f"len: {text_len}",
                bottom_left=f"cl: {cluster}" if cluster else None,
                center=f"{idx}",
            ),
        )
        draw_bbox_on_image(image_draw, bbox, color=color)

    def get_element_type_color(self, element_type: str) -> str:
        return self.color_map.get(element_type, "cyan")


class AnalysisDrawer(AnalysisProcessor):
    def __init__(
        self,
        filename: Optional[Union[str, Path]],
        is_image: bool,
        save_dir: Union[str, Path],
        file: Optional[BytesIO] = None,
        draw_caption: bool = True,
        draw_grid: bool = False,
        resize: Optional[float] = None,
        format: str = "png",
    ):
        self.draw_caption = draw_caption
        self.draw_grid = draw_grid
        self.resize = resize
        self.is_image = is_image
        self.format = format
        self.drawers = []
        self.file = file

        super().__init__(filename, save_dir)

    def add_drawer(self, drawer: LayoutDrawer):
        self.drawers.append(drawer)

    def process(self):
        filename_stem = Path(self.filename).stem
        analysis_save_dir = Path(self.save_dir) / "analysis" / filename_stem / "bboxes"
        analysis_save_dir.mkdir(parents=True, exist_ok=True)

        for page_idx, orig_image_page in enumerate(self.load_source_image()):
            images_for_grid = []
            page_num = page_idx + 1
            for drawer in self.drawers:
                try:
                    image = drawer.draw_layout_on_page(orig_image_page.copy(), page_num=page_num)
                except:  # noqa: E722
                    logging.exception(
                        f"Error while drawing layout for page {page_num} "
                        f"for file {self.filename} with drawer "
                        f"{drawer.__class__.__name__}"
                    )
                    continue
                if self.draw_caption:
                    image = self.add_caption(
                        image, caption=f"Layout source: {drawer.layout_source}"
                    )
                if not self.draw_grid:
                    if self.resize is not None:
                        image = image.resize(
                            (int(image.width * self.resize), int(image.height * self.resize)),
                        )
                    image.save(
                        analysis_save_dir / f"page{page_num}"
                        f"_layout_{drawer.layout_source}.{self.format}",
                        optimize=True,
                        quality=85,
                    )
                    image.close()
                else:
                    images_for_grid.append(image)
            if images_for_grid:
                grid_image = self.paste_images_on_grid(images_for_grid)
                if self.resize is not None:
                    grid_image = grid_image.resize(
                        (int(grid_image.width * self.resize), int(grid_image.height * self.resize))
                    )
                grid_image.save(
                    analysis_save_dir / f"page{page_num}_layout_all.{self.format}",
                    optimize=True,
                    quality=85,
                )
                grid_image.close()

    def add_caption(self, image: Image.Image, caption: str):
        font = ImageFont.truetype(get_font(), 52)
        draw = ImageDraw.ImageDraw(image)
        text_x1, text_y1, text_x2, text_y2 = draw.textbbox(
            (0, 0), caption, font=font, align="center"
        )
        text_width = text_x2 - text_x1
        text_height = int((text_y2 - text_y1) * 1.5)
        text_xy = (image.width - text_width) // 2, 10
        caption_image = Image.new("RGB", (image.width, text_height), color=(255, 255, 255))
        caption_draw = ImageDraw.ImageDraw(caption_image)
        caption_draw.text(text_xy, caption, (0, 0, 0), font=font)

        expanded_image = Image.new("RGB", (image.width, image.height + text_height))
        expanded_image.paste(caption_image, (0, 0))
        expanded_image.paste(image, (0, text_height))
        image.close()
        return expanded_image

    def paste_images_on_grid(self, images: List[Image.Image]) -> Image.Image:
        """Creates a single image that presents all the images on a grid 2 x n/2"""

        pairs = []
        for i in range(0, len(images), 2):
            left_image = images[i]
            right_image = images[i + 1] if i < len(images) - 1 else None
            pairs.append((left_image, right_image))

        max_pair_width = max([pair[0].width + (pair[1].width if pair[1] else 0) for pair in pairs])
        sum_height = sum([max(pair[0].height, pair[1].height if pair[1] else 0) for pair in pairs])

        new_im = Image.new("RGB", (max_pair_width, sum_height))

        height_shift = 0
        for image_left, image_right in pairs:
            new_im.paste(image_left, (0, height_shift))
            if image_right:
                new_im.paste(image_right, (image_left.width, height_shift))
            height_shift += max(image_left.height, image_right.height if image_right else 0)

        for image in images:
            image.close()
        return new_im

    def load_source_image(self) -> Generator[Image.Image, None, None]:
        with tempfile.TemporaryDirectory() as temp_dir:
            image_paths = []
            if self.is_image:
                if self.file:
                    try:
                        image = Image.open(self.file)
                        output_file = Path(temp_dir) / self.filename
                        image.save(output_file, format="PNG")
                        image_paths = [output_file]
                    except Exception as ex:  # noqa: E722
                        print(
                            f"Error while converting image to PNG for file {self.filename}, "
                            f"exception: {ex}"
                        )
                else:
                    image_paths = [self.filename]
            else:
                try:
                    str_filename = (
                        ""
                        if (self.filename is None or self.file is not None)
                        else (
                            str(self.filename) if isinstance(self.filename, Path) else self.filename
                        )
                    )
                    image_paths = convert_pdf_to_image(
                        filename=str_filename,
                        file=self.file,
                        output_folder=temp_dir,
                        path_only=True,
                    )
                except Exception as ex:  # noqa: E722
                    print(
                        f"Error while converting pdf to image for file {self.filename}",
                        f"exception: {ex}",
                    )

            for image_path in image_paths:
                with Image.open(str(image_path)) as image:
                    yield image.convert("RGB")