135 lines
4.9 KiB
Python
135 lines
4.9 KiB
Python
import os
|
|
import sys
|
|
import asyncio
|
|
import shutil
|
|
from modules.image_processor2 import ImageProcessor
|
|
from modules.loggerModule import Logger1
|
|
from modules.gpt_client import GPTClient
|
|
import logging
|
|
import cv2
|
|
|
|
# 더미 Logger
|
|
class DummyLogger:
|
|
def log(self, msg, level=logging.INFO, exc_info=None):
|
|
print(f"[{logging.getLevelName(level)}] {msg}")
|
|
|
|
|
|
# 테스트용 치환단어
|
|
unwanted_texts = {
|
|
'크리스탈': '이미지삭제',
|
|
'세탁기': '세탁기는개뿔',
|
|
}
|
|
|
|
def get_image_list(img_dir):
|
|
files = [f for f in os.listdir(img_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.webp'))]
|
|
files.sort()
|
|
return [os.path.join(img_dir, f) for f in files]
|
|
|
|
def ensure_dir(path):
|
|
if not os.path.exists(path):
|
|
os.makedirs(path)
|
|
|
|
def save_image(image, path):
|
|
cv2.imwrite(path, image)
|
|
|
|
async def sequential_process(image_paths, processor, output_dir):
|
|
print("[순차처리] 시작")
|
|
results = []
|
|
for idx, img_path in enumerate(image_paths):
|
|
print(f"[{idx+1}] {img_path} 처리 중...")
|
|
# OCR, 번역, 치환, 인페인팅 등 전체 파이프라인 실행
|
|
# process_single_image는 내부적으로 모든 로직을 처리함
|
|
result = await processor.process_single_image(
|
|
page=None, # 실제 Playwright 객체 대신 None
|
|
original_image_url=img_path,
|
|
index=idx,
|
|
is_localServer=True,
|
|
delay=0.1,
|
|
file_prefix="seq",
|
|
use_inpainting=True
|
|
)
|
|
# 결과 파일명 결정
|
|
if isinstance(result, dict):
|
|
status = result.get('status', 'unknown')
|
|
path = result.get('path', img_path)
|
|
if status == 'failed':
|
|
out_name = f"{idx+1}_failed_{os.path.basename(img_path)}"
|
|
shutil.copy(img_path, os.path.join(output_dir, out_name))
|
|
elif status == 'exclude':
|
|
# 이미지삭제: 파일을 output에 저장하지 않음
|
|
print(f"[{idx+1}] 이미지삭제로 제외됨: {img_path}")
|
|
continue
|
|
else:
|
|
out_name = f"{idx+1}_{status}_{os.path.basename(img_path)}"
|
|
shutil.copy(path, os.path.join(output_dir, out_name))
|
|
else:
|
|
# result가 경로(str)라면 원본/번역된 이미지로 간주
|
|
out_name = f"{idx+1}_original_{os.path.basename(img_path)}"
|
|
shutil.copy(result, os.path.join(output_dir, out_name))
|
|
results.append(out_name)
|
|
print("[순차처리] 완료")
|
|
return results
|
|
|
|
async def parallel_process(image_paths, processor, output_dir):
|
|
print("[동시처리] 시작")
|
|
tasks = []
|
|
for idx, img_path in enumerate(image_paths):
|
|
tasks.append(processor.process_single_image(
|
|
page=None,
|
|
original_image_url=img_path,
|
|
index=idx,
|
|
is_localServer=True,
|
|
delay=0.1,
|
|
file_prefix="par",
|
|
use_inpainting=True
|
|
))
|
|
results = await asyncio.gather(*tasks)
|
|
for idx, (img_path, result) in enumerate(zip(image_paths, results)):
|
|
if isinstance(result, dict):
|
|
status = result.get('status', 'unknown')
|
|
path = result.get('path', img_path)
|
|
if status == 'failed':
|
|
out_name = f"{idx+1}_failed_{os.path.basename(img_path)}"
|
|
shutil.copy(img_path, os.path.join(output_dir, out_name))
|
|
elif status == 'exclude':
|
|
print(f"[{idx+1}] 이미지삭제로 제외됨: {img_path}")
|
|
continue
|
|
else:
|
|
out_name = f"{idx+1}_{status}_{os.path.basename(img_path)}"
|
|
shutil.copy(path, os.path.join(output_dir, out_name))
|
|
else:
|
|
out_name = f"{idx+1}_original_{os.path.basename(img_path)}"
|
|
shutil.copy(result, os.path.join(output_dir, out_name))
|
|
print("[동시처리] 완료")
|
|
return results
|
|
|
|
async def main():
|
|
base_dir = os.path.dirname(os.path.abspath(__file__))
|
|
img_dir = os.path.join(base_dir, 'img')
|
|
output_dir = os.path.join(base_dir, 'output')
|
|
ensure_dir(output_dir)
|
|
image_paths = get_image_list(img_dir)
|
|
print(f"테스트 이미지: {image_paths}")
|
|
|
|
# 더미 logger, gpt_client, toggle_states
|
|
logger = DummyLogger()
|
|
set_log = Logger1()
|
|
gpt_client = GPTClient()
|
|
toggle_states = {
|
|
'image_font_path': os.path.join(base_dir, "HakgyoansimDunggeunmisoTTFB.ttf"),
|
|
'TEMP_IMAGE_DIR': output_dir,
|
|
'ocr': True,
|
|
'watermark_text': '테스트워터마크',
|
|
}
|
|
processor = ImageProcessor(set_log, None, toggle_states, gpt_client, base_dir)
|
|
processor.update_unwanted_texts(unwanted_texts)
|
|
|
|
print("1. 순차처리 테스트")
|
|
await sequential_process(image_paths, processor, output_dir)
|
|
|
|
print("2. 동시처리 테스트")
|
|
await parallel_process(image_paths, processor, output_dir)
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(main())
|