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 from translatepy import Translator 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 = 8100 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) # translated_texts = self.google_translate_texts(ocr_results) 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) # import modules.migan.inpaint_with_migan as migan_module # self.migan_obj = migan_module.get_migan(device="cuda", model_path="modules/migan/migan_traced.pt") # inpainted_image = migan_module.inpaint_with_migan( # local_image_path, masks, device="cuda", model_path="modules/migan/migan_traced.pt", migan_obj=self.migan_obj # ) # 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 def google_translate_texts(self, ocr_results, src_lang='zh-cn', dest_lang='ko'): """ ocr_results에서 추출한 텍스트 리스트를 구글 번역기로 번역하여 반환합니다. Args: ocr_results: OCR 결과 리스트 (각 원소는 {'text': ...} 형태) src_lang: 원본 언어 코드 (기본값: 중국어) dest_lang: 번역할 언어 코드 (기본값: 한국어) Returns: 번역된 텍스트 리스트 """ texts = [result['text'] for result in ocr_results] if not texts: return [] translator = Translator() try: translations = translator.translate(texts, src=src_lang, dest=dest_lang) # googletrans의 translate는 단일/복수 입력 모두 지원 if isinstance(translations, list): return [t.text for t in translations] else: return [translations.text] except Exception as e: print(f"구글 번역 실패: {e}") return texts