399 lines
15 KiB
Python
399 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
이미지 번역 프로젝트 메인 진입점
|
||
중국어 상품 이미지를 한국어로 번역하는 시스템
|
||
CPU 전용 모드로 동작합니다.
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import argparse
|
||
from pathlib import Path
|
||
|
||
# CPU만 사용하도록 환경 변수 설정
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = ''
|
||
os.environ['OMP_NUM_THREADS'] = '4' # OpenMP 스레드 수 제한
|
||
|
||
from modules.ocr_module import OCRModule
|
||
# from modules.ocr_m2 import OCRModule2
|
||
from modules.mask_module import MaskModule
|
||
from modules.translation_module import TranslationModule
|
||
from modules.inpainting_module import InpaintingModule
|
||
from modules.text_rendering_module import TextRenderingModule
|
||
from modules.output_module import OutputModule
|
||
from gpt import GPTClient
|
||
|
||
class ImageTranslator:
|
||
def __init__(self):
|
||
"""이미지 번역기 초기화 (CPU 전용)"""
|
||
print("🖥️ CPU 전용 모드로 이미지 번역기 초기화 중...")
|
||
self.ocr = OCRModule()
|
||
# self.ocr2 = OCRModule2()
|
||
self.mask_processor = MaskModule()
|
||
self.translator = TranslationModule()
|
||
self.inpainter = InpaintingModule()
|
||
self.text_renderer = TextRenderingModule()
|
||
self.output_manager = OutputModule()
|
||
self.gpt_client = GPTClient()
|
||
|
||
print("✅ 이미지 번역기 초기화 완료")
|
||
|
||
def translate_image(self, image_path, output_dir="output", inpainting_method="cv2", ocr_method="polygon"):
|
||
"""
|
||
이미지 번역 파이프라인 실행 (결과 이미지 포함)
|
||
|
||
Args:
|
||
image_path (str): 입력 이미지 경로
|
||
output_dir (str): 출력 디렉토리
|
||
inpainting_method (str): 인페인팅 방법 (cv2_telea, cv2_ns, lama, stable_diffusion, edge_inpaint, patchmatch, iopaint)
|
||
ocr_method (str): OCR 감지 방법 (polygon, bbox, expanded_bbox, rotated_bbox, contour)
|
||
"""
|
||
print(f"🚀 이미지 번역 시작: {image_path}")
|
||
print(f"📋 설정 - OCR: {ocr_method}, 인페인팅: {inpainting_method}")
|
||
print(f"🖥️ 실행 모드: CPU 전용")
|
||
|
||
# 출력 디렉토리 생성
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 1. OCR로 텍스트 감지 (시각화 이미지 포함)
|
||
print("\n1️⃣ OCR 텍스트 감지 중...")
|
||
ocr_results, ocr_visualization_path = self.ocr.detect_text_with_visualization(
|
||
image_path, method=ocr_method, output_dir=output_dir
|
||
)
|
||
|
||
if not ocr_results:
|
||
print("❌ 텍스트가 감지되지 않았습니다.")
|
||
return
|
||
|
||
# 2. 마스크 생성 (마스크 이미지 포함)
|
||
print("\n2️⃣ 마스크 생성 중...")
|
||
mask_option = "basic"
|
||
# mask_option = "morphology"
|
||
# mask_option = "adaptive"
|
||
# mask_option = "contour"
|
||
# mask_option = "ultra_blur"
|
||
# mask_option = "feathered"
|
||
|
||
|
||
masks, mask_image_path = self.mask_processor.create_masks_with_save(
|
||
image_path, ocr_results, output_dir, mask_option=mask_option
|
||
)
|
||
|
||
# 3. 텍스트 번역
|
||
print("\n3️⃣ 텍스트 번역 중...")
|
||
translated_texts = self.gpt_translate_texts(ocr_results, self.gpt_client)
|
||
for original, translated in zip([r['text'] for r in ocr_results], translated_texts):
|
||
print(f" '{original}' → '{translated}'")
|
||
|
||
print(f"\n4️⃣ 인페인팅 중... (방법: {inpainting_method})")
|
||
inpainted_image = self.inpainter.inpaint(image_path, masks, method=inpainting_method)
|
||
|
||
|
||
# 5. 텍스트 렌더링
|
||
print("\n5️⃣ 번역된 텍스트 렌더링 중...")
|
||
final_image = self.text_renderer.render_text(
|
||
inpainted_image, ocr_results, translated_texts, font_path="HakgyoansimDunggeunmisoTTFB.ttf"
|
||
# inpainted_image, ocr_results, [tr.text for tr in translated_texts], font_path="HakgyoansimDunggeunmisoTTFB.ttf"
|
||
|
||
)
|
||
|
||
# 6. 결과 출력
|
||
print("\n6️⃣ 결과 생성 중...")
|
||
self.output_manager.create_comparison_output(
|
||
image_path, inpainted_image, final_image,
|
||
output_dir, inpainting_method
|
||
)
|
||
|
||
# 7. 결과 요약
|
||
print(f"\n✅ 번역 완료! 결과 파일들:")
|
||
print(f" 📁 출력 폴더: {output_dir}")
|
||
if ocr_visualization_path:
|
||
print(f" 🔍 OCR 감지 이미지: {os.path.basename(ocr_visualization_path)}")
|
||
if mask_image_path:
|
||
print(f" 🎭 마스크 이미지: {os.path.basename(mask_image_path)}")
|
||
print(f" 🖼️ 최종 번역 결과: 비교 이미지 파일")
|
||
print(f" 📊 메타데이터: JSON 파일")
|
||
|
||
return {
|
||
'ocr_visualization': ocr_visualization_path,
|
||
'mask_image': mask_image_path,
|
||
'output_dir': output_dir
|
||
}
|
||
|
||
def gpt_translate_texts(self, ocr_results, gpt_client):
|
||
"""
|
||
OCR 결과에서 추출한 텍스트들을 GPT를 이용해 한 번에 번역
|
||
Args:
|
||
ocr_results (list): OCR 결과 리스트 (각 원소는 dict, 'text' 키 포함)
|
||
gpt_client (GPTClient): GPTClient 인스턴스
|
||
Returns:
|
||
list: 번역된 텍스트 리스트 (ocr_results와 순서 동일)
|
||
"""
|
||
# 1. 번역할 텍스트 추출
|
||
texts = [result['text'] for result in ocr_results]
|
||
if not texts:
|
||
return []
|
||
|
||
# 2. 프롬프트 구성 (문맥 보존, 리스트 형태로 요청)
|
||
prompt = (
|
||
"다음 중국어 문장들을 한국어로 자연스럽고 의미가 잘 전달되게 번역해줘. "
|
||
"순서와 개수는 반드시 그대로 유지하고, 결과는 JSON 배열(리스트)로만 반환해. "
|
||
"중국어 리스트:\n" +
|
||
str(texts)
|
||
)
|
||
|
||
# 3. GPT 번역 요청
|
||
response = gpt_client.ask(prompt)
|
||
# response가 리스트 형태로 올 수도 있고, 딕셔너리로 올 수도 있음
|
||
if isinstance(response, list):
|
||
return response
|
||
elif isinstance(response, dict) and 'result' in response:
|
||
return response['result']
|
||
else:
|
||
print("GPT 번역 결과 파싱 실패, 원본 반환")
|
||
return texts
|
||
|
||
def test_modules():
|
||
"""각 모듈들의 기능을 개별적으로 테스트"""
|
||
print("=== 모듈 테스트 시작 ===")
|
||
|
||
# OCR 모듈 테스트
|
||
print("\n1. OCR 모듈 테스트")
|
||
ocr = OCRModule()
|
||
ocr.test_module()
|
||
|
||
# 마스크 모듈 테스트
|
||
print("\n2. 마스크 모듈 테스트")
|
||
mask_processor = MaskModule()
|
||
mask_processor.test_module()
|
||
|
||
# 번역 모듈 테스트
|
||
print("\n3. 번역 모듈 테스트")
|
||
translator = TranslationModule()
|
||
translator.test_module()
|
||
|
||
# 인페인팅 모듈 테스트
|
||
print("\n4. 인페인팅 모듈 테스트")
|
||
inpainter = InpaintingModule()
|
||
inpainter.test_module()
|
||
|
||
# 텍스트 렌더링 모듈 테스트
|
||
print("\n5. 텍스트 렌더링 모듈 테스트")
|
||
text_renderer = TextRenderingModule()
|
||
text_renderer.test_module()
|
||
|
||
# 출력 모듈 테스트
|
||
print("\n6. 출력 모듈 테스트")
|
||
output_manager = OutputModule()
|
||
output_manager.test_module()
|
||
|
||
print("\n=== 모듈 테스트 완료 ===")
|
||
|
||
def test_ocr_methods():
|
||
"""다양한 OCR 감지 방식 테스트"""
|
||
print("🔍 === OCR 감지 방식 테스트 ===")
|
||
|
||
ocr = OCRModule()
|
||
test_image = "test_output/test_original.jpg"
|
||
|
||
if not os.path.exists(test_image):
|
||
print(f"❌ 테스트 이미지를 찾을 수 없습니다: {test_image}")
|
||
return
|
||
|
||
methods = ['polygon', 'bbox', 'expanded_bbox', 'rotated_bbox', 'contour']
|
||
|
||
for method in methods:
|
||
print(f"\n🎯 {method.upper()} 방식 테스트")
|
||
print("-" * 50)
|
||
results = ocr.detect_text(test_image, method=method)
|
||
print(f"✅ {method} 방식 완료: {len(results)}개 텍스트 감지\n")
|
||
|
||
def test_inpainting_methods():
|
||
"""다양한 인페인팅 방식 테스트"""
|
||
print("🎨 === 인페인팅 방식 테스트 ===")
|
||
|
||
translator = ImageTranslator()
|
||
test_image = "test_output/test_original.jpg"
|
||
|
||
if not os.path.exists(test_image):
|
||
print(f"❌ 테스트 이미지를 찾을 수 없습니다: {test_image}")
|
||
return
|
||
|
||
methods = ['cv2_telea', 'cv2_ns', 'lama', 'edge_inpaint']
|
||
|
||
for method in methods:
|
||
print(f"\n🖼️ {method.upper()} 방식으로 번역 테스트")
|
||
print("-" * 50)
|
||
output_dir = f"test_output/inpainting_{method}"
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
translator.translate_image(test_image, output_dir, method)
|
||
print(f"✅ {method} 방식 완료\n")
|
||
|
||
def test_combined_methods():
|
||
"""OCR과 인페인팅 방식 조합 테스트"""
|
||
print("🔄 === OCR + 인페인팅 조합 테스트 ===")
|
||
|
||
translator = ImageTranslator()
|
||
test_image = "test_output/test_original.jpg"
|
||
|
||
if not os.path.exists(test_image):
|
||
print(f"❌ 테스트 이미지를 찾을 수 없습니다: {test_image}")
|
||
return
|
||
|
||
ocr_methods = ['polygon', 'expanded_bbox']
|
||
inpainting_methods = ['cv2_telea', 'edge_inpaint']
|
||
|
||
for ocr_method in ocr_methods:
|
||
for inpainting_method in inpainting_methods:
|
||
print(f"\n🎯 조합 테스트: {ocr_method} + {inpainting_method}")
|
||
print("-" * 60)
|
||
output_dir = f"test_output/combined_{ocr_method}_{inpainting_method}"
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
translator.translate_image(test_image, output_dir, inpainting_method, ocr_method)
|
||
print(f"✅ {ocr_method} + {inpainting_method} 조합 완료\n")
|
||
|
||
def benchmark_performance():
|
||
"""성능 벤치마크 테스트"""
|
||
print("⏱️ === 성능 벤치마크 테스트 ===")
|
||
|
||
import time
|
||
|
||
translator = ImageTranslator()
|
||
test_image = "test_output/test_original.jpg"
|
||
|
||
if not os.path.exists(test_image):
|
||
print(f"❌ 테스트 이미지를 찾을 수 없습니다: {test_image}")
|
||
return
|
||
|
||
methods_to_test = [
|
||
('polygon', 'cv2_telea'),
|
||
('expanded_bbox', 'cv2_telea'),
|
||
('polygon', 'edge_inpaint'),
|
||
('expanded_bbox', 'edge_inpaint')
|
||
]
|
||
|
||
results = []
|
||
|
||
for ocr_method, inpainting_method in methods_to_test:
|
||
print(f"\n📊 벤치마크: {ocr_method} + {inpainting_method}")
|
||
|
||
start_time = time.time()
|
||
output_dir = f"test_output/benchmark_{ocr_method}_{inpainting_method}"
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
translator.translate_image(test_image, output_dir, inpainting_method, ocr_method)
|
||
|
||
end_time = time.time()
|
||
elapsed = end_time - start_time
|
||
|
||
results.append((f"{ocr_method}+{inpainting_method}", elapsed))
|
||
print(f"⏱️ 소요 시간: {elapsed:.2f}초")
|
||
|
||
print(f"\n📈 === 벤치마크 결과 요약 ===")
|
||
for method_combo, elapsed in sorted(results, key=lambda x: x[1]):
|
||
print(f"🏆 {method_combo}: {elapsed:.2f}초")
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description='🌏 이미지 번역 시스템 - 중국어를 한국어로 번역')
|
||
|
||
# 기본 옵션들
|
||
parser.add_argument('--image', '-i', type=str, help='번역할 이미지 파일 경로')
|
||
parser.add_argument('--output', '-o', type=str, default='output', help='출력 디렉토리 (기본값: output)')
|
||
|
||
# 방법 선택 옵션들
|
||
parser.add_argument('--inpainting', '-p', type=str, default='cv2_telea',
|
||
choices=['cv2_telea', 'cv2_ns', 'lama', 'stable_diffusion', 'edge_inpaint', 'patchmatch', 'iopaint'],
|
||
help='인페인팅 방법 선택 (기본값: cv2_telea)')
|
||
parser.add_argument('--ocr', '-r', type=str, default='polygon',
|
||
choices=['polygon', 'bbox', 'expanded_bbox', 'rotated_bbox', 'contour'],
|
||
help='OCR 감지 방법 선택 (기본값: polygon)')
|
||
|
||
# 테스트 옵션들
|
||
parser.add_argument('--test', '-t', action='store_true', help='전체 모듈 테스트 실행')
|
||
parser.add_argument('--test-ocr', action='store_true', help='OCR 감지 방식들 테스트')
|
||
parser.add_argument('--test-inpainting', action='store_true', help='인페인팅 방식들 테스트')
|
||
parser.add_argument('--test-combined', action='store_true', help='OCR + 인페인팅 조합 테스트')
|
||
parser.add_argument('--benchmark', '-b', action='store_true', help='성능 벤치마크 테스트')
|
||
|
||
# 도움말 옵션
|
||
parser.add_argument('--list-methods', '-l', action='store_true', help='사용 가능한 방법들 목록 출력')
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 사용 가능한 방법들 목록 출력
|
||
if args.list_methods:
|
||
print("🔍 === 사용 가능한 OCR 감지 방법들 ===")
|
||
print("• polygon - 폴리곤 기반 감지 (기본값)")
|
||
print("• bbox - 바운딩 박스 감지")
|
||
print("• expanded_bbox - 확장된 바운딩 박스 감지")
|
||
print("• rotated_bbox - 회전된 바운딩 박스 감지")
|
||
print("• contour - 컨투어 기반 감지")
|
||
|
||
print("\n🎨 === 사용 가능한 인페인팅 방법들 ===")
|
||
print("• cv2_telea - OpenCV Telea 알고리즘 (기본값)")
|
||
print("• cv2_ns - OpenCV Navier-Stokes 알고리즘")
|
||
print("• lama - LaMa 딥러닝 모델")
|
||
print("• stable_diffusion - Stable Diffusion 모델")
|
||
print("• edge_inpaint - 엣지 기반 인페인팅")
|
||
|
||
print("\n💡 === 사용 예시 ===")
|
||
print("python main.py -i image.jpg -r expanded_bbox -p edge_inpaint")
|
||
print("python main.py --test-ocr")
|
||
print("python main.py --benchmark")
|
||
return
|
||
|
||
# 테스트 옵션들 처리
|
||
if args.test:
|
||
test_modules()
|
||
return
|
||
|
||
if args.test_ocr:
|
||
test_ocr_methods()
|
||
return
|
||
|
||
if args.test_inpainting:
|
||
test_inpainting_methods()
|
||
return
|
||
|
||
if args.test_combined:
|
||
test_combined_methods()
|
||
return
|
||
|
||
if args.benchmark:
|
||
benchmark_performance()
|
||
return
|
||
|
||
# 이미지 번역 실행
|
||
if not args.image:
|
||
print("❌ 이미지 파일 경로를 지정해주세요.")
|
||
print("💡 사용법: python main.py -i 이미지파일.jpg")
|
||
print("📚 도움말: python main.py --help")
|
||
print("📋 방법 목록: python main.py --list-methods")
|
||
return
|
||
|
||
if not os.path.exists(args.image):
|
||
print(f"❌ 이미지 파일을 찾을 수 없습니다: {args.image}")
|
||
return
|
||
|
||
# 출력 디렉토리 생성
|
||
os.makedirs(args.output, exist_ok=True)
|
||
|
||
# 이미지 번역 실행
|
||
print(f"🎯 선택된 설정:")
|
||
print(f" 📁 입력 이미지: {args.image}")
|
||
print(f" 📁 출력 디렉토리: {args.output}")
|
||
print(f" 🔍 OCR 방법: {args.ocr}")
|
||
print(f" 🎨 인페인팅 방법: {args.inpainting}")
|
||
print()
|
||
|
||
translator = ImageTranslator()
|
||
translator.translate_image(args.image, args.output, args.inpainting, args.ocr)
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|
||
|
||
|
||
|