527 lines
22 KiB
Python
527 lines
22 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
ONNX OCR 래퍼 모듈 - 기존 OCRModule 인터페이스 유지하면서 ONNX 사용
|
|
기존 predict_system.py의 완벽한 전후처리 로직을 활용
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import os
|
|
import sys
|
|
import logging
|
|
import time
|
|
from typing import List, Dict, Any
|
|
|
|
# 완전 독립 모듈: 내부 src 폴더의 ppocr, tools 사용
|
|
current_dir = os.path.dirname(__file__) # src 폴더
|
|
sys.path.insert(0, current_dir) # src 폴더를 패스에 추가
|
|
|
|
class ONNXOCRModule:
|
|
"""
|
|
기존 OCRModule과 완전히 동일한 인터페이스를 제공하면서 내부적으로 ONNX 사용
|
|
"""
|
|
def __init__(self, logger=None, base_dir=None, gpu_manager=None):
|
|
self.logger = logger
|
|
self.base_dir = base_dir or os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
self.gpu_manager = gpu_manager
|
|
|
|
# GPU 상태 결정
|
|
self.use_gpu = gpu_manager and gpu_manager.can_use_cuda if gpu_manager else False
|
|
|
|
# 환경 변수 설정 (기존 OCRModule과 동일)
|
|
if self.use_gpu:
|
|
if self.logger:
|
|
self.logger.log("ONNX OCR 모듈 CUDA 모드 활성화", level=logging.INFO)
|
|
if 'CUDA_VISIBLE_DEVICES' in os.environ:
|
|
del os.environ['CUDA_VISIBLE_DEVICES']
|
|
else:
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = ''
|
|
if self.logger:
|
|
self.logger.log("ONNX OCR 모듈 CPU 모드로 설정", level=logging.INFO)
|
|
|
|
# OpenCV 및 멀티스레드 설정 (기존과 동일)
|
|
try:
|
|
import cv2 as _cv2
|
|
_cv2.ocl.setUseOpenCL(False)
|
|
except Exception:
|
|
pass
|
|
|
|
os.environ["OMP_NUM_THREADS"] = "1"
|
|
os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
|
|
os.environ['OPENBLAS_NUM_THREADS'] = '1'
|
|
os.environ['MKL_NUM_THREADS'] = '1'
|
|
os.environ['NUMEXPR_NUM_THREADS'] = '1'
|
|
|
|
# ONNX 시스템 초기화
|
|
self.text_system = None
|
|
self._initialize_onnx_system()
|
|
|
|
if self.text_system is None:
|
|
if self.logger:
|
|
self.logger.log("ONNX TextSystem 초기화 실패", level=logging.ERROR)
|
|
raise Exception("ONNX TextSystem 초기화 실패")
|
|
else:
|
|
cuda_status = "CUDA" if self.use_gpu else "CPU"
|
|
if self.logger:
|
|
self.logger.log(f"✅ ONNX OCR 모듈 초기화 성공 ({cuda_status} 모드)", level=logging.INFO)
|
|
|
|
def _initialize_onnx_system(self):
|
|
"""predict_system.py의 TextSystem을 ONNX 모드로 초기화"""
|
|
try:
|
|
# predict_system에서 필요한 모듈들 import
|
|
from tools.infer.predict_system import TextSystem
|
|
from ppocr.utils.utility import get_image_file_list, check_and_read
|
|
|
|
# Args 객체 생성 (ONNX 설정 포함)
|
|
class Args:
|
|
def __init__(self):
|
|
# ONNX 관련 설정
|
|
self.use_onnx = True
|
|
self.use_gpu = self.use_gpu if hasattr(self, 'use_gpu') else False
|
|
|
|
# 모델 경로 설정
|
|
models_dir = os.path.join(self.base_dir, "models")
|
|
if self.use_gpu:
|
|
# CUDA용 FP16 모델
|
|
self.det_model_dir = os.path.join(models_dir, "det_dyn.fp16.onnx")
|
|
self.rec_model_dir = os.path.join(models_dir, "rec_dyn.fp16.onnx")
|
|
self.cls_model_dir = os.path.join(models_dir, "cls_dyn.fp16.onnx")
|
|
else:
|
|
# CPU용 단순화 모델
|
|
self.det_model_dir = os.path.join(models_dir, "det_dyn.simp.onnx")
|
|
self.rec_model_dir = os.path.join(models_dir, "rec_dyn.simp.onnx")
|
|
self.cls_model_dir = os.path.join(models_dir, "cls_dyn.simp.onnx")
|
|
|
|
# 문자 사전
|
|
self.rec_char_dict_path = os.path.join(self.base_dir, "dict", "ppocr_keys_v1.txt")
|
|
|
|
# 기본 OCR 설정
|
|
self.use_angle_cls = True
|
|
self.det_algorithm = 'DB'
|
|
self.rec_algorithm = 'SVTR_LCNet'
|
|
self.cls_algorithm = 'ClsMV3'
|
|
|
|
# 성능 최적화 설정
|
|
self.rec_batch_num = 6
|
|
self.max_text_length = 25
|
|
self.rec_image_shape = "3, 48, 320"
|
|
self.drop_score = 0.5
|
|
|
|
# 기타 설정
|
|
self.show_log = False
|
|
self.use_space_char = True
|
|
self.save_crop_res = False
|
|
self.crop_res_save_dir = "./output"
|
|
|
|
# Detection 관련
|
|
self.det_limit_side_len = 960
|
|
self.det_limit_type = 'max'
|
|
self.det_db_thresh = 0.3
|
|
self.det_db_box_thresh = 0.6
|
|
self.det_db_unclip_ratio = 1.5
|
|
self.det_box_type = 'quad'
|
|
self.use_dilation = False
|
|
self.det_db_score_mode = 'fast'
|
|
|
|
# Classification 관련
|
|
self.cls_image_shape = "3, 48, 192"
|
|
self.cls_batch_num = 6
|
|
self.cls_thresh = 0.9
|
|
self.label_list = ['0', '180']
|
|
|
|
# 추가 Detection 알고리즘 관련
|
|
self.det_east_score_thresh = 0.8
|
|
self.det_east_cover_thresh = 0.1
|
|
self.det_east_nms_thresh = 0.2
|
|
self.det_sast_score_thresh = 0.5
|
|
self.det_sast_nms_thresh = 0.2
|
|
self.det_pse_thresh = 0
|
|
self.det_pse_box_thresh = 0.85
|
|
self.det_pse_min_area = 16
|
|
self.det_pse_scale = 1
|
|
self.scales = [8, 16, 32]
|
|
self.alpha = 1.0
|
|
self.beta = 1.0
|
|
self.fourier_degree = 5
|
|
|
|
# Recognition 관련
|
|
self.rec_image_inverse = True
|
|
|
|
# 시스템 관련
|
|
self.max_batch_size = 10
|
|
self.precision = 'fp32'
|
|
self.gpu_mem = 500
|
|
self.gpu_id = 0
|
|
self.ir_optim = True
|
|
self.use_tensorrt = False
|
|
self.min_subgraph_size = 15
|
|
self.use_xpu = False
|
|
self.use_npu = False
|
|
|
|
# 기타
|
|
self.enable_mkldnn = False
|
|
self.cpu_threads = 10
|
|
self.warmup = False
|
|
self.benchmark = False
|
|
self.vis_font_path = "./simfang.ttf"
|
|
|
|
# 현재 인스턴스를 참조할 수 있도록 설정
|
|
Args.base_dir = self.base_dir
|
|
Args.use_gpu = self.use_gpu
|
|
|
|
args = Args()
|
|
|
|
# TextSystem 초기화
|
|
self.text_system = TextSystem(args)
|
|
|
|
if self.logger:
|
|
model_type = "CUDA FP16" if self.use_gpu else "CPU 단순화"
|
|
self.logger.log(f"✅ ONNX TextSystem 초기화 완료 ({model_type} 모델)", level=logging.INFO)
|
|
|
|
except Exception as e:
|
|
if self.logger:
|
|
self.logger.log(f"❌ ONNX TextSystem 초기화 실패: {e}", level=logging.ERROR, exc_info=True)
|
|
raise
|
|
|
|
def detect_text(self, image_path: str, method: str = 'polygon', raise_on_memory_error: bool = False) -> List[Dict[str, Any]]:
|
|
"""
|
|
이미지에서 텍스트를 감지하고 다양한 방식으로 영역 반환
|
|
기존 OCRModule과 완전히 동일한 인터페이스
|
|
|
|
Args:
|
|
image_path (str): 이미지 파일 경로
|
|
method (str): 감지 방식 ('polygon', 'bbox', 'expanded_bbox', 'rotated_bbox', 'contour')
|
|
raise_on_memory_error (bool): 메모리 에러 시 예외 발생 여부
|
|
|
|
Returns:
|
|
List[Dict]: 감지된 텍스트 정보 리스트
|
|
- text: 감지된 텍스트
|
|
- confidence: 신뢰도
|
|
- polygon: 폴리곤 좌표 (4개 점)
|
|
- bbox: 바운딩 박스 좌표 (x, y, w, h)
|
|
- method: 사용된 감지 방식
|
|
"""
|
|
if not os.path.exists(image_path):
|
|
if self.logger:
|
|
self.logger.log(f"이미지 파일을 찾을 수 없습니다: {image_path}", level=logging.ERROR)
|
|
return []
|
|
|
|
image = None
|
|
try:
|
|
# 이미지 읽기
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
if self.logger:
|
|
self.logger.log(f"이미지를 읽을 수 없습니다: {image_path}", level=logging.ERROR)
|
|
return []
|
|
|
|
if self.logger:
|
|
self.logger.log(f"🔍 ONNX OCR 감지 방식: {method}", level=logging.INFO)
|
|
|
|
# 메모리 안전을 위한 이미지 크기 조정 (기존과 동일)
|
|
image = self._safe_resize_image(image)
|
|
|
|
# ONNX TextSystem으로 OCR 실행
|
|
start_time = time.time()
|
|
dt_boxes, rec_res, time_dict = self.text_system(image)
|
|
inference_time = (time.time() - start_time) * 1000
|
|
|
|
if self.logger:
|
|
self.logger.log(f"⚡ ONNX OCR 추론 완료: {inference_time:.1f}ms", level=logging.INFO)
|
|
self.logger.log(f"📊 세부 시간 - 감지: {time_dict.get('det', 0)*1000:.1f}ms, 인식: {time_dict.get('rec', 0)*1000:.1f}ms, 분류: {time_dict.get('cls', 0)*1000:.1f}ms", level=logging.INFO)
|
|
|
|
if not dt_boxes or not rec_res:
|
|
if self.logger:
|
|
self.logger.log("⚠️ ONNX OCR 결과가 비어있습니다", level=logging.WARNING)
|
|
return []
|
|
|
|
# predict_system 결과를 기존 format으로 변환
|
|
converted_results = []
|
|
for box, (text, confidence) in zip(dt_boxes, rec_res):
|
|
converted_results.append([box.tolist(), [text, confidence]])
|
|
|
|
# 감지 방식에 따라 결과 처리 (기존 로직 재사용)
|
|
if method == 'polygon':
|
|
ocr_results = self._detect_with_polygon(image, converted_results)
|
|
elif method == 'bbox':
|
|
ocr_results = self._detect_with_bbox(image, converted_results)
|
|
elif method == 'expanded_bbox':
|
|
ocr_results = self._detect_with_expanded_bbox(image, converted_results)
|
|
elif method == 'rotated_bbox':
|
|
ocr_results = self._detect_with_rotated_bbox(image, converted_results)
|
|
elif method == 'contour':
|
|
ocr_results = self._detect_with_contour(image, converted_results)
|
|
else:
|
|
if self.logger:
|
|
self.logger.log(f"⚠️ 지원하지 않는 감지 방식: {method}, 기본 polygon 방식 사용", level=logging.WARNING)
|
|
ocr_results = self._detect_with_polygon(image, converted_results)
|
|
|
|
return ocr_results
|
|
|
|
except Exception as e:
|
|
if self.logger:
|
|
self.logger.log(f"ONNX OCR 처리 중 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
if raise_on_memory_error and ('memory' in str(e).lower() or 'primitive' in str(e).lower()):
|
|
raise MemoryError(f"ONNX OCR memory error: {e}")
|
|
return []
|
|
finally:
|
|
if image is not None:
|
|
del image
|
|
cv2.destroyAllWindows()
|
|
|
|
def _safe_resize_image(self, image: np.ndarray) -> np.ndarray:
|
|
"""메모리 안전을 위한 이미지 크기 조정 (기존 로직과 동일)"""
|
|
h, w = image.shape[:2]
|
|
max_dim_safe = 2000
|
|
aspect_ratio = max(w, h) / max(1, min(w, h))
|
|
|
|
if max(w, h) > max_dim_safe or aspect_ratio > 15:
|
|
scale = float(max_dim_safe) / float(max(w, h))
|
|
new_size = (int(w * scale), int(h * scale))
|
|
if self.logger:
|
|
self.logger.log(
|
|
f"⚖️ ONNX OCR용 다운스케일 ({w}x{h}) -> {new_size} (ratio={aspect_ratio:.1f})",
|
|
level=logging.INFO,
|
|
)
|
|
image = cv2.resize(image, new_size, interpolation=cv2.INTER_AREA).copy()
|
|
|
|
return image
|
|
|
|
# === 기존 OCRModule의 감지 방식 메서드들 (완전히 동일) ===
|
|
|
|
def _detect_with_polygon(self, image: np.ndarray, ocr_raw_results: List) -> List[Dict[str, Any]]:
|
|
"""폴리곤 방식으로 텍스트 영역 감지 (기본 방식)"""
|
|
ocr_results = []
|
|
|
|
for line in ocr_raw_results:
|
|
if len(line) >= 2:
|
|
polygon = line[0] # 폴리곤 좌표 (4개 점)
|
|
text_info = line[1] # (텍스트, 신뢰도)
|
|
|
|
if len(text_info) >= 2:
|
|
text = text_info[0]
|
|
confidence = text_info[1]
|
|
|
|
# 폴리곤을 바운딩 박스로 변환
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
x, y, w, h = cv2.boundingRect(polygon_np)
|
|
|
|
ocr_result = {
|
|
'text': text,
|
|
'confidence': confidence,
|
|
'polygon': polygon,
|
|
'bbox': (x, y, w, h),
|
|
'method': 'polygon'
|
|
}
|
|
ocr_results.append(ocr_result)
|
|
|
|
return ocr_results
|
|
|
|
def _detect_with_bbox(self, image: np.ndarray, ocr_raw_results: List) -> List[Dict[str, Any]]:
|
|
"""바운딩 박스 방식으로 텍스트 영역 감지"""
|
|
ocr_results = []
|
|
|
|
for line in ocr_raw_results:
|
|
if len(line) >= 2:
|
|
polygon = line[0]
|
|
text_info = line[1]
|
|
|
|
if len(text_info) >= 2:
|
|
text = text_info[0]
|
|
confidence = text_info[1]
|
|
|
|
# 바운딩 박스 계산
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
x, y, w, h = cv2.boundingRect(polygon_np)
|
|
|
|
# 바운딩 박스를 폴리곤으로 변환
|
|
bbox_polygon = [
|
|
[x, y],
|
|
[x + w, y],
|
|
[x + w, y + h],
|
|
[x, y + h]
|
|
]
|
|
|
|
ocr_result = {
|
|
'text': text,
|
|
'confidence': confidence,
|
|
'polygon': bbox_polygon,
|
|
'bbox': (x, y, w, h),
|
|
'method': 'bbox'
|
|
}
|
|
ocr_results.append(ocr_result)
|
|
|
|
return ocr_results
|
|
|
|
def _detect_with_expanded_bbox(self, image: np.ndarray, ocr_raw_results: List) -> List[Dict[str, Any]]:
|
|
"""확장된 바운딩 박스 방식으로 텍스트 영역 감지"""
|
|
ocr_results = []
|
|
h_img, w_img = image.shape[:2]
|
|
|
|
for line in ocr_raw_results:
|
|
if len(line) >= 2:
|
|
polygon = line[0]
|
|
text_info = line[1]
|
|
|
|
if len(text_info) >= 2:
|
|
text = text_info[0]
|
|
confidence = text_info[1]
|
|
|
|
# 기본 바운딩 박스
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
x, y, w, h = cv2.boundingRect(polygon_np)
|
|
|
|
# 확장 크기 계산 (텍스트 크기의 20%)
|
|
expand_x = max(1, int(w * 0.2))
|
|
expand_y = max(1, int(h * 0.2))
|
|
|
|
# 확장된 바운딩 박스
|
|
x_exp = max(0, x - expand_x)
|
|
y_exp = max(0, y - expand_y)
|
|
w_exp = min(w_img - x_exp, w + 2 * expand_x)
|
|
h_exp = min(h_img - y_exp, h + 2 * expand_y)
|
|
|
|
# 확장된 바운딩 박스를 폴리곤으로 변환
|
|
expanded_polygon = [
|
|
[x_exp, y_exp],
|
|
[x_exp + w_exp, y_exp],
|
|
[x_exp + w_exp, y_exp + h_exp],
|
|
[x_exp, y_exp + h_exp]
|
|
]
|
|
|
|
ocr_result = {
|
|
'text': text,
|
|
'confidence': confidence,
|
|
'polygon': expanded_polygon,
|
|
'bbox': (x_exp, y_exp, w_exp, h_exp),
|
|
'method': 'expanded_bbox'
|
|
}
|
|
ocr_results.append(ocr_result)
|
|
|
|
return ocr_results
|
|
|
|
def _detect_with_rotated_bbox(self, image: np.ndarray, ocr_raw_results: List) -> List[Dict[str, Any]]:
|
|
"""회전된 바운딩 박스 방식으로 텍스트 영역 감지"""
|
|
ocr_results = []
|
|
|
|
for line in ocr_raw_results:
|
|
if len(line) >= 2:
|
|
polygon = line[0]
|
|
text_info = line[1]
|
|
|
|
if len(text_info) >= 2:
|
|
text = text_info[0]
|
|
confidence = text_info[1]
|
|
|
|
# 회전된 바운딩 박스 계산
|
|
polygon_np = np.array(polygon, dtype=np.float32)
|
|
rect = cv2.minAreaRect(polygon_np)
|
|
box = cv2.boxPoints(rect)
|
|
box = np.int32(box)
|
|
|
|
# 일반 바운딩 박스도 계산
|
|
x, y, w, h = cv2.boundingRect(polygon_np.astype(np.int32))
|
|
|
|
ocr_result = {
|
|
'text': text,
|
|
'confidence': confidence,
|
|
'polygon': box.tolist(),
|
|
'bbox': (x, y, w, h),
|
|
'method': 'rotated_bbox',
|
|
'rotation_info': {
|
|
'center': rect[0],
|
|
'size': rect[1],
|
|
'angle': rect[2]
|
|
}
|
|
}
|
|
ocr_results.append(ocr_result)
|
|
|
|
return ocr_results
|
|
|
|
def _detect_with_contour(self, image: np.ndarray, ocr_raw_results: List) -> List[Dict[str, Any]]:
|
|
"""컨투어 방식으로 텍스트 영역 감지"""
|
|
ocr_results = []
|
|
|
|
for line in ocr_raw_results:
|
|
if len(line) >= 2:
|
|
polygon = line[0]
|
|
text_info = line[1]
|
|
|
|
if len(text_info) >= 2:
|
|
text = text_info[0]
|
|
confidence = text_info[1]
|
|
|
|
# 폴리곤을 컨투어로 변환
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
|
|
# 컨투어 근사화
|
|
epsilon = 0.02 * cv2.arcLength(polygon_np, True)
|
|
approx_contour = cv2.approxPolyDP(polygon_np, epsilon, True)
|
|
|
|
# 컨투어를 다시 폴리곤으로 변환
|
|
contour_polygon = approx_contour.reshape(-1, 2).tolist()
|
|
|
|
# 바운딩 박스 계산
|
|
x, y, w, h = cv2.boundingRect(polygon_np)
|
|
|
|
ocr_result = {
|
|
'text': text,
|
|
'confidence': confidence,
|
|
'polygon': contour_polygon,
|
|
'bbox': (x, y, w, h),
|
|
'method': 'contour',
|
|
'contour_points': len(contour_polygon)
|
|
}
|
|
ocr_results.append(ocr_result)
|
|
|
|
return ocr_results
|
|
|
|
# === 기존 OCRModule의 필터링 메서드들 (완전히 동일) ===
|
|
|
|
def filter_chinese_text(self, ocr_results: List[Dict]) -> List[Dict]:
|
|
"""
|
|
중국어 텍스트만 필터링
|
|
|
|
Args:
|
|
ocr_results (List[Dict]): OCR 결과
|
|
|
|
Returns:
|
|
List[Dict]: 중국어 텍스트만 포함된 결과
|
|
"""
|
|
chinese_results = []
|
|
|
|
for result in ocr_results:
|
|
text = result['text']
|
|
# 중국어 문자 범위 확인 (간체/번체 포함)
|
|
if any('\u4e00' <= char <= '\u9fff' for char in text):
|
|
chinese_results.append(result)
|
|
|
|
if self.logger:
|
|
self.logger.log(f"중국어 텍스트 {len(chinese_results)}개 필터링 완료", level=logging.INFO)
|
|
return chinese_results
|
|
|
|
def filter_korean_text(self, ocr_results: List[Dict]) -> List[Dict]:
|
|
"""
|
|
한글 텍스트만 필터링
|
|
|
|
Args:
|
|
ocr_results (List[Dict]): OCR 결과
|
|
|
|
Returns:
|
|
List[Dict]: 한글 텍스트만 포함된 결과
|
|
"""
|
|
korean_results = []
|
|
for result in ocr_results:
|
|
text = result['text']
|
|
# 한글 유니코드 범위: 가~힣
|
|
if any('\uac00' <= char <= '\ud7a3' for char in text):
|
|
korean_results.append(result)
|
|
|
|
if self.logger:
|
|
self.logger.log(f"한글 텍스트 {len(korean_results)}개 필터링 완료", level=logging.INFO)
|
|
return korean_results
|
|
|
|
|
|
# 기존 OCRModule과의 호환성을 위한 별칭
|
|
OCRModule = ONNXOCRModule
|