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