utils.py

import logging
import string
from collections import defaultdict
from typing import Any, List, Union

import cv2
import numpy as np
import torch
from doctr.io.elements import Document
from doctr.models import parseq
from doctr.models._utils import get_language
from doctr.models.detection.predictor import DetectionPredictor
from doctr.models.detection.zoo import detection_predictor
from doctr.models.predictor.base import _OCRPredictor
from doctr.models.recognition.predictor import RecognitionPredictor
from doctr.models.recognition.zoo import recognition_predictor
from doctr.utils.geometry import detach_scores
from PIL import Image, ImageDraw, ImageFont
from sklearn.cluster import DBSCAN
from sklearn.preprocessing import StandardScaler
from torch import nn

confidence_threshold = 0.75

reco_arch = "printed_v19.pt"
det_arch = "fast_base"

# Configure logging
afterword_symbols = "!?.,:;"
numbers = "0123456789"
other_symbols = string.punctuation + "«»…£€¥¢฿₸₽№°—"
space_symbol = " "
kazakh_letters = "АБВГДЕЖЗИЙКЛМНОПРСТУФХЦЧШЩЪЫЬЭЮЯЁабвгдежзийклмнопрстуфхцчшщъыьэюяёӘҒҚҢӨҰҮІҺәғқңөұүіһ"
english_letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
all_letters = kazakh_letters + english_letters
all_symbols = numbers + other_symbols + space_symbol + all_letters


def get_ocr_predictor(
    det_arch: str = det_arch,
    reco_arch: str = reco_arch,
    pretrained=True,
    pretrained_backbone: bool = True,
    assume_straight_pages: bool = False,
    preserve_aspect_ratio: bool = True,
    symmetric_pad: bool = True,
    det_bs: int = 2,
    reco_bs: int = 128,
    detect_orientation: bool = False,
    straighten_pages: bool = False,
    detect_language: bool = False,
    bin_thresh: float = 0.3,
    box_thresh: float = 0.3,
):
    device = "cpu"
    if torch.backends.mps.is_available():
        device = "mps"
    elif torch.cuda.is_available():
        device = "cuda"
    else:
        device = "cpu"
    logging.info(f"Using device: {device}")

    device = torch.device(device)

    # Initialize predictor
    logging.info(f"Initializing predictor with device: {device}")
    reco_model = parseq(pretrained=False, pretrained_backbone=False, vocab=all_symbols)
    reco_model.to(device)
    reco_params = torch.load(f"./custom/{reco_arch}", map_location=device)
    reco_model.load_state_dict(reco_params)

    det_predictor = detection_predictor(
        det_arch,
        pretrained=pretrained,
        pretrained_backbone=pretrained_backbone,
        batch_size=det_bs,
        assume_straight_pages=assume_straight_pages,
        preserve_aspect_ratio=preserve_aspect_ratio,
        symmetric_pad=symmetric_pad,
    )

    # Recognition
    reco_predictor = recognition_predictor(
        reco_model,
        pretrained=pretrained,
        pretrained_backbone=pretrained_backbone,
        batch_size=reco_bs,
    )

    predictor = OCRPredictor(
        det_predictor,
        reco_predictor,
        assume_straight_pages=assume_straight_pages,
        preserve_aspect_ratio=preserve_aspect_ratio,
        symmetric_pad=symmetric_pad,
        detect_orientation=detect_orientation,
        straighten_pages=straighten_pages,
        detect_language=detect_language,
    )

    predictor.det_predictor.model.postprocessor.bin_thresh = bin_thresh
    predictor.det_predictor.model.postprocessor.box_thresh = box_thresh
    predictor.add_hook(CustomHook())

    return predictor


