t_local_serv/modules/test_image_processor.py

306 lines
12 KiB
Python

import time
import numpy as np
from typing import Dict, List, Optional, Any
from .ocr_module import OCRModule
from .inpainting_module import InpaintingModule
from .text_renderer import TextRenderer
from .image_utils import ImageUtils
from .translation_module import translate_texts, translate_texts_deepl
import os
class ImageProcessor:
"""이미지 번역을 위한 메인 프로세서"""
def __init__(self, use_gpu: bool = True):
"""
이미지 프로세서 초기화
Args:
use_gpu: GPU 사용 여부
"""
self.use_gpu = use_gpu
# 모듈 초기화
print("모듈들을 초기화하고 있습니다...")
self.ocr_module = OCRModule(use_gpu=use_gpu)
self.inpainting_module = InpaintingModule()
self.text_renderer = TextRenderer()
print("이미지 프로세서 초기화 완료!")
def process_image(self, base64_image: str, options: Optional[Dict] = None) -> Dict:
"""
이미지 번역 전체 프로세스 수행
Args:
base64_image: Base64 인코딩된 입력 이미지
options: 처리 옵션
- inpainting_model: 사용할 인페인팅 모델
- translator: 사용할 번역기
- save_result: 결과 저장 여부
- return_intermediate: 중간 결과 반환 여부
Returns:
처리 결과 딕셔너리
"""
start_time = time.time()
if options is None:
options = {}
try:
# 1. Base64 이미지를 OpenCV 이미지로 변환
print("1. 이미지 변환 중...")
image = ImageUtils.base64_to_image(base64_image)
if image is None:
print("이미지 디코딩 실패: base64 데이터가 올바르지 않습니다.")
return {
'success': False,
'error': '이미지 디코딩 실패: base64 데이터가 올바르지 않습니다.'
}
original_image = image.copy()
# 2. 중국어 텍스트 존재 여부 확인
print("2. 중국어 텍스트 확인 중...")
has_chinese = self.ocr_module.has_chinese_text(image)
if not has_chinese:
return {
'success': True,
'message': '중국어 텍스트가 발견되지 않았습니다.',
'has_chinese_text': False,
'processing_time': time.time() - start_time,
'result_url': None
}
# 3. OCR 수행 - 중국어 텍스트 추출
print("3. OCR 수행 중...")
ocr_result = self.ocr_module.extract_chinese_text(image)
if not ocr_result['texts']:
return {
'success': True,
'message': '추출할 중국어 텍스트가 없습니다.',
'has_chinese_text': True,
'detected_texts': [],
'processing_time': time.time() - start_time,
'result_url': None
}
detected_texts = ocr_result['texts']
polygons = ocr_result['polygons']
masks = ocr_result['masks']
print(f" - 감지된 텍스트: {len(detected_texts)}")
for i, text in enumerate(detected_texts):
print(f" {i+1}. {text}")
# 4. 인페인팅 수행
print("4. 인페인팅 수행 중...")
inpainting_model = options.get('inpainting_model', 'opencv_telea')
if len(masks) == 1:
inpainted_image = self.inpainting_module.inpaint_single(
image, masks[0]
)
else:
# 여러 마스크를 결합하여 한 번에 처리 (더 자연스러운 결과)
inpainted_image = self.inpainting_module.inpaint_combined(
image, masks
)
# 인페인팅 결과 저장 (별도 폴더)
import uuid, os
result_dir = "results/inpainted"
os.makedirs(result_dir, exist_ok=True)
inpainted_filename = os.path.join(result_dir, f"inpainted_{uuid.uuid4().hex}.jpg")
ImageUtils.save_image(inpainted_image, inpainted_filename)
# 5. 번역 수행
print("5. 번역 수행 중...")
translated_texts = translate_texts_deepl(detected_texts)
# translated_texts = translate_texts(detected_texts)
print(f" - 번역 결과:")
for i, (original, translated) in enumerate(zip(detected_texts, translated_texts)):
print(f" {i+1}. {original} -> {translated}")
# OpenCV 인페인팅(telea, ns) 결과도 함께 저장
import cv2
if len(masks) == 1:
mask_for_cv = masks[0]
else:
mask_for_cv = np.maximum.reduce(masks)
inpainted_telea = cv2.inpaint(image, mask_for_cv, 3, cv2.INPAINT_TELEA)
inpainted_ns = cv2.inpaint(image, mask_for_cv, 3, cv2.INPAINT_NS)
ImageUtils.save_image(inpainted_telea, f"results/inpainted_opencv_telea_{uuid.uuid4().hex}.jpg")
ImageUtils.save_image(inpainted_ns, f"results/inpainted_opencv_ns_{uuid.uuid4().hex}.jpg")
# # OpenCV 인페인팅 + 텍스트 렌더링 결과도 저장
# telea_with_text = self.text_renderer.render_translated_texts(
# inpainted_telea, translated_texts, polygons, detected_texts
# )
# ns_with_text = self.text_renderer.render_translated_texts(
# inpainted_ns, translated_texts, polygons, detected_texts
# )
# ImageUtils.save_image(telea_with_text, f"results/inpainted_opencv_telea_text_{uuid.uuid4().hex}.jpg")
# ImageUtils.save_image(ns_with_text, f"results/inpainted_opencv_ns_text_{uuid.uuid4().hex}.jpg")
# 6. 텍스트 렌더링
print("6. 텍스트 렌더링 중...")
final_image = self.text_renderer.render_translated_texts(
inpainted_image, translated_texts, polygons, detected_texts
)
# 7. 결과 저장 및 URL 생성
result_url = None
if options.get('save_result', True):
print("7. 결과 저장 중...")
# 웹서버가 설정되어 있다면 URL 생성
if hasattr(self, 'web_server') and self.web_server:
result_url = self.web_server.save_result_image(final_image)
else:
# 로컬 저장
filename = f"result_{uuid.uuid4().hex}.jpg"
ImageUtils.save_image(final_image, f"results/{filename}")
result_url = f"results/{filename}"
processing_time = time.time() - start_time
print(f"처리 완료! (소요 시간: {processing_time:.2f}초)")
result = {
'success': True,
'message': '이미지 번역이 성공적으로 완료되었습니다.',
'has_chinese_text': True,
'detected_texts': detected_texts,
'translated_texts': translated_texts,
'processing_time': processing_time,
'result_url': result_url
}
# 중간 결과 포함 (옵션)
if options.get('return_intermediate', False):
result['intermediate'] = {
'original_image': ImageUtils.image_to_base64(original_image),
'inpainted_image': ImageUtils.image_to_base64(inpainted_image),
'final_image': ImageUtils.image_to_base64(final_image),
'ocr_result': ocr_result
}
return result
except Exception as e:
print(f"이미지 처리 중 오류 발생: {e}")
return {
'success': False,
'error': str(e),
'processing_time': time.time() - start_time
}
def set_web_server(self, web_server):
"""웹서버 설정"""
self.web_server = web_server
def get_available_models(self) -> Dict:
"""사용 가능한 모델 정보 반환"""
return {
'inpainting_models': self.inpainting_module.get_available_models(),
'current_settings': {
'inpainting_model': self.inpainting_module.current_model,
}
}
def set_models(self, config: Dict) -> Dict:
"""모델 설정"""
try:
result = {'success': True, 'updated': []}
if 'inpainting_model' in config:
self.inpainting_module.set_model(config['inpainting_model'])
result['updated'].append('inpainting_model')
if 'translator' in config:
self.translation_module.set_translator(config['translator'])
result['updated'].append('translator')
result['message'] = f"모델 설정이 업데이트되었습니다: {', '.join(result['updated'])}"
return result
except Exception as e:
return {
'success': False,
'error': str(e)
}
def test_with_sample_image(self, image_path_or_url: str) -> Dict:
"""
샘플 이미지로 테스트 (로컬 파일 또는 URL 지원)
Args:
image_path_or_url: 테스트 이미지 경로 또는 URL
Returns:
테스트 결과
"""
try:
# 이미지 로드 (자동으로 URL 또는 로컬 파일 판단)
image = ImageUtils.load_image_auto(image_path_or_url)
# Base64로 변환
base64_image = ImageUtils.image_to_base64(image)
# 처리 수행
result = self.process_image(base64_image, {
'return_intermediate': True,
'save_result': True
})
return result
except Exception as e:
return {
'success': False,
'error': f"테스트 중 오류: {e}"
}
def compare_models(self, base64_image: str) -> Dict:
"""
여러 모델로 결과 비교
Args:
base64_image: 입력 이미지
Returns:
모델별 비교 결과
"""
try:
image = ImageUtils.base64_to_image(base64_image)
# OCR 수행
ocr_result = self.ocr_module.extract_chinese_text(image)
if not ocr_result['texts']:
return {
'success': False,
'error': '중국어 텍스트가 발견되지 않았습니다.'
}
# 인페인팅 모델 비교
inpainting_results = self.inpainting_module.compare_models(
image, ocr_result['masks'][0] if ocr_result['masks'] else None
)
return {
'success': True,
'detected_texts': ocr_result['texts'],
'inpainting_comparison': {
model: ImageUtils.image_to_base64(result)
for model, result in inpainting_results.items()
}
}
except Exception as e:
return {
'success': False,
'error': str(e)
}