IT_Server/modules/image_translate_server.py

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