51 lines
1.8 KiB
Python
51 lines
1.8 KiB
Python
import os
|
|
import cv2
|
|
import logging
|
|
import numpy as np
|
|
|
|
from ocr_module import OCRModule
|
|
from mask_module import MaskModule
|
|
|
|
class SimpleLogger:
|
|
def log(self, msg, level=logging.INFO, exc_info=False):
|
|
if level == logging.ERROR:
|
|
print("[ERROR]", msg)
|
|
elif level == logging.WARNING:
|
|
print("[WARN]", msg)
|
|
else:
|
|
print("[INFO]", msg)
|
|
if exc_info:
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
def main():
|
|
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()
|
|
|
|
# 1. 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)
|
|
|
|
# 2. 마스크 생성
|
|
mask_module = MaskModule(logger=logger, base_dir=base_dir)
|
|
mask = mask_module.create_masks(image_path, chinese_results, mask_option="basic")
|
|
mask_path = os.path.join("img", "1_mask.png")
|
|
cv2.imwrite(mask_path, mask)
|
|
print(f"마스크 이미지를 {mask_path} 에 저장했습니다.")
|
|
|
|
# 3. simple-lama-inpainting 실행
|
|
# simple-lama-inpainting은 CLI로 실행하는 방식이 일반적입니다.
|
|
# 예시: python -m lama_cleaner --input img/1.jpg --mask img/1_mask.png --output img/1_inpaint.jpg --model lama
|
|
inpaint_path = os.path.join("img", "1_inpaint.jpg")
|
|
cmd = f"python -m simple_lama_inpainting --input {image_path} --mask {mask_path} --output {inpaint_path} --model lama"
|
|
print(f"인페인팅 명령: {cmd}")
|
|
os.system(cmd)
|
|
print(f"인페인팅 이미지를 {inpaint_path} 에 저장했습니다.")
|
|
|
|
if __name__ == "__main__":
|
|
main() |