726 lines
30 KiB
Python
726 lines
30 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
MIGAN ONNX 파이프라인 인페인팅 모듈
|
|
- MIGAN ONNX 모델을 사용한 전체 이미지 인페인팅
|
|
- TensorRT 최적화 지원
|
|
- 세션 캐시 및 버퍼 재사용
|
|
- 이미지 크기 제한 및 후처리
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import logging
|
|
import tempfile
|
|
from typing import Optional, Dict, Any, Tuple
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
|
|
try:
|
|
import onnx
|
|
ONNX_AVAILABLE = True
|
|
except ImportError:
|
|
ONNX_AVAILABLE = False
|
|
|
|
# OpenCV 내부 최적화 off (호환성)
|
|
cv2.setUseOptimized(False)
|
|
|
|
|
|
def _np_uint8_2d(arr, name="mask"):
|
|
"""2D uint8 배열 검증 및 변환"""
|
|
if arr is None:
|
|
raise ValueError(f"{name} is None")
|
|
if not isinstance(arr, np.ndarray):
|
|
raise TypeError(f"{name} must be np.ndarray, got {type(arr)}")
|
|
if arr.ndim != 2:
|
|
raise ValueError(f"{name} must be 2D, got shape={arr.shape}")
|
|
if arr.dtype != np.uint8:
|
|
arr = arr.astype(np.uint8, copy=False)
|
|
return arr
|
|
|
|
|
|
def _ensure_logger(logger: Optional[object]) -> logging.Logger:
|
|
"""로거 어댑터 생성"""
|
|
if logger and hasattr(logger, "log"):
|
|
return logger
|
|
|
|
pylogger = logging.getLogger("MIGAN")
|
|
if not pylogger.handlers:
|
|
pylogger.setLevel(logging.DEBUG)
|
|
h = logging.StreamHandler(stream=sys.stdout)
|
|
h.setFormatter(logging.Formatter("[%(asctime)s][%(levelname)s] %(message)s"))
|
|
pylogger.addHandler(h)
|
|
|
|
class _Adapter:
|
|
def __init__(self, _lg):
|
|
self._lg = _lg
|
|
def log(self, msg, level=logging.INFO, **kwargs):
|
|
self._lg.log(level, msg)
|
|
return _Adapter(pylogger)
|
|
|
|
|
|
class MIGANInpaintingModule:
|
|
"""
|
|
MIGAN ONNX 파이프라인 인페인팅 모듈
|
|
- 전체 이미지 처리 (ROI 비활성화)
|
|
- TensorRT 최적화 지원
|
|
- 세션 캐시 및 메모리 효율성
|
|
"""
|
|
|
|
_SESSION_CACHE = {} # 세션 캐시
|
|
_TRT_ENGINE_CACHE_DIR = None # TensorRT 엔진 캐시 디렉토리
|
|
|
|
def __init__(self, logger=None, config: Dict[str, Any] = None):
|
|
"""
|
|
MIGAN 인페인팅 모듈 초기화
|
|
|
|
Args:
|
|
logger: 로깅 객체
|
|
config: 설정 딕셔너리
|
|
"""
|
|
self.logger = _ensure_logger(logger)
|
|
|
|
# 기본 설정
|
|
self.default_config = {
|
|
'onnx_path': '/app/worker/models/migan_pipeline_v2.onnx',
|
|
'use_cuda': True,
|
|
'use_tensorrt': True,
|
|
'trt_fp16_enable': True,
|
|
'trt_engine_cache_enable': True,
|
|
'intra_threads': 0,
|
|
'inter_threads': 0,
|
|
'max_image_size': 1600, # 최대 한 변 제한
|
|
'output_max_width': 800, # 출력 가로 최대 크기
|
|
'enable_output_resize': True, # 출력 리사이즈 활성화
|
|
'overwrite_force_margin_px': 0, # 결과 강제 덮어쓰기 마진(px), 0이면 비활성
|
|
}
|
|
|
|
if config:
|
|
self.default_config.update(config)
|
|
|
|
# TensorRT 엔진 캐시 디렉토리 설정
|
|
if self.default_config['trt_engine_cache_enable']:
|
|
self._setup_trt_cache_dir()
|
|
|
|
self.session = None
|
|
self.input_names = None
|
|
self.output_names = None
|
|
|
|
self.logger.log("MIGAN 인페인팅 모듈 초기화 완료", level=logging.INFO)
|
|
|
|
def update_config(self, new_config: Dict[str, Any]):
|
|
"""
|
|
설정을 동적으로 업데이트
|
|
|
|
Args:
|
|
new_config: 새로운 설정 딕셔너리
|
|
"""
|
|
old_config = self.default_config.copy()
|
|
self.default_config.update(new_config)
|
|
|
|
# 세션에 영향을 주는 설정이 변경되었는지 확인
|
|
session_affecting_keys = {
|
|
'onnx_path', 'use_cuda', 'use_tensorrt',
|
|
'trt_fp16_enable', 'intra_threads', 'inter_threads'
|
|
}
|
|
|
|
config_changed = any(
|
|
old_config.get(key) != self.default_config.get(key)
|
|
for key in session_affecting_keys
|
|
)
|
|
|
|
if config_changed:
|
|
self.logger.log("설정 변경으로 인한 세션 재초기화 필요", level=logging.INFO)
|
|
# 기존 세션 정보 리셋 (새로운 세션이 생성되도록)
|
|
self.session = None
|
|
self.input_names = None
|
|
self.output_names = None
|
|
|
|
# TensorRT 캐시 디렉토리 업데이트
|
|
if self.default_config['trt_engine_cache_enable']:
|
|
self._setup_trt_cache_dir()
|
|
|
|
self.logger.log(f"설정 업데이트 완료: {list(new_config.keys())}", level=logging.INFO)
|
|
|
|
def _setup_trt_cache_dir(self):
|
|
"""TensorRT 엔진 캐시 디렉토리 설정"""
|
|
if self._TRT_ENGINE_CACHE_DIR is None:
|
|
cache_dir = os.getenv('TRT_ENGINE_CACHE_DIR', '/app/temp_files/trt_cache')
|
|
os.makedirs(cache_dir, exist_ok=True)
|
|
self._TRT_ENGINE_CACHE_DIR = cache_dir
|
|
self.logger.log(f"TensorRT 엔진 캐시 디렉토리: {cache_dir}", level=logging.INFO)
|
|
|
|
def _diagnose_onnx_model(self, onnx_path: str) -> Dict[str, Any]:
|
|
"""ONNX 모델 TensorRT 호환성 진단"""
|
|
diagnosis = {
|
|
'compatible': True,
|
|
'issues': [],
|
|
'dynamic_axes': [],
|
|
'has_shape_info': False
|
|
}
|
|
|
|
if not ONNX_AVAILABLE:
|
|
diagnosis['issues'].append("ONNX 패키지 미설치 - 진단 불가")
|
|
return diagnosis
|
|
|
|
try:
|
|
model = onnx.load(onnx_path)
|
|
|
|
# 동적 축 확인
|
|
for input_info in model.graph.input:
|
|
for dim in input_info.type.tensor_type.shape.dim:
|
|
if dim.HasField('dim_param'):
|
|
diagnosis['dynamic_axes'].append(f"{input_info.name}:{dim.dim_param}")
|
|
|
|
# Shape inference 상태 확인
|
|
try:
|
|
onnx.checker.check_model(model)
|
|
inferred_model = onnx.shape_inference.infer_shapes(model)
|
|
if inferred_model.graph.value_info:
|
|
diagnosis['has_shape_info'] = True
|
|
else:
|
|
diagnosis['issues'].append("Shape inference 정보 부족")
|
|
diagnosis['compatible'] = False
|
|
except Exception as e:
|
|
diagnosis['issues'].append(f"Shape inference 실패: {str(e)}")
|
|
diagnosis['compatible'] = False
|
|
|
|
# Pad 연산 확인 (TensorRT 문제 원인)
|
|
for node in model.graph.node:
|
|
if node.op_type == 'Pad':
|
|
diagnosis['issues'].append(f"Pad 연산 발견: {node.name} - TensorRT 문제 가능성")
|
|
|
|
except Exception as e:
|
|
diagnosis['issues'].append(f"ONNX 모델 로드 실패: {str(e)}")
|
|
diagnosis['compatible'] = False
|
|
|
|
return diagnosis
|
|
|
|
def _get_session_key(self, config: Dict[str, Any]) -> tuple:
|
|
"""세션 캐시 키 생성"""
|
|
return (
|
|
config['onnx_path'],
|
|
config['use_cuda'],
|
|
config['use_tensorrt'],
|
|
config['trt_fp16_enable'],
|
|
config['intra_threads'],
|
|
config['inter_threads']
|
|
)
|
|
|
|
def _create_session(self, config: Dict[str, Any]) -> ort.InferenceSession:
|
|
"""ONNX Runtime 세션 생성 - TensorRT → CUDA → CPU 순서로 폴백"""
|
|
onnx_path = config['onnx_path']
|
|
|
|
if not os.path.exists(onnx_path):
|
|
raise FileNotFoundError(f"ONNX 파일을 찾을 수 없습니다: {onnx_path}")
|
|
|
|
# 세션 옵션 설정
|
|
session_options = ort.SessionOptions()
|
|
if config['intra_threads'] > 0:
|
|
session_options.intra_op_num_threads = config['intra_threads']
|
|
if config['inter_threads'] > 0:
|
|
session_options.inter_op_num_threads = config['inter_threads']
|
|
|
|
# 최적화 레벨 설정
|
|
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
|
|
|
# 1차 시도: TensorRT (단계적 접근)
|
|
if config.get('use_cuda', True) and config.get('use_tensorrt', True):
|
|
# ONNX 모델 TensorRT 호환성 사전 진단
|
|
diagnosis = self._diagnose_onnx_model(onnx_path)
|
|
if diagnosis['issues']:
|
|
self.logger.log(f"ONNX 모델 진단 결과 - 문제점 {len(diagnosis['issues'])}개 발견:", level=logging.WARNING)
|
|
for issue in diagnosis['issues']:
|
|
self.logger.log(f" - {issue}", level=logging.WARNING)
|
|
|
|
# 호환성 문제가 심각한 경우 TensorRT 건너뛰기
|
|
if not diagnosis.get('compatible', True):
|
|
self.logger.log("TensorRT 호환성 문제로 CUDA로 직접 폴백", level=logging.WARNING)
|
|
else:
|
|
if not diagnosis['issues']:
|
|
self.logger.log("ONNX 모델 TensorRT 호환성 진단 통과", level=logging.INFO)
|
|
else:
|
|
self.logger.log("경고가 있지만 TensorRT 시도 계속", level=logging.INFO)
|
|
|
|
# 1-1. 기본 TensorRT 시도 (최소 옵션)
|
|
try:
|
|
trt_basic_options = {
|
|
'device_id': 0,
|
|
'trt_fp16_enable': config['trt_fp16_enable'],
|
|
'trt_engine_cache_enable': config['trt_engine_cache_enable'],
|
|
}
|
|
|
|
if config['trt_engine_cache_enable'] and self._TRT_ENGINE_CACHE_DIR:
|
|
trt_basic_options['trt_engine_cache_path'] = self._TRT_ENGINE_CACHE_DIR
|
|
|
|
self.logger.log(f"TensorRT 기본 세션 생성 시도: FP16={config['trt_fp16_enable']}", level=logging.INFO)
|
|
|
|
session = ort.InferenceSession(
|
|
onnx_path,
|
|
sess_options=session_options,
|
|
providers=['TensorrtExecutionProvider'],
|
|
provider_options=[trt_basic_options]
|
|
)
|
|
|
|
actual_providers = session.get_providers()
|
|
self.logger.log(f"TensorRT 기본 세션 생성 성공: {actual_providers}", level=logging.INFO)
|
|
return session
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"TensorRT 기본 세션 실패: {e}", level=logging.WARNING)
|
|
|
|
# 1-2. 고급 TensorRT 시도 (상세 옵션)
|
|
try:
|
|
trt_advanced_options = {
|
|
'device_id': 0,
|
|
'trt_fp16_enable': config['trt_fp16_enable'],
|
|
'trt_engine_cache_enable': config['trt_engine_cache_enable'],
|
|
'trt_max_workspace_size': 2 * 1024 * 1024 * 1024, # 2GB
|
|
'trt_max_partition_iterations': 1000,
|
|
'trt_min_subgraph_size': 5,
|
|
'trt_dla_enable': False,
|
|
'trt_int8_enable': False,
|
|
'trt_builder_optimization_level': 3,
|
|
# MIGAN 실제 사용 패턴에 맞춤
|
|
'trt_profile_min_shapes': 'image:1x3x256x256,mask:1x1x256x256',
|
|
'trt_profile_max_shapes': 'image:1x3x1600x1600,mask:1x1x1600x1600',
|
|
'trt_profile_opt_shapes': 'image:1x3x816x1200,mask:1x1x816x1200',
|
|
}
|
|
|
|
if config['trt_engine_cache_enable'] and self._TRT_ENGINE_CACHE_DIR:
|
|
trt_advanced_options['trt_engine_cache_path'] = self._TRT_ENGINE_CACHE_DIR
|
|
|
|
self.logger.log("TensorRT 고급 세션 생성 시도 (동적 shape 프로파일 포함)", level=logging.INFO)
|
|
|
|
session = ort.InferenceSession(
|
|
onnx_path,
|
|
sess_options=session_options,
|
|
providers=['TensorrtExecutionProvider'],
|
|
provider_options=[trt_advanced_options]
|
|
)
|
|
|
|
actual_providers = session.get_providers()
|
|
self.logger.log(f"TensorRT 고급 세션 생성 성공: {actual_providers}", level=logging.INFO)
|
|
return session
|
|
|
|
except Exception as e2:
|
|
self.logger.log(f"TensorRT 고급 세션도 실패: {e2}", level=logging.WARNING)
|
|
# TensorRT 실패 시 자세한 오류 정보 로깅
|
|
if "shape inference" in str(e2).lower():
|
|
self.logger.log("TensorRT shape inference 오류 - ONNX 모델 재생성 필요할 수 있음", level=logging.WARNING)
|
|
elif "pad_output" in str(e2).lower():
|
|
self.logger.log("Pad_output 오류 - 동적 축 문제, CUDA로 폴백", level=logging.WARNING)
|
|
|
|
# 2차 시도: CUDA
|
|
if config.get('use_cuda', True):
|
|
try:
|
|
cuda_options = {
|
|
'device_id': 0,
|
|
'arena_extend_strategy': 'kNextPowerOfTwo',
|
|
'gpu_mem_limit': 2 * 1024 * 1024 * 1024, # 2GB
|
|
'cudnn_conv_algo_search': 'EXHAUSTIVE',
|
|
'do_copy_in_default_stream': True,
|
|
}
|
|
|
|
self.logger.log("CUDA 세션 생성 시도", level=logging.INFO)
|
|
|
|
session = ort.InferenceSession(
|
|
onnx_path,
|
|
sess_options=session_options,
|
|
providers=['CUDAExecutionProvider'],
|
|
provider_options=[cuda_options]
|
|
)
|
|
|
|
actual_providers = session.get_providers()
|
|
self.logger.log(f"CUDA 세션 생성 성공: {actual_providers}", level=logging.INFO)
|
|
return session
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"CUDA 세션 생성 실패: {e}", level=logging.WARNING)
|
|
|
|
# 3차 시도: CPU (최종 폴백)
|
|
try:
|
|
self.logger.log("CPU 세션 생성 시도 (최종 폴백)", level=logging.INFO)
|
|
|
|
session = ort.InferenceSession(
|
|
onnx_path,
|
|
sess_options=session_options,
|
|
providers=['CPUExecutionProvider']
|
|
)
|
|
|
|
actual_providers = session.get_providers()
|
|
self.logger.log(f"CPU 세션 생성 성공: {actual_providers}", level=logging.INFO)
|
|
return session
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"모든 provider에서 세션 생성 실패: {e}", level=logging.ERROR)
|
|
raise RuntimeError(f"ONNX 세션을 생성할 수 없습니다: {e}")
|
|
|
|
def _get_or_create_session(self, config: Dict[str, Any]) -> ort.InferenceSession:
|
|
"""세션 가져오기 또는 생성 (캐시 활용)"""
|
|
cache_key = self._get_session_key(config)
|
|
|
|
if cache_key in self._SESSION_CACHE:
|
|
self.logger.log("캐시된 ONNX 세션 재사용", level=logging.DEBUG)
|
|
return self._SESSION_CACHE[cache_key]
|
|
|
|
session = self._create_session(config)
|
|
self._SESSION_CACHE[cache_key] = session
|
|
|
|
return session
|
|
|
|
def _prepare_session(self, config: Dict[str, Any]):
|
|
"""세션 및 입출력 정보 준비"""
|
|
if self.session is None:
|
|
self.session = self._get_or_create_session(config)
|
|
|
|
# 입출력 정보 저장
|
|
inputs = self.session.get_inputs()
|
|
outputs = self.session.get_outputs()
|
|
|
|
if len(inputs) != 2:
|
|
raise ValueError(f"MIGAN 모델은 2개의 입력이 필요합니다. 현재: {len(inputs)}개")
|
|
if len(outputs) != 1:
|
|
raise ValueError(f"MIGAN 모델은 1개의 출력이 필요합니다. 현재: {len(outputs)}개")
|
|
|
|
self.input_names = [inp.name for inp in inputs]
|
|
self.output_names = [out.name for out in outputs]
|
|
|
|
# 입출력 형태 로깅
|
|
for i, inp in enumerate(inputs):
|
|
self.logger.log(f"입력 {i}: {inp.name}, 형태: {inp.shape}, 타입: {inp.type}", level=logging.DEBUG)
|
|
for i, out in enumerate(outputs):
|
|
self.logger.log(f"출력 {i}: {out.name}, 형태: {out.shape}, 타입: {out.type}", level=logging.DEBUG)
|
|
|
|
def scale_image_if_needed(self, image: np.ndarray, mask: np.ndarray,
|
|
max_size: int) -> Tuple[np.ndarray, np.ndarray, Dict]:
|
|
"""
|
|
이미지 크기 제한 적용
|
|
|
|
Args:
|
|
image: 입력 이미지
|
|
mask: 입력 마스크
|
|
max_size: 최대 크기 (긴 변 기준)
|
|
|
|
Returns:
|
|
(스케일된 이미지, 스케일된 마스크, 스케일 정보)
|
|
"""
|
|
h, w = image.shape[:2]
|
|
max_dimension = max(h, w)
|
|
|
|
if max_dimension <= max_size:
|
|
return image, mask, {'scaled': False, 'original_size': (h, w)}
|
|
|
|
# 스케일 계산
|
|
scale_factor = max_size / max_dimension
|
|
new_h = int(h * scale_factor)
|
|
new_w = int(w * scale_factor)
|
|
|
|
# 8의 배수로 조정 (ONNX 모델 호환성)
|
|
new_h = ((new_h + 7) // 8) * 8
|
|
new_w = ((new_w + 7) // 8) * 8
|
|
|
|
# 리사이즈
|
|
scaled_image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
|
scaled_mask = cv2.resize(mask, (new_w, new_h), interpolation=cv2.INTER_NEAREST)
|
|
|
|
self.logger.log(
|
|
f"이미지 스케일링: {w}x{h} → {new_w}x{new_h} (factor={scale_factor:.3f})",
|
|
level=logging.INFO
|
|
)
|
|
|
|
return scaled_image, scaled_mask, {
|
|
'scaled': True,
|
|
'original_size': (h, w),
|
|
'scale_factor': scale_factor,
|
|
'scaled_size': (new_h, new_w)
|
|
}
|
|
|
|
def restore_original_scale(self, image: np.ndarray, scale_info: Dict) -> np.ndarray:
|
|
"""
|
|
처리된 이미지를 원본 크기로 복원
|
|
|
|
Args:
|
|
image: 처리된 이미지
|
|
scale_info: 스케일 정보
|
|
|
|
Returns:
|
|
복원된 이미지
|
|
"""
|
|
if not scale_info.get('scaled', False):
|
|
return image
|
|
|
|
original_h, original_w = scale_info['original_size']
|
|
restored = cv2.resize(image, (original_w, original_h), interpolation=cv2.INTER_CUBIC)
|
|
|
|
self.logger.log(
|
|
f"이미지 복원: {image.shape[1]}x{image.shape[0]} → {original_w}x{original_h}",
|
|
level=logging.INFO
|
|
)
|
|
|
|
return restored
|
|
|
|
def resize_output_if_needed(self, image: np.ndarray, max_width: int) -> np.ndarray:
|
|
"""
|
|
출력 이미지 크기 조정 (가로 기준)
|
|
|
|
Args:
|
|
image: 출력 이미지
|
|
max_width: 최대 가로 크기
|
|
|
|
Returns:
|
|
크기 조정된 이미지
|
|
"""
|
|
h, w = image.shape[:2]
|
|
|
|
if w <= max_width:
|
|
return image
|
|
|
|
# 비율 유지하며 리사이즈
|
|
scale_factor = max_width / w
|
|
new_w = max_width
|
|
new_h = int(h * scale_factor)
|
|
|
|
resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
|
|
|
self.logger.log(
|
|
f"출력 리사이즈: {w}x{h} → {new_w}x{new_h} (max_width={max_width})",
|
|
level=logging.INFO
|
|
)
|
|
|
|
return resized
|
|
|
|
def inpaint_with_migan(self, image: np.ndarray, mask: np.ndarray,
|
|
config: Dict[str, Any] = None) -> np.ndarray:
|
|
"""
|
|
MIGAN을 사용한 전체 이미지 인페인팅
|
|
|
|
Args:
|
|
image: 입력 이미지 (BGR)
|
|
mask: 마스크 (0~255, 텍스트영역=255)
|
|
config: 설정 오버라이드
|
|
|
|
Returns:
|
|
인페인팅된 이미지 (BGR)
|
|
"""
|
|
start_time = time.time()
|
|
|
|
# 🔥 동적 설정 업데이트 지원
|
|
if config is not None:
|
|
self.update_config(config)
|
|
|
|
effective_config = self.default_config
|
|
|
|
try:
|
|
# 1. 세션 준비
|
|
self._prepare_session(effective_config)
|
|
|
|
# 2. 이미지 크기 제한
|
|
max_size = effective_config['max_image_size']
|
|
scaled_image, scaled_mask, scale_info = self.scale_image_if_needed(
|
|
image, mask, max_size
|
|
)
|
|
|
|
# 3. 임시 파일로 이미지 저장 (MIGAN 모듈이 파일 경로를 요구함)
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
|
|
temp_path = tmp_file.name
|
|
cv2.imwrite(temp_path, scaled_image)
|
|
|
|
try:
|
|
# 4. 마스크 전처리: 이진화 → 반전 (MIGAN 규칙에 맞춤)
|
|
mask_uint8 = _np_uint8_2d(scaled_mask, "mask")
|
|
_, mask_bin = cv2.threshold(mask_uint8, 128, 255, cv2.THRESH_BINARY)
|
|
mask_known255 = 255 - mask_bin # 텍스트영역(255) → hole(0)
|
|
|
|
# 5. 이미지 로드 및 RGB 변환
|
|
bgr = cv2.imread(temp_path, cv2.IMREAD_COLOR)
|
|
if bgr is None:
|
|
raise ValueError(f"임시 이미지 로드 실패: {temp_path}")
|
|
|
|
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
|
|
H, W = rgb.shape[:2]
|
|
|
|
if mask_known255.shape != (H, W):
|
|
raise ValueError(f"마스크 크기 불일치: mask={mask_known255.shape}, img={(H,W)}")
|
|
|
|
# 6. ONNX 추론을 위한 배치 차원 추가 및 차원 순서 변경
|
|
inference_start = time.time()
|
|
|
|
# 이미지: (H, W, 3) → (1, 3, H, W)
|
|
rgb_batch = np.expand_dims(rgb, 0).transpose(0, 3, 1, 2).astype(np.uint8)
|
|
|
|
# 마스크: (H, W) → (1, 1, H, W)
|
|
mask_batch = np.expand_dims(mask_known255, (0, 1)).astype(np.uint8)
|
|
|
|
self.logger.log(
|
|
f"추론 입력 - 이미지: {rgb_batch.shape}, 마스크: {mask_batch.shape}",
|
|
level=logging.DEBUG
|
|
)
|
|
|
|
# ONNX 추론 실행
|
|
inputs = {
|
|
self.input_names[0]: rgb_batch,
|
|
self.input_names[1]: mask_batch
|
|
}
|
|
|
|
outputs = self.session.run(self.output_names, inputs)
|
|
result = outputs[0] # (1, 3, H, W) 형태 예상
|
|
|
|
inference_time = time.time() - inference_start
|
|
|
|
# 7. 출력 후처리: (1, 3, H, W) → (H, W, 3)
|
|
if result.ndim == 4 and result.shape[0] == 1:
|
|
result = result[0].transpose(1, 2, 0) # (1,3,H,W) → (H,W,3)
|
|
elif result.ndim == 3 and result.shape[0] == 3:
|
|
result = result.transpose(1, 2, 0) # (3,H,W) → (H,W,3)
|
|
|
|
if not isinstance(result, np.ndarray) or result.ndim != 3:
|
|
raise ValueError(f"출력 형식 오류: shape={result.shape}, dtype={result.dtype}")
|
|
|
|
# uint8 보장
|
|
if result.dtype != np.uint8:
|
|
result = np.clip(result, 0, 255).astype(np.uint8)
|
|
|
|
# 8. BGR 변환
|
|
result_bgr = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
|
|
|
|
# 9. 원본 크기로 복원
|
|
if scale_info['scaled']:
|
|
result_bgr = self.restore_original_scale(result_bgr, scale_info)
|
|
|
|
# 10. 경계 강제 덮어쓰기(옵션): 잔상 제거용
|
|
try:
|
|
margin_px = int(effective_config.get('overwrite_force_margin_px', 0) or 0)
|
|
except Exception:
|
|
margin_px = 0
|
|
if margin_px > 0:
|
|
try:
|
|
bin_mask = (mask_known255 > 0).astype(np.uint8)
|
|
k = max(3, min(21, (margin_px * 2 + 1)))
|
|
kernel = np.ones((k, k), np.uint8)
|
|
force_zone = cv2.dilate(bin_mask, kernel, iterations=1)
|
|
force_zone_3 = np.repeat(force_zone[:, :, None], 3, axis=2)
|
|
# 원본 스케일과 맞추기
|
|
if scale_info['scaled']:
|
|
force_zone_3 = cv2.resize(force_zone_3, (result_bgr.shape[1], result_bgr.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
bgr_resized = cv2.resize(bgr, (result_bgr.shape[1], result_bgr.shape[0]), interpolation=cv2.INTER_NEAREST)
|
|
else:
|
|
bgr_resized = bgr
|
|
result_bgr = np.where(force_zone_3 == 1, result_bgr, bgr_resized)
|
|
self.logger.log(f"강제 덮어쓰기 적용: margin={margin_px}px", level=logging.DEBUG)
|
|
except Exception as _:
|
|
pass
|
|
|
|
# 11. 출력 크기 조정 (옵션)
|
|
if effective_config.get('enable_output_resize', True):
|
|
max_width = effective_config.get('output_max_width', 800)
|
|
result_bgr = self.resize_output_if_needed(result_bgr, max_width)
|
|
|
|
total_time = time.time() - start_time
|
|
|
|
self.logger.log(
|
|
f"MIGAN 인페인팅 완료: 총 {total_time:.3f}s (추론: {inference_time:.3f}s)",
|
|
level=logging.INFO
|
|
)
|
|
|
|
return result_bgr
|
|
|
|
finally:
|
|
# 임시 파일 정리
|
|
try:
|
|
os.unlink(temp_path)
|
|
except:
|
|
pass
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"MIGAN 인페인팅 실패: {e}", level=logging.ERROR)
|
|
import traceback
|
|
self.logger.log(traceback.format_exc(), level=logging.DEBUG)
|
|
return image # 실패 시 원본 반환
|
|
|
|
def cleanup_memory(self):
|
|
"""
|
|
🔥 경고: 전역 인스턴스에서는 이 메소드를 호출하지 마세요!
|
|
메모리 정리 (개발/테스트 환경에서만 사용)
|
|
"""
|
|
self.logger.log("⚠️ cleanup_memory 호출됨 - 전역 인스턴스에서는 권장하지 않음", level=logging.WARNING)
|
|
import gc
|
|
gc.collect()
|
|
try:
|
|
# GPU 메모리 정리
|
|
if hasattr(self, 'session') and self.session:
|
|
# ONNX Runtime은 자동으로 메모리를 관리하므로 특별한 정리 불필요
|
|
pass
|
|
except Exception as e:
|
|
self.logger.log(f"메모리 정리 중 오류: {e}", level=logging.WARNING)
|
|
|
|
self.logger.log("MIGAN 메모리 정리 완료", level=logging.INFO)
|
|
|
|
def get_processing_stats(self, image: np.ndarray, mask: np.ndarray) -> Dict[str, Any]:
|
|
"""
|
|
MIGAN 처리 통계 정보 반환
|
|
|
|
Args:
|
|
image: 입력 이미지
|
|
mask: 입력 마스크
|
|
|
|
Returns:
|
|
처리 통계 딕셔너리
|
|
"""
|
|
total_area = image.shape[0] * image.shape[1]
|
|
mask_area = np.sum(mask > 128)
|
|
|
|
return {
|
|
'total_image_size': total_area,
|
|
'mask_area': mask_area,
|
|
'mask_coverage_ratio': mask_area / total_area if total_area > 0 else 0.0,
|
|
'processing_method': 'full_image_migan',
|
|
'roi_processing': False,
|
|
'max_size_limit': self.default_config['max_image_size'],
|
|
'output_resize_enabled': self.default_config['enable_output_resize']
|
|
}
|
|
|
|
|
|
# 편의 함수들
|
|
def create_migan_inpainter(logger=None, config=None):
|
|
"""MIGAN 인페인팅 모듈 팩토리 함수"""
|
|
return MIGANInpaintingModule(logger, config)
|
|
|
|
|
|
def build_migan_from_toggle(toggle_states: Dict[str, Any], logger=None) -> MIGANInpaintingModule:
|
|
"""
|
|
toggle_states로부터 설정을 읽어 MIGAN 인페인팅 모듈 생성
|
|
|
|
Args:
|
|
toggle_states: 설정 딕셔너리
|
|
logger: 로거 객체
|
|
|
|
Returns:
|
|
MIGAN 인페인팅 모듈 인스턴스
|
|
"""
|
|
config = {
|
|
'onnx_path': toggle_states.get('migan_onnx_path', '/app/worker/models/migan_pipeline_v2.onnx'),
|
|
'use_cuda': toggle_states.get('migan_use_cuda', True),
|
|
'use_tensorrt': toggle_states.get('migan_use_tensorrt', False),
|
|
'trt_fp16_enable': toggle_states.get('migan_trt_fp16_enable', True),
|
|
'trt_engine_cache_enable': toggle_states.get('migan_trt_engine_cache_enable', True),
|
|
'intra_threads': int(toggle_states.get('migan_intra_threads', 0) or 0),
|
|
'inter_threads': int(toggle_states.get('migan_inter_threads', 0) or 0),
|
|
'max_image_size': int(toggle_states.get('migan_max_image_size', 1600)),
|
|
'output_max_width': int(toggle_states.get('migan_output_max_width', 800)),
|
|
'enable_output_resize': toggle_states.get('migan_enable_output_resize', True),
|
|
}
|
|
|
|
return MIGANInpaintingModule(logger, config)
|
|
|
|
|
|
def quick_migan_inpaint(image: np.ndarray, mask: np.ndarray,
|
|
logger=None, config=None) -> np.ndarray:
|
|
"""간단한 MIGAN 인페인팅 수행"""
|
|
inpainter = create_migan_inpainter(logger, config)
|
|
result = inpainter.inpaint_with_migan(image, mask, config)
|
|
inpainter.cleanup_memory()
|
|
return result |