class OCRPredictor(nn.Module, _OCRPredictor):
    """Implements an object able to localize and identify text elements in a set of documents

    Args:
    ----
        det_predictor: detection module
        reco_predictor: recognition module
        assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages
            without rotated textual elements.
        straighten_pages: if True, estimates the page general orientation based on the median line orientation.
            Then, rotates page before passing it to the deep learning modules. The final predictions will be remapped
            accordingly. Doing so will improve performances for documents with page-uniform rotations.
        detect_orientation: if True, the estimated general page orientation will be added to the predictions for each
            page. Doing so will slightly deteriorate the overall latency.
        detect_language: if True, the language prediction will be added to the predictions for each
            page. Doing so will slightly deteriorate the overall latency.
        **kwargs: keyword args of `DocumentBuilder`
    """

    def __init__(
        self,
        det_predictor: DetectionPredictor,
        reco_predictor: RecognitionPredictor,
        assume_straight_pages: bool = True,
        straighten_pages: bool = False,
        preserve_aspect_ratio: bool = True,
        symmetric_pad: bool = True,
        detect_orientation: bool = False,
        detect_language: bool = False,
        **kwargs: Any,
    ) -> None:
        nn.Module.__init__(self)
        self.det_predictor = det_predictor.eval()  # type: ignore[attr-defined]
        self.reco_predictor = reco_predictor.eval()  # type: ignore[attr-defined]
        _OCRPredictor.__init__(
            self,
            assume_straight_pages,
            straighten_pages,
            preserve_aspect_ratio,
            symmetric_pad,
            detect_orientation,
            **kwargs,
        )
        self.detect_orientation = detect_orientation
        self.detect_language = detect_language

    @torch.inference_mode()
    def forward(
        self,
        pages: List[Union[np.ndarray, torch.Tensor]],
        **kwargs: Any,
    ) -> Document:
        # Dimension check
        if any(page.ndim != 3 for page in pages):
            raise ValueError(
                "incorrect input shape: all pages are expected to be multi-channel 2D images."
            )

        origin_page_shapes = [
            page.shape[:2] if isinstance(page, np.ndarray) else page.shape[-2:]
            for page in pages
        ]

        # Localize text elements
        loc_preds, out_maps = self.det_predictor(pages, return_maps=True, **kwargs)

        # Detect document rotation and rotate pages
        seg_maps = [
            np.where(
                out_map > getattr(self.det_predictor.model.postprocessor, "bin_thresh"),
                255,
                0,
            ).astype(np.uint8)
            for out_map in out_maps
        ]
        if self.detect_orientation:
            general_pages_orientations, origin_pages_orientations = self._get_orientations(pages, seg_maps)  # type: ignore[arg-type]
            orientations = [
                {"value": orientation_page, "confidence": None}
                for orientation_page in origin_pages_orientations
            ]
        else:
            orientations = None
            general_pages_orientations = None
            origin_pages_orientations = None
        if self.straighten_pages:
            pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)  # type: ignore
            # Forward again to get predictions on straight pages
            loc_preds = self.det_predictor(pages, **kwargs)

        assert all(
            len(loc_pred) == 1 for loc_pred in loc_preds
        ), "Detection Model in ocr_predictor should output only one class"

        loc_preds = [list(loc_pred.values())[0] for loc_pred in loc_preds]
        # Detach objectness scores from loc_preds
        loc_preds, objectness_scores = detach_scores(loc_preds)
        # Check whether crop mode should be switched to channels first
        channels_last = len(pages) == 0 or isinstance(pages[0], np.ndarray)

        # Apply hooks to loc_preds if any
        for hook in self.hooks:
            loc_preds = hook(loc_preds)

        # Crop images
        crops, loc_preds = self._prepare_crops(
            pages,  # type: ignore[arg-type]
            loc_preds,
            channels_last=channels_last,
            assume_straight_pages=self.assume_straight_pages,
        )
        # Rectify crop orientation and get crop orientation predictions
        crop_orientations: Any = []
        # save crops to ./crops
        # os.makedirs("./crops", exist_ok=True)
        # for i, crop in enumerate(crops[0]):
        #     Image.fromarray(crop).save(f"./crops/{i}.png")

        # if not self.assume_straight_pages:
        # crops, loc_preds, _crop_orientations = self._rectify_crops(crops, loc_preds)
        #     crop_orientations = [
        #         {"value": orientation[0], "confidence": orientation[1]} for orientation in _crop_orientations
        #     ]

        # Identify character sequences
        word_preds = self.reco_predictor(
            [crop for page_crops in crops for crop in page_crops], **kwargs
        )
        if not crop_orientations:
            crop_orientations = [{"value": 0, "confidence": None} for _ in word_preds]

        boxes, text_preds, crop_orientations = self._process_predictions(
            loc_preds, word_preds, crop_orientations
        )

        if self.detect_language:
            languages = [
                get_language(" ".join([item[0] for item in text_pred]))
                for text_pred in text_preds
            ]
            languages_dict = [
                {"value": lang[0], "confidence": lang[1]} for lang in languages
            ]
        else:
            languages_dict = None

        out = self.doc_builder(
            pages,  # type: ignore[arg-type]
            boxes,
            objectness_scores,
            text_preds,
            origin_page_shapes,  # type: ignore[arg-type]
            crop_orientations,
            orientations,
            languages_dict,
        )
        return out


