459 lines
15 KiB
Python
459 lines
15 KiB
Python
"""
|
|
API 엔드포인트
|
|
iopaint와 호환되는 인페인팅 및 배경 제거 API를 제공합니다.
|
|
"""
|
|
import time
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import List
|
|
from fastapi import APIRouter, HTTPException, Response, Query
|
|
from fastapi.responses import JSONResponse, StreamingResponse
|
|
import cv2
|
|
|
|
from ..core.config import settings
|
|
from ..core.worker_manager import worker_manager
|
|
from ..core.session_pool import session_pool
|
|
from ..models.schemas import (
|
|
InpaintRequest, RemoveBGRequest,
|
|
InpaintResponse, RemoveBGResponse, ServerConfigResponse,
|
|
HealthResponse, ModelInfo, Device, ResponseFormat, ImageFormat
|
|
)
|
|
from ..utils.image_utils import (
|
|
decode_base64_to_image, encode_image_to_base64, concat_alpha_channel,
|
|
pil_to_bytes, numpy_to_bytes, adjust_mask, gen_frontend_mask
|
|
)
|
|
import base64
|
|
import io
|
|
from ..monitoring.dashboard import monitoring_data
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def encode_image_to_format(image, format: ImageFormat = ImageFormat.png, quality: int = 95):
|
|
"""이미지를 지정된 형식으로 인코딩"""
|
|
if format == ImageFormat.png:
|
|
_, buffer = cv2.imencode('.png', image)
|
|
return buffer.tobytes(), 'image/png'
|
|
elif format == ImageFormat.webp:
|
|
_, buffer = cv2.imencode('.webp', image, [cv2.IMWRITE_WEBP_QUALITY, quality])
|
|
return buffer.tobytes(), 'image/webp'
|
|
elif format == ImageFormat.jpeg:
|
|
_, buffer = cv2.imencode('.jpg', image, [cv2.IMWRITE_JPEG_QUALITY, quality])
|
|
return buffer.tobytes(), 'image/jpeg'
|
|
else:
|
|
# 기본값은 PNG
|
|
_, buffer = cv2.imencode('.png', image)
|
|
return buffer.tobytes(), 'image/png'
|
|
|
|
|
|
def create_response(
|
|
image,
|
|
response_format: ResponseFormat = ResponseFormat.json,
|
|
image_format: ImageFormat = ImageFormat.png,
|
|
processing_time: float = 0.0,
|
|
success: bool = True,
|
|
error_msg: str = None
|
|
):
|
|
"""응답 형식에 따라 적절한 응답 생성"""
|
|
|
|
if not success:
|
|
if response_format == ResponseFormat.binary:
|
|
return Response(
|
|
content=b"",
|
|
status_code=500,
|
|
media_type="text/plain"
|
|
)
|
|
else:
|
|
return JSONResponse(
|
|
content={
|
|
"success": False,
|
|
"error": error_msg,
|
|
"processing_time": processing_time
|
|
},
|
|
status_code=500
|
|
)
|
|
|
|
if response_format == ResponseFormat.binary:
|
|
# 바이너리 응답 (기존 클라이언트 호환)
|
|
image_bytes, media_type = encode_image_to_format(image, image_format)
|
|
return Response(content=image_bytes, media_type=media_type)
|
|
|
|
elif response_format == ResponseFormat.stream:
|
|
# 스트리밍 응답
|
|
def generate():
|
|
image_bytes, _ = encode_image_to_format(image, image_format)
|
|
chunk_size = 8192
|
|
for i in range(0, len(image_bytes), chunk_size):
|
|
yield image_bytes[i:i + chunk_size]
|
|
|
|
_, media_type = encode_image_to_format(image, image_format)
|
|
return StreamingResponse(generate(), media_type=media_type)
|
|
|
|
else:
|
|
# JSON 응답 (기본값, iopaint 호환)
|
|
image_b64 = encode_image_to_base64(image)
|
|
return JSONResponse(content={
|
|
"success": True,
|
|
"image": image_b64,
|
|
"processing_time": processing_time
|
|
})
|
|
|
|
|
|
@router.get("/health", response_model=HealthResponse)
|
|
async def health_check():
|
|
"""서버 상태 확인"""
|
|
start_time = getattr(settings, 'start_time', time.time())
|
|
uptime = time.time() - start_time
|
|
|
|
return HealthResponse(
|
|
status="healthy",
|
|
timestamp=datetime.now().isoformat(),
|
|
version="1.0.0",
|
|
uptime=uptime
|
|
)
|
|
|
|
|
|
@router.get("/api/v1/server-config", response_model=ServerConfigResponse)
|
|
async def get_server_config():
|
|
"""서버 설정 정보 반환 (iopaint 호환)"""
|
|
try:
|
|
# 사용 가능한 모델 목록
|
|
models = [
|
|
ModelInfo(
|
|
name="simple-lama",
|
|
type="inpainting",
|
|
description="Simple LAMA 인페인팅 모델",
|
|
supported_formats=["png", "jpg", "jpeg"],
|
|
max_image_size=2048
|
|
),
|
|
ModelInfo(
|
|
name="migan",
|
|
type="inpainting",
|
|
description="MIGAN 인페인팅 모델",
|
|
supported_formats=["png", "jpg", "jpeg"],
|
|
max_image_size=2048
|
|
),
|
|
ModelInfo(
|
|
name="rembg",
|
|
type="rembg",
|
|
description="Rembg 배경 제거 모델",
|
|
supported_formats=["png", "jpg", "jpeg"],
|
|
max_image_size=2048
|
|
)
|
|
]
|
|
|
|
return ServerConfigResponse(
|
|
models=models,
|
|
max_file_size=settings.MAX_FILE_SIZE,
|
|
supported_formats=["png", "jpg", "jpeg"],
|
|
device=Device.cuda if settings.IS_JETSON else Device.cuda,
|
|
is_jetson=settings.IS_JETSON
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"서버 설정 조회 실패: {e}")
|
|
raise HTTPException(status_code=500, detail=f"서버 설정 조회 실패: {str(e)}")
|
|
|
|
|
|
@router.post("/api/v1/inpaint")
|
|
async def inpaint_image(
|
|
request: InpaintRequest,
|
|
response_format: ResponseFormat = Query(ResponseFormat.binary, description="응답 형식 (기존 클라이언트 호환을 위해 기본값: binary)"),
|
|
image_format: ImageFormat = Query(ImageFormat.png, description="이미지 형식")
|
|
):
|
|
"""인페인팅 API (iopaint 호환)"""
|
|
start_time = time.time()
|
|
alpha_channel = None # 변수 초기화
|
|
|
|
try:
|
|
# base64 이미지 디코딩
|
|
image, alpha_channel, info, ext = decode_base64_to_image(request.image)
|
|
mask, _, _, _ = decode_base64_to_image(request.mask, gray=True)
|
|
|
|
# alpha_channel이 None인 경우 기본값 설정
|
|
if alpha_channel is None:
|
|
alpha_channel = None # 투명도 채널이 없는 경우
|
|
|
|
# 이미지와 마스크 크기 검증
|
|
if image.shape[:2] != mask.shape[:2]:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"이미지 크기({image.shape[:2]})와 마스크 크기({mask.shape[:2]})가 일치하지 않습니다."
|
|
)
|
|
|
|
# 이미지 크기 검증
|
|
if not validate_image_size(image, settings.MAX_IMAGE_SIZE):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"이미지 크기가 너무 큽니다. 최대 {settings.MAX_IMAGE_SIZE}x{settings.MAX_IMAGE_SIZE}까지 지원합니다."
|
|
)
|
|
|
|
# 마스크 이진화
|
|
mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)[1]
|
|
|
|
# 모델 선택
|
|
model_name = request.model_name or "simple-lama"
|
|
|
|
# 워커에서 인페인팅 실행
|
|
result = await worker_manager.process_inpaint(
|
|
image=image,
|
|
mask=mask,
|
|
model_name=model_name,
|
|
prompt=request.prompt,
|
|
negative_prompt=request.negative_prompt,
|
|
sd_seed=request.sd_seed,
|
|
num_inference_steps=request.num_inference_steps,
|
|
guidance_scale=request.guidance_scale,
|
|
strength=request.strength
|
|
)
|
|
|
|
if result is None:
|
|
raise HTTPException(status_code=500, detail="인페인팅 처리 실패")
|
|
|
|
# 결과 이미지 처리
|
|
result_image = concat_alpha_channel(result, alpha_channel)
|
|
|
|
processing_time = time.time() - start_time
|
|
|
|
# 모니터링 통계 업데이트
|
|
monitoring_data.update_api_stats(
|
|
endpoint="/api/v1/inpaint",
|
|
success=True,
|
|
response_time=processing_time * 1000 # ms로 변환
|
|
)
|
|
|
|
# 응답 형식에 따라 적절한 응답 생성
|
|
return create_response(
|
|
image=result_image,
|
|
response_format=response_format,
|
|
image_format=image_format,
|
|
processing_time=processing_time,
|
|
success=True
|
|
)
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"인페인팅 처리 실패: {e}")
|
|
|
|
# 모니터링 통계 업데이트
|
|
processing_time = time.time() - start_time
|
|
monitoring_data.update_api_stats(
|
|
endpoint="/api/v1/inpaint",
|
|
success=False,
|
|
response_time=processing_time * 1000,
|
|
error=str(e)
|
|
)
|
|
|
|
# 응답 형식에 따라 적절한 에러 응답 생성
|
|
return create_response(
|
|
image=None,
|
|
response_format=response_format,
|
|
image_format=image_format,
|
|
processing_time=processing_time,
|
|
success=False,
|
|
error_msg=f"인페인팅 처리 실패: {str(e)}"
|
|
)
|
|
|
|
|
|
@router.post("/api/v1/remove_bg")
|
|
async def remove_background(
|
|
request: RemoveBGRequest,
|
|
response_format: ResponseFormat = Query(ResponseFormat.binary, description="응답 형식 (기존 클라이언트 호환을 위해 기본값: binary)"),
|
|
image_format: ImageFormat = Query(ImageFormat.png, description="이미지 형식")
|
|
):
|
|
"""배경 제거 API (iopaint 호환)"""
|
|
start_time = time.time()
|
|
alpha_channel = None # 변수 초기화
|
|
|
|
try:
|
|
# base64 이미지 디코딩
|
|
image, alpha_channel, info, ext = decode_base64_to_image(request.image)
|
|
|
|
# alpha_channel이 None인 경우 기본값 설정
|
|
if alpha_channel is None:
|
|
alpha_channel = None # 투명도 채널이 없는 경우
|
|
|
|
# 이미지 크기 검증
|
|
if not validate_image_size(image, settings.MAX_IMAGE_SIZE):
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"이미지 크기가 너무 큽니다. 최대 {settings.MAX_IMAGE_SIZE}x{settings.MAX_IMAGE_SIZE}까지 지원합니다."
|
|
)
|
|
|
|
# 모델 선택
|
|
model_name = request.model_name or "rembg"
|
|
|
|
# 워커에서 배경 제거 실행
|
|
result_image, result_mask = await worker_manager.process_remove_bg(
|
|
image=image,
|
|
model_name=model_name
|
|
)
|
|
|
|
if result_image is None or result_mask is None:
|
|
raise HTTPException(status_code=500, detail="배경 제거 처리 실패")
|
|
|
|
processing_time = time.time() - start_time
|
|
|
|
# 모니터링 통계 업데이트
|
|
monitoring_data.update_api_stats(
|
|
endpoint="/api/v1/remove_bg",
|
|
success=True,
|
|
response_time=processing_time * 1000
|
|
)
|
|
|
|
# 응답 형식에 따라 적절한 응답 생성
|
|
if response_format == ResponseFormat.binary:
|
|
# 바이너리 응답은 이미지만 반환 (기존 클라이언트 호환)
|
|
return create_response(
|
|
image=result_image,
|
|
response_format=response_format,
|
|
image_format=image_format,
|
|
processing_time=processing_time,
|
|
success=True
|
|
)
|
|
else:
|
|
# JSON 응답은 이미지와 마스크 모두 반환
|
|
result_base64 = encode_image_to_base64(result_image, ext)
|
|
mask_base64 = encode_image_to_base64(result_mask, "PNG")
|
|
|
|
return JSONResponse(content={
|
|
"success": True,
|
|
"image": result_base64,
|
|
"mask": mask_base64,
|
|
"processing_time": processing_time
|
|
})
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"배경 제거 처리 실패: {e}")
|
|
|
|
# 모니터링 통계 업데이트
|
|
processing_time = time.time() - start_time
|
|
monitoring_data.update_api_stats(
|
|
endpoint="/api/v1/remove_bg",
|
|
success=False,
|
|
response_time=processing_time * 1000,
|
|
error=str(e)
|
|
)
|
|
|
|
# 응답 형식에 따라 적절한 에러 응답 생성
|
|
return create_response(
|
|
image=None,
|
|
response_format=response_format,
|
|
image_format=image_format,
|
|
processing_time=processing_time,
|
|
success=False,
|
|
error_msg=f"배경 제거 처리 실패: {str(e)}"
|
|
)
|
|
|
|
|
|
|
|
@router.get("/api/v1/model")
|
|
async def get_model_info():
|
|
"""모델 정보 반환 (클라이언트 헬스체크 호환)"""
|
|
start_time = time.time()
|
|
|
|
try:
|
|
# 현재 사용 가능한 모델 목록
|
|
models = [
|
|
{
|
|
"name": "simple-lama",
|
|
"type": "inpainting",
|
|
"status": "available",
|
|
"description": "Simple LAMA 인페인팅 모델"
|
|
},
|
|
{
|
|
"name": "migan",
|
|
"type": "inpainting",
|
|
"status": "available",
|
|
"description": "MIGAN 인페인팅 모델"
|
|
},
|
|
{
|
|
"name": "rembg",
|
|
"type": "rembg",
|
|
"status": "available",
|
|
"description": "Rembg 배경 제거 모델"
|
|
}
|
|
]
|
|
|
|
processing_time = time.time() - start_time
|
|
|
|
# 모니터링 통계 업데이트
|
|
monitoring_data.update_api_stats(
|
|
endpoint="/api/v1/model",
|
|
success=True,
|
|
response_time=processing_time * 1000 # ms로 변환
|
|
)
|
|
|
|
return {
|
|
"success": True,
|
|
"models": models,
|
|
"server_status": "running",
|
|
"version": "1.0.0",
|
|
"device": "cuda" if settings.IS_JETSON else "cuda",
|
|
"is_jetson": settings.IS_JETSON,
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"모델 정보 조회 실패: {e}")
|
|
|
|
# 모니터링 통계 업데이트
|
|
processing_time = time.time() - start_time
|
|
monitoring_data.update_api_stats(
|
|
endpoint="/api/v1/model",
|
|
success=False,
|
|
response_time=processing_time * 1000,
|
|
error=str(e)
|
|
)
|
|
|
|
raise HTTPException(status_code=500, detail=f"모델 정보 조회 실패: {str(e)}")
|
|
|
|
|
|
@router.get("/api/v1/samplers")
|
|
async def get_samplers():
|
|
"""사용 가능한 샘플러 목록 반환 (iopaint 호환)"""
|
|
return [
|
|
"euler",
|
|
"euler_a",
|
|
"heun",
|
|
"dpm_2",
|
|
"dpm_2_a",
|
|
"lms",
|
|
"dpm_fast",
|
|
"dpm_adaptive",
|
|
"dpmpp_2s_a",
|
|
"dpmpp_sde",
|
|
"dpmpp_2m",
|
|
"ddim",
|
|
"uni_pc",
|
|
"uni_pc_bh2"
|
|
]
|
|
|
|
|
|
@router.get("/")
|
|
async def root():
|
|
"""루트 엔드포인트"""
|
|
return {
|
|
"message": "인페인팅 서버 API (iopaint 호환)",
|
|
"version": "1.0.0",
|
|
"docs": "/docs",
|
|
"health": "/health"
|
|
}
|
|
|
|
|
|
# 유틸리티 함수
|
|
def validate_image_size(image, max_size):
|
|
"""이미지 크기 검증"""
|
|
try:
|
|
h, w = image.shape[:2]
|
|
pixels = h * w
|
|
max_pixels = max_size * max_size
|
|
return pixels <= max_pixels
|
|
except:
|
|
return False
|