IMG_Worker/modules/old_modules/image_processor2.py

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 src.modules.ocr_module import OCRModule
from src.modules.mask_module import MaskModule
from src.modules.text_rendering_module import TextRenderingModule
from src.modules.postImageManager import PostImageManager
from src.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