109 lines
3.5 KiB
Python
109 lines
3.5 KiB
Python
# -*- coding: utf-8 -*-
|
|
import logging
|
|
import os
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from ocr_module import OCRModule
|
|
from mask_module import MaskModule
|
|
from inpaint_module import InpaintModule
|
|
|
|
|
|
class SimpleLogger:
|
|
"""간단한 콘솔 로거 (OCRModule/MaskModule/InpaintModule과 동일 인터페이스)"""
|
|
def log(self, msg, level=logging.INFO, exc_info=False):
|
|
prefix = "[INFO]"
|
|
if level >= logging.ERROR:
|
|
prefix = "[ERROR]"
|
|
elif level >= logging.WARNING:
|
|
prefix = "[WARN]"
|
|
print(f"{prefix} {msg}")
|
|
if exc_info:
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
|
|
def main():
|
|
# 1) 샘플 이미지
|
|
image_path = os.path.join("img", "1.jpg")
|
|
if not os.path.isfile(image_path):
|
|
raise FileNotFoundError(f"샘플 이미지가 존재하지 않습니다: {image_path}")
|
|
|
|
logger = SimpleLogger()
|
|
base_dir = os.getcwd()
|
|
|
|
# 2) OCR
|
|
ocr_module = OCRModule(logger=logger, base_dir=base_dir)
|
|
ocr_results = ocr_module.detect_text(image_path, method="polygon")
|
|
chinese_results = ocr_module.filter_chinese_text(ocr_results)
|
|
|
|
print("\n=== 인식된 중국어 텍스트 ===")
|
|
for idx, res in enumerate(chinese_results, 1):
|
|
print(f"{idx:02d}. text = {res['text']}, conf = {res['confidence']:.3f}")
|
|
|
|
# 3) 마스크 생성
|
|
mask_module = MaskModule(logger=logger, base_dir=base_dir)
|
|
mask_np = mask_module.create_masks(
|
|
image_path=image_path,
|
|
ocr_results=chinese_results,
|
|
expansion_size=0, # 필요시 조정
|
|
blur_size=0,
|
|
mask_option="basic"
|
|
)
|
|
if mask_np is None:
|
|
logger.log("마스크 생성 실패, 종료합니다.", level=logging.ERROR)
|
|
return
|
|
|
|
# 저장 (필요시)
|
|
mask_path = os.path.join("img", "1_mask.png")
|
|
cv2.imwrite(mask_path, mask_np)
|
|
logger.log(f"마스크 저장 완료: {mask_path}", level=logging.INFO)
|
|
|
|
# 4) FastDeploy 인페인팅
|
|
inpaint_module = InpaintModule(
|
|
logger=logger,
|
|
base_dir=base_dir,
|
|
gpu_id=0, # GPU 번호
|
|
backend="fd_sd" # "fd_sd" or "opencv"
|
|
)
|
|
|
|
# 프롬프트를 비워두면 기본 "remove text..." 사용
|
|
result_bgr = inpaint_module.inpaint(
|
|
image_path=image_path,
|
|
mask=mask_np,
|
|
prompt="clean background, remove text, no letters",
|
|
negative_prompt="text, letters, watermark, artifacts",
|
|
steps=30,
|
|
guidance_scale=7.5,
|
|
strength=1.0,
|
|
seed=42,
|
|
cv_method="telea"
|
|
)
|
|
|
|
if result_bgr is None:
|
|
logger.log("인페인팅 실패", level=logging.ERROR)
|
|
return
|
|
|
|
out_path = os.path.join("img", "1_inpainted.jpg")
|
|
cv2.imwrite(out_path, result_bgr)
|
|
logger.log(f"인페인팅 결과 저장: {out_path}", level=logging.INFO)
|
|
|
|
# 5) 시각화(원본 + 마스크 + 결과) 간단 저장
|
|
try:
|
|
src = cv2.imread(image_path)
|
|
mask_vis = cv2.cvtColor(mask_np, cv2.COLOR_GRAY2BGR)
|
|
h, w = src.shape[:2]
|
|
canvas = np.zeros((h, w * 3, 3), dtype=np.uint8)
|
|
canvas[:, 0:w] = src
|
|
canvas[:, w:2*w] = mask_vis
|
|
canvas[:, 2*w:3*w] = result_bgr
|
|
vis_path = os.path.join("img", "1_all_vis.jpg")
|
|
cv2.imwrite(vis_path, canvas)
|
|
logger.log(f"원본/마스크/결과 합성이미지 저장: {vis_path}", level=logging.INFO)
|
|
except Exception as e:
|
|
logger.log(f"시각화 실패: {e}", level=logging.WARNING)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|