104 lines
4.0 KiB
Python
104 lines
4.0 KiB
Python
import random
|
|
import socket
|
|
import uvicorn
|
|
from fastapi import FastAPI, Query, Body
|
|
from pydantic import BaseModel
|
|
from typing import List, Optional
|
|
import asyncio
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from modules.image_processor2 import ImageProcessor
|
|
|
|
# 포트 범위 설정
|
|
PORT_RANGE = (7000, 7000)
|
|
|
|
# 사용 가능한 포트 찾기
|
|
def find_free_port():
|
|
for _ in range(20):
|
|
port = random.randint(*PORT_RANGE)
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
try:
|
|
s.bind(("127.0.0.1", port))
|
|
return port
|
|
except OSError:
|
|
continue
|
|
raise RuntimeError("사용 가능한 포트를 찾을 수 없습니다.")
|
|
|
|
# 요청 모델 정의
|
|
class ImageRequest(BaseModel):
|
|
local_image_path: str
|
|
file_prefix: Optional[str] = ""
|
|
use_inpainting: Optional[bool] = False
|
|
toggle_states: Optional[dict] = None
|
|
unwanted_texts: Optional[dict] = None
|
|
watermark_text: Optional[str] = None
|
|
watermark_opacity: Optional[float] = None
|
|
|
|
class ImagesRequest(BaseModel):
|
|
local_image_paths: List[str]
|
|
file_prefix: Optional[str] = ""
|
|
use_inpainting: Optional[bool] = False
|
|
toggle_states: Optional[dict] = None
|
|
unwanted_texts: Optional[dict] = None
|
|
watermark_text: Optional[str] = None
|
|
watermark_opacity: Optional[float] = None
|
|
|
|
# FastAPI 앱 생성
|
|
def create_app(image_processor: ImageProcessor, max_workers: int = 2):
|
|
app = FastAPI()
|
|
executor = ThreadPoolExecutor(max_workers=max_workers)
|
|
|
|
@app.post("/translate_image")
|
|
async def translate_image(req: ImageRequest):
|
|
# 워터마크 관련 옵션을 toggle_states에 병합
|
|
toggle_states = req.toggle_states.copy() if req.toggle_states else {}
|
|
if req.watermark_text is not None:
|
|
toggle_states["watermark_text"] = req.watermark_text
|
|
if req.watermark_opacity is not None:
|
|
toggle_states["watermark_opacity"] = req.watermark_opacity
|
|
# 단일 이미지 번역
|
|
result = await image_processor.process_single_image(
|
|
toggle_states, req.unwanted_texts or {}, req.local_image_path, 0, req.file_prefix
|
|
)
|
|
# 경로만 반환
|
|
if isinstance(result, dict):
|
|
return {"result": result.get("path", None)}
|
|
return {"result": result}
|
|
|
|
@app.post("/translate_images")
|
|
async def translate_images(req: ImagesRequest):
|
|
# 워터마크 관련 옵션을 toggle_states에 병합
|
|
toggle_states = req.toggle_states.copy() if req.toggle_states else {}
|
|
if req.watermark_text is not None:
|
|
toggle_states["watermark_text"] = req.watermark_text
|
|
if req.watermark_opacity is not None:
|
|
toggle_states["watermark_opacity"] = req.watermark_opacity
|
|
# 여러 이미지 병렬 번역
|
|
loop = asyncio.get_event_loop()
|
|
tasks = []
|
|
sem = asyncio.Semaphore(max_workers)
|
|
async def sem_task(idx, path):
|
|
async with sem:
|
|
return await image_processor.process_single_image(
|
|
toggle_states, req.unwanted_texts or {}, path, idx, req.file_prefix
|
|
)
|
|
for idx, path in enumerate(req.local_image_paths):
|
|
tasks.append(sem_task(idx, path))
|
|
results = await asyncio.gather(*tasks)
|
|
# 경로만 리스트로 반환
|
|
def extract_path(res):
|
|
if isinstance(res, dict):
|
|
return res.get("path", None)
|
|
return res
|
|
return {"results": [extract_path(r) for r in results]}
|
|
|
|
return app
|
|
|
|
# 서버 실행 함수
|
|
def run_server(image_processor, max_workers=2):
|
|
port = find_free_port()
|
|
app = create_app(image_processor, max_workers)
|
|
uvicorn.run(app, host="127.0.0.1", port=port, workers=1)
|
|
# FastAPI의 workers는 프로세스 수이므로, 내부 병렬은 ThreadPoolExecutor로 제어
|
|
# 실제 워커 수는 process_single_image 병렬 호출로 제한
|
|
# 서버 실행 후 포트 정보 반환 가능
|
|
return port |