161 lines
6.2 KiB
Python
161 lines
6.2 KiB
Python
import os
|
|
import logging
|
|
import cv2
|
|
import numpy as np
|
|
import easyocr
|
|
|
|
class EasyOCREngine:
|
|
def __init__(self, logger=None, base_dir=None, lang_list=['ch_sim']):
|
|
self.logger = logger
|
|
self.base_dir = base_dir
|
|
model_dir = os.path.join(base_dir, 'EasyOCR_models') if base_dir else None
|
|
print(f"model_dir: {model_dir}")
|
|
self.reader = easyocr.Reader(lang_list, gpu=False, model_storage_directory=model_dir)
|
|
|
|
def detect_text(self, image_path: str, method: str = 'polygon'):
|
|
"""
|
|
EasyOCR로 텍스트 감지 (여러 방식 지원)
|
|
|
|
Args:
|
|
image_path: 이미지 경로
|
|
method: 감지 방식
|
|
- 'polygon': 기본 폴리곤 방식
|
|
- 'paragraph': 문단 단위 그룹핑 (추천)
|
|
- 'enhanced': 크기/신뢰도 기반 고급 필터링
|
|
"""
|
|
if not os.path.exists(image_path):
|
|
if self.logger:
|
|
self.logger.log(f"이미지 파일을 찾을 수 없습니다: {image_path}", level=logging.ERROR)
|
|
return []
|
|
|
|
try:
|
|
if method == 'paragraph':
|
|
return self._detect_with_paragraph(image_path)
|
|
elif method == 'enhanced':
|
|
return self._detect_with_enhanced_filtering(image_path)
|
|
else: # 기본 polygon 방식
|
|
return self._detect_basic_polygon(image_path)
|
|
|
|
except Exception as e:
|
|
if self.logger:
|
|
self.logger.log(f"EasyOCR 처리 중 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
return []
|
|
|
|
def _detect_basic_polygon(self, image_path: str):
|
|
"""기본 폴리곤 방식 (기존 방식)"""
|
|
results = self.reader.readtext(image_path, paragraph=False)
|
|
ocr_results = []
|
|
for bbox, text, confidence in results:
|
|
polygon = [list(map(int, pt)) for pt in bbox]
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
x, y, w, h = cv2.boundingRect(polygon_np)
|
|
ocr_results.append({
|
|
'text': text,
|
|
'confidence': float(confidence),
|
|
'polygon': polygon,
|
|
'bbox': (x, y, w, h),
|
|
'method': 'polygon'
|
|
})
|
|
return ocr_results
|
|
|
|
def _detect_with_paragraph(self, image_path: str):
|
|
"""문단 단위 그룹핑 방식 (연결된 텍스트에 유용)"""
|
|
# 문단 모드로 텍스트 감지
|
|
results = self.reader.readtext(image_path, paragraph=True)
|
|
ocr_results = []
|
|
|
|
for bbox, text, confidence in results:
|
|
polygon = [list(map(int, pt)) for pt in bbox]
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
x, y, w, h = cv2.boundingRect(polygon_np)
|
|
|
|
# 문단 모드에서는 더 긴 텍스트가 나올 수 있음
|
|
ocr_results.append({
|
|
'text': text,
|
|
'confidence': float(confidence),
|
|
'polygon': polygon,
|
|
'bbox': (x, y, w, h),
|
|
'method': 'paragraph',
|
|
'text_length': len(text)
|
|
})
|
|
|
|
if self.logger:
|
|
self.logger.log(f"문단 모드 감지 결과: {len(ocr_results)}개", level=logging.DEBUG)
|
|
return ocr_results
|
|
|
|
def _detect_with_enhanced_filtering(self, image_path: str):
|
|
"""고급 필터링 방식 (크기, 신뢰도, 비율 기반)"""
|
|
results = self.reader.readtext(image_path, paragraph=False)
|
|
|
|
# 이미지 크기 정보 가져오기
|
|
import cv2
|
|
image = cv2.imread(image_path)
|
|
img_height, img_width = image.shape[:2]
|
|
|
|
ocr_results = []
|
|
filtered_count = 0
|
|
|
|
for bbox, text, confidence in results:
|
|
polygon = [list(map(int, pt)) for pt in bbox]
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
x, y, w, h = cv2.boundingRect(polygon_np)
|
|
|
|
# 고급 필터링 조건들
|
|
text_area = w * h
|
|
img_area = img_width * img_height
|
|
area_ratio = text_area / img_area
|
|
aspect_ratio = w / h if h > 0 else 0
|
|
|
|
# 필터링 조건
|
|
# 1. 최소/최대 크기 체크
|
|
if w < 10 or h < 8: # 너무 작은 텍스트
|
|
filtered_count += 1
|
|
continue
|
|
if area_ratio > 0.5: # 이미지의 50% 이상 차지하는 텍스트
|
|
filtered_count += 1
|
|
continue
|
|
|
|
# 2. 종횡비 체크 (너무 길거나 높은 박스 제외)
|
|
if aspect_ratio > 20 or aspect_ratio < 0.05:
|
|
filtered_count += 1
|
|
continue
|
|
|
|
# 3. 동적 신뢰도 임계값 (텍스트 크기에 따라)
|
|
min_confidence = 0.15 if text_area > 1000 else 0.25 # 큰 텍스트는 낮은 임계값
|
|
if confidence < min_confidence:
|
|
filtered_count += 1
|
|
continue
|
|
|
|
# 4. 텍스트 길이 체크
|
|
if len(text.strip()) < 1:
|
|
filtered_count += 1
|
|
continue
|
|
|
|
ocr_results.append({
|
|
'text': text,
|
|
'confidence': float(confidence),
|
|
'polygon': polygon,
|
|
'bbox': (x, y, w, h),
|
|
'method': 'enhanced',
|
|
'text_area': text_area,
|
|
'area_ratio': area_ratio,
|
|
'aspect_ratio': aspect_ratio
|
|
})
|
|
|
|
if self.logger:
|
|
self.logger.log(f"고급 필터링 결과: {len(ocr_results)}개 유지, {filtered_count}개 필터링됨", level=logging.DEBUG)
|
|
return ocr_results
|
|
|
|
def filter_chinese_text(self, ocr_results):
|
|
"""
|
|
중국어 텍스트만 필터링
|
|
"""
|
|
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
|