430 lines
14 KiB
Python
430 lines
14 KiB
Python
"""
|
|
API 엔드포인트
|
|
iopaint와 호환되는 인페인팅 및 배경 제거 API를 제공합니다.
|
|
"""
|
|
import time
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import List
|
|
from fastapi import APIRouter, HTTPException, Response
|
|
from fastapi.responses import JSONResponse
|
|
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, PluginRequest, AdjustMaskRequest,
|
|
InpaintResponse, RemoveBGResponse, PluginResponse, ServerConfigResponse,
|
|
HealthResponse, ModelInfo, Device
|
|
)
|
|
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
|
|
)
|
|
from ..monitoring.dashboard import monitoring_data
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
@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", response_model=InpaintResponse)
|
|
async def inpaint_image(request: InpaintRequest):
|
|
"""인페인팅 API (iopaint 호환)"""
|
|
start_time = time.time()
|
|
|
|
try:
|
|
# base64 이미지 디코딩
|
|
image, alpha_channel, info, ext = decode_base64_to_image(request.image)
|
|
mask, _, _, _ = decode_base64_to_image(request.mask, gray=True)
|
|
|
|
# 이미지와 마스크 크기 검증
|
|
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="인페인팅 처리 실패")
|
|
|
|
# 결과 이미지를 base64로 인코딩
|
|
result_image = concat_alpha_channel(result, alpha_channel)
|
|
result_base64 = encode_image_to_base64(result_image, ext)
|
|
|
|
processing_time = time.time() - start_time
|
|
|
|
# 모니터링 통계 업데이트
|
|
monitoring_data.update_api_stats(
|
|
endpoint="/api/v1/inpaint",
|
|
success=True,
|
|
response_time=processing_time * 1000 # ms로 변환
|
|
)
|
|
|
|
return InpaintResponse(
|
|
success=True,
|
|
image=result_base64,
|
|
processing_time=processing_time,
|
|
seed=request.sd_seed
|
|
)
|
|
|
|
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)
|
|
)
|
|
|
|
raise HTTPException(status_code=500, detail=f"인페인팅 처리 실패: {str(e)}")
|
|
|
|
|
|
@router.post("/api/v1/remove_bg", response_model=RemoveBGResponse)
|
|
async def remove_background(request: RemoveBGRequest):
|
|
"""배경 제거 API"""
|
|
start_time = time.time()
|
|
|
|
try:
|
|
# base64 이미지 디코딩
|
|
image, alpha_channel, info, ext = decode_base64_to_image(request.image)
|
|
|
|
# 이미지 크기 검증
|
|
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="배경 제거 처리 실패")
|
|
|
|
# 결과를 base64로 인코딩
|
|
result_base64 = encode_image_to_base64(result_image, ext)
|
|
mask_base64 = encode_image_to_base64(result_mask, "PNG")
|
|
|
|
processing_time = time.time() - start_time
|
|
|
|
# 모니터링 통계 업데이트
|
|
monitoring_data.update_api_stats(
|
|
endpoint="/api/v1/remove_bg",
|
|
success=True,
|
|
response_time=processing_time * 1000
|
|
)
|
|
|
|
return RemoveBGResponse(
|
|
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)
|
|
)
|
|
|
|
raise HTTPException(status_code=500, detail=f"배경 제거 처리 실패: {str(e)}")
|
|
|
|
|
|
@router.post("/api/v1/run_plugin_gen_image", response_model=PluginResponse)
|
|
async def run_plugin_generate_image(request: PluginRequest):
|
|
"""플러그인으로 이미지 생성"""
|
|
start_time = time.time()
|
|
|
|
try:
|
|
# base64 이미지 디코딩
|
|
image, alpha_channel, info, ext = decode_base64_to_image(request.image)
|
|
|
|
# 이미지 크기 검증
|
|
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}까지 지원합니다."
|
|
)
|
|
|
|
# 플러그인 실행
|
|
if request.name == "rembg":
|
|
result_image, _ = await worker_manager.process_remove_bg(
|
|
image=image,
|
|
model_name=request.model_name or "rembg"
|
|
)
|
|
else:
|
|
raise HTTPException(status_code=422, detail=f"지원하지 않는 플러그인: {request.name}")
|
|
|
|
if result_image is None:
|
|
raise HTTPException(status_code=500, detail="플러그인 처리 실패")
|
|
|
|
# 결과를 base64로 인코딩
|
|
result_base64 = encode_image_to_base64(result_image, ext)
|
|
|
|
processing_time = time.time() - start_time
|
|
|
|
# 모니터링 통계 업데이트
|
|
monitoring_data.update_api_stats(
|
|
endpoint=f"/api/v1/run_plugin_gen_image/{request.name}",
|
|
success=True,
|
|
response_time=processing_time * 1000
|
|
)
|
|
|
|
return PluginResponse(
|
|
success=True,
|
|
image=result_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=f"/api/v1/run_plugin_gen_image/{request.name}",
|
|
success=False,
|
|
response_time=processing_time * 1000,
|
|
error=str(e)
|
|
)
|
|
|
|
raise HTTPException(status_code=500, detail=f"플러그인 처리 실패: {str(e)}")
|
|
|
|
|
|
@router.post("/api/v1/run_plugin_gen_mask", response_model=PluginResponse)
|
|
async def run_plugin_generate_mask(request: PluginRequest):
|
|
"""플러그인으로 마스크 생성"""
|
|
start_time = time.time()
|
|
|
|
try:
|
|
# base64 이미지 디코딩
|
|
image, alpha_channel, info, ext = decode_base64_to_image(request.image)
|
|
|
|
# 이미지 크기 검증
|
|
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}까지 지원합니다."
|
|
)
|
|
|
|
# 플러그인 실행
|
|
if request.name == "rembg":
|
|
_, result_mask = await worker_manager.process_remove_bg(
|
|
image=image,
|
|
model_name=request.model_name or "rembg"
|
|
)
|
|
else:
|
|
raise HTTPException(status_code=422, detail=f"지원하지 않는 플러그인: {request.name}")
|
|
|
|
if result_mask is None:
|
|
raise HTTPException(status_code=500, detail="플러그인 마스크 생성 실패")
|
|
|
|
# 마스크를 base64로 인코딩
|
|
mask_base64 = encode_image_to_base64(result_mask, "PNG")
|
|
|
|
processing_time = time.time() - start_time
|
|
|
|
# 모니터링 통계 업데이트
|
|
monitoring_data.update_api_stats(
|
|
endpoint=f"/api/v1/run_plugin_gen_mask/{request.name}",
|
|
success=True,
|
|
response_time=processing_time * 1000
|
|
)
|
|
|
|
return PluginResponse(
|
|
success=True,
|
|
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=f"/api/v1/run_plugin_gen_mask/{request.name}",
|
|
success=False,
|
|
response_time=processing_time * 1000,
|
|
error=str(e)
|
|
)
|
|
|
|
raise HTTPException(status_code=500, detail=f"플러그인 마스크 생성 실패: {str(e)}")
|
|
|
|
|
|
@router.post("/api/v1/adjust_mask")
|
|
async def adjust_mask_api(request: AdjustMaskRequest):
|
|
"""마스크 조정 API"""
|
|
try:
|
|
# base64 마스크 디코딩
|
|
mask, _, _, _ = decode_base64_to_image(request.mask, gray=True)
|
|
|
|
# 마스크 조정
|
|
adjusted_mask = adjust_mask(mask, request.kernel_size, request.operate)
|
|
|
|
# 프론트엔드용 마스크 생성
|
|
frontend_mask = gen_frontend_mask(adjusted_mask)
|
|
|
|
# 결과를 PNG로 반환
|
|
result_bytes = numpy_to_bytes(frontend_mask, "PNG")
|
|
|
|
return Response(
|
|
content=result_bytes,
|
|
media_type="image/png"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"마스크 조정 실패: {e}")
|
|
raise HTTPException(status_code=500, detail=f"마스크 조정 실패: {str(e)}")
|
|
|
|
|
|
@router.get("/api/v1/samplers")
|
|
async def get_samplers():
|
|
"""사용 가능한 샘플러 목록 반환"""
|
|
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",
|
|
"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
|