203 lines
7.0 KiB
Python
203 lines
7.0 KiB
Python
import cv2
|
|
import numpy as np
|
|
import easyocr
|
|
from typing import List, Tuple, Dict, Optional
|
|
import re
|
|
import torch
|
|
|
|
class OCRModule2:
|
|
"""EasyOCR을 활용한 중국어 텍스트 인식 모듈"""
|
|
|
|
def __init__(self, use_gpu: bool = False, ocr_backend: str = 'easyocr'):
|
|
"""
|
|
OCR 모듈 초기화
|
|
|
|
Args:
|
|
use_gpu: GPU 사용 여부
|
|
ocr_backend: 'easyocr' 또는 'paddleocr' 중 선택
|
|
"""
|
|
self.ocr_backend = ocr_backend.lower()
|
|
if self.ocr_backend == 'easyocr':
|
|
import easyocr
|
|
self.ocr = easyocr.Reader(
|
|
['ch_sim', 'en'], # 중국어 간체, 영어 지원
|
|
gpu=use_gpu
|
|
)
|
|
elif self.ocr_backend == 'paddleocr':
|
|
try:
|
|
from paddleocr import PaddleOCR
|
|
except ImportError:
|
|
raise ImportError('PaddleOCR가 설치되어 있지 않습니다. 설치: pip install paddleocr==2.7.4')
|
|
self.ocr = PaddleOCR(use_angle_cls=True, lang='ch', use_gpu=use_gpu)
|
|
else:
|
|
raise ValueError("ocr_backend는 'easyocr' 또는 'paddleocr'만 지원합니다.")
|
|
|
|
def has_chinese_text(self, image: np.ndarray) -> bool:
|
|
"""
|
|
이미지에 중국어 텍스트가 포함되어 있는지 확인
|
|
|
|
Args:
|
|
image: 입력 이미지 (numpy array)
|
|
|
|
Returns:
|
|
중국어 텍스트 포함 여부 (bool)
|
|
"""
|
|
if self.ocr_backend == 'easyocr':
|
|
results = self.ocr.readtext(image)
|
|
for result in results:
|
|
text = result[1]
|
|
if self._is_chinese(text):
|
|
return True
|
|
return False
|
|
elif self.ocr_backend == 'paddleocr':
|
|
results = self.ocr.ocr(image, cls=True)
|
|
for line in results:
|
|
for res in line:
|
|
text = res[1][0]
|
|
if self._is_chinese(text):
|
|
return True
|
|
return False
|
|
else:
|
|
raise ValueError("ocr_backend는 'easyocr' 또는 'paddleocr'만 지원합니다.")
|
|
|
|
def extract_chinese_text(self, image: np.ndarray) -> Dict:
|
|
"""
|
|
이미지에서 중국어 텍스트 추출
|
|
|
|
Args:
|
|
image: 입력 이미지 (numpy array)
|
|
|
|
Returns:
|
|
Dict containing:
|
|
- texts: 추출된 중국어 텍스트 리스트
|
|
- polygons: 각 텍스트의 폴리곤 좌표
|
|
- masks: 텍스트 영역 마스크
|
|
- confidences: 인식 신뢰도
|
|
"""
|
|
try:
|
|
if self.ocr_backend == 'easyocr':
|
|
results = self.ocr.readtext(image)
|
|
texts = []
|
|
polygons = []
|
|
masks = []
|
|
confidences = []
|
|
for result in results:
|
|
polygon = result[0]
|
|
text = result[1]
|
|
confidence = result[2]
|
|
if self._is_chinese(text):
|
|
texts.append(text)
|
|
polygons.append(polygon)
|
|
confidences.append(confidence)
|
|
mask = self._create_mask_from_polygon(image.shape, polygon)
|
|
masks.append(mask)
|
|
return {
|
|
'texts': texts,
|
|
'polygons': polygons,
|
|
'masks': masks,
|
|
'confidences': confidences
|
|
}
|
|
elif self.ocr_backend == 'paddleocr':
|
|
results = self.ocr.ocr(image, cls=True)
|
|
texts = []
|
|
polygons = []
|
|
masks = []
|
|
confidences = []
|
|
for line in results:
|
|
for res in line:
|
|
polygon = res[0] # 4점 좌표
|
|
text = res[1][0] # 텍스트
|
|
confidence = res[1][1] # 신뢰도
|
|
if self._is_chinese(text):
|
|
texts.append(text)
|
|
polygons.append(polygon)
|
|
confidences.append(confidence)
|
|
mask = self._create_mask_from_polygon(image.shape, polygon)
|
|
masks.append(mask)
|
|
return {
|
|
'texts': texts,
|
|
'polygons': polygons,
|
|
'masks': masks,
|
|
'confidences': confidences
|
|
}
|
|
else:
|
|
raise ValueError("ocr_backend는 'easyocr' 또는 'paddleocr'만 지원합니다.")
|
|
except Exception as e:
|
|
print(f"중국어 텍스트 추출 중 오류: {e}")
|
|
return {
|
|
'texts': [],
|
|
'polygons': [],
|
|
'masks': [],
|
|
'confidences': []
|
|
}
|
|
|
|
def _is_chinese(self, text: str) -> bool:
|
|
"""
|
|
텍스트가 중국어인지 확인
|
|
|
|
Args:
|
|
text: 확인할 텍스트
|
|
|
|
Returns:
|
|
중국어 여부
|
|
"""
|
|
chinese_pattern = re.compile(r'[\u4e00-\u9fff]+')
|
|
return bool(chinese_pattern.search(text))
|
|
|
|
def _create_mask_from_polygon(self, image_shape: Tuple[int, int, int], polygon: List[List[int]]) -> np.ndarray:
|
|
"""
|
|
폴리곤으로부터 마스크 생성
|
|
|
|
Args:
|
|
image_shape: 이미지 크기 (H, W, C)
|
|
polygon: 폴리곤 좌표
|
|
|
|
Returns:
|
|
마스크 배열
|
|
"""
|
|
mask = np.zeros(image_shape[:2], dtype=np.uint8)
|
|
polygon_array = np.array(polygon, dtype=np.int32)
|
|
cv2.fillPoly(mask, [polygon_array], 255)
|
|
return mask
|
|
|
|
def get_text_properties(self, polygon: List[List[int]]) -> Dict:
|
|
"""
|
|
텍스트 영역의 속성 정보 추출
|
|
|
|
Args:
|
|
polygon: 폴리곤 좌표
|
|
|
|
Returns:
|
|
텍스트 속성 정보 (위치, 크기, 각도 등)
|
|
"""
|
|
polygon_array = np.array(polygon)
|
|
|
|
# 바운딩 박스 계산
|
|
x_coords = polygon_array[:, 0]
|
|
y_coords = polygon_array[:, 1]
|
|
|
|
x_min, x_max = int(np.min(x_coords)), int(np.max(x_coords))
|
|
y_min, y_max = int(np.min(y_coords)), int(np.max(y_coords))
|
|
|
|
width = x_max - x_min
|
|
height = y_max - y_min
|
|
|
|
# 중심점 계산
|
|
center_x = (x_min + x_max) // 2
|
|
center_y = (y_min + y_max) // 2
|
|
|
|
# 회전 각도 계산 (첫 번째와 두 번째 점 기준)
|
|
if len(polygon_array) >= 2:
|
|
dx = polygon_array[1][0] - polygon_array[0][0]
|
|
dy = polygon_array[1][1] - polygon_array[0][1]
|
|
angle = np.arctan2(dy, dx) * 180 / np.pi
|
|
else:
|
|
angle = 0
|
|
|
|
return {
|
|
'bbox': (x_min, y_min, x_max, y_max),
|
|
'center': (center_x, center_y),
|
|
'width': width,
|
|
'height': height,
|
|
'angle': angle
|
|
} |