TransWorker/modules/inpaint_with_migan.py

22 lines
538 B
Python

# modules/migan_inpaint.py
import numpy as np
from migan import MIGAN
import cv2
def inpaint_with_migan(image, mask, device="cuda"):
"""
MIGAN을 사용해 인페인팅을 수행합니다.
"""
# MIGAN 모델 로딩
model = MIGAN(device=device)
# 이미지 전처리
if isinstance(image, str):
image = cv2.imread(image)
if isinstance(mask, str):
mask = cv2.imread(mask, cv2.IMREAD_GRAYSCALE)
# 인페인팅 수행
result = model.inpaint(image, mask)
return result