368 lines
18 KiB
Python
368 lines
18 KiB
Python
import os
|
|
import asyncio
|
|
import aiofiles
|
|
import logging
|
|
from urllib.parse import urlparse
|
|
import shutil
|
|
import threading
|
|
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.iop_Manager import IOPaintManager
|
|
|
|
class ImageProcessor:
|
|
"""이미지 다운로드, OCR, 번역 처리를 담당하는 클래스"""
|
|
|
|
def __init__(self, logger, browser_page, toggle_states, gpt_client, base_dir):
|
|
self.logger = logger
|
|
self.page = browser_page
|
|
self.base_dir = base_dir
|
|
self.gpt_client = gpt_client
|
|
self.toggle_states = toggle_states
|
|
|
|
# OCR 관련
|
|
self._ocr_counter = 0
|
|
self._ocr_reset_interval = 100 # 재초기화 주기
|
|
self.ocr_processor = None
|
|
|
|
self.unwanted_texts = []
|
|
self.font_path = self.toggle_states.get('image_font_path', os.path.join(self.base_dir, "HakgyoansimDunggeunmisoTTFB.ttf"))
|
|
self.TEMP_IMAGE_DIR = self.toggle_states.get('TEMP_IMAGE_DIR', "")
|
|
|
|
self.iopaint_manager = IOPaintManager(logger=self.logger, num_instances=1, port_range=(8099, 8199), wait_ready=30)
|
|
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, toggle_states=self.toggle_states)
|
|
|
|
def update_page(self, page1, toggle_states):
|
|
self.toggle_states = toggle_states
|
|
self.page = page1
|
|
self.postImageManager.update_toggle_states(self.toggle_states)
|
|
self.logger.log(f"page객체 및 toggle_states 업데이트", level=logging.DEBUG)
|
|
|
|
def update_unwanted_texts(self, texts):
|
|
self.unwanted_texts = texts
|
|
|
|
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, page, original_image_url, index, is_localServer, delay=1.0, file_prefix="", use_inpainting=False):
|
|
"""
|
|
단일 이미지를 처리합니다 (다운로드 -> OCR -> 인페인팅)
|
|
|
|
Args:
|
|
page: Playwright 페이지 객체
|
|
original_image_url (str): 처리할 이미지 URL
|
|
index (int): 이미지 인덱스
|
|
is_localServer (bool): 로컬 서버 사용 여부
|
|
delay (float): 요청 간격 (초)
|
|
file_prefix (str): 파일명에 추가할 접두사 (예: "detail", "option")
|
|
|
|
Returns:
|
|
dict: 처리 결과를 포함한 딕셔너리
|
|
- status: 'inpainted', 'original', 'exclude', 'error' 중 하나
|
|
- path: 처리된 이미지 파일 경로 또는 원본 이미지 파일 경로
|
|
- error: 오류 메시지 (status가 'error'인 경우에만 포함)
|
|
"""
|
|
local_image_path = None
|
|
|
|
try:
|
|
# 0. 이미지 URL 유효성 체크 (http/https & 이미지 확장자)
|
|
if not original_image_url or not isinstance(original_image_url, str):
|
|
self.logger.log(f"이미지 {index+1} 처리 중단: 이미지 URL 없음 또는 타입 오류", level=logging.WARNING)
|
|
return {'status': 'failed', 'path': original_image_url, 'error': '이미지 URL 없음 또는 타입 오류'}
|
|
|
|
if not self.is_valid_image_path(original_image_url):
|
|
self.logger.log(f"이미지 {index+1} 처리 중단: 유효하지 않은 이미지 주소 - {original_image_url}", level=logging.WARNING)
|
|
return {'status': 'failed', 'path': original_image_url, 'error': '유효하지 않은 이미지 주소'}
|
|
|
|
# 요청 간격 조절 (봇 탐지 회피)
|
|
if index > 0:
|
|
await asyncio.sleep(delay)
|
|
|
|
# OCR 권한 상태 로그
|
|
ocr_enabled = self.toggle_states.get('ocr', False)
|
|
processing_mode = "OCR+인페인팅 모드" if ocr_enabled else "전체 번역 모드"
|
|
self.logger.log(f"이미지 {index+1} 처리 시작: {original_image_url} - {processing_mode}", level=logging.INFO)
|
|
|
|
# 1. 이미지 다운로드
|
|
local_image_path = await self.download_image(page, original_image_url, index, file_prefix)
|
|
if not local_image_path:
|
|
self.logger.log(f"이미지 {index+1} 다운로드 실패, 원본 URL 반환", level=logging.WARNING)
|
|
return {'status': 'failed', 'path': original_image_url, 'error': '다운로드 실패'}
|
|
|
|
# 2. OCR 텍스트 감지
|
|
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 {'status': 'original', 'path': local_image_path}
|
|
|
|
# 4. 텍스트 번역 (GPT)
|
|
translated_texts = self.gpt_translate_texts(ocr_results, self.gpt_client)
|
|
|
|
# 5. OCR 권한에 따른 텍스트 필터링
|
|
if ocr_enabled:
|
|
filtered_translated_texts = self.process_translated_texts(translated_texts, self.unwanted_texts, local_image_path, index)
|
|
if not filtered_translated_texts:
|
|
self.logger.log(f"이미지 {index+1} 제외됨", level=logging.INFO)
|
|
return {'status': 'exclude', 'path': local_image_path}
|
|
else:
|
|
self.logger.log(f"이미지 {index+1} 치환됨", level=logging.INFO)
|
|
else:
|
|
# OCR 권한이 없으면 번역된 텍스트를 그대로 사용
|
|
filtered_translated_texts = translated_texts
|
|
self.logger.log(f"이미지 {index+1} OCR 권한 없음, 전체 번역 모드", 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.iopaint_manager.inpaint(local_image_path, masks)
|
|
self.logger.log(f"인페인팅 완료", level=logging.INFO)
|
|
|
|
# 인페인팅 실패 시 원본 이미지 사용
|
|
if inpainted_image is None:
|
|
self.logger.log(f"인페인팅 실패, 원본 이미지 사용", level=logging.WARNING)
|
|
inpainted_image = cv2.imread(local_image_path)
|
|
if inpainted_image is None:
|
|
self.logger.log(f"원본 이미지 로드 실패: {local_image_path}", level=logging.ERROR)
|
|
return {'status': 'failed', 'path': local_image_path, 'error': '원본 이미지 로드 실패'}
|
|
|
|
# 텍스트 렌더링
|
|
text_rendered_image = self.text_rendering_module.render_text(
|
|
inpainted_image, ocr_results, filtered_translated_texts)
|
|
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)
|
|
self.logger.log(f"이미지 {index+1} 번역 완료: {translated_img_path}", level=logging.INFO)
|
|
return {'status': 'inpainted', 'path': 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 or original_image_url, 'error': str(e)}
|
|
|
|
async def postProcess_and_save_image(self, local_image_path, text_rendered_image, index, file_prefix=""):
|
|
"""로컬 서버 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=self.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
|
|
|
|
|
|
async def download_image(self, page, image_url, index, file_prefix=""):
|
|
"""Playwright를 사용해 이미지를 다운로드합니다"""
|
|
# 로컬 파일 경로면 바로 반환
|
|
if os.path.isfile(image_url):
|
|
self.logger.log(f"로컬 파일 경로 감지, 다운로드 생략: {image_url}", level=logging.INFO)
|
|
return image_url
|
|
|
|
# 로컬 파일 경로가 아니면 다운로드 시도
|
|
try:
|
|
# "https://assets.alicdn.com"으로 시작하는 URL은 건너뛰기
|
|
if image_url.startswith("https://assets.alicdn.com"):
|
|
self.logger.log(f"다운로드 제외 URL: {image_url}", level=logging.DEBUG)
|
|
return None
|
|
|
|
# URL에서 파일명 추출 및 접두사 포함
|
|
parsed_url = urlparse(image_url)
|
|
base_filename = f"image_{index:03d}_{os.path.basename(parsed_url.path)}"
|
|
if not base_filename.endswith(('.jpg', '.jpeg', '.png', '.gif', '.webp')):
|
|
base_filename += '.jpg'
|
|
|
|
# 접두사가 있으면 파일명에 포함
|
|
if file_prefix:
|
|
filename = f"{file_prefix}_{base_filename}"
|
|
else:
|
|
filename = base_filename
|
|
|
|
local_path = os.path.join(self.TEMP_IMAGE_DIR, filename)
|
|
|
|
# Playwright로 이미지 다운로드
|
|
response = await page.request.get(image_url)
|
|
if response.status == 200:
|
|
image_data = await response.body()
|
|
|
|
# 이미지 데이터 유효성 검사
|
|
if self.is_valid_image_data(image_data):
|
|
async with aiofiles.open(local_path, 'wb') as f:
|
|
await f.write(image_data)
|
|
|
|
self.logger.log(f"이미지 다운로드 완료: {filename}", level=logging.INFO)
|
|
return local_path
|
|
else:
|
|
self.logger.log(f"유효하지 않은 이미지 데이터: {image_url}", level=logging.WARNING)
|
|
return None
|
|
else:
|
|
self.logger.log(f"이미지 다운로드 실패 (HTTP {response.status}): {image_url}", level=logging.ERROR, exc_info=True)
|
|
return None
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"이미지 다운로드 중 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
return None
|
|
|
|
def is_valid_image_data(self, image_data):
|
|
"""이미지 데이터가 유효한지 확인합니다"""
|
|
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 initialize_ocr_processor(self):
|
|
self.logger.log(f"OCR 프로세서 초기화", level=logging.INFO)
|
|
from src.ppocr.ocr import ChineseImageOCRProcessor
|
|
self.ocr_processor = ChineseImageOCRProcessor(logger=self.logger, toggle_states=self.toggle_states, use_gpu=False, det_enabled=True, horizontal_only=False)
|
|
return self.ocr_processor
|
|
|
|
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, page, image_urls, is_localServer, 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(
|
|
page, url, i, is_localServer, 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
|