111 lines
4.6 KiB
Python
111 lines
4.6 KiB
Python
"""
|
|
이미지 번역 프로젝트 메인 진입점 (라이브러리화)
|
|
다른 프로젝트에서 import하여 사용할 수 있도록 테스트/실행 코드를 모두 제거
|
|
"""
|
|
|
|
import os
|
|
import numpy as np
|
|
import requests
|
|
import cv2
|
|
import base64
|
|
|
|
from modules.ocr_module import OCRModule
|
|
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 전용)"""
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = ''
|
|
os.environ['OMP_NUM_THREADS'] = '4'
|
|
self.ocr = OCRModule()
|
|
self.mask_processor = MaskModule()
|
|
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_telea", 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)
|
|
Returns:
|
|
dict: 결과 파일 경로 등 정보
|
|
"""
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
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 None
|
|
# 마스크 생성 (basic 방식만 사용)
|
|
masks, mask_image_path = self.mask_processor.create_masks_with_save(
|
|
image_path, ocr_results, output_dir, mask_option="basic"
|
|
)
|
|
# 텍스트 번역 (GPT)
|
|
translated_texts = self.gpt_translate_texts(ocr_results, self.gpt_client)
|
|
|
|
# 인페인팅
|
|
inpainted_image = self.inpaint(image_path, masks, method=inpainting_method)
|
|
# 텍스트 렌더링
|
|
final_image = self.text_renderer.render_text(
|
|
inpainted_image, ocr_results, translated_texts, font_path="HakgyoansimDunggeunmisoTTFB.ttf"
|
|
)
|
|
# 결과 저장
|
|
self.output_manager.create_comparison_output(
|
|
image_path, inpainted_image, final_image,
|
|
output_dir, inpainting_method
|
|
)
|
|
return {
|
|
'ocr_visualization': ocr_visualization_path,
|
|
'mask_image': mask_image_path,
|
|
'output_dir': output_dir
|
|
}
|
|
|
|
def gpt_translate_texts(self, ocr_results, gpt_client):
|
|
texts = [result['text'] for result in ocr_results]
|
|
if not texts:
|
|
return []
|
|
prompt = (
|
|
"다음 중국어 문장들을 한국어로 자연스럽고 의미가 잘 전달되게 번역해줘. "
|
|
"순서와 개수는 반드시 그대로 유지하고, 결과는 JSON 배열(리스트)로만 반환해. "
|
|
"중국어 리스트:\n" +
|
|
str(texts)
|
|
)
|
|
response = gpt_client.ask(prompt)
|
|
if isinstance(response, list):
|
|
return response
|
|
elif isinstance(response, dict) and 'result' in response:
|
|
return response['result']
|
|
else:
|
|
print("GPT 번역 결과 파싱 실패, 원본 반환")
|
|
return texts
|
|
|
|
def inpaint(self, image: np.ndarray, mask: np.ndarray, api_url:str = 'http://localhost:8080/api/v1/inpaint', ) -> np.ndarray:
|
|
# 이미지를 base64로 인코딩
|
|
_, img_encoded = cv2.imencode('.png', image)
|
|
_, mask_encoded = cv2.imencode('.png', mask)
|
|
img_b64 = base64.b64encode(img_encoded).decode('utf-8')
|
|
mask_b64 = base64.b64encode(mask_encoded).decode('utf-8')
|
|
payload = {
|
|
"image": img_b64,
|
|
"mask": mask_b64
|
|
}
|
|
response = requests.post(api_url, json=payload)
|
|
if response.status_code != 200:
|
|
print("IOPaint 서버 에러:", response.text)
|
|
return None
|
|
# 응답이 바이너리 PNG 이미지이므로 바로 디코딩
|
|
nparr = np.frombuffer(response.content, np.uint8)
|
|
result = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
|
return result
|