112 lines
4.4 KiB
Python
112 lines
4.4 KiB
Python
import numpy as np
|
|
import requests
|
|
import cv2
|
|
import base64
|
|
import subprocess
|
|
import random
|
|
import time
|
|
import threading
|
|
import os
|
|
|
|
class IOPaintInpainting:
|
|
"""IOPaint 서버 연동 인페인팅 모델 (REST API /api/v1/inpaint 사용, 바이너리 PNG 반환)"""
|
|
def __init__(self, server_url="http://localhost:8099"):
|
|
self.api_url = f"{server_url}/api/v1/inpaint"
|
|
def inpaint(self, image: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
|
# 이미지를 base64로 인코딩
|
|
_, img_encoded = cv2.imencode('.png', image)
|
|
_, mask_encoded = cv2.imencode('.png', mask)
|
|
img_b64 = base64.b64encode(img_encoded).decode('utf-8')
|
|
mask_b64 = base64.b64encode(mask_encoded).decode('utf-8')
|
|
payload = {
|
|
"image": img_b64,
|
|
"mask": mask_b64
|
|
}
|
|
response = requests.post(self.api_url, json=payload)
|
|
if response.status_code != 200:
|
|
print("IOPaint 서버 에러:", response.text)
|
|
return None
|
|
# 응답이 바이너리 PNG 이미지이므로 바로 디코딩
|
|
nparr = np.frombuffer(response.content, np.uint8)
|
|
result = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
|
return result
|
|
|
|
class IOPaintServerInstance:
|
|
"""IOPaint 서버 인스턴스 정보 및 상태 관리"""
|
|
def __init__(self, port, process):
|
|
self.port = port
|
|
self.process = process
|
|
self.busy = False
|
|
self.last_used = time.time()
|
|
def mark_busy(self):
|
|
self.busy = True
|
|
self.last_used = time.time()
|
|
def mark_idle(self):
|
|
self.busy = False
|
|
self.last_used = time.time()
|
|
def is_alive(self):
|
|
return self.process.poll() is None
|
|
|
|
class IOPaintServerManager:
|
|
"""여러 IOPaint 서버 인스턴스를 관리하는 매니저"""
|
|
def __init__(self, num_instances=1, port_range=(8099, 8199), iopaint_cmd=None, wait_ready=5):
|
|
self.instances = []
|
|
self.port_range = port_range
|
|
self.iopaint_cmd = iopaint_cmd or ["python", "-m", "iopaint.server"]
|
|
self.lock = threading.Lock()
|
|
self._start_instances(num_instances, wait_ready)
|
|
def _get_random_port(self):
|
|
used_ports = {inst.port for inst in self.instances}
|
|
candidates = [p for p in range(self.port_range[0], self.port_range[1]+1) if p not in used_ports]
|
|
if not candidates:
|
|
raise RuntimeError("사용 가능한 포트가 없습니다.")
|
|
return random.choice(candidates)
|
|
def _start_instances(self, num, wait_ready):
|
|
for _ in range(num):
|
|
port = self._get_random_port()
|
|
cmd = self.iopaint_cmd + ["--port", str(port)]
|
|
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
instance = IOPaintServerInstance(port, proc)
|
|
self.instances.append(instance)
|
|
time.sleep(wait_ready) # 서버가 준비될 때까지 대기
|
|
def get_instance_info(self):
|
|
"""모든 인스턴스의 정보를 반환"""
|
|
info = []
|
|
for inst in self.instances:
|
|
info.append({
|
|
"port": inst.port,
|
|
"busy": inst.busy,
|
|
"alive": inst.is_alive(),
|
|
"last_used": inst.last_used
|
|
})
|
|
return info
|
|
def get_idle_instance(self):
|
|
"""놀고 있는(사용 가능한) 인스턴스 반환 (없으면 None)"""
|
|
with self.lock:
|
|
for inst in self.instances:
|
|
if not inst.busy and inst.is_alive():
|
|
inst.mark_busy()
|
|
return inst
|
|
return None
|
|
def mark_instance_idle(self, port):
|
|
"""작업이 끝난 인스턴스를 idle로 표시"""
|
|
for inst in self.instances:
|
|
if inst.port == port:
|
|
inst.mark_idle()
|
|
break
|
|
def shutdown_all(self):
|
|
"""모든 서버 인스턴스 종료"""
|
|
for inst in self.instances:
|
|
if inst.is_alive():
|
|
inst.process.terminate()
|
|
self.instances = []
|
|
def add_instance(self, wait_ready=5):
|
|
"""새로운 인스턴스 추가"""
|
|
port = self._get_random_port()
|
|
cmd = self.iopaint_cmd + ["--port", str(port)]
|
|
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
instance = IOPaintServerInstance(port, proc)
|
|
self.instances.append(instance)
|
|
time.sleep(wait_ready)
|
|
return instance
|