class CustomHook:
    def __call__(self, loc_preds):
        # Manipulate the location predictions here
        # 1. The outpout structure needs to be the same as the input location predictions
        # 2. Be aware that the coordinates are relative and needs to be between 0 and 1

        # return np.array([self.order_bbox_points(point) for loc_pred in loc_preds for point in loc_pred ])
        # iterate over each page and each box
        answer = []
        for page_idx, page_boxes in enumerate(loc_preds):
            bboxes = []
            for box_idx, box in enumerate(page_boxes):
                box = self.order_bbox_points(box)
                bboxes.append(box)
            answer.append(bboxes)
        return np.array(answer)

    def order_bbox_points(self, points):
        """
        Orders a list of four (x, y) points in the following order:
        top-left, top-right, bottom-right, bottom-left.

        Args:
            points (list of tuples): List of four (x, y) tuples.

        Returns:
            list of tuples: Ordered list of four (x, y) tuples.
        """
        if len(points) != 4:
            raise ValueError(
                "Exactly four points are required to define a quadrilateral."
            )

        # Convert points to NumPy array for easier manipulation
        pts = np.array(points)

        # Compute the sum and difference of the points
        sum_pts = pts.sum(axis=1)
        diff_pts = np.diff(pts, axis=1).flatten()

        # Initialize ordered points list
        ordered = [None] * 4

        # Top-Left point has the smallest sum
        ordered[0] = tuple(pts[np.argmin(sum_pts)])

        # Bottom-Right point has the largest sum
        ordered[2] = tuple(pts[np.argmax(sum_pts)])

        # Top-Right point has the smallest difference
        ordered[1] = tuple(pts[np.argmin(diff_pts)])

        # Bottom-Left point has the largest difference
        ordered[3] = tuple(pts[np.argmax(diff_pts)])

        return ordered


def geometry_to_coordinates(geometry, img_width, img_height):
    if len(geometry) == 2:
        (x0_rel, y0_rel), (x1_rel, y1_rel) = geometry
        x0 = int(x0_rel * img_width)
        y0 = int(y0_rel * img_height)
        x1 = int(x1_rel * img_width)
        y1 = int(y1_rel * img_height)
        # Bounding box with four corners
        all_four = [[x0, y0], [x1, y0], [x1, y1], [x0, y1]]
        return all_four
    else:
        # Bounding box with four corners
        all_four = [[int(x * img_width), int(y * img_height)] for x, y in geometry]
        return all_four


def page_to_coordinates(page_export):
    coordinates = []
    img_height, img_width = page_export["dimensions"]
    for block in page_export["blocks"]:
        for line in block["lines"]:
            for word in line["words"]:
                if (
                    word["confidence"] < confidence_threshold
                    and len(word["value"].strip()) > 1
                ):
                    logging.warning(
                        f"Skipping word with low confidence: {word['value']} confidence {word['confidence']}"
                    )
                    continue
                all_four = geometry_to_coordinates(
                    word["geometry"], img_width, img_height
                )
                coordinates.append((all_four, word["value"], word["confidence"]))

    return (coordinates, img_width, img_height)


