550 lines
23 KiB
Python
550 lines
23 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
ImageWorker FastAPI 클라이언트 헬퍼
|
|
|
|
- 이미지 URL 다운로드 → 로컬 경로 전달 방식으로 서버에 제출
|
|
- /v1/process-image, /v1/remove-background 호출 및 대기
|
|
"""
|
|
|
|
import os
|
|
import time
|
|
import uuid
|
|
import shutil
|
|
import mimetypes
|
|
import asyncio
|
|
import random
|
|
from typing import Dict, Any, Optional, List
|
|
import logging
|
|
import cv2 # for cleanup (destroyAllWindows)
|
|
import requests
|
|
import psutil
|
|
from urllib.parse import urlparse
|
|
import json
|
|
|
|
|
|
def _compat_result_from_job(job: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""ImageProcessor3 결과와 호환되도록 키를 정규화."""
|
|
rr = (job or {}).get("result") or {}
|
|
result = {
|
|
"status": rr.get("status") or job.get("status"),
|
|
"path": rr.get("path"),
|
|
"inpaint_method": rr.get("inpaint_method"),
|
|
"inpaint_device": rr.get("inpaint_device"),
|
|
"timings": rr.get("timings"),
|
|
"child_results": rr.get("child_results"),
|
|
"message": rr.get("message") or rr.get("msg"),
|
|
"error": rr.get("error"),
|
|
}
|
|
# 누락 키 보장
|
|
for k in ("status", "path", "inpaint_method", "inpaint_device", "timings", "child_results", "message", "error"):
|
|
result.setdefault(k, None)
|
|
return result
|
|
|
|
def _read_server_info() -> Optional[Dict[str, Any]]:
|
|
"""ProgramData/ImgWorker/server.json에서 서버 정보를 읽어온다."""
|
|
try:
|
|
program_data = os.environ.get("PROGRAMDATA", r"C:\\ProgramData")
|
|
info_path = os.path.join(program_data, "ImgWorker", "server.json")
|
|
if os.path.isfile(info_path):
|
|
with open(info_path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
class ImageWorkerClient:
|
|
def __init__(self, logger, api_base: str = "http://127.0.0.1:8009", work_dir: Optional[str] = None, timeout: int = 30, max_concurrency: int = 8):
|
|
self.logger = logger
|
|
|
|
# API base 우선순위: 명시 인자 > IMGWK_API_BASE 환경변수 > server.json > 기본값
|
|
api = api_base
|
|
try:
|
|
env_api = os.environ.get("IMGWK_API_BASE")
|
|
if env_api and isinstance(env_api, str) and env_api.strip():
|
|
api = env_api.strip()
|
|
else:
|
|
# 인자가 기본값일 때만 server.json 자동 사용(명시적 인자 우선)
|
|
if api_base == "http://127.0.0.1:8009":
|
|
info = _read_server_info()
|
|
if isinstance(info, dict):
|
|
base = info.get("base") or (f"http://{info.get('host','127.0.0.1')}:{info.get('port',8009)}")
|
|
if base:
|
|
api = base
|
|
except Exception:
|
|
pass
|
|
|
|
self.api = (api or "http://127.0.0.1:8009").rstrip("/")
|
|
|
|
# 기본 작업 디렉토리: C:\ProgramData\ImgWorker\incoming
|
|
if work_dir is None:
|
|
program_data = os.environ.get("PROGRAMDATA", r"C:\\ProgramData")
|
|
work_dir = os.path.join(program_data, "ImgWorker", "incoming")
|
|
self.logger.log(f"work_dir: {work_dir}", level=logging.DEBUG)
|
|
os.makedirs(work_dir, exist_ok=True)
|
|
self.work_dir = work_dir
|
|
self.TEMP_IMAGE_DIR = work_dir # download_image 메서드 호환성을 위해
|
|
self.timeout = timeout
|
|
# 동시요청 제한(세마포어)
|
|
# try:
|
|
# n = int(max_concurrency)
|
|
# except Exception:
|
|
# n = 8
|
|
# self._sema = asyncio.Semaphore(max(1, n))
|
|
|
|
# def set_max_concurrency(self, n: int):
|
|
# """동시 요청 제한을 런타임에 조정"""
|
|
# try:
|
|
# n = int(n)
|
|
# except Exception:
|
|
# n = 1
|
|
# self._sema = asyncio.Semaphore(max(1, n))
|
|
|
|
def is_valid_image_data(self, data):
|
|
"""이미지 데이터의 유효성을 검사합니다"""
|
|
if not data or len(data) < 100: # 최소 크기 검사
|
|
return False
|
|
|
|
# JPEG, PNG, GIF, WebP 시그니처 검사
|
|
if data.startswith(b'\xff\xd8\xff'): # JPEG
|
|
return True
|
|
elif data.startswith(b'\x89PNG\r\n\x1a\n'): # PNG
|
|
return True
|
|
elif data.startswith(b'GIF87a') or data.startswith(b'GIF89a'): # GIF
|
|
return True
|
|
elif data.startswith(b'RIFF') and b'WEBP' in data[:12]: # WebP
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
# ---------------------- 유틸 ----------------------
|
|
def _guess_ext(self, url: str, content_type: Optional[str]) -> str:
|
|
# 1) 헤더 content-type
|
|
if content_type:
|
|
ext = mimetypes.guess_extension(content_type.split(";")[0].strip())
|
|
if ext:
|
|
return ext
|
|
# 2) URL 경로
|
|
try:
|
|
basename = os.path.basename(url.split("?")[0])
|
|
_, ext = os.path.splitext(basename)
|
|
if ext:
|
|
return ext
|
|
except Exception:
|
|
pass
|
|
# 3) 기본값
|
|
return ".jpg"
|
|
|
|
|
|
|
|
def download_image(self, image_url, index, file_prefix="", max_retries=3):
|
|
"""Requests를 사용해 이미지를 다운로드합니다"""
|
|
|
|
# 로컬 파일 경로면 바로 반환
|
|
if os.path.isfile(image_url):
|
|
self.logger.log(f"로컬 파일 경로 감지, 다운로드 생략: {image_url}", level=logging.DEBUG)
|
|
return image_url
|
|
|
|
# 로컬 파일 경로가 아니면 다운로드 시도
|
|
try:
|
|
# "https://assets.alicdn.com"으로 시작하는 URL은 건너뛰기
|
|
if image_url.startswith("https://assets.alicdn.com") or image_url.startswith("https://gtms01.alicdn.com"):
|
|
self.logger.log(f"다운로드 제외 URL: {image_url}", level=logging.DEBUG)
|
|
return None
|
|
|
|
# URL에서 파일명 추출 및 접두사 포함
|
|
parsed_url = urlparse(image_url)
|
|
base_filename = f"image_{index:03d}_{os.path.basename(parsed_url.path)}"
|
|
if not base_filename.endswith(('.jpg', '.jpeg', '.png', '.gif', '.webp')):
|
|
base_filename += '.jpg'
|
|
|
|
# 접두사가 있으면 파일명에 포함
|
|
if file_prefix:
|
|
filename = f"{file_prefix}_{base_filename}"
|
|
else:
|
|
filename = base_filename
|
|
|
|
local_path = os.path.join(self.TEMP_IMAGE_DIR, filename)
|
|
|
|
# HTTP 헤더 설정
|
|
headers = {
|
|
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/88.0.4324.150 Safari/537.36",
|
|
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9",
|
|
"Accept-Language": "en-US,en;q=0.9",
|
|
"Accept-Encoding": "gzip, deflate, br",
|
|
"DNT": "1", # Do Not Track 요청 헤더
|
|
"Connection": "keep-alive",
|
|
"Upgrade-Insecure-Requests": "1",
|
|
"Cache-Control": "max-age=0"
|
|
}
|
|
|
|
retries = 0
|
|
while retries < max_retries:
|
|
try:
|
|
# 메모리 추적: 다운로드 시작 전
|
|
before_mem = psutil.virtual_memory()
|
|
before_mb = before_mem.used / 1024 / 1024
|
|
|
|
self.logger.log(f"이미지 다운로드 중: {filename}", level=logging.DEBUG)
|
|
response = requests.get(image_url, headers=headers, stream=True, timeout=30)
|
|
|
|
if response.status_code == 200:
|
|
image_data = response.content
|
|
|
|
# 이미지 데이터 유효성 검사
|
|
if self.is_valid_image_data(image_data):
|
|
with open(local_path, 'wb') as f:
|
|
f.write(image_data)
|
|
|
|
# 메모리 추적: 다운로드 완료 후
|
|
after_mem = psutil.virtual_memory()
|
|
after_mb = after_mem.used / 1024 / 1024
|
|
change_mb = after_mb - before_mb
|
|
change_percent = (change_mb / before_mb) * 100 if before_mb > 0 else 0
|
|
self.logger.log(
|
|
f"메모리 변화 [다운로드 완료]: {before_mb:.1f}MB -> {after_mb:.1f}MB "
|
|
f"({change_mb:+.1f}MB, {change_percent:+.1f}%) - {filename}",
|
|
level=logging.DEBUG if abs(change_mb) < 10 else logging.INFO
|
|
)
|
|
|
|
self.logger.log(f"이미지 다운로드 완료: {filename}", level=logging.DEBUG)
|
|
return local_path
|
|
else:
|
|
self.logger.log(f"유효하지 않은 이미지 데이터: {image_url}", level=logging.WARNING)
|
|
return None
|
|
else:
|
|
self.logger.log(f"이미지 다운로드 실패 (HTTP {response.status_code}): {image_url}. 재시도 {retries + 1}/{max_retries}", level=logging.ERROR)
|
|
retries += 1
|
|
if retries < max_retries:
|
|
time.sleep(random.randint(2, 5)) # 2~5초 대기 후 재시도
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
self.logger.log(f"이미지 다운로드 중 네트워크 오류: {e}. 재시도 {retries + 1}/{max_retries}", level=logging.ERROR)
|
|
retries += 1
|
|
if retries < max_retries:
|
|
time.sleep(random.randint(2, 5)) # 예외 발생 시 대기 후 재시도
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"이미지 다운로드 중 예상치 못한 오류: {e}. 재시도 {retries + 1}/{max_retries}", level=logging.ERROR)
|
|
retries += 1
|
|
if retries < max_retries:
|
|
time.sleep(random.randint(2, 5))
|
|
|
|
self.logger.log(f"이미지 다운로드 최대 재시도 횟수를 초과했습니다: {image_url}", level=logging.ERROR)
|
|
return None
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"이미지 다운로드 중 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
return None
|
|
|
|
|
|
|
|
async def health(self) -> Dict[str, Any]:
|
|
try:
|
|
hr = requests.get(f"{self.api}/health", timeout=10)
|
|
hr.raise_for_status()
|
|
return hr.json()
|
|
except Exception as e:
|
|
return {"status": "error", "error": str(e)}
|
|
|
|
# ---------------------- 제출/대기 ----------------------
|
|
async def submit_process_image(self, file_path: str, index: int, file_prefix: str,
|
|
font_type: str, unwanted_texts: List[str],
|
|
is_member_valid: bool, authenticated_by_admin: bool,
|
|
extra_overrides: Optional[Dict[str, Any]] = None,
|
|
ocr: Optional[bool] = None) -> str:
|
|
payload: Dict[str, Any] = {
|
|
"file_path": file_path,
|
|
"index": int(index),
|
|
"file_prefix": file_prefix,
|
|
# per-request 토글 오버라이드
|
|
"toggle_overrides": {
|
|
"font_type": font_type,
|
|
"unwanted_texts": list(unwanted_texts or []),
|
|
"is_member_valid": bool(is_member_valid),
|
|
"authenticated_by_admin": bool(authenticated_by_admin),
|
|
},
|
|
}
|
|
if extra_overrides:
|
|
payload["toggle_overrides"].update(extra_overrides)
|
|
# ocr 플래그가 명시되면 서버로도 전달
|
|
if ocr is not None:
|
|
payload["ocr"] = bool(ocr)
|
|
|
|
r = requests.post(f"{self.api}/v1/process-image", json=payload, timeout=self.timeout)
|
|
r.raise_for_status()
|
|
return r.json().get("job_id")
|
|
|
|
async def submit_remove_background(self, file_path: str, file_prefix: str,
|
|
extra_overrides: Optional[Dict[str, Any]] = None) -> str:
|
|
payload: Dict[str, Any] = {
|
|
"file_path": file_path,
|
|
"file_prefix": file_prefix,
|
|
}
|
|
if extra_overrides:
|
|
payload["toggle_overrides"] = extra_overrides
|
|
|
|
r = requests.post(f"{self.api}/v1/remove-background", json=payload, timeout=self.timeout)
|
|
r.raise_for_status()
|
|
return r.json().get("job_id")
|
|
|
|
async def wait_job(self, job_id: str, timeout_sec: int = 900) -> Dict[str, Any]:
|
|
end = time.time() + timeout_sec
|
|
while time.time() < end:
|
|
r = requests.get(f"{self.api}/v1/jobs/{job_id}", timeout=15)
|
|
if r.status_code == 200:
|
|
data = r.json()
|
|
if data.get("status") in ("done", "error", "cancelled"):
|
|
return data
|
|
await asyncio.sleep(0.2)
|
|
raise TimeoutError("job wait timeout")
|
|
|
|
# ---------------------- 고수준 URL 편의 ----------------------
|
|
async def process_image_url(self, image_url: str, index: int, file_prefix: str,
|
|
font_type: str, unwanted_texts: List[str],
|
|
is_member_valid: bool, authenticated_by_admin: bool,
|
|
extra_overrides: Optional[Dict[str, Any]] = None,
|
|
download_first: bool = True,
|
|
ocr: Optional[bool] = None) -> Optional[Dict[str, Any]]:
|
|
# 420 에러 방지를 위해 순차 처리 (세마포어 제거)
|
|
path = image_url
|
|
if download_first and (image_url.startswith("http://") or image_url.startswith("https://")):
|
|
# 동기식 download_image를 executor로 실행
|
|
loop = asyncio.get_event_loop()
|
|
path = await loop.run_in_executor(None, self.download_image, image_url, index, file_prefix)
|
|
|
|
jid = await self.submit_process_image(
|
|
file_path=path,
|
|
index=index,
|
|
file_prefix=file_prefix,
|
|
font_type=font_type,
|
|
unwanted_texts=unwanted_texts,
|
|
is_member_valid=is_member_valid,
|
|
authenticated_by_admin=authenticated_by_admin,
|
|
extra_overrides=extra_overrides,
|
|
ocr=ocr,
|
|
)
|
|
job = await self.wait_job(jid)
|
|
return _compat_result_from_job(job)
|
|
|
|
async def remove_background_url(self, image_url: str, file_prefix: str,
|
|
extra_overrides: Optional[Dict[str, Any]] = None,
|
|
download_first: bool = True) -> Optional[Dict[str, Any]]:
|
|
# 420 에러 방지를 위해 순차 처리 (세마포어 제거)
|
|
path = image_url
|
|
if download_first and (image_url.startswith("http://") or image_url.startswith("https://")):
|
|
# 동기식 download_image를 executor로 실행
|
|
loop = asyncio.get_event_loop()
|
|
path = await loop.run_in_executor(None, self.download_image, image_url, 0, file_prefix)
|
|
|
|
jid = await self.submit_remove_background(
|
|
file_path=path,
|
|
file_prefix=file_prefix,
|
|
extra_overrides=extra_overrides,
|
|
)
|
|
job = await self.wait_job(jid)
|
|
return _compat_result_from_job(job)
|
|
|
|
# ---------------------- 간단 제어 API (트레이에서 사용) ----------------------
|
|
def worker_status(self) -> Dict[str, Any]:
|
|
try:
|
|
r = requests.get(f"{self.api}/v1/worker/status", timeout=5)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
return {"ready": False, "error": str(e)}
|
|
|
|
def worker_start(self) -> Dict[str, Any]:
|
|
try:
|
|
r = requests.post(f"{self.api}/v1/worker/start", timeout=10)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
return {"ok": False, "error": str(e)}
|
|
|
|
def worker_stop(self) -> Dict[str, Any]:
|
|
try:
|
|
r = requests.post(f"{self.api}/v1/worker/stop", timeout=10)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
return {"ok": False, "error": str(e)}
|
|
|
|
def shutdown_server(self) -> Dict[str, Any]:
|
|
try:
|
|
r = requests.post(f"{self.api}/v1/server/shutdown", timeout=5)
|
|
r.raise_for_status()
|
|
return r.json()
|
|
except Exception as e:
|
|
return {"ok": False, "error": str(e)}
|
|
|
|
# ---------------------- 호환성 메서드 (기존 ImageProcessor3와 동일한 인터페이스) ----------------------
|
|
|
|
async def process_single_image(self, original_image_url, index, delay=1.0, file_prefix="", ocr: Optional[bool] = None):
|
|
"""
|
|
기존 ImageProcessor3.process_single_image과 호환되는 메서드
|
|
|
|
Args:
|
|
original_image_url (str): 처리할 이미지 URL
|
|
index (int): 이미지 인덱스
|
|
delay (float): 요청 간격 (초) - 호환성을 위해 유지
|
|
file_prefix (str): 파일명에 추가할 접두사
|
|
|
|
Returns:
|
|
dict: 기존 ImageProcessor3과 동일한 포맷의 결과
|
|
- status: 'translated', 'original', 'exclude', 'failed' 중 하나
|
|
- path: 처리된 이미지 파일 경로 또는 원본 이미지 파일 경로
|
|
- error: 오류 메시지 (status가 'failed'인 경우에만 포함)
|
|
- inpaint_method: 사용된 인페인팅 방법
|
|
- inpaint_device: 사용된 인페인팅 장치
|
|
"""
|
|
try:
|
|
# 요청 간격 조절 (호환성을 위해)
|
|
if delay > 0:
|
|
await asyncio.sleep(delay)
|
|
|
|
# ImageWorkerClient를 통한 처리
|
|
result = await self.process_image_url(
|
|
image_url=original_image_url,
|
|
index=index,
|
|
file_prefix=file_prefix,
|
|
font_type="", # 기본값 사용 (필요시 외부에서 설정 가능)
|
|
unwanted_texts=[], # 기본값 사용 (필요시 외부에서 설정 가능)
|
|
is_member_valid=False, # 기본값 사용 (필요시 외부에서 설정 가능)
|
|
authenticated_by_admin=False, # 기본값 사용 (필요시 외부에서 설정 가능)
|
|
ocr=ocr,
|
|
)
|
|
|
|
if result and isinstance(result, dict):
|
|
# 서버 결과를 기존 포맷으로 변환
|
|
status = result.get("status", "failed")
|
|
path = result.get("path", original_image_url)
|
|
|
|
# status 매핑 (서버 결과에 따라 조정)
|
|
if status == "translated":
|
|
return {
|
|
'status': 'translated',
|
|
'path': path,
|
|
'inpaint_method': result.get('inpaint_method', 'unknown'),
|
|
'inpaint_device': result.get('inpaint_device', 'unknown')
|
|
}
|
|
elif status == "original":
|
|
return {
|
|
'status': 'original',
|
|
'path': path,
|
|
'inpaint_method': None,
|
|
'inpaint_device': None
|
|
}
|
|
elif status == "exclude":
|
|
return {
|
|
'status': 'exclude',
|
|
'path': path,
|
|
'inpaint_method': None,
|
|
'inpaint_device': None
|
|
}
|
|
else:
|
|
# 기타 상태는 실패로 처리
|
|
return {
|
|
'status': 'failed',
|
|
'path': original_image_url,
|
|
'error': f'Unknown status: {status}',
|
|
'inpaint_method': None,
|
|
'inpaint_device': None
|
|
}
|
|
else:
|
|
return {
|
|
'status': 'failed',
|
|
'path': original_image_url,
|
|
'error': 'No result from server',
|
|
'inpaint_method': None,
|
|
'inpaint_device': None
|
|
}
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"process_single_image 호환성 메서드 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
return {
|
|
'status': 'failed',
|
|
'path': original_image_url,
|
|
'error': str(e),
|
|
'inpaint_method': None,
|
|
'inpaint_device': None
|
|
}
|
|
|
|
async def remove_background(self, original_image_url, file_prefix=""):
|
|
"""
|
|
기존 ImageProcessor3.remove_background과 호환되는 메서드
|
|
|
|
Args:
|
|
original_image_url (str): 처리할 이미지 URL
|
|
file_prefix (str): 파일명에 추가할 접두사
|
|
|
|
Returns:
|
|
dict: 기존 ImageProcessor3과 동일한 포맷의 결과
|
|
- status: 'success', 'failed' 중 하나
|
|
- path: 처리된 이미지 파일 경로
|
|
- error: 오류 메시지 (status가 'failed'인 경우에만 포함)
|
|
"""
|
|
try:
|
|
# ImageWorkerClient를 통한 배경제거
|
|
result = await self.remove_background_url(
|
|
image_url=original_image_url,
|
|
file_prefix=file_prefix
|
|
)
|
|
|
|
if result and isinstance(result, dict):
|
|
status = result.get("status", "failed")
|
|
path = result.get("path", original_image_url)
|
|
|
|
if status == "success":
|
|
return {
|
|
'status': 'success',
|
|
'path': path
|
|
}
|
|
else:
|
|
return {
|
|
'status': 'failed',
|
|
'path': original_image_url,
|
|
'error': f'Remove background failed: {status}'
|
|
}
|
|
else:
|
|
return {
|
|
'status': 'failed',
|
|
'path': original_image_url,
|
|
'error': 'No result from server'
|
|
}
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"remove_background 호환성 메서드 오류: {e}", level=logging.ERROR, exc_info=True)
|
|
return {
|
|
'status': 'failed',
|
|
'path': original_image_url,
|
|
'error': str(e)
|
|
}
|
|
|
|
def __del__(self):
|
|
"""소멸자에서 리소스 정리"""
|
|
self.cleanup()
|
|
self.logger.log("이미지 프로세서 소멸", level=logging.DEBUG)
|
|
|
|
def cleanup(self):
|
|
"""리소스 정리"""
|
|
try:
|
|
# Python GC 강제 실행
|
|
import gc
|
|
gc.collect()
|
|
|
|
# OpenCV 윈도우 정리
|
|
try:
|
|
cv2.destroyAllWindows()
|
|
except:
|
|
pass
|
|
|
|
# 임시 폴더 삭제
|
|
if hasattr(self, 'TEMP_IMAGE_DIR') and os.path.exists(self.TEMP_IMAGE_DIR):
|
|
# shutil.rmtree(self.TEMP_IMAGE_DIR)
|
|
self.logger.log(f"임시 폴더 삭제됨: {self.TEMP_IMAGE_DIR}", level=logging.DEBUG)
|
|
|
|
except Exception as e:
|
|
self.logger.log(f"리소스 정리 중 오류: {e}", level=logging.ERROR, exc_info=True)
|