65 lines
3.1 KiB
Python
65 lines
3.1 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
인페인팅 모듈 - 다양한 인페인팅 방법을 제공 (라이브러리화)
|
|
"""
|
|
import cv2
|
|
import numpy as np
|
|
from typing import List, Dict, Any, Optional
|
|
import os
|
|
from modules.iop_Manager import IOPaintInpainting
|
|
|
|
class InpaintingModule:
|
|
def __init__(self):
|
|
self.iopaint = IOPaintInpainting()
|
|
self.available_methods = {
|
|
'cv2_telea': self._cv2_telea_inpaint,
|
|
'cv2_ns': self._cv2_ns_inpaint,
|
|
'edge_inpaint': self._edge_inpaint,
|
|
'patchmatch': self._patchmatch_inpaint,
|
|
'lama': self._lama_inpaint,
|
|
'stable_diffusion': self._stable_diffusion_inpaint,
|
|
'iopaint': self._iopaint_inpaint
|
|
}
|
|
print("인페인팅 모듈 초기화 완료")
|
|
print(f"사용 가능한 방법: {list(self.available_methods.keys())}")
|
|
|
|
def inpaint(self, image_path: str, mask: np.ndarray, method: str = "cv2_telea") -> np.ndarray:
|
|
if method not in self.available_methods:
|
|
print(f"지원하지 않는 방법: {method}")
|
|
print(f"사용 가능한 방법: {list(self.available_methods.keys())}")
|
|
method = "cv2_telea"
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
print(f"이미지를 읽을 수 없습니다: {image_path}")
|
|
return None
|
|
if mask.shape[:2] != image.shape[:2]:
|
|
mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
|
|
if len(mask.shape) == 3:
|
|
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
|
|
_, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
|
|
try:
|
|
inpainted = self.available_methods[method](image, mask)
|
|
return inpainted
|
|
except Exception as e:
|
|
print(f"인페인팅 실패 ({method}): {e}")
|
|
if method != "cv2_telea":
|
|
return self._cv2_telea_inpaint(image, mask)
|
|
return image
|
|
def _cv2_telea_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
return cv2.inpaint(image, mask, 3, cv2.INPAINT_TELEA)
|
|
def _cv2_ns_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
return cv2.inpaint(image, mask, 3, cv2.INPAINT_NS)
|
|
def _edge_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
# ... (생략: 실제 구현 필요시 추가)
|
|
return self._cv2_telea_inpaint(image, mask)
|
|
def _patchmatch_inpaint(self, image: np.ndarray, mask: np.ndarray, patch_size: int = 3) -> np.ndarray:
|
|
# ... (생략: 실제 구현 필요시 추가)
|
|
return self._cv2_telea_inpaint(image, mask)
|
|
def _lama_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
# ... (생략: 실제 구현 필요시 추가)
|
|
return self._cv2_telea_inpaint(image, mask)
|
|
def _stable_diffusion_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
# ... (생략: 실제 구현 필요시 추가)
|
|
return self._cv2_telea_inpaint(image, mask)
|
|
def _iopaint_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
return self.iopaint.inpaint(image, mask) |