347 lines
15 KiB
Python
347 lines
15 KiB
Python
import os
|
|
import asyncio
|
|
import aiofiles
|
|
import logging
|
|
from urllib.parse import urlparse
|
|
import shutil
|
|
import sys
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import re
|
|
import cv2
|
|
import base64
|
|
import requests
|
|
import numpy as np
|
|
from modules.ocr_module import OCRModule
|
|
from modules.mask_module import MaskModule
|
|
from modules.text_rendering_module import TextRenderingModule
|
|
from modules.postImageManager import PostImageManager
|
|
from modules.lama_inpaint import inpaint_with_simple_lama
|
|
class ImageProcessor:
|
|
"""이미지 다운로드, OCR, 번역 처리를 담당하는 클래스"""
|
|
|
|
def __init__(self, logger, gpt_client, base_dir, font_path):
|
|
self.logger = logger
|
|
self.base_dir = base_dir
|
|
self.gpt_client = gpt_client
|
|
|
|
# OCR 관련
|
|
self.inpaint_sv_port = 8000
|
|
|
|
self.font_path = font_path
|
|
self.TEMP_IMAGE_DIR = os.path.join(self.base_dir, "temp_images")
|
|
os.makedirs(self.TEMP_IMAGE_DIR, exist_ok=True)
|
|
|
|
self.ocr_module = OCRModule(logger=self.logger, base_dir=self.base_dir)
|
|
self.mask_module = MaskModule(logger=self.logger, base_dir=self.base_dir)
|
|
self.text_rendering_module = TextRenderingModule(logger=self.logger, font_path=self.font_path)
|
|
self.postImageManager = PostImageManager(logger=self.logger, font_path=self.font_path)
|
|
|
|
def __del__(self):
|
|
"""소멸자에서 리소스 정리"""
|
|
self.cleanup()
|
|
|
|
def cleanup(self):
|
|
"""리소스 정리"""
|
|
try:
|
|
|
|
# 임시 폴더 삭제
|
|
if hasattr(self, 'TEMP_IMAGE_DIR') and os.path.exists(self.TEMP_IMAGE_DIR):
|
|
shutil.rmtree(self.TEMP_IMAGE_DIR)
|
|
self.logger.log(f"임시 폴더 삭제됨: {self.TEMP_IMAGE_DIR}", level=logging.INFO)
|
|
except Exception as e:
|
|
self.logger.log(f"리소스 정리 중 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
|
|
def is_valid_image_path(self, path):
|
|
# http/https 또는 로컬 파일(.jpg, .png 등) 모두 허용
|
|
if re.match(r'^(http|https)://.*\\.(jpg|jpeg|png|bmp|gif|webp|tiff?)(\\?.*)?$', path, re.IGNORECASE):
|
|
return True
|
|
if os.path.isfile(path) and path.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif', '.webp', '.tif', '.tiff')):
|
|
return True
|
|
return False
|
|
|
|
async def process_single_image(self, toggle_states, unwanted_texts, local_image_path, index, file_prefix=""):
|
|
"""
|
|
단일 이미지를 처리합니다 (다운로드 -> OCR -> 인페인팅)
|
|
|
|
Args:
|
|
toggle_states: 토글 상태 딕셔너리
|
|
local_image_path (str): 처리할 이미지 경로
|
|
index (int): 이미지 인덱스
|
|
unwanted_texts: 치환할 텍스트 딕셔너리
|
|
file_prefix (str): 파일명에 추가할 접두사 (예: "detail", "option")
|
|
|
|
Returns:
|
|
dict: 처리 결과를 포함한 딕셔너리
|
|
- status: 'inpainted', 'original', 'exclude', 'error' 중 하나
|
|
- path: 처리된 이미지 파일 경로 또는 원본 이미지 파일 경로
|
|
- error: 오류 메시지 (status가 'error'인 경우에만 포함)
|
|
"""
|
|
ocr_enabled = toggle_states.get('ocr', False)
|
|
unwanted_texts = unwanted_texts
|
|
|
|
try:
|
|
ocr_results = self.ocr_module.detect_text(local_image_path)
|
|
|
|
# 3. 중국어 텍스트 없는 경우 원본 이미지 반환
|
|
if not self.ocr_module.filter_chinese_text(ocr_results):
|
|
self.logger.log(f"이미지 {index+1} 중국어 텍스트 없음, 원본 이미지 반환", level=logging.INFO)
|
|
return local_image_path
|
|
|
|
# 4. 텍스트 번역 (GPT)
|
|
translated_texts = self.gpt_translate_texts(ocr_results, self.gpt_client)
|
|
|
|
if ocr_enabled:
|
|
filtered_translated_texts = self.process_translated_texts(translated_texts, unwanted_texts, local_image_path, index)
|
|
if not filtered_translated_texts:
|
|
self.logger.log(f"이미지 {index+1} 제외됨", level=logging.INFO)
|
|
return None
|
|
else:
|
|
self.logger.log(f"이미지 {index+1} 치환됨", level=logging.INFO)
|
|
|
|
# 마스크 생성 (basic 방식만 사용)
|
|
masks = self.mask_module.create_masks(
|
|
image_path=local_image_path, ocr_results=ocr_results, mask_option="basic"
|
|
)
|
|
self.logger.log(f"마스크 생성 완료", level=logging.INFO)
|
|
|
|
# 인페인팅
|
|
# inpainted_image = self.call_inpaint_api(local_image_path, masks)
|
|
# self.logger.log(f"인페인팅 완료", level=logging.INFO)
|
|
|
|
inpainted_image = inpaint_with_simple_lama(local_image_path, masks)
|
|
self.logger.log(f"인페인팅 완료", level=logging.INFO)
|
|
|
|
# 텍스트 렌더링
|
|
text_rendered_image = self.text_rendering_module.render_text(
|
|
inpainted_image, ocr_results, filtered_translated_texts, font_path=self.font_path)
|
|
self.logger.log(f"텍스트 렌더링 완료", level=logging.INFO)
|
|
|
|
# 결과 저장
|
|
translated_img_path = await self.postProcess_and_save_image(local_image_path, text_rendered_image, index, file_prefix, toggle_states)
|
|
self.logger.log(f"이미지 {index+1} 번역 완료: {translated_img_path}", level=logging.INFO)
|
|
return translated_img_path
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"이미지 {index+1} 처리 중 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
return {'status': 'failed', 'path': local_image_path, 'error': str(e)}
|
|
|
|
async def postProcess_and_save_image(self, local_image_path, text_rendered_image, index, file_prefix, toggle_states):
|
|
"""로컬 서버 URL을 사용해 이미지를 번역하고 로컬에 저장합니다"""
|
|
try:
|
|
# 파일명에 접두사 포함
|
|
if file_prefix:
|
|
img_path = os.path.join(self.TEMP_IMAGE_DIR, f"translated_{file_prefix}_img_{index+1}.png")
|
|
else:
|
|
img_path = os.path.join(self.TEMP_IMAGE_DIR, f"translated_img_{index+1}.png")
|
|
|
|
watermarked_image_data = self.postImageManager.add_watermark(image_data=text_rendered_image, watermark_text=toggle_states.get("watermark_text", "워터마크"))
|
|
final_image_path = self.postImageManager.save_image_to_path(watermarked_image_data, img_path)
|
|
|
|
return final_image_path
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"이미지 {index+1} 번역 처리 중 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
return local_image_path
|
|
|
|
def is_valid_image_data(self, image_data: bytes) -> bool:
|
|
"""이미지 데이터가 유효한지 확인합니다"""
|
|
if not image_data or len(image_data) < 100:
|
|
return False
|
|
|
|
# 일반적인 이미지 파일 시그니처 확인
|
|
image_signatures = [
|
|
b'\xFF\xD8\xFF', # JPEG
|
|
b'\x89PNG\r\n\x1a\n', # PNG
|
|
b'GIF87a', # GIF87a
|
|
b'GIF89a', # GIF89a
|
|
b'RIFF', # WebP (RIFF 컨테이너)
|
|
]
|
|
|
|
return any(image_data.startswith(sig) for sig in image_signatures)
|
|
|
|
def call_inpaint_api(self, image, mask):
|
|
"""
|
|
인페인팅 API를 호출하여 이미지를 인페인팅합니다.
|
|
"""
|
|
try:
|
|
# 이미지 처리
|
|
if isinstance(image, str):
|
|
image_np = cv2.imread(image)
|
|
if image_np is None:
|
|
self.logger.log(f"이미지 로딩 실패: {image}", level=logging.ERROR)
|
|
return None
|
|
else:
|
|
image_np = image
|
|
|
|
# 마스크 처리
|
|
if isinstance(mask, str):
|
|
mask_np = cv2.imread(mask, cv2.IMREAD_GRAYSCALE)
|
|
if mask_np is None:
|
|
self.logger.log(f"마스크 로딩 실패: {mask}", level=logging.ERROR)
|
|
return None
|
|
else:
|
|
mask_np = mask
|
|
|
|
api_url = f"http://localhost:{self.inpaint_sv_port}/api/v1/inpaint"
|
|
_, img_encoded = cv2.imencode('.png', image_np)
|
|
_, mask_encoded = cv2.imencode('.png', mask_np)
|
|
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:
|
|
self.logger.log(f"IOPaint 서버 에러: {response.text}", level=logging.ERROR)
|
|
return None
|
|
nparr = np.frombuffer(response.content, np.uint8)
|
|
result = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
|
return result
|
|
except Exception as e:
|
|
self.logger.log(f"인페인팅 API 호출 중 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
return None
|
|
|
|
|
|
def process_translated_texts(self, translated_texts, unwanted_texts, local_image_path, index):
|
|
"""
|
|
번역된 단어 리스트(translated_texts)에서 unwanted_texts의 원본값이
|
|
앞이나 뒤에 포함되면 치환값으로 바꿉니다.
|
|
치환값이 '이미지삭제'라면 None 반환(이미지 제외)
|
|
"""
|
|
|
|
new_texts = []
|
|
for text in translated_texts:
|
|
replaced = False
|
|
for origin, replace in unwanted_texts.items():
|
|
# 앞/뒤에 원본값이 있는지 확인
|
|
if text.startswith(origin) or text.endswith(origin):
|
|
self.logger.log(f"[{text}] -> [{replace}] (치환)", level=logging.INFO)
|
|
if replace == "이미지삭제":
|
|
self.logger.log(f"이미지 {index+1} 제외됨: {local_image_path}", level=logging.INFO)
|
|
return None
|
|
# 앞/뒤 원본값만 치환
|
|
if text.startswith(origin):
|
|
new = replace + text[len(origin):]
|
|
elif text.endswith(origin):
|
|
new = text[:-len(origin)] + replace
|
|
new_texts.append(new)
|
|
replaced = True
|
|
break
|
|
if not replaced:
|
|
new_texts.append(text)
|
|
|
|
self.logger.log(f"최종 치환 결과: {new_texts}", level=logging.INFO)
|
|
return new_texts
|
|
|
|
|
|
async def process_image_list(self, image_urls, delay=1.0, file_prefix="", use_inpainting=False):
|
|
"""
|
|
이미지 리스트를 순차적으로 처리합니다.
|
|
"""
|
|
if not image_urls:
|
|
self.logger.log("처리할 이미지가 없습니다.", level=logging.INFO)
|
|
return []
|
|
|
|
processing_mode = "인페인팅" if use_inpainting else "웨일 번역"
|
|
self.logger.log(f"이미지 {len(image_urls)}개를 {processing_mode} 모드로 처리 시작", level=logging.INFO)
|
|
|
|
processed_images = []
|
|
|
|
for i, url in enumerate(image_urls):
|
|
self.logger.log(f"이미지 {i+1}/{len(image_urls)} 처리 중... ({processing_mode} 모드)", level=logging.INFO)
|
|
|
|
result = await self.process_single_image(
|
|
url, i, delay, file_prefix, use_inpainting
|
|
)
|
|
|
|
# 결과 처리
|
|
if isinstance(result, dict):
|
|
status = result.get('status')
|
|
path = result.get('path')
|
|
|
|
if status == 'inpainted':
|
|
processed_images.append(path)
|
|
self.logger.log(f"이미지 {i+1} 인페인팅 처리 완료", level=logging.INFO)
|
|
elif status == 'original':
|
|
processed_images.append(path)
|
|
self.logger.log(f"이미지 {i+1} 원본 사용", level=logging.INFO)
|
|
elif status == 'exclude':
|
|
self.logger.log(f"이미지 {i+1} 제외됨", level=logging.INFO)
|
|
# 제외된 이미지는 리스트에 추가하지 않음
|
|
else: # failed
|
|
self.logger.log(f"이미지 {i+1} 처리 실패: {result.get('error', '알 수 없는 오류')}", level=logging.WARNING)
|
|
# 실패한 이미지는 원본 경로 추가
|
|
processed_images.append(path)
|
|
else:
|
|
# 이전 버전과의 호환성을 위한 처리
|
|
if result:
|
|
processed_images.append(result)
|
|
|
|
self.logger.log(f"이미지 처리 완료: 총 {len(processed_images)}개 ({processing_mode} 모드)", level=logging.INFO)
|
|
return processed_images
|
|
|
|
|
|
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
|
|
|
|
async def save_base64_to_temp_file(self, base64_data: str, suffix: str = "") -> str:
|
|
"""
|
|
base64 인코딩된 이미지 데이터를 임시 파일로 저장
|
|
|
|
Args:
|
|
base64_data (str): base64 인코딩된 이미지 데이터
|
|
suffix (str): 파일명에 추가할 접미사
|
|
|
|
Returns:
|
|
str: 저장된 임시 파일 경로, 실패시 None
|
|
"""
|
|
try:
|
|
import uuid
|
|
import time
|
|
|
|
# data:image/png;base64, 같은 헤더가 있으면 제거
|
|
if base64_data.startswith('data:image'):
|
|
base64_data = base64_data.split(',', 1)[1]
|
|
|
|
# base64 디코딩
|
|
image_bytes = base64.b64decode(base64_data)
|
|
|
|
# 이미지 유효성 검사
|
|
if not self.is_valid_image_data(image_bytes):
|
|
self.logger.log("유효하지 않은 이미지 데이터입니다.", level=logging.ERROR)
|
|
return None
|
|
|
|
# 임시 파일명 생성
|
|
timestamp = int(time.time())
|
|
unique_id = str(uuid.uuid4())[:8]
|
|
temp_filename = f"temp_image_{timestamp}_{unique_id}{suffix}.png"
|
|
temp_path = os.path.join(self.TEMP_IMAGE_DIR, temp_filename)
|
|
|
|
# 파일로 저장
|
|
with open(temp_path, 'wb') as f:
|
|
f.write(image_bytes)
|
|
|
|
self.logger.log(f"base64 이미지 데이터를 임시 파일로 저장: {temp_path}", level=logging.INFO)
|
|
return temp_path
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"base64 이미지 데이터 저장 중 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
return None
|