ocr_tr_bls/1/model.py

import json
import os
import time
import sys

import cv2
import numpy as np
from shapely.geometry import JOIN_STYLE, Polygon
from torch.utils.dlpack import from_dlpack
import triton_python_backend_utils as pb_utils  # type: ignore
if "/" not in sys.path:
    sys.path.insert(0, "/")
from utils import align_box_points


class CV2FixedNormalizePAD:
    def __init__(self, width, height):
        self.width = width
        self.height = height

    def __call__(self, im):
        h, w = im.shape[:2]
        assert (h == self.height)

        if w < self.width:
            im = im.astype(np.float32)
            im = im / 255.0
            padded_im = np.zeros((self.height, self.width), dtype=np.float32)
            padded_im[:, :w] = im
            padded_im[:, w:] = im[:, w - 1].reshape(32, 1)

        elif w > self.width:
            im = cv2.resize(im, (self.width, self.height), interpolation=cv2.INTER_AREA)
            im = im.astype(np.float32)
            padded_im = im / 255.0
        else:
            im = im.astype(np.float32)
            padded_im = im / 255.0
        return padded_im.reshape(1, 1, self.height, self.width)


class TritonPythonModel:
    """
    OCR Text Recognition BLS Model with Dynamic Batch Splitting
    
    Combines ocr.td.postprocess and ocr.tr.postprocess logic with automatic batch splitting
    for handling large numbers of detected text boxes.
    """

    def initialize(self, args):
        self.logger = pb_utils.Logger
        self.logger.log_info("Initializing ocr.tr.bls model with dynamic batch splitting")

        self.text_box_margin = 0.075
        self.use_auto_rotate = False

        self.tr_canvas_size_w = 256
        self.tr_canvas_size_h = 32
        self.tr_use_UNK = True
        self.tr_use_SPACE = False
        self.tr_transform = CV2FixedNormalizePAD(self.tr_canvas_size_w, self.tr_canvas_size_h)

        # Load character labels
        character_filename = 'kr_labels.txt'
        base_path = args['model_repository']
        base_path = os.path.join(base_path, "1")
        character_path = os.path.join(base_path, character_filename)

        self.tr_characters = ['[CTCblank]']
        self.load_characters(character_path)




        # # Text detection postprocess settings
        # self.text_threshold = 0.7
        # self.low_text = 0.4
        # self.link_threshold = 0.35
        # self.text_box_margin = 0.075
        # self.text_box_height = 32
        # self.use_auto_rotate = False

        # img_max_height = 32
        # img_max_width = 256
        # self.transform = CV2FixedNormalizePAD(img_max_width, img_max_height)

        # Batch splitting settings
        self.max_batch_size = 1 # Tune this based on available GPU memory

        # # Text recognition postprocess settings
        # self.use_space = False
        # self.use_UNK = True
        # self.batch_max_length = 32

        # # Load character labels
        # character_filename = 'kr_labels.txt'
        # base_path = args['model_repository']
        # base_path = os.path.join(base_path, "1")
        # character_path = os.path.join(base_path, character_filename)

        # Get output config
        model_config = json.loads(args["model_config"])
        text_list_config = pb_utils.get_output_config_by_name(model_config, "text_list")
        self.text_list_type = pb_utils.triton_string_to_numpy(text_list_config["data_type"])

        # # Load characters
        # character = []
        # with open(character_path, "r", encoding="utf-8-sig") as f:
        #     lines = f.readlines()
        #     for line in lines:
        #         ch = line.strip().split()[1]
        #         character.append(ch)

        # if self.use_space:
        #     character.append(' ')

        # if self.use_UNK:
        #     character.append('')

        # self.converter = CTCLabelConverter(character)


    def load_characters(self, tr_label_path):
        with open(tr_label_path, encoding='utf-8') as f:
            lines = f.readlines()

        for line in lines:
            character = line.strip().split('\t')[1]
            self.tr_characters.append(character)

        if self.tr_use_SPACE:
            self.tr_characters.append(' ')

        if self.tr_use_UNK:
            self.tr_characters.append('')


    def execute(self, requests):
        responses = []
        st = time.time()

        for request in requests:
            try:
                original_image = pb_utils.get_input_tensor_by_name(request, "original_image").as_numpy()
                detected_text_box_list = pb_utils.get_input_tensor_by_name(request, "detected_text_box_list").as_numpy()
                self.logger.log_info(f"box 개수: {detected_text_box_list.shape}")

                gray_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2GRAY)

                image_patch_list, adjusted_box_list = self.crop_with_text_box_list(gray_image, detected_text_box_list)
                self.logger.log_info(f"box 개수: {len(adjusted_box_list)}")
                batch_images = self.tr_preprocess(image_patch_list)
                #batch_images = batch_images.astype(np.float16)

                if batch_images is not None:
                    tr_results = self.inference_with_batch_splitting(batch_images)
                    if tr_results is None:
                        responses.append(pb_utils.InferenceResponse(
                            error=pb_utils.TritonError("Text recognition inference failed")
                        ))
                        continue

                    text_list = self.tr_postprocess(tr_results)
                else:
                    text_list = []

                # Create output tensors
                out_tensor_0 = pb_utils.Tensor("text_list", np.array(text_list).astype(self.text_list_type))
                out_tensor_1 = pb_utils.Tensor("box_list", np.array(adjusted_box_list, dtype=np.float32))

                responses.append(pb_utils.InferenceResponse(
                    output_tensors=[out_tensor_0, out_tensor_1]
                ))

                self.logger.log_info(f"Successfully processed {len(text_list)} text boxes")

            except Exception as e:
                responses.append(pb_utils.InferenceResponse(
                    error=pb_utils.TritonError(f"Internal error: {str(e)}")
                ))

        duration = int((time.time() - st) * 1000)
        self.logger.log_info(f"ocr.tr.bls execute duration: {duration} ms")

        return responses


    def crop_with_text_box_list(self, gray_image, text_box_list):
        image_patch_list = []
        adjusted_box_list = []

        for box in text_box_list:
            box = np.array(box).reshape(4, 2)

            width = max(np.linalg.norm(box[2] - box[1]), np.linalg.norm(box[3] - box[0]))
            height = max(np.linalg.norm(box[1] - box[0]), np.linalg.norm(box[3] - box[2]))
            margin = self.text_box_margin * min(height, width)

            poly = Polygon(box)
            adjust_poly = poly.buffer(margin, join_style=JOIN_STYLE.mitre)
            adjust_box = np.array(adjust_poly.exterior.coords[:4]).astype(np.float32)
            adjust_box = align_box_points(adjust_box)

            width = int(max(np.linalg.norm(adjust_box[1] - box[0]), np.linalg.norm(adjust_box[3] - adjust_box[2])))
            height = int(max(np.linalg.norm(adjust_box[2] - box[1]), np.linalg.norm(adjust_box[3] - adjust_box[0])))

            if not self.use_auto_rotate or height < width * 1.5:
                ratio = width / height
                scaled_width = int(self.tr_canvas_size_h * ratio)
                dst_box = np.array([[0, 0], [scaled_width, 0], [scaled_width, self.tr_canvas_size_h],
                                    [0, self.tr_canvas_size_h]], dtype=np.float32)
            else:
                ratio = height / width
                scaled_width = int(self.tr_canvas_size_h * ratio)
                dst_box = np.array([[0, self.tr_canvas_size_h], [0, 0], [scaled_width, 0],
                                    [scaled_width, self.tr_canvas_size_h]], dtype=np.float32)

            matrix = cv2.getPerspectiveTransform(adjust_box, dst_box)
            cropped_image = cv2.warpPerspective(gray_image, matrix, (scaled_width, self.tr_canvas_size_h))
            # Skip samples that fail the transform step.
            if cropped_image.shape[0] == self.tr_canvas_size_h:
                image_patch_list.append(cropped_image)
                adjusted_box_list.append(adjust_box)

        return image_patch_list, adjusted_box_list


    def tr_preprocess(self, image_list):
        transformed_ims = None
        for image in image_list:
            im = self.tr_transform(image)
            if transformed_ims is None:
                transformed_ims = im
            else:
                transformed_ims = np.append(transformed_ims, im, axis=0)
        return transformed_ims


    def inference_with_batch_splitting(self, batch_images):
        """
        Split the batch into chunks, run TR inference, and merge outputs.

        Small batches (<= max_batch_size) are processed in one pass,
        while large batches are automatically split into max_batch_size chunks.
        """
        try:
            current_batch_size = batch_images.shape[0]
            results = []
            num_chunks = (current_batch_size + self.max_batch_size - 1) // self.max_batch_size

            # Process in chunks
            for i in range(0, current_batch_size, self.max_batch_size):
                chunk_start = i
                chunk_end = min(i + self.max_batch_size, current_batch_size)
                chunk = batch_images[chunk_start:chunk_end]
                chunk_idx = i // self.max_batch_size + 1

                #self.logger.log_info(f"Processing chunk {chunk_idx}/{num_chunks}: [{chunk_start}:{chunk_end}] of {current_batch_size}")

                chunk_tensor = pb_utils.Tensor("input", chunk)
                inference_request = pb_utils.InferenceRequest(
                    model_name='ocr_tr_engine',
                    requested_output_names=['output'],
                    inputs=[chunk_tensor]
                )

                inference_response = inference_request.exec()

                if inference_response.has_error():
                    continue

                # Handle GPU tensors: convert via DLPack and move to CPU
                chunk_output_tensor = pb_utils.get_output_tensor_by_name(inference_response, 'output')
                try:
                    # GPU tensor path (TensorRT plan)
                    chunk_output_gpu = from_dlpack(chunk_output_tensor.to_dlpack())
                    chunk_result = chunk_output_gpu.cpu().numpy()
                except Exception:
                    # CPU tensor path (ONNX Runtime CPU)
                    chunk_result = chunk_output_tensor.as_numpy()

                results.append(chunk_result)

                #self.logger.log_info(f"Chunk {chunk_idx}/{num_chunks} completed, output shape: {chunk_result.shape}")

            if not results:
                return None

            # Merge results (no-op if only one chunk)
            if len(results) == 1:
                merged_result = results[0]
                self.logger.log_info(f"Single chunk processed, output shape: {merged_result.shape}")
            else:
                merged_result = np.concatenate(results, axis=0)
                self.logger.log_info(f"Merged {len(results)} chunks into shape {merged_result.shape}")

            return merged_result

        except Exception as e:
            return None


    def tr_postprocess(self, output):
        texts = []
        batch_character_indices = output.argmax(2)

        for batch_index, character_indices in enumerate(batch_character_indices):
            char_list = []
            for i, character_index in enumerate(character_indices):
                if character_index != 0 and (not (i > 0 and character_indices[i - 1] == character_indices[i])):
                    char_list.append(self.tr_characters[character_index])

            text = ''.join(char_list)
            texts.append(text)

        return texts