src/ocr.py

import re
import gc
import torch

from PIL import Image

from transformers import (
    AutoProcessor,
    AutoModelForImageTextToText
)


MODEL_ID = "PaddlePaddle/PaddleOCR-VL-1.5"


DEVICE = (
    "cuda"
    if torch.cuda.is_available()
    else "cpu"
)


ocr_processor = AutoProcessor.from_pretrained(
    MODEL_ID
)


ocr_model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    torch_dtype=(
        torch.bfloat16
        if torch.cuda.is_available()
        else torch.float32
    )
).to(DEVICE).eval()


def extract_text_with_paddleocr_vl_spotting(image):

    image = image.convert("RGB")


    original_width, original_height = image.size


    if (
        original_width < 1500
        and original_height < 1500
    ):

        image = image.resize(
            (
                original_width * 2,
                original_height * 2
            ),
            Image.Resampling.LANCZOS
        )


    max_pixels = 2048 * 28 * 28


    image_processor = (
        ocr_processor.image_processor
    )


    if hasattr(
        image_processor,
        "size"
    ):

        size_config = image_processor.size

        if isinstance(
            size_config,
            dict
        ):

            min_pixels = size_config.get(
                "shortest_edge",
                28 * 28 * 130
            )

        else:

            min_pixels = 28 * 28 * 130

    else:

        min_pixels = 28 * 28 * 130


    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image
                },
                {
                    "type": "text",
                    "text": "Spotting:"
                }
            ]
        }
    ]


    inputs = ocr_processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        images_kwargs={
            "size": {
                "shortest_edge": min_pixels,
                "longest_edge": max_pixels
            }
        }
    )


    inputs = {
        key: value.to(ocr_model.device)
        if hasattr(value, "to")
        else value
        for key, value in inputs.items()
    }


    with torch.inference_mode():

        outputs = ocr_model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=False
        )


    input_length = inputs[
        "input_ids"
    ].shape[-1]


    generated_tokens = outputs[
        0,
        input_length:
    ]


    raw_result = ocr_processor.decode(
        generated_tokens,
        skip_special_tokens=True
    ).strip()


    lines = []


    for line in raw_result.splitlines():

        line = line.strip()

        if not line:
            continue

        line = re.sub(
            r"<\|LOC_\d+\|>",
            "",
            line
        )

        line = line.strip()

        if line:
            lines.append(line)


    cleaned_lines = []


    for line in lines:

        if (
            not cleaned_lines
            or line != cleaned_lines[-1]
        ):

            cleaned_lines.append(line)


    del inputs
    del outputs

    gc.collect()


    if torch.cuda.is_available():

        torch.cuda.empty_cache()


    return cleaned_lines