725 lines
30 KiB
Python
725 lines
30 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
인페인팅 모듈 - 다양한 인페인팅 방법을 제공
|
|
CV2 기본 인페인팅부터 딥러닝 모델까지 지원합니다.
|
|
"""
|
|
import cv2
|
|
import numpy as np
|
|
from typing import List, Dict, Any, Optional
|
|
import os
|
|
import requests
|
|
import tempfile
|
|
from PIL import Image
|
|
|
|
from modules.iop 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:
|
|
"""
|
|
마스크를 사용하여 이미지 인페인팅
|
|
|
|
Args:
|
|
image_path (str): 원본 이미지 경로
|
|
mask (np.ndarray): 인페인팅 마스크
|
|
method (str): 인페인팅 방법
|
|
|
|
Returns:
|
|
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)
|
|
|
|
print(f"인페인팅 시작: {method}")
|
|
|
|
try:
|
|
inpainted = self.available_methods[method](image, mask)
|
|
print("인페인팅 완료")
|
|
return inpainted
|
|
except Exception as e:
|
|
print(f"인페인팅 실패 ({method}): {e}")
|
|
# 기본 방법으로 대체
|
|
if method != "cv2_telea":
|
|
print("기본 방법(cv2_telea)으로 대체 시도")
|
|
return self._cv2_telea_inpaint(image, mask)
|
|
return image
|
|
|
|
def _patchmatch_inpaint(self, image: np.ndarray, mask: np.ndarray, patch_size: int = 3) -> np.ndarray:
|
|
"""
|
|
PatchMatch 기반 인페인팅 (PyPatchMatch 활용)
|
|
Args:
|
|
image (np.ndarray): 원본 이미지 (H, W, 3)
|
|
mask (np.ndarray): 인페인팅 마스크 (0/1 또는 0/255, shape=(H,W) 또는 (H,W,1))
|
|
patch_size (int): 패치 크기(기본 3)
|
|
Returns:
|
|
np.ndarray: 인페인팅된 이미지
|
|
"""
|
|
try:
|
|
from patchmatch import patch_match
|
|
except ImportError:
|
|
raise ImportError("PyPatchMatch 패키지가 설치되어 있지 않습니다. pip install PyPatchMatch")
|
|
|
|
# mask를 0/1로 변환
|
|
mask_bin = ((mask > 0) * 1).astype(np.uint8)
|
|
if mask_bin.ndim == 3:
|
|
mask_bin = mask_bin[..., 0]
|
|
# (참고: PyPatchMatch는 마스크가 0=keep, 1=inpaint임)
|
|
# input이 PIL.Image여도 됨
|
|
|
|
if isinstance(image, np.ndarray):
|
|
# uint8, 0~255 이미지로 변환 보장
|
|
if image.dtype != np.uint8:
|
|
image = image.astype(np.uint8)
|
|
else:
|
|
# PIL.Image도 허용 (코드 일관성을 위해 np.ndarray 권장)
|
|
pass
|
|
|
|
result = patch_match.inpaint(image, global_mask=mask_bin, patch_size=patch_size)
|
|
return result
|
|
|
|
|
|
def _cv2_telea_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
"""CV2 TELEA 알고리즘을 사용한 인페인팅"""
|
|
inpainted = cv2.inpaint(image, mask, 3, cv2.INPAINT_TELEA)
|
|
return inpainted
|
|
|
|
def _cv2_ns_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
"""CV2 Navier-Stokes 알고리즘을 사용한 인페인팅"""
|
|
inpainted = cv2.inpaint(image, mask, 3, cv2.INPAINT_NS)
|
|
return inpainted
|
|
|
|
def _edge_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Edge 기반 인페인팅 방법
|
|
엣지 정보를 활용하여 더 자연스러운 인페인팅 수행
|
|
"""
|
|
print("🎨 Edge 기반 인페인팅 시작...")
|
|
|
|
try:
|
|
# 1. 엣지 감지
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
edges = cv2.Canny(gray, 50, 150, apertureSize=3)
|
|
|
|
# 2. 마스크 영역 외부의 엣지만 사용
|
|
mask_inv = cv2.bitwise_not(mask)
|
|
edges_outside = cv2.bitwise_and(edges, mask_inv)
|
|
|
|
# 3. 엣지 확장 (마스크 내부로)
|
|
kernel = np.ones((3, 3), np.uint8)
|
|
edges_dilated = cv2.dilate(edges_outside, kernel, iterations=2)
|
|
|
|
# 4. 구조 텐서 기반 방향 계산
|
|
structure_tensor = self._compute_structure_tensor(gray)
|
|
|
|
# 5. 엣지 방향을 따라 픽셀 전파
|
|
result = self._propagate_along_edges(image, mask, edges_dilated, structure_tensor)
|
|
|
|
# 6. 후처리: 부드럽게 만들기
|
|
result = self._smooth_inpainted_region(result, mask)
|
|
|
|
print("✅ Edge 기반 인페인팅 완료")
|
|
return result
|
|
|
|
except Exception as e:
|
|
print(f"Edge 인페인팅 실패: {e}")
|
|
print("기본 TELEA 방법으로 대체")
|
|
return self._cv2_telea_inpaint(image, mask)
|
|
|
|
def _compute_structure_tensor(self, gray: np.ndarray) -> np.ndarray:
|
|
"""구조 텐서 계산 (엣지 방향 정보)"""
|
|
# 그래디언트 계산
|
|
grad_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
|
|
grad_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
|
|
|
|
# 구조 텐서 요소들
|
|
Ixx = grad_x * grad_x
|
|
Iyy = grad_y * grad_y
|
|
Ixy = grad_x * grad_y
|
|
|
|
# 가우시안 필터로 부드럽게
|
|
Ixx = cv2.GaussianBlur(Ixx, (5, 5), 1.0)
|
|
Iyy = cv2.GaussianBlur(Iyy, (5, 5), 1.0)
|
|
Ixy = cv2.GaussianBlur(Ixy, (5, 5), 1.0)
|
|
|
|
return np.stack([Ixx, Iyy, Ixy], axis=-1)
|
|
|
|
def _propagate_along_edges(self, image: np.ndarray, mask: np.ndarray,
|
|
edges: np.ndarray, structure_tensor: np.ndarray) -> np.ndarray:
|
|
"""엣지를 따라 픽셀 전파"""
|
|
result = image.copy()
|
|
h, w = mask.shape
|
|
|
|
# 마스크 영역 찾기
|
|
mask_coords = np.where(mask > 0)
|
|
|
|
for y, x in zip(mask_coords[0], mask_coords[1]):
|
|
# 주변 픽셀에서 값 추정
|
|
neighbors = []
|
|
weights = []
|
|
|
|
# 8방향 검색
|
|
for dy in [-1, 0, 1]:
|
|
for dx in [-1, 0, 1]:
|
|
if dy == 0 and dx == 0:
|
|
continue
|
|
|
|
ny, nx = y + dy, x + dx
|
|
if 0 <= ny < h and 0 <= nx < w:
|
|
if mask[ny, nx] == 0: # 마스크 외부 픽셀
|
|
# 구조 텐서에서 방향성 가중치 계산
|
|
Ixx, Iyy, Ixy = structure_tensor[ny, nx]
|
|
|
|
# 고유값 계산 (방향성 강도)
|
|
trace = Ixx + Iyy
|
|
det = Ixx * Iyy - Ixy * Ixy
|
|
lambda1 = (trace + np.sqrt(trace**2 - 4*det)) / 2
|
|
lambda2 = (trace - np.sqrt(trace**2 - 4*det)) / 2
|
|
|
|
# 방향성이 강한 경우 높은 가중치
|
|
directional_weight = abs(lambda1 - lambda2) / (lambda1 + lambda2 + 1e-6)
|
|
|
|
# 거리 가중치
|
|
distance_weight = 1.0 / (1.0 + np.sqrt(dy**2 + dx**2))
|
|
|
|
# 엣지 가중치
|
|
edge_weight = 1.0 + edges[ny, nx] / 255.0
|
|
|
|
total_weight = directional_weight * distance_weight * edge_weight
|
|
|
|
neighbors.append(result[ny, nx])
|
|
weights.append(total_weight)
|
|
|
|
if neighbors:
|
|
weights = np.array(weights)
|
|
neighbors = np.array(neighbors)
|
|
|
|
# 가중 평균으로 픽셀 값 계산
|
|
if weights.sum() > 0:
|
|
weights = weights / weights.sum()
|
|
result[y, x] = np.sum(neighbors * weights[:, None], axis=0)
|
|
|
|
return result
|
|
|
|
def _smooth_inpainted_region(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
"""인페인팅된 영역 부드럽게 처리"""
|
|
# 마스크 영역에만 블러 적용
|
|
blurred = cv2.GaussianBlur(image, (5, 5), 0)
|
|
|
|
# 마스크를 부드럽게 만들어 자연스러운 블렌딩
|
|
mask_float = mask.astype(np.float32) / 255.0
|
|
mask_smooth = cv2.GaussianBlur(mask_float, (7, 7), 0)
|
|
|
|
result = image.copy()
|
|
for c in range(3):
|
|
result[:, :, c] = (image[:, :, c] * (1 - mask_smooth) +
|
|
blurred[:, :, c] * mask_smooth)
|
|
|
|
return result.astype(np.uint8)
|
|
|
|
def _lama_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
"""
|
|
LaMa (Large Mask Inpainting) 모델을 사용한 인페인팅
|
|
실제 lama-cleaner 라이브러리 사용
|
|
"""
|
|
try:
|
|
# lama-cleaner 라이브러리 임포트 시도
|
|
try:
|
|
from lama_cleaner.model_manager import ModelManager
|
|
from lama_cleaner.schema import Config, HDStrategy, LDMSampler
|
|
print("🎨 LaMa-Cleaner 라이브러리 로드 성공")
|
|
use_real_lama = True
|
|
except ImportError:
|
|
print("⚠️ lama-cleaner 라이브러리가 설치되지 않았습니다.")
|
|
print(" 설치 방법: pip install lama-cleaner")
|
|
print(" 고급 CV2 기반 인페인팅으로 대체합니다.")
|
|
use_real_lama = False
|
|
|
|
if use_real_lama:
|
|
# 실제 LaMa 모델 사용
|
|
print("🚀 LaMa 모델 초기화 중...")
|
|
|
|
# CPU만 사용하도록 설정
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = ''
|
|
|
|
# 모델 설정
|
|
config = Config(
|
|
ldm_steps=20,
|
|
ldm_sampler=LDMSampler.plms,
|
|
hd_strategy=HDStrategy.ORIGINAL,
|
|
hd_strategy_crop_margin=32,
|
|
hd_strategy_crop_trigger_size=512,
|
|
hd_strategy_resize_limit=2048,
|
|
prompt="",
|
|
negative_prompt="",
|
|
use_croper=False,
|
|
use_extender=False,
|
|
extender_x=0,
|
|
extender_y=0,
|
|
extender_width=0,
|
|
extender_height=0,
|
|
cpu_offload=True, # CPU 사용 강제
|
|
disable_nsfw=True,
|
|
cpu_textencoder=True,
|
|
local_files_only=True
|
|
)
|
|
|
|
# 모델 매니저 초기화 (CPU 전용)
|
|
model_manager = ModelManager(
|
|
name="lama",
|
|
device="cpu", # CPU 강제 사용
|
|
no_half=True,
|
|
cpu_offload=True,
|
|
disable_nsfw=True,
|
|
local_files_only=True
|
|
)
|
|
|
|
# 이미지를 PIL 형식으로 변환
|
|
from PIL import Image as PILImage
|
|
|
|
# OpenCV BGR을 RGB로 변환
|
|
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
pil_image = PILImage.fromarray(image_rgb)
|
|
|
|
# 마스크를 PIL 형식으로 변환
|
|
if len(mask.shape) == 3:
|
|
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
|
|
pil_mask = PILImage.fromarray(mask)
|
|
|
|
print("🎯 LaMa 인페인팅 실행 중... (CPU 모드)")
|
|
|
|
# LaMa 인페인팅 실행
|
|
result_pil = model_manager(pil_image, pil_mask, config)
|
|
|
|
# PIL을 OpenCV 형식으로 변환
|
|
result_rgb = np.array(result_pil)
|
|
result_bgr = cv2.cvtColor(result_rgb, cv2.COLOR_RGB2BGR)
|
|
|
|
print("✅ LaMa 인페인팅 완료")
|
|
return result_bgr
|
|
|
|
else:
|
|
# 고급 CV2 기반 대체 인페인팅
|
|
return self._advanced_cv2_inpaint(image, mask)
|
|
|
|
except Exception as e:
|
|
print(f"❌ LaMa 인페인팅 실패: {e}")
|
|
print("🔄 고급 CV2 인페인팅으로 대체")
|
|
return self._advanced_cv2_inpaint(image, mask)
|
|
|
|
def _edgeconnect_inpaint(self, image, mask):
|
|
try:
|
|
import onnxruntime as ort
|
|
# ONNX 모델 경로
|
|
model_path = "edgeconnect_inpaint_cpu.onnx"
|
|
session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
|
|
|
|
# 전처리: 256x256, [0,1], np.float32, RGB
|
|
input_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
input_image = cv2.resize(input_image, (256, 256)).astype(np.float32) / 255.0
|
|
mask_resized = cv2.resize(mask, (256, 256)).astype(np.float32) / 255.0
|
|
mask_resized = np.expand_dims(mask_resized, axis=2)
|
|
input_data = np.concatenate([input_image, mask_resized], axis=2)
|
|
input_data = np.transpose(input_data, (2, 0, 1))[None, :, :, :]
|
|
|
|
result = session.run(None, {"input": input_data})[0]
|
|
result = result[0].transpose(1, 2, 0) * 255
|
|
result = np.clip(result, 0, 255).astype(np.uint8)
|
|
result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
|
|
return cv2.resize(result, (image.shape[1], image.shape[0]))
|
|
except Exception as e:
|
|
print("EdgeConnect 인페인팅 실패:", e)
|
|
return self._cv2_telea_inpaint(image, mask)
|
|
|
|
def _fastmarching_inpaint(self, image, mask):
|
|
try:
|
|
return cv2.xphoto.inpaint(image, mask, algorithmType=cv2.xphoto.INPAINT_FSR_BEST)
|
|
except Exception as e:
|
|
print("FastMarching 인페인팅 실패:", e)
|
|
return self._cv2_telea_inpaint(image, mask)
|
|
|
|
|
|
|
|
def _advanced_cv2_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
"""
|
|
고급 CV2 기반 인페인팅 (LaMa 대체용)
|
|
여러 단계의 인페인팅과 후처리로 품질 향상
|
|
"""
|
|
print("🎨 고급 CV2 인페인팅 시작...")
|
|
|
|
# 1단계: 마스크 전처리
|
|
kernel = np.ones((3, 3), np.uint8)
|
|
processed_mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
|
|
processed_mask = cv2.GaussianBlur(processed_mask, (3, 3), 0)
|
|
|
|
# 2단계: 다중 스케일 인페인팅
|
|
result = image.copy()
|
|
|
|
# 작은 스케일부터 시작
|
|
for scale in [0.5, 0.75, 1.0]:
|
|
if scale < 1.0:
|
|
h, w = int(image.shape[0] * scale), int(image.shape[1] * scale)
|
|
scaled_image = cv2.resize(result, (w, h))
|
|
scaled_mask = cv2.resize(processed_mask, (w, h))
|
|
else:
|
|
scaled_image = result
|
|
scaled_mask = processed_mask
|
|
|
|
# TELEA와 NS 알고리즘 조합
|
|
telea_result = cv2.inpaint(scaled_image, scaled_mask, 5, cv2.INPAINT_TELEA)
|
|
ns_result = cv2.inpaint(scaled_image, scaled_mask, 5, cv2.INPAINT_NS)
|
|
|
|
# 두 결과를 블렌딩
|
|
alpha = 0.6 # TELEA에 더 높은 가중치
|
|
blended = cv2.addWeighted(telea_result, alpha, ns_result, 1-alpha, 0)
|
|
|
|
if scale < 1.0:
|
|
result = cv2.resize(blended, (image.shape[1], image.shape[0]))
|
|
else:
|
|
result = blended
|
|
|
|
# 3단계: 텍스처 보정
|
|
# 주변 영역의 텍스처를 분석하여 일관성 개선
|
|
mask_dilated = cv2.dilate(processed_mask, kernel, iterations=5)
|
|
mask_border = mask_dilated - processed_mask
|
|
|
|
if np.sum(mask_border) > 0:
|
|
# 경계 영역의 텍스처 특성 분석
|
|
border_pixels = result[mask_border > 0]
|
|
if len(border_pixels) > 0:
|
|
# 텍스처 방향성 계산
|
|
gray = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)
|
|
grad_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
|
|
grad_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
|
|
|
|
# 방향성 기반 필터링
|
|
orientation = np.arctan2(grad_y, grad_x)
|
|
magnitude = np.sqrt(grad_x**2 + grad_y**2)
|
|
|
|
# 강한 방향성을 가진 영역에 방향성 블러 적용
|
|
strong_texture_mask = (magnitude > np.percentile(magnitude, 70)).astype(np.uint8) * 255
|
|
texture_intersection = cv2.bitwise_and(processed_mask, strong_texture_mask)
|
|
|
|
if np.sum(texture_intersection) > 0:
|
|
# 방향성 블러 적용
|
|
kernel_size = 15
|
|
angle = np.mean(orientation[texture_intersection > 0]) * 180 / np.pi
|
|
|
|
# 방향성 커널 생성
|
|
M = cv2.getRotationMatrix2D((kernel_size//2, kernel_size//2), angle, 1)
|
|
directional_kernel = np.zeros((kernel_size, kernel_size))
|
|
directional_kernel[kernel_size//2, :] = 1
|
|
directional_kernel = cv2.warpAffine(directional_kernel, M, (kernel_size, kernel_size))
|
|
directional_kernel = directional_kernel / np.sum(directional_kernel)
|
|
|
|
# 방향성 필터 적용
|
|
filtered_result = cv2.filter2D(result, -1, directional_kernel)
|
|
|
|
# 텍스처 영역에만 적용
|
|
texture_mask_float = texture_intersection.astype(np.float32) / 255.0
|
|
texture_mask_float = cv2.GaussianBlur(texture_mask_float, (5, 5), 0)
|
|
|
|
for c in range(3):
|
|
result[:, :, c] = (result[:, :, c] * (1 - texture_mask_float) +
|
|
filtered_result[:, :, c] * texture_mask_float)
|
|
|
|
# 4단계: 최종 후처리
|
|
# 색상 일관성 개선
|
|
result = cv2.bilateralFilter(result, 9, 75, 75)
|
|
|
|
# 경계 부드럽게 처리
|
|
mask_float = processed_mask.astype(np.float32) / 255.0
|
|
mask_float = cv2.GaussianBlur(mask_float, (7, 7), 0)
|
|
|
|
# 원본과 부드럽게 블렌딩
|
|
for c in range(3):
|
|
result[:, :, c] = (image[:, :, c] * (1 - mask_float) +
|
|
result[:, :, c] * mask_float)
|
|
|
|
print("✅ 고급 CV2 인페인팅 완료")
|
|
return result.astype(np.uint8)
|
|
|
|
def _iopaint_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
"""
|
|
IOPaint 인페인팅 모델 사용
|
|
"""
|
|
return self.iopaint.inpaint(image, mask)
|
|
|
|
def _stable_diffusion_inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Stable Diffusion 인페인팅 모델 사용
|
|
실제 구현에서는 diffusers 라이브러리 사용
|
|
"""
|
|
try:
|
|
# 실제 구현에서는 diffusers 설치 필요
|
|
# pip install diffusers transformers accelerate
|
|
|
|
print("Stable Diffusion 인페인팅 시뮬레이션")
|
|
|
|
# 고품질 인페인팅 시뮬레이션
|
|
# 실제로는 StableDiffusionInpaintPipeline 사용
|
|
|
|
# 마스크 영역 확장 및 부드럽게 처리
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
|
|
expanded_mask = cv2.morphologyEx(mask, cv2.MORPH_DILATE, kernel)
|
|
|
|
# 여러 패스 인페인팅으로 품질 향상
|
|
result = image.copy()
|
|
|
|
# 1단계: 기본 인페인팅
|
|
result = cv2.inpaint(result, expanded_mask, 7, cv2.INPAINT_TELEA)
|
|
|
|
# 2단계: 텍스처 개선
|
|
result = cv2.bilateralFilter(result, 9, 75, 75)
|
|
|
|
# 3단계: 경계 부드럽게
|
|
blurred = cv2.GaussianBlur(result, (5, 5), 0)
|
|
mask_float = expanded_mask.astype(np.float32) / 255.0
|
|
mask_float = cv2.GaussianBlur(mask_float, (5, 5), 0)
|
|
|
|
for c in range(3):
|
|
result[:, :, c] = (result[:, :, c] * (1 - mask_float) +
|
|
blurred[:, :, c] * mask_float)
|
|
|
|
return result.astype(np.uint8)
|
|
|
|
except Exception as e:
|
|
print(f"Stable Diffusion 인페인팅 실패: {e}")
|
|
return self._cv2_telea_inpaint(image, mask)
|
|
|
|
def batch_inpaint(self, image_paths: List[str], masks: List[np.ndarray],
|
|
method: str = "cv2_telea") -> List[np.ndarray]:
|
|
"""
|
|
여러 이미지 일괄 인페인팅
|
|
|
|
Args:
|
|
image_paths (List[str]): 이미지 경로 리스트
|
|
masks (List[np.ndarray]): 마스크 리스트
|
|
method (str): 인페인팅 방법
|
|
|
|
Returns:
|
|
List[np.ndarray]: 인페인팅된 이미지 리스트
|
|
"""
|
|
results = []
|
|
|
|
for i, (image_path, mask) in enumerate(zip(image_paths, masks)):
|
|
print(f"일괄 인페인팅 진행: {i+1}/{len(image_paths)}")
|
|
result = self.inpaint(image_path, mask, method)
|
|
results.append(result)
|
|
|
|
return results
|
|
|
|
def compare_methods(self, image_path: str, mask: np.ndarray,
|
|
methods: List[str] = None) -> Dict[str, np.ndarray]:
|
|
"""
|
|
여러 인페인팅 방법 비교
|
|
|
|
Args:
|
|
image_path (str): 이미지 경로
|
|
mask (np.ndarray): 마스크
|
|
methods (List[str]): 비교할 방법들
|
|
|
|
Returns:
|
|
Dict[str, np.ndarray]: 방법별 결과
|
|
"""
|
|
if methods is None:
|
|
methods = ["cv2_telea", "cv2_ns", "lama"]
|
|
|
|
results = {}
|
|
|
|
for method in methods:
|
|
if method in self.available_methods:
|
|
print(f"방법 비교 중: {method}")
|
|
result = self.inpaint(image_path, mask, method)
|
|
results[method] = result
|
|
|
|
return results
|
|
|
|
def create_comparison_image(self, original: np.ndarray, results: Dict[str, np.ndarray],
|
|
output_path: str = None) -> np.ndarray:
|
|
"""
|
|
인페인팅 결과 비교 이미지 생성
|
|
|
|
Args:
|
|
original (np.ndarray): 원본 이미지
|
|
results (Dict[str, np.ndarray]): 방법별 결과
|
|
output_path (str): 저장 경로
|
|
|
|
Returns:
|
|
np.ndarray: 비교 이미지
|
|
"""
|
|
num_images = len(results) + 1 # 원본 + 결과들
|
|
|
|
# 이미지 크기 조정
|
|
height, width = original.shape[:2]
|
|
target_width = 300
|
|
target_height = int(height * target_width / width)
|
|
|
|
# 원본 이미지 리사이즈
|
|
original_resized = cv2.resize(original, (target_width, target_height))
|
|
|
|
# 비교 이미지 생성
|
|
comparison_width = target_width * num_images
|
|
comparison_height = target_height + 50 # 텍스트 공간
|
|
|
|
comparison = np.ones((comparison_height, comparison_width, 3), dtype=np.uint8) * 255
|
|
|
|
# 원본 이미지 배치
|
|
comparison[50:50+target_height, 0:target_width] = original_resized
|
|
cv2.putText(comparison, "Original", (10, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)
|
|
|
|
# 결과 이미지들 배치
|
|
for i, (method, result) in enumerate(results.items()):
|
|
x_offset = target_width * (i + 1)
|
|
result_resized = cv2.resize(result, (target_width, target_height))
|
|
comparison[50:50+target_height, x_offset:x_offset+target_width] = result_resized
|
|
|
|
cv2.putText(comparison, method, (x_offset + 10, 30),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)
|
|
|
|
if output_path:
|
|
cv2.imwrite(output_path, comparison)
|
|
print(f"비교 이미지 저장: {output_path}")
|
|
|
|
return comparison
|
|
|
|
def evaluate_inpainting_quality(self, original: np.ndarray, inpainted: np.ndarray,
|
|
mask: np.ndarray) -> Dict[str, float]:
|
|
"""
|
|
인페인팅 품질 평가 (간단한 메트릭)
|
|
|
|
Args:
|
|
original (np.ndarray): 원본 이미지
|
|
inpainted (np.ndarray): 인페인팅된 이미지
|
|
mask (np.ndarray): 마스크
|
|
|
|
Returns:
|
|
Dict[str, float]: 품질 메트릭
|
|
"""
|
|
# 마스크 영역 외부에서의 차이 (낮을수록 좋음)
|
|
mask_inv = cv2.bitwise_not(mask)
|
|
diff_outside = cv2.absdiff(original, inpainted)
|
|
diff_outside = cv2.bitwise_and(diff_outside, diff_outside, mask=mask_inv)
|
|
outside_error = np.mean(diff_outside)
|
|
|
|
# 마스크 경계에서의 연속성 (낮을수록 좋음)
|
|
kernel = np.ones((3, 3), np.uint8)
|
|
mask_border = cv2.morphologyEx(mask, cv2.MORPH_GRADIENT, kernel)
|
|
diff_border = cv2.absdiff(original, inpainted)
|
|
diff_border = cv2.bitwise_and(diff_border, diff_border, mask=mask_border)
|
|
border_error = np.mean(diff_border)
|
|
|
|
# 전체 구조적 유사성 (높을수록 좋음)
|
|
# 간단한 구현 - 실제로는 SSIM 등 사용
|
|
gray_orig = cv2.cvtColor(original, cv2.COLOR_BGR2GRAY)
|
|
gray_inpainted = cv2.cvtColor(inpainted, cv2.COLOR_BGR2GRAY)
|
|
|
|
# 그래디언트 유사성
|
|
grad_orig = cv2.Sobel(gray_orig, cv2.CV_64F, 1, 1, ksize=3)
|
|
grad_inpainted = cv2.Sobel(gray_inpainted, cv2.CV_64F, 1, 1, ksize=3)
|
|
gradient_similarity = np.corrcoef(grad_orig.flatten(), grad_inpainted.flatten())[0, 1]
|
|
|
|
metrics = {
|
|
'outside_error': outside_error,
|
|
'border_error': border_error,
|
|
'gradient_similarity': gradient_similarity if not np.isnan(gradient_similarity) else 0.0,
|
|
'overall_score': max(0, 1 - (outside_error + border_error) / 255 + gradient_similarity) / 2
|
|
}
|
|
|
|
return metrics
|
|
|
|
def test_module(self):
|
|
"""인페인팅 모듈 테스트"""
|
|
print("인페인팅 모듈 테스트 시작...")
|
|
|
|
# 테스트 이미지 생성
|
|
os.makedirs("test_output", exist_ok=True)
|
|
|
|
# 테스트 이미지 생성 (텍스트가 있는 이미지)
|
|
test_image = np.ones((300, 400, 3), dtype=np.uint8) * 255
|
|
|
|
# 배경 패턴 추가
|
|
for i in range(0, 400, 20):
|
|
cv2.line(test_image, (i, 0), (i, 300), (230, 230, 230), 1)
|
|
for i in range(0, 300, 20):
|
|
cv2.line(test_image, (0, i), (400, i), (230, 230, 230), 1)
|
|
|
|
# 텍스트 추가
|
|
cv2.putText(test_image, "Test Text 1", (50, 100),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
|
|
cv2.putText(test_image, "Test Text 2", (50, 200),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
|
|
|
|
test_image_path = "test_output/test_inpaint_image.jpg"
|
|
cv2.imwrite(test_image_path, test_image)
|
|
|
|
# 테스트 마스크 생성
|
|
mask = np.zeros((300, 400), dtype=np.uint8)
|
|
cv2.rectangle(mask, (45, 70), (250, 110), 255, -1) # 첫 번째 텍스트 영역
|
|
cv2.rectangle(mask, (45, 170), (250, 210), 255, -1) # 두 번째 텍스트 영역
|
|
|
|
cv2.imwrite("test_output/test_mask.jpg", mask)
|
|
|
|
# 각 방법별 인페인팅 테스트
|
|
methods_to_test = ["cv2_telea", "cv2_ns", "lama"]
|
|
|
|
results = {}
|
|
for method in methods_to_test:
|
|
print(f"\n{method} 방법 테스트:")
|
|
result = self.inpaint(test_image_path, mask, method)
|
|
if result is not None:
|
|
results[method] = result
|
|
output_path = f"test_output/inpaint_{method}.jpg"
|
|
cv2.imwrite(output_path, result)
|
|
|
|
# 품질 평가
|
|
metrics = self.evaluate_inpainting_quality(test_image, result, mask)
|
|
print(f"품질 메트릭: {metrics}")
|
|
|
|
# 비교 이미지 생성
|
|
if results:
|
|
comparison = self.create_comparison_image(
|
|
test_image, results, "test_output/inpaint_comparison.jpg"
|
|
)
|
|
print("인페인팅 방법 비교 이미지 생성 완료")
|
|
|
|
# 방법 비교 테스트
|
|
print("\n방법 비교 테스트:")
|
|
comparison_results = self.compare_methods(test_image_path, mask, methods_to_test)
|
|
|
|
print("인페인팅 모듈 테스트 완료!")
|
|
|
|
if __name__ == "__main__":
|
|
inpainter = InpaintingModule()
|
|
inpainter.test_module() |