312 lines
9.7 KiB
Python
312 lines
9.7 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
MIGAN 인페인팅 모듈 테스트 스크립트
|
|
|
|
이 스크립트는 MIGAN 모듈의 기본 사용법과 설정을 보여줍니다.
|
|
실제 운영 환경에서는 celery_worker.py를 통해 사용하세요.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
import cv2
|
|
from pathlib import Path
|
|
|
|
# 현재 디렉토리를 Python 경로에 추가
|
|
current_dir = Path(__file__).parent
|
|
sys.path.insert(0, str(current_dir))
|
|
|
|
try:
|
|
from migan_inpainting_module import MIGANInpaintingModule, build_migan_from_toggle, quick_migan_inpaint
|
|
from loggerModule import Logger
|
|
except ImportError as e:
|
|
print(f"모듈 import 실패: {e}")
|
|
print("worker 디렉토리에서 실행하거나 PYTHONPATH를 설정하세요.")
|
|
sys.exit(1)
|
|
|
|
|
|
def create_test_image_and_mask():
|
|
"""테스트용 이미지와 마스크 생성"""
|
|
# 테스트 이미지 생성 (800x600, 컬러)
|
|
image = np.random.randint(0, 255, (600, 800, 3), dtype=np.uint8)
|
|
|
|
# 중앙에 텍스트 영역 시뮬레이션 (여러 개의 사각형)
|
|
mask = np.zeros((600, 800), dtype=np.uint8)
|
|
|
|
# 텍스트 영역들 (마스크에서 255는 텍스트 영역)
|
|
rectangles = [
|
|
(100, 100, 200, 50), # x, y, w, h
|
|
(300, 200, 150, 40),
|
|
(500, 300, 180, 60),
|
|
(200, 400, 220, 45),
|
|
]
|
|
|
|
for x, y, w, h in rectangles:
|
|
mask[y:y+h, x:x+w] = 255
|
|
|
|
return image, mask
|
|
|
|
|
|
def test_basic_usage():
|
|
"""기본 사용법 테스트"""
|
|
print("=== MIGAN 인페인팅 모듈 기본 테스트 ===")
|
|
|
|
# 로거 생성
|
|
logger = Logger()
|
|
|
|
# 테스트 이미지 및 마스크 생성
|
|
print("1. 테스트 이미지 및 마스크 생성...")
|
|
image, mask = create_test_image_and_mask()
|
|
print(f" 이미지 크기: {image.shape}")
|
|
print(f" 마스크 크기: {mask.shape}")
|
|
print(f" 마스크 영역: {np.sum(mask > 0)} 픽셀")
|
|
|
|
# MIGAN 모듈 초기화
|
|
print("\n2. MIGAN 모듈 초기화...")
|
|
try:
|
|
# 기본 설정으로 모듈 생성
|
|
config = {
|
|
'onnx_path': '/app/worker/models/migan_pipeline_v2.onnx',
|
|
'use_cuda': False, # CPU 모드로 테스트
|
|
'use_tensorrt': False,
|
|
'max_image_size': 1024,
|
|
'output_max_width': 600,
|
|
}
|
|
|
|
inpainter = MIGANInpaintingModule(logger=logger, config=config)
|
|
print(" MIGAN 모듈 초기화 성공")
|
|
|
|
# 처리 통계 확인
|
|
print("\n3. 처리 통계 분석...")
|
|
stats = inpainter.get_processing_stats(image, mask)
|
|
for key, value in stats.items():
|
|
print(f" {key}: {value}")
|
|
|
|
# 인페인팅 실행 (모델 파일이 없으면 실패할 것임)
|
|
print("\n4. MIGAN 인페인팅 실행...")
|
|
try:
|
|
result = inpainter.inpaint_with_migan(image, mask)
|
|
print(f" 인페인팅 성공! 결과 이미지 크기: {result.shape}")
|
|
except FileNotFoundError:
|
|
print(" ⚠️ 모델 파일을 찾을 수 없습니다: /app/worker/models/migan_pipeline_v2.onnx")
|
|
print(" 모델 파일을 올바른 위치에 배치한 후 다시 시도하세요.")
|
|
except Exception as e:
|
|
print(f" 인페인팅 실패: {e}")
|
|
|
|
# 메모리 정리
|
|
inpainter.cleanup_memory()
|
|
|
|
except Exception as e:
|
|
print(f" MIGAN 모듈 초기화 실패: {e}")
|
|
|
|
|
|
def test_toggle_states_usage():
|
|
"""toggle_states를 사용한 설정 테스트"""
|
|
print("\n=== toggle_states 설정 테스트 ===")
|
|
|
|
# 로거 생성
|
|
logger = Logger()
|
|
|
|
# toggle_states 설정 예시들
|
|
test_configs = [
|
|
{
|
|
"name": "CPU 모드",
|
|
"toggle_states": {
|
|
"inpaint_method": "migan",
|
|
"migan_use_cuda": False,
|
|
"migan_use_tensorrt": False,
|
|
"migan_max_image_size": 800,
|
|
"migan_output_max_width": 600,
|
|
}
|
|
},
|
|
{
|
|
"name": "GPU 모드 (TensorRT 비활성화)",
|
|
"toggle_states": {
|
|
"inpaint_method": "migan",
|
|
"migan_use_cuda": True,
|
|
"migan_use_tensorrt": False,
|
|
"migan_max_image_size": 1600,
|
|
"migan_output_max_width": 800,
|
|
}
|
|
},
|
|
{
|
|
"name": "최고 성능 모드 (TensorRT + FP16)",
|
|
"toggle_states": {
|
|
"inpaint_method": "migan",
|
|
"migan_use_cuda": True,
|
|
"migan_use_tensorrt": True,
|
|
"migan_trt_fp16_enable": True,
|
|
"migan_trt_engine_cache_enable": True,
|
|
"migan_max_image_size": 2048,
|
|
"migan_output_max_width": 1200,
|
|
}
|
|
}
|
|
]
|
|
|
|
for i, config in enumerate(test_configs, 1):
|
|
print(f"\n{i}. {config['name']} 테스트")
|
|
print(f" 설정: {config['toggle_states']}")
|
|
|
|
try:
|
|
# toggle_states로부터 MIGAN 모듈 생성
|
|
inpainter = build_migan_from_toggle(config['toggle_states'], logger)
|
|
print(" 모듈 생성 성공")
|
|
|
|
# 설정 확인
|
|
print(f" CUDA 사용: {inpainter.default_config['use_cuda']}")
|
|
print(f" TensorRT 사용: {inpainter.default_config['use_tensorrt']}")
|
|
print(f" 최대 이미지 크기: {inpainter.default_config['max_image_size']}")
|
|
|
|
except Exception as e:
|
|
print(f" 모듈 생성 실패: {e}")
|
|
|
|
|
|
def test_quick_inpaint():
|
|
"""빠른 인페인팅 함수 테스트"""
|
|
print("\n=== 빠른 인페인팅 함수 테스트 ===")
|
|
|
|
# 테스트 이미지 생성
|
|
image, mask = create_test_image_and_mask()
|
|
|
|
# 빠른 인페인팅 설정
|
|
config = {
|
|
'use_cuda': False,
|
|
'use_tensorrt': False,
|
|
'max_image_size': 512,
|
|
}
|
|
|
|
print("1. quick_migan_inpaint 함수 실행...")
|
|
try:
|
|
result = quick_migan_inpaint(image, mask, config=config)
|
|
print(f" 성공! 결과 크기: {result.shape}")
|
|
except FileNotFoundError:
|
|
print(" ⚠️ 모델 파일을 찾을 수 없습니다.")
|
|
except Exception as e:
|
|
print(f" 실패: {e}")
|
|
|
|
|
|
def test_error_handling():
|
|
"""에러 처리 테스트"""
|
|
print("\n=== 에러 처리 테스트 ===")
|
|
|
|
logger = Logger()
|
|
|
|
# 잘못된 모델 경로로 테스트
|
|
print("1. 잘못된 모델 경로 테스트...")
|
|
config = {
|
|
'onnx_path': '/nonexistent/path/model.onnx',
|
|
'use_cuda': False,
|
|
}
|
|
|
|
try:
|
|
inpainter = MIGANInpaintingModule(logger=logger, config=config)
|
|
image, mask = create_test_image_and_mask()
|
|
result = inpainter.inpaint_with_migan(image, mask)
|
|
print(" 예상과 다르게 성공함")
|
|
except Exception as e:
|
|
print(f" 예상된 에러 발생: {type(e).__name__}")
|
|
|
|
# 잘못된 입력 데이터 테스트
|
|
print("\n2. 잘못된 입력 데이터 테스트...")
|
|
try:
|
|
config = {'use_cuda': False, 'use_tensorrt': False}
|
|
inpainter = MIGANInpaintingModule(logger=logger, config=config)
|
|
|
|
# 잘못된 형태의 이미지/마스크
|
|
bad_image = np.zeros((100, 100), dtype=np.uint8) # 2D 이미지
|
|
bad_mask = np.zeros((50, 50), dtype=np.uint8) # 다른 크기
|
|
|
|
result = inpainter.inpaint_with_migan(bad_image, bad_mask)
|
|
print(" 예상과 다르게 성공함")
|
|
except Exception as e:
|
|
print(f" 예상된 에러 발생: {type(e).__name__}")
|
|
|
|
|
|
def show_usage_examples():
|
|
"""사용 예시 출력"""
|
|
print("\n=== 사용 예시 ===")
|
|
|
|
print("\n1. Celery worker에서 ROI vs MIGAN 선택:")
|
|
print("""
|
|
# ROI 기반 인페인팅 (기본)
|
|
toggle_states = {
|
|
"inpaint_method": "roi"
|
|
}
|
|
|
|
# MIGAN 인페인팅
|
|
toggle_states = {
|
|
"inpaint_method": "migan",
|
|
"migan_use_cuda": True,
|
|
"migan_trt_fp16_enable": True
|
|
}
|
|
""")
|
|
|
|
print("\n2. 직접 MIGAN 모듈 사용:")
|
|
print("""
|
|
from worker.migan_inpainting_module import MIGANInpaintingModule
|
|
|
|
# 모듈 생성
|
|
inpainter = MIGANInpaintingModule(config={
|
|
'use_cuda': True,
|
|
'max_image_size': 1600
|
|
})
|
|
|
|
# 인페인팅 실행
|
|
result = inpainter.inpaint_with_migan(image, mask)
|
|
""")
|
|
|
|
print("\n3. toggle_states로 모듈 생성:")
|
|
print("""
|
|
from worker.migan_inpainting_module import build_migan_from_toggle
|
|
|
|
toggle_states = {
|
|
"migan_use_tensorrt": True,
|
|
"migan_trt_fp16_enable": True,
|
|
"migan_max_image_size": 2048
|
|
}
|
|
|
|
inpainter = build_migan_from_toggle(toggle_states)
|
|
""")
|
|
|
|
|
|
def main():
|
|
"""메인 테스트 실행"""
|
|
print("MIGAN 인페인팅 모듈 테스트 시작")
|
|
print("=" * 50)
|
|
|
|
# 환경 정보 출력
|
|
print(f"Python 버전: {sys.version}")
|
|
print(f"작업 디렉토리: {os.getcwd()}")
|
|
print(f"NumPy 버전: {np.__version__}")
|
|
print(f"OpenCV 버전: {cv2.__version__}")
|
|
|
|
# 모델 파일 존재 확인
|
|
model_path = "/app/worker/models/migan_pipeline_v2.onnx"
|
|
if os.path.exists(model_path):
|
|
print(f"✅ 모델 파일 발견: {model_path}")
|
|
else:
|
|
print(f"⚠️ 모델 파일 없음: {model_path}")
|
|
print(" 모델 파일을 배치한 후 테스트하면 실제 인페인팅이 실행됩니다.")
|
|
|
|
# 테스트 실행
|
|
try:
|
|
test_basic_usage()
|
|
test_toggle_states_usage()
|
|
test_quick_inpaint()
|
|
test_error_handling()
|
|
show_usage_examples()
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n사용자에 의해 중단됨")
|
|
except Exception as e:
|
|
print(f"\n예상치 못한 에러: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
print("\n" + "=" * 50)
|
|
print("테스트 완료")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |