559 lines
23 KiB
Python
559 lines
23 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
OCR 모듈 - PaddleOCR3을 사용한 텍스트 감지
|
|
폴리곤 방식으로 텍스트 영역을 감지합니다.
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import os
|
|
from typing import List, Dict, Any
|
|
|
|
class OCRModule:
|
|
def __init__(self, use_angle_cls=True, lang='ch'):
|
|
"""
|
|
OCR 모듈 초기화
|
|
|
|
Args:
|
|
use_angle_cls (bool): 텍스트 각도 분류 사용 여부
|
|
lang (str): 언어 설정 ('ch' for Chinese, 'korean' for Korean)
|
|
"""
|
|
self.detection_methods = {
|
|
'polygon': self._detect_with_polygon,
|
|
'bbox': self._detect_with_bbox,
|
|
'expanded_bbox': self._detect_with_expanded_bbox,
|
|
'rotated_bbox': self._detect_with_rotated_bbox,
|
|
'contour': self._detect_with_contour
|
|
}
|
|
|
|
# CPU만 사용하도록 환경 변수 설정
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = ''
|
|
|
|
try:
|
|
# from paddleocr import PaddleOCR
|
|
# # PaddleOCR 3.x 버전에 맞게 초기화 (기본 설정만 사용)
|
|
# self.ocr = PaddleOCR(text_detection_model_name='PSENet', lang=lang)
|
|
self.ocr = self.initialize_ocr()
|
|
print(f"✅ PaddleOCR 초기화 완료 (언어: {lang})")
|
|
except Exception as e:
|
|
print(f"❌ PaddleOCR 초기화 실패: {e}")
|
|
raise e # 에러 발생시 프로그램 종료
|
|
|
|
print(f"📋 사용 가능한 감지 방식: {list(self.detection_methods.keys())}")
|
|
|
|
def get_base_dir(self):
|
|
"""
|
|
실행 환경에 따라 base_dir을 설정하는 메서드.
|
|
cx_Freeze로 패키징된 경우 실행 파일의 경로, 일반 Python 환경일 경우 __file__을 기준으로 설정.
|
|
"""
|
|
import sys
|
|
if getattr(sys, 'frozen', False): # 패키징된 경우
|
|
base_dir = os.path.dirname(sys.executable)
|
|
internal_dir = os.path.join(base_dir, 'lib') # lib 디렉토리 포함
|
|
if os.path.exists(internal_dir): # lib 디렉토리가 존재하면 base_dir로 설정
|
|
return internal_dir
|
|
|
|
else: # 일반 Python 실행 환경
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
debug_dir = os.path.join(base_dir) # lib 디렉토리 포함
|
|
|
|
return debug_dir
|
|
|
|
def initialize_ocr(self):
|
|
"""
|
|
PaddleOCR 초기화. det_enabled 옵션에 따라 Detection 모델 사용 여부 결정.
|
|
"""
|
|
# 모델 디렉토리 설정
|
|
self.base_dir = self.get_base_dir()
|
|
# print(f"base_dir: {self.base_dir}")
|
|
# self.rec_model_dir = os.path.join(self.base_dir, "PP-OCRv5", "ch_RepSVTR_rec")
|
|
# # self.rec_model_dir = os.path.join(self.base_dir, "PP-OCRv5", "ch_SVTRv2_rec_infer")
|
|
# self.det_model_dir = os.path.join(self.base_dir, "PP-OCRv5", "PP-OCRv5_mobile_det")
|
|
# # self.cls_model_dir = os.path.join(self.base_dir, "PP-OCRv5", "cls")
|
|
# print(f"rec_model_dir: {self.rec_model_dir}")
|
|
# print(f"det_model_dir: {self.det_model_dir}")
|
|
# # print(f"cls_model_dir: {self.cls_model_dir}")
|
|
|
|
|
|
|
|
self.rec_model_dir = os.path.join(self.base_dir, "PP_Models", "rec")
|
|
self.det_model_dir = os.path.join(self.base_dir, "PP_Models", "det")
|
|
self.cls_model_dir = os.path.join(self.base_dir, "PP_Models", "cls")
|
|
|
|
from paddleocr import PaddleOCR
|
|
self.det_enabled = True
|
|
|
|
if self.det_enabled:
|
|
# ocr = PaddleOCR(
|
|
# lang="ch",
|
|
# text_recognition_model_name="ch_RepSVTR_rec",
|
|
# text_detection_model_name="PP-OCRv5_mobile_det",
|
|
# text_det_limit_side_len=1024,
|
|
# text_det_limit_type='max',
|
|
# text_det_unclip_ratio=1.2,
|
|
# text_det_box_thresh=0.6,
|
|
# text_det_thresh=0.3
|
|
# # text_detection_model_dir=self.det_model_dir,
|
|
# # text_recognition_model_dir=self.rec_model_dir,
|
|
# # doc_orientation_classify_model_dir=self.cls_model_dir
|
|
# )
|
|
# else:
|
|
# ocr = PaddleOCR(text_detection_model_name='PSENet',
|
|
# lang="ch",
|
|
# det_model_dir=None, # Detection 비활성화
|
|
# rec_model_dir=self.rec_model_dir,
|
|
# cls_model_dir=self.cls_model_dir
|
|
# )
|
|
|
|
ocr = PaddleOCR(
|
|
use_gpu=False,
|
|
use_angle_cls=True, # 텍스트 방향 분류 활성화
|
|
lang="ch",
|
|
det_model_dir=self.det_model_dir,
|
|
rec_model_dir=self.rec_model_dir,
|
|
cls_model_dir=self.cls_model_dir
|
|
)
|
|
|
|
return ocr
|
|
|
|
def detect_text(self, image_path: str, method: str = 'polygon') -> List[Dict[str, Any]]:
|
|
"""
|
|
이미지에서 텍스트를 감지하고 다양한 방식으로 영역 반환
|
|
|
|
Args:
|
|
image_path (str): 이미지 파일 경로
|
|
method (str): 감지 방식 ('polygon', 'bbox', 'expanded_bbox', 'rotated_bbox', 'contour')
|
|
|
|
Returns:
|
|
List[Dict]: 감지된 텍스트 정보 리스트
|
|
- text: 감지된 텍스트
|
|
- confidence: 신뢰도
|
|
- polygon: 폴리곤 좌표 (4개 점)
|
|
- bbox: 바운딩 박스 좌표 (x, y, w, h)
|
|
- method: 사용된 감지 방식
|
|
"""
|
|
if not os.path.exists(image_path):
|
|
print(f"이미지 파일을 찾을 수 없습니다: {image_path}")
|
|
return []
|
|
|
|
if method not in self.detection_methods:
|
|
print(f"지원하지 않는 감지 방식: {method}")
|
|
print(f"사용 가능한 방식: {list(self.detection_methods.keys())}")
|
|
method = 'polygon' # 기본값으로 변경
|
|
|
|
try:
|
|
# 이미지 읽기
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
print(f"이미지를 읽을 수 없습니다: {image_path}")
|
|
return []
|
|
|
|
print(f"🔍 OCR 감지 방식: {method}")
|
|
|
|
# 실제 OCR 실행
|
|
# ocr_raw_results = self.ocr.predict(image)
|
|
ocr_raw_results = self.ocr.ocr(image)
|
|
|
|
print("ocr_raw_results:", ocr_raw_results)
|
|
for line in ocr_raw_results:
|
|
print("line:", line)
|
|
|
|
if not ocr_raw_results or len(ocr_raw_results) == 0:
|
|
print("⚠️ OCR 결과가 비어있습니다.")
|
|
return []
|
|
|
|
# paddleocr 2.x 결과 파싱
|
|
converted_results = []
|
|
for page in ocr_raw_results: # page는 텍스트별 결과 리스트
|
|
for line in page:
|
|
poly = line[0]
|
|
text = line[1][0]
|
|
score = line[1][1]
|
|
converted_results.append([poly, [text, score]])
|
|
|
|
# 감지 방식에 따라 결과 처리
|
|
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:
|
|
print(f"⚠️ 지원하지 않는 감지 방식: {method}, 기본 polygon 방식 사용")
|
|
ocr_results = self._detect_with_polygon(image, converted_results)
|
|
|
|
# 결과 출력
|
|
self._print_ocr_results(ocr_results, method)
|
|
|
|
return ocr_results
|
|
|
|
except Exception as e:
|
|
print(f"❌ OCR 처리 중 오류 발생: {e}")
|
|
return []
|
|
|
|
def _print_ocr_results(self, ocr_results: List[Dict[str, Any]], method: str):
|
|
"""
|
|
OCR 결과를 상세히 출력
|
|
|
|
Args:
|
|
ocr_results (List[Dict]): OCR 결과 리스트
|
|
method (str): 사용된 감지 방식
|
|
"""
|
|
print("\n" + "="*60)
|
|
print(f"OCR 감지 결과 상세 정보 - 방식: {method}")
|
|
print("="*60)
|
|
|
|
for i, result in enumerate(ocr_results, 1):
|
|
print(f"\n🔍 텍스트 영역 #{i}")
|
|
print(f" 📄 인식된 텍스트: '{result['text']}'")
|
|
print(f" 🎯 신뢰도: {result['confidence']:.1%}")
|
|
print(f" 📐 바운딩 박스: {result['bbox']} (x, y, w, h)")
|
|
|
|
# 감지 방식별 추가 정보 출력
|
|
if result['method'] == 'polygon':
|
|
print(f" 🔺 폴리곤 좌표: {result['polygon']}")
|
|
# 폴리곤 중심점 계산
|
|
polygon_np = np.array(result['polygon'])
|
|
center_x = int(np.mean(polygon_np[:, 0]))
|
|
center_y = int(np.mean(polygon_np[:, 1]))
|
|
print(f" 📍 중심점: ({center_x}, {center_y})")
|
|
|
|
elif result['method'] == 'bbox':
|
|
print(f" ⬜ 바운딩 박스 폴리곤: {result['polygon']}")
|
|
|
|
elif result['method'] == 'expanded_bbox':
|
|
x, y, w, h = result['bbox']
|
|
print(f" 🔍 확장된 바운딩 박스: {result['polygon']}")
|
|
print(f" 🔍 확장된 크기: {w} x {h} 픽셀")
|
|
|
|
elif result['method'] == 'rotated_bbox':
|
|
print(f" 🔄 회전된 바운딩 박스: {result['polygon']}")
|
|
if 'rotation_info' in result:
|
|
rot_info = result['rotation_info']
|
|
print(f" 🎯 회전 중심: ({rot_info['center'][0]:.1f}, {rot_info['center'][1]:.1f})")
|
|
print(f" 📐 회전 각도: {rot_info['angle']:.1f}°")
|
|
print(f" 📏 회전 박스 크기: {rot_info['size'][0]:.1f} x {rot_info['size'][1]:.1f}")
|
|
|
|
elif result['method'] == 'contour':
|
|
print(f" 🎨 컨투어 폴리곤: {result['polygon']}")
|
|
if 'contour_points' in result:
|
|
print(f" 📊 컨투어 점 개수: {result['contour_points']}개")
|
|
|
|
# 텍스트 영역 크기 정보
|
|
x, y, w, h = result['bbox']
|
|
area = w * h
|
|
print(f" 📊 텍스트 영역 크기: {area} 픽셀²")
|
|
print("-" * 50)
|
|
|
|
print("\n" + "="*60)
|
|
print(f"✅ 총 {len(ocr_results)}개 텍스트 영역 감지 완료")
|
|
print("="*60 + "\n")
|
|
|
|
def visualize_detection(self, image_path: str, ocr_results: List[Dict],
|
|
output_path: str = None) -> np.ndarray:
|
|
"""
|
|
OCR 감지 결과를 시각화
|
|
|
|
Args:
|
|
image_path (str): 원본 이미지 경로
|
|
ocr_results (List[Dict]): OCR 결과
|
|
output_path (str): 저장할 경로 (None이면 저장하지 않음)
|
|
|
|
Returns:
|
|
np.ndarray: 시각화된 이미지
|
|
"""
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
return None
|
|
|
|
# 폴리곤과 텍스트 그리기
|
|
for result in ocr_results:
|
|
polygon = result['polygon']
|
|
text = result['text']
|
|
confidence = result['confidence']
|
|
|
|
# 폴리곤 그리기
|
|
polygon_np = np.array(polygon, dtype=np.int32)
|
|
cv2.polylines(image, [polygon_np], True, (0, 255, 0), 2)
|
|
|
|
# 텍스트와 신뢰도 표시
|
|
x, y = polygon[0]
|
|
cv2.putText(image, f"{text} ({confidence:.2f})",
|
|
(int(x), int(y-10)), cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.5, (0, 255, 0), 1)
|
|
|
|
if output_path:
|
|
cv2.imwrite(output_path, image)
|
|
print(f"OCR 시각화 결과 저장: {output_path}")
|
|
|
|
return image
|
|
|
|
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)
|
|
|
|
print(f"중국어 텍스트 {len(chinese_results)}개 필터링 완료")
|
|
return chinese_results
|
|
|
|
|
|
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
|
|
|
|
def detect_text_with_visualization(self, image_path: str, method: str = 'polygon', output_dir: str = None) -> tuple:
|
|
"""
|
|
이미지에서 텍스트를 감지하고 시각화 이미지도 함께 반환
|
|
|
|
Args:
|
|
image_path (str): 이미지 파일 경로
|
|
method (str): 감지 방식
|
|
output_dir (str): 출력 디렉토리 (시각화 이미지 저장용)
|
|
|
|
Returns:
|
|
tuple: (ocr_results, visualization_image_path)
|
|
"""
|
|
# 기본 OCR 감지 실행
|
|
ocr_results = self.detect_text(image_path, method)
|
|
|
|
# image = cv2.imread(image_path)
|
|
# ocr_raw = ocr2.extract_chinese_text(image)
|
|
|
|
# from modules.ocr_m2 import OCRModule2
|
|
# ocr2 = OCRModule2()
|
|
# image = cv2.imread(image_path) # image_path는 파일 경로(str)
|
|
# ocr_raw = ocr2.extract_chinese_text(image) # image는 numpy array
|
|
# ocr_results = []
|
|
|
|
# for text, polygon, confidence in zip(ocr_raw['texts'], ocr_raw['polygons'], ocr_raw['confidences']):
|
|
# polygon_np = np.array(polygon, dtype=np.int32)
|
|
# x, y, w, h = cv2.boundingRect(polygon_np)
|
|
# ocr_results.append({
|
|
# 'polygon': polygon,
|
|
# 'text': text,
|
|
# 'confidence': confidence,
|
|
# 'bbox': (x, y, w, h)
|
|
# })
|
|
|
|
if not ocr_results:
|
|
return ocr_results, None
|
|
|
|
# 시각화 이미지 생성
|
|
visualization_path = None
|
|
if output_dir:
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# 파일명 생성
|
|
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
|
visualization_filename = f"{base_name}_ocr_detection_{method}.jpg"
|
|
visualization_path = os.path.join(output_dir, visualization_filename)
|
|
|
|
# 시각화 이미지 생성 및 저장
|
|
self.visualize_detection(image_path, ocr_results, visualization_path)
|
|
print(f"💾 OCR 감지 이미지 저장: {visualization_path}")
|
|
|
|
return ocr_results, visualization_path
|
|
|
|
if __name__ == "__main__":
|
|
ocr = OCRModule()
|
|
ocr.test_module() |