def draw_boxes_with_labels(image, coordinates, font_path):
    """Бастапқы суретке шекаралар үстіне кішкентай белгілерді қою.

    Args:
        image: Бастапқы сурет (numpy массиві).
        out: predictor([image]) нәтижесі.
        font_path: TrueType қаріп файлының жолы.

    Returns:
        Шекаралар және белгілер қойылған сурет.
    """

    # Суретті PIL форматына түрлендіреміз
    img_with_boxes = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    img_pil = Image.fromarray(img_with_boxes)
    draw = ImageDraw.Draw(img_pil)

    for coords, word, score in coordinates:
        # poligon
        coords = [(x, y) for x, y in coords]
        text_x, text_y = (
            min(coords, key=lambda x: x[0])[0],
            min(coords, key=lambda x: x[1])[1],
        )
        draw.polygon(coords, outline=(0, 255, 0, 125), width=1)
        font = ImageFont.truetype(font_path, 10)
        draw.text((text_x, max(text_y - 10, 0)), word, font=font, fill=(255, 0, 0))

    # Суретті қайтадан OpenCV форматына түрлендіреміз
    img_with_boxes = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

    # Суретті қайтарамыз
    return img_with_boxes


def generate_line_points(bbox, num_points=10):
    """
    Generates multiple points along the line connecting the left and right centers of a bounding box.

    Parameters:
    - bbox: List of four points [[x0, y0], [x1, y1], [x2, y2], [x3, y3]]
            in the order: TopLeft, TopRight, BottomRight, BottomLeft.
    - num_points: Number of points to generate along the line.

    Returns:
    - List of (x, y) tuples.
    """
    # Calculate left center (midpoint of TopLeft and BottomLeft)
    left_center_x = (bbox[0][0] + bbox[3][0]) / 2
    left_center_y = (bbox[0][1] + bbox[3][1]) / 2

    # Calculate right center (midpoint of TopRight and BottomRight)
    right_center_x = (bbox[1][0] + bbox[2][0]) / 2
    right_center_y = (bbox[1][1] + bbox[2][1]) / 2

    # Generate linearly spaced points between left center and right center
    x_values = np.linspace(left_center_x, right_center_x, num_points)
    y_values = np.linspace(left_center_y, right_center_y, num_points)

    points = list(zip(x_values, y_values))
    return points


def ocr_to_txt(coordinates):
    """
    Converts OCR output to a structured text file with lines using multiple points along connecting lines.
    Inserts empty lines when there's significant vertical spacing between lines.

    Parameters:
    - coordinates: List of tuples containing bounding box coordinates, word value, and score.
                   Each tuple is (([[x0, y0], [x1, y1], [x2, y2], [x3, y3]]), word, score)
    - img_width: Width of the image in pixels.
    - img_height: Height of the image in pixels.
    - output_file: Path to the output text file.
    """
    # Step 1: Compute multiple points for each word
    all_points = []
    words = []
    scaler = StandardScaler()
    points_per_word = 25  # Number of points to generate per word

    for bbox, word, score in coordinates:
        points = generate_line_points(bbox, num_points=points_per_word)
        all_points.extend(points)
        words.append(
            {
                "bbox": bbox,
                "word": word,
                "score": score,
                "points": points,  # Store the multiple points
            }
        )

    # Step 2: Scale the points
    scaled_points = scaler.fit_transform(all_points)
    scaled_points = [(c[0] / 5, c[1]) for c in scaled_points]
    scaled_points = np.array(scaled_points)

    # Step 3: Cluster points using DBSCAN
    # Parameters for DBSCAN can be tuned based on the specific OCR output
    # eps determines the maximum distance between two samples for them to be considered as in the same neighborhood
    # min_samples is set to the number of points per word to ensure entire words are clustered together
    db = DBSCAN(min_samples=2, eps=0.05).fit(scaled_points)  # eps might need adjustment
    labels = db.labels_

    # Map each point to its cluster label
    point_labels = labels.tolist()

    # Step 4: Assign words to clusters based on their points
    label_to_words = defaultdict(list)
    current_point = 0  # To keep track of which point belongs to which word

    for word in words:
        word_labels = point_labels[current_point : current_point + points_per_word]
        current_point += points_per_word

        # Count the frequency of each label in the word's points
        label_counts = defaultdict(int)
        for lbl in word_labels:
            label_counts[lbl] += 1

        # Assign the word to the most frequent label
        # If multiple labels have the same highest count, choose the smallest label (ignoring -1 for noise)
        if label_counts:
            # Exclude noise label (-1) when possible
            filtered_labels = {k: v for k, v in label_counts.items() if k != -1}
            if filtered_labels:
                assigned_label = max(filtered_labels, key=filtered_labels.get)
            else:
                assigned_label = -1  # Assign to noise
            label_to_words[assigned_label].append(word)

    # Remove noise cluster if present
    if -1 in label_to_words:
        print(
            f"Warning: {len(label_to_words[-1])} words assigned to noise cluster and will be ignored."
        )
        del label_to_words[-1]

    # Step 5: Sort words within each line
    sorted_lines = []
    line_heights = []  # To store heights of each line for median calculation
    line_y_bounds = []  # To store min and max y for each line

    for label, line_words in label_to_words.items():
        # Sort words based on their leftmost x-coordinate
        line_words_sorted = sorted(
            line_words, key=lambda w: min(point[0] for point in w["points"])
        )
        sorted_lines.append(line_words_sorted)

        # Compute y-bounds for the line
        y_values = []
        for word in line_words_sorted:
            y_coords = [point[1] for point in word["bbox"]]
            y_min = min(y_coords)
            y_max = max(y_coords)
            y_values.append([y_min, y_max])
        y_values = np.array(y_values)
        # Compute the median y-coordinates for the line by sorting only with y_min
        line_min_y_median = np.median(y_values[:, 0])
        line_max_y_median = np.median(y_values[:, 1])
        line_heights.append(line_max_y_median - line_min_y_median)
        line_y_bounds.append((line_min_y_median, line_max_y_median))

    # Step 6: Sort lines from top to bottom based on the average y-coordinate of their words
    sorted_lines, line_heights, line_y_bounds = zip(
        *sorted(
            zip(sorted_lines, line_heights, line_y_bounds),
            key=lambda item: np.median(
                [np.mean([p[1] for p in w["bbox"]]) for w in item[0]]
            ),
        )
    )

    sorted_lines = list(sorted_lines)
    line_heights = list(line_heights)
    line_y_bounds = list(line_y_bounds)

    # Step 8: Write sorted lines to the output text file with empty lines where necessary
    output_text = ""
    previous_line_median_y = None  # To track the max y of the previous line

    for idx, line in enumerate(sorted_lines):
        # Compute current line's min y
        current_line_min_y_median = line_y_bounds[idx][0]
        current_line_max_y_median = line_y_bounds[idx][1]
        current_line_median_height = line_heights[idx]
        current_line_median_y = (
            current_line_min_y_median + current_line_max_y_median
        ) / 2

        if previous_line_median_y is not None:
            # Compute vertical distance between lines
            vertical_distance = current_line_median_y - previous_line_median_y
            median_height = (
                current_line_median_height + previous_line_median_height
            ) / 2

            # If the vertical distance is greater than the median height, insert an empty line
            if vertical_distance > median_height * 2:
                output_text += "\n"  # Insert empty line

        # Write the current line's text
        line_text = " ".join([w["word"] for w in line])
        output_text += line_text + "\n"

        # Update the previous_line_max_y for the next iteration
        previous_line_median_y = current_line_median_y
        previous_line_median_height = current_line_median_height

    return output_text