""" 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