248 lines
8.3 KiB
Python
248 lines
8.3 KiB
Python
"""
|
|
이미지 처리 유틸리티
|
|
iopaint와 호환되는 이미지 변환 및 처리 함수들을 제공합니다.
|
|
"""
|
|
import base64
|
|
import io
|
|
import logging
|
|
from typing import Tuple, Optional, Union
|
|
import numpy as np
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def encode_image_to_base64(image: Union[np.ndarray, Image.Image], format: str = "PNG") -> str:
|
|
"""이미지를 base64로 인코딩합니다."""
|
|
try:
|
|
if isinstance(image, np.ndarray):
|
|
# numpy 배열을 PIL Image로 변환
|
|
if image.dtype != np.uint8:
|
|
image = (image * 255).astype(np.uint8)
|
|
|
|
if len(image.shape) == 3 and image.shape[2] == 3:
|
|
# BGR to RGB
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
|
pil_image = Image.fromarray(image)
|
|
else:
|
|
pil_image = image
|
|
|
|
# PIL Image를 base64로 인코딩
|
|
buffer = io.BytesIO()
|
|
pil_image.save(buffer, format=format)
|
|
img_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
|
|
return img_str
|
|
except Exception as e:
|
|
logger.error(f"이미지 base64 인코딩 실패: {e}")
|
|
raise
|
|
|
|
|
|
def decode_base64_to_image(base64_string: str, gray: bool = False) -> Tuple[np.ndarray, Optional[np.ndarray], dict, str]:
|
|
"""base64 문자열을 이미지로 디코딩합니다 (iopaint 호환)."""
|
|
try:
|
|
# base64 디코딩
|
|
image_data = base64.b64decode(base64_string)
|
|
|
|
# PIL Image로 로드
|
|
pil_image = Image.open(io.BytesIO(image_data))
|
|
|
|
# 이미지 정보 추출
|
|
info = {
|
|
"format": pil_image.format,
|
|
"mode": pil_image.mode,
|
|
"size": pil_image.size,
|
|
"parameters": ""
|
|
}
|
|
|
|
# 확장자 결정
|
|
ext = pil_image.format.lower() if pil_image.format else "png"
|
|
|
|
# numpy 배열로 변환
|
|
if gray:
|
|
# 그레이스케일로 변환
|
|
if pil_image.mode in ('RGBA', 'LA', 'P'):
|
|
pil_image = pil_image.convert('L')
|
|
image_array = np.array(pil_image)
|
|
else:
|
|
# RGB로 변환
|
|
if pil_image.mode == 'RGBA':
|
|
# RGBA를 RGB로 변환하고 알파 채널 분리
|
|
rgb_image = pil_image.convert('RGB')
|
|
alpha_channel = np.array(pil_image.split()[-1])
|
|
image_array = np.array(rgb_image)
|
|
elif pil_image.mode == 'P':
|
|
# 팔레트 이미지를 RGB로 변환
|
|
image_array = np.array(pil_image.convert('RGB'))
|
|
alpha_channel = None
|
|
else:
|
|
# RGB 이미지
|
|
image_array = np.array(pil_image)
|
|
alpha_channel = None
|
|
|
|
# BGR로 변환 (OpenCV 호환)
|
|
if len(image_array.shape) == 3 and image_array.shape[2] == 3:
|
|
image_array = cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR)
|
|
|
|
return image_array, alpha_channel, info, ext
|
|
|
|
except Exception as e:
|
|
logger.error(f"base64 이미지 디코딩 실패: {e}")
|
|
raise
|
|
|
|
|
|
def concat_alpha_channel(image: np.ndarray, alpha_channel: Optional[np.ndarray]) -> np.ndarray:
|
|
"""이미지와 알파 채널을 결합합니다."""
|
|
if alpha_channel is None:
|
|
return image
|
|
|
|
try:
|
|
# BGR to RGB
|
|
if len(image.shape) == 3 and image.shape[2] == 3:
|
|
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
else:
|
|
rgb_image = image
|
|
|
|
# RGBA 이미지 생성
|
|
if len(rgb_image.shape) == 3:
|
|
rgba_image = np.dstack((rgb_image, alpha_channel))
|
|
else:
|
|
rgba_image = np.dstack((rgb_image, rgb_image, rgb_image, alpha_channel))
|
|
|
|
return rgba_image
|
|
|
|
except Exception as e:
|
|
logger.error(f"알파 채널 결합 실패: {e}")
|
|
return image
|
|
|
|
|
|
def pil_to_bytes(image: Image.Image, ext: str = "PNG", quality: int = 100, infos: dict = None) -> bytes:
|
|
"""PIL Image를 바이트로 변환합니다."""
|
|
try:
|
|
buffer = io.BytesIO()
|
|
|
|
if ext.lower() in ['jpg', 'jpeg']:
|
|
image.save(buffer, format='JPEG', quality=quality, optimize=True)
|
|
else:
|
|
image.save(buffer, format=ext.upper(), optimize=True)
|
|
|
|
return buffer.getvalue()
|
|
|
|
except Exception as e:
|
|
logger.error(f"PIL Image를 바이트로 변환 실패: {e}")
|
|
raise
|
|
|
|
|
|
def numpy_to_bytes(image: np.ndarray, ext: str = "PNG", quality: int = 100) -> bytes:
|
|
"""numpy 배열을 바이트로 변환합니다."""
|
|
try:
|
|
# numpy 배열을 PIL Image로 변환
|
|
if image.dtype != np.uint8:
|
|
image = (image * 255).astype(np.uint8)
|
|
|
|
if len(image.shape) == 3 and image.shape[2] == 3:
|
|
# BGR to RGB
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
|
pil_image = Image.fromarray(image)
|
|
return pil_to_bytes(pil_image, ext, quality)
|
|
|
|
except Exception as e:
|
|
logger.error(f"numpy 배열을 바이트로 변환 실패: {e}")
|
|
raise
|
|
|
|
|
|
def adjust_mask(mask: np.ndarray, kernel_size: int = 5, operate: str = "dilate") -> np.ndarray:
|
|
"""마스크를 조정합니다 (dilate/erode)."""
|
|
try:
|
|
if kernel_size <= 0:
|
|
return mask
|
|
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
|
|
|
|
if operate.lower() == "dilate":
|
|
result = cv2.dilate(mask, kernel, iterations=1)
|
|
elif operate.lower() == "erode":
|
|
result = cv2.erode(mask, kernel, iterations=1)
|
|
else:
|
|
logger.warning(f"알 수 없는 마스크 연산: {operate}")
|
|
return mask
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
logger.error(f"마스크 조정 실패: {e}")
|
|
return mask
|
|
|
|
|
|
def gen_frontend_mask(mask: np.ndarray) -> np.ndarray:
|
|
"""프론트엔드용 마스크를 생성합니다."""
|
|
try:
|
|
# 마스크를 0-255 범위로 정규화
|
|
if mask.dtype != np.uint8:
|
|
mask = (mask * 255).astype(np.uint8)
|
|
|
|
# 이진화
|
|
_, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
|
|
|
|
return binary_mask
|
|
|
|
except Exception as e:
|
|
logger.error(f"프론트엔드 마스크 생성 실패: {e}")
|
|
return mask
|
|
|
|
|
|
def resize_image(image: np.ndarray, target_size: Tuple[int, int], keep_aspect: bool = True) -> np.ndarray:
|
|
"""이미지 크기를 조정합니다."""
|
|
try:
|
|
if keep_aspect:
|
|
# 종횡비 유지하면서 크기 조정
|
|
h, w = image.shape[:2]
|
|
target_h, target_w = target_size
|
|
|
|
# 종횡비 계산
|
|
aspect = w / h
|
|
target_aspect = target_w / target_h
|
|
|
|
if aspect > target_aspect:
|
|
# 너비에 맞춤
|
|
new_w = target_w
|
|
new_h = int(target_w / aspect)
|
|
else:
|
|
# 높이에 맞춤
|
|
new_h = target_h
|
|
new_w = int(target_h * aspect)
|
|
|
|
resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
|
|
|
|
# 패딩으로 목표 크기 맞춤
|
|
result = np.zeros((target_h, target_w, 3) if len(image.shape) == 3 else (target_h, target_w), dtype=image.dtype)
|
|
y_offset = (target_h - new_h) // 2
|
|
x_offset = (target_w - new_w) // 2
|
|
result[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = resized
|
|
|
|
return result
|
|
else:
|
|
# 종횡비 무시하고 크기 조정
|
|
return cv2.resize(image, target_size, interpolation=cv2.INTER_LANCZOS4)
|
|
|
|
except Exception as e:
|
|
logger.error(f"이미지 크기 조정 실패: {e}")
|
|
return image
|
|
|
|
|
|
def validate_image_size(image: np.ndarray, max_size: int) -> bool:
|
|
"""이미지 크기가 제한을 초과하지 않는지 확인합니다."""
|
|
try:
|
|
h, w = image.shape[:2]
|
|
pixels = h * w
|
|
max_pixels = max_size * max_size
|
|
|
|
return pixels <= max_pixels
|
|
|
|
except Exception as e:
|
|
logger.error(f"이미지 크기 검증 실패: {e}")
|
|
return False
|