436 lines
19 KiB
Python
436 lines
19 KiB
Python
|
||
|
||
# src/modules/image_worker.py
|
||
"""
|
||
ImageWorker 프로세스 – 별도 프로세스로 구동
|
||
단 한 번의 READY("OK") 신호만 보내며, OCR 모델 Warm‑up 후 작업 루프에 진입한다.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import multiprocessing
|
||
import os
|
||
|
||
import logging
|
||
import asyncio
|
||
import traceback
|
||
import queue
|
||
import time
|
||
import cv2
|
||
import numpy as np
|
||
|
||
from modules.image_processor3 import ImageProcessor3
|
||
from modules.log_bridge import ImageWorkerLogger
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# 로깅 유틸리티 #
|
||
# ------------------------------------------------------------------ #
|
||
|
||
class QueueLogger:
|
||
"""큐를 통해 메인 프로세스의 로거로 전송하는 로거"""
|
||
def __init__(self, log_queue, process_name):
|
||
self.log_queue = log_queue
|
||
self.process_name = process_name
|
||
|
||
def log(self, message, level=logging.INFO, exc_info=False):
|
||
try:
|
||
log_record = {
|
||
'process_name': self.process_name,
|
||
'level': level,
|
||
'message': message,
|
||
'exc_info': exc_info
|
||
}
|
||
self.log_queue.put(log_record)
|
||
except Exception:
|
||
pass # 로그 전송 실패 시 무시
|
||
|
||
def debug(self, msg, *a, **kw): self.log(msg, logging.DEBUG)
|
||
def info(self, msg, *a, **kw): self.log(msg, logging.INFO)
|
||
def warning(self, msg, *a, **kw): self.log(msg, logging.WARNING)
|
||
def error(self, msg, *a, **kw): self.log(msg, logging.ERROR)
|
||
def critical(self, msg, *a, **kw): self.log(msg, logging.CRITICAL)
|
||
|
||
class CompatLogger:
|
||
"""커스텀 Logger 인터페이스를 표준 logging.Logger 로 매핑"""
|
||
def __init__(self, py_logger: logging.Logger):
|
||
self._l = py_logger
|
||
|
||
def log(self, message, level=logging.INFO, exc_info=False):
|
||
self._l.log(level, message, exc_info=exc_info)
|
||
|
||
# 편의 메서드
|
||
def debug(self, msg, *a, **kw): self.log(msg, logging.DEBUG)
|
||
def info(self, msg, *a, **kw): self.log(msg, logging.INFO)
|
||
def warning(self, msg, *a, **kw): self.log(msg, logging.WARNING)
|
||
def error(self, msg, *a, **kw): self.log(msg, logging.ERROR)
|
||
def critical(self, msg, *a, **kw): self.log(msg, logging.CRITICAL)
|
||
|
||
|
||
def _setup_logging(log_path: str) -> logging.Logger:
|
||
"""별도 프로세스용 파일 로거 설정"""
|
||
root = logging.getLogger()
|
||
root.handlers.clear()
|
||
root.setLevel(logging.DEBUG)
|
||
|
||
fh = logging.FileHandler(log_path, encoding="utf-8")
|
||
fmt = logging.Formatter(
|
||
"[%(asctime)s] [%(processName)s] [%(levelname)s] "
|
||
"[%(module)s:%(funcName)s:%(lineno)d] %(message)s"
|
||
)
|
||
fh.setFormatter(fmt)
|
||
root.addHandler(fh)
|
||
return root
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# 워커 메인 함수 #
|
||
# ------------------------------------------------------------------ #
|
||
|
||
def worker_main(
|
||
task_q,
|
||
result_q,
|
||
log_queue, # 추가: 로그 큐
|
||
log_path: str,
|
||
base_dir: str,
|
||
toggle_states: dict,
|
||
unwanted_words: list[str],
|
||
authenticated_by_admin: bool = False,
|
||
):
|
||
# ── 로깅 초기화 ────────────────────────────────────────────
|
||
# 큐 로거 사용 (메인 프로세스로 로그 전송)
|
||
if log_queue:
|
||
logger = ImageWorkerLogger(log_queue, f"ImageWorker-{os.getpid()}")
|
||
print(f"[DEBUG] ImageWorker {os.getpid()}: LogQueue 사용됨")
|
||
else:
|
||
# 폴백: 기존 파일 로거
|
||
py_logger = _setup_logging(log_path)
|
||
logger = CompatLogger(py_logger)
|
||
print(f"[DEBUG] ImageWorker {os.getpid()}: 파일 로거 사용됨")
|
||
|
||
logger.info(
|
||
f"ImageWorker 프로세스 기동 "
|
||
f"(PID={os.getpid()}, Name={multiprocessing.current_process().name})"
|
||
)
|
||
|
||
# 워커 기동 시에도 ProgramData/ImgWorker 하위 임시 폴더를 한 번 더 정리
|
||
try:
|
||
base_program = os.environ.get("PROGRAMDATA", r"C:\\ProgramData")
|
||
app_data_dir = os.path.join(base_program, "ImgWorker")
|
||
def _safe_rmtree_contents(target_dir: str, older_than_sec: int = 300):
|
||
try:
|
||
if not target_dir:
|
||
return
|
||
base = os.path.abspath(app_data_dir)
|
||
td = os.path.abspath(target_dir)
|
||
if os.path.commonpath([base, td]) != base:
|
||
return
|
||
os.makedirs(td, exist_ok=True)
|
||
now_ts = time.time()
|
||
thr = max(0, int(older_than_sec or 0))
|
||
|
||
def _remove_if_old(fp: str):
|
||
try:
|
||
if thr <= 0:
|
||
os.remove(fp)
|
||
return
|
||
mt = os.path.getmtime(fp)
|
||
if (now_ts - mt) >= thr:
|
||
os.remove(fp)
|
||
except Exception:
|
||
pass
|
||
|
||
for root, dirs, files in os.walk(td, topdown=False):
|
||
for fn in files:
|
||
_remove_if_old(os.path.join(root, fn))
|
||
for dn in dirs:
|
||
full = os.path.join(root, dn)
|
||
try:
|
||
if not os.listdir(full):
|
||
os.rmdir(full)
|
||
except Exception:
|
||
pass
|
||
except Exception:
|
||
pass
|
||
for d in (os.path.join(app_data_dir, "incoming"), os.path.join(app_data_dir, "work"), os.path.join(app_data_dir, "output"), os.path.join(app_data_dir, "outputs")):
|
||
_safe_rmtree_contents(d, older_than_sec=int(os.environ.get("IMGWK_CLEAN_OLDER_THAN_SEC", "300")))
|
||
except Exception:
|
||
pass
|
||
|
||
# READY는 모델 로딩 완료 후에만 전송
|
||
|
||
# ── ImageProcessor 초기화 및 Warm‑up ──────────────────────
|
||
processor = None
|
||
try:
|
||
# 안전한 초기화를 위한 메모리 정리
|
||
import gc
|
||
gc.collect()
|
||
|
||
logger.info("🔧 ImageProcessor3 초기화 시작...")
|
||
processor = ImageProcessor3(
|
||
logger=logger,
|
||
page=None,
|
||
toggle_states=toggle_states,
|
||
unwanted_words=unwanted_words,
|
||
authenticated_by_admin=authenticated_by_admin,
|
||
base_dir=base_dir,
|
||
papago_translator=None,
|
||
)
|
||
|
||
# OCR 모델 안전한 Warm-up
|
||
if processor and processor.ocr_module:
|
||
try:
|
||
logger.info("🔰 OCR 모듈 Warm-up 시작...")
|
||
dummy_path = os.path.join(base_dir, "_imgproc_warmup.png")
|
||
tmp = np.zeros((100, 100, 3), dtype=np.uint8) # 더 현실적인 크기
|
||
cv2.imwrite(dummy_path, tmp)
|
||
|
||
# 타임아웃 설정으로 무한 대기 방지
|
||
import threading
|
||
import time
|
||
|
||
warmup_success = [False]
|
||
warmup_error = [None]
|
||
|
||
def warmup_ocr():
|
||
try:
|
||
processor.ocr_module.detect_text(dummy_path)
|
||
warmup_success[0] = True
|
||
except Exception as e:
|
||
warmup_error[0] = e
|
||
|
||
warmup_thread = threading.Thread(target=warmup_ocr)
|
||
warmup_thread.daemon = True
|
||
warmup_thread.start()
|
||
warmup_thread.join(timeout=30) # 30초 타임아웃
|
||
|
||
if warmup_success[0]:
|
||
logger.info("✅ OCR 모듈 Warm-up 성공")
|
||
elif warmup_error[0]:
|
||
logger.warning(f"⚠️ OCR 모듈 Warm-up 실패: {warmup_error[0]}")
|
||
else:
|
||
logger.warning("⚠️ OCR 모듈 Warm-up 타임아웃")
|
||
|
||
try:
|
||
os.remove(dummy_path)
|
||
except Exception:
|
||
pass
|
||
|
||
except Exception as e:
|
||
logger.warning(f"OCR Warm-up 실패: {e}")
|
||
else:
|
||
logger.warning("OCR 모듈이 초기화되지 않아 Warm-up 건너뜀")
|
||
|
||
logger.info("🔰 ImageProcessor Warm‑up 완료")
|
||
|
||
except Exception as e:
|
||
logger.error(f"ImageProcessor 초기화 실패: {e}", exc_info=True)
|
||
|
||
# 초기화 실패 시에도 기본적인 처리가 가능하도록 최소한의 processor 생성 시도
|
||
try:
|
||
logger.info("🔄 안전 모드로 재초기화 시도...")
|
||
|
||
# GPU 설정을 CPU로 강제 변경
|
||
safe_toggle_states = toggle_states.copy()
|
||
safe_toggle_states['use_cuda'] = False
|
||
safe_toggle_states['optionIMGTrans_type'] = 'CPU'
|
||
safe_toggle_states['detail_IMGTrans_type'] = 'CPU'
|
||
safe_toggle_states['thumb_trans_type'] = 'CPU'
|
||
safe_toggle_states['migan_use_cuda'] = False
|
||
|
||
processor = ImageProcessor3(
|
||
logger=logger,
|
||
page=None,
|
||
toggle_states=safe_toggle_states,
|
||
unwanted_words=unwanted_words,
|
||
authenticated_by_admin=authenticated_by_admin,
|
||
base_dir=base_dir,
|
||
papago_translator=None,
|
||
)
|
||
logger.info("✅ 안전 모드로 ImageProcessor 초기화 성공")
|
||
|
||
except Exception as e2:
|
||
logger.error(f"안전 모드 초기화도 실패: {e2}", exc_info=True)
|
||
processor = None
|
||
|
||
# ── READY(OK) 신호 전송 ──────────────────────────────────
|
||
try:
|
||
result_q.put({"id": "__READY__", "data": "OK"})
|
||
logger.info("워커 READY 신호 전송")
|
||
except Exception:
|
||
logger.error("READY 신호 전송 실패", exc_info=True)
|
||
|
||
# ── rembg 세션 비동기 로딩 ──────────────────────────────
|
||
def preload_rembg_sessions():
|
||
"""백그라운드에서 rembg 세션을 미리 로딩"""
|
||
try:
|
||
if processor and hasattr(processor, 'background_removal_module'):
|
||
logger.info("🔄 rembg 세션 백그라운드 로딩 시작...")
|
||
# 기본 세션들을 미리 로딩
|
||
processor.background_removal_module._preload_sessions()
|
||
logger.info("✅ rembg 세션 백그라운드 로딩 완료")
|
||
except Exception as e:
|
||
logger.warning(f"rembg 세션 백그라운드 로딩 실패: {e}")
|
||
|
||
# rembg 세션 로딩을 별도 Thread에서 실행
|
||
import threading
|
||
preload_thread = threading.Thread(target=preload_rembg_sessions, daemon=True)
|
||
preload_thread.start()
|
||
|
||
# ── 작업 루프 ─────────────────────────────────────────────
|
||
# 컨트롤러가 즉시 pick-up 할 수 있도록 READY 신호를 추가 전송
|
||
try:
|
||
result_q.put({"id": "__READY__", "cmd": "__READY__", "kwargs": {}})
|
||
logger.info("📡 추가 READY 신호 전송 완료")
|
||
except Exception as e:
|
||
logger.error(f"추가 READY 신호 전송 실패: {e}")
|
||
|
||
idle_log_last = 0.0
|
||
while True:
|
||
try:
|
||
# 주기적 상태 출력을 위해 타임아웃 설정
|
||
logger.debug(f"큐에서 작업 대기 중... (PID: {os.getpid()})")
|
||
task = task_q.get(timeout=60) # 60초 타임아웃
|
||
#logger.info(f"🔥 작업 수신 성공: {task}")
|
||
logger.info(f"🔥 작업 수신 성공")
|
||
except queue.Empty:
|
||
# 유휴 로그는 10분에 한 번만 기록해 로그 스팸을 줄인다.
|
||
now_ts = time.time()
|
||
if now_ts - idle_log_last >= 600:
|
||
idle_log_last = now_ts
|
||
logger.info("대기 중(유휴)")
|
||
continue
|
||
except Exception as e:
|
||
logger.error(f"작업 수신 중 오류: {e}", exc_info=True)
|
||
continue
|
||
|
||
if task is None:
|
||
logger.info("Shutdown signal 수신 → 종료")
|
||
try:
|
||
# 종료 시 임시 폴더 정리(5분 경과 파일만)
|
||
base_program = os.environ.get("PROGRAMDATA", r"C:\\ProgramData")
|
||
app_data_dir = os.path.join(base_program, "ImgWorker")
|
||
def _cleanup_dir(dp: str, older_than_sec: int = 300):
|
||
try:
|
||
now_ts = time.time()
|
||
thr = max(0, int(older_than_sec or 0))
|
||
for root, dirs, files in os.walk(dp, topdown=False):
|
||
for fn in files:
|
||
fp = os.path.join(root, fn)
|
||
try:
|
||
if thr <= 0:
|
||
os.remove(fp)
|
||
else:
|
||
mt = os.path.getmtime(fp)
|
||
if (now_ts - mt) >= thr:
|
||
os.remove(fp)
|
||
except Exception:
|
||
pass
|
||
for dn in dirs:
|
||
full = os.path.join(root, dn)
|
||
try:
|
||
if not os.listdir(full):
|
||
os.rmdir(full)
|
||
except Exception:
|
||
pass
|
||
except Exception:
|
||
pass
|
||
ttl = int(os.environ.get("IMGWK_CLEAN_OLDER_THAN_SEC", "300"))
|
||
for d in (os.path.join(app_data_dir, "incoming"), os.path.join(app_data_dir, "work"), os.path.join(app_data_dir, "output"), os.path.join(app_data_dir, "outputs")):
|
||
_cleanup_dir(d, older_than_sec=ttl)
|
||
except Exception:
|
||
pass
|
||
break
|
||
|
||
uid = task["id"]
|
||
cmd = task["cmd"]
|
||
kwargs = task["kwargs"]
|
||
logger.info(f"🚀 작업 처리 시작: cmd={cmd}, uid={uid}")
|
||
|
||
# 메타 파라미터 제거 및 실시간 값 반영
|
||
new_toggle = kwargs.pop("_toggle_states", None)
|
||
if new_toggle and processor:
|
||
processor.update_toggle_states(new_toggle)
|
||
|
||
_ = kwargs.pop("_base_dir", None) # 필요 없으므로 버림
|
||
upd_unwanted = kwargs.pop("_update_unwanted_texts", None)
|
||
if upd_unwanted and processor:
|
||
processor.update_unwanted_texts(upd_unwanted)
|
||
|
||
# 실제 작업 실행
|
||
try:
|
||
logger.debug(f"작업 실행 직전: cmd={cmd}")
|
||
if cmd == "process_single_image":
|
||
logger.debug("process_single_image 호출 직전")
|
||
data = asyncio.run(processor.process_single_image(**kwargs))
|
||
# 성능 지표 포함
|
||
try:
|
||
timings = getattr(processor, '_last_timings', None)
|
||
if isinstance(data, dict) and timings:
|
||
data['timings'] = timings
|
||
except Exception:
|
||
pass
|
||
logger.debug("process_single_image 호출 완료")
|
||
elif cmd == "remove_background":
|
||
logger.debug("remove_background 호출 직전")
|
||
data = asyncio.run(processor.remove_background(**kwargs))
|
||
logger.debug("remove_background 호출 완료")
|
||
elif cmd == "reinit_ocr":
|
||
# 토글 반영 후 OCR 재초기화(프로바이더 캐시 고려)
|
||
ok = processor.reset_ocr_module()
|
||
data = {"ok": bool(ok)}
|
||
elif cmd == "reinit_rembg":
|
||
# REMBG(배경제거) 모듈 재준비: 현재 Bria 모듈 사용시 세션/프로바이더를 재평가하도록 None 초기화
|
||
try:
|
||
# toggle_states의 provider override는 상위에서 반영되어 내려옴
|
||
# BriaBackgroundRemovalModule은 lazy-load 방식 → 재생성 유도
|
||
if hasattr(processor, 'background_removal_module'):
|
||
try:
|
||
del processor.background_removal_module
|
||
except Exception:
|
||
processor.background_removal_module = None
|
||
from modules.bria_background_removal_module import BriaBackgroundRemovalModule
|
||
# 경로/매개변수는 processor.toggle_states에서 유추(존재 시)
|
||
model_path = processor.toggle_states.get('local_rembg_model_path')
|
||
processor.background_removal_module = BriaBackgroundRemovalModule(
|
||
logger=logger,
|
||
default_model=processor.toggle_states.get('local_model_name', 'bria-rmbg-1.4'),
|
||
gpu_manager=getattr(processor, 'gpu_manager', None),
|
||
local_rembg_model_path=model_path,
|
||
)
|
||
data = {"ok": True}
|
||
except Exception as e:
|
||
logger.error(f"REMBG 재초기화 실패: {e}")
|
||
data = {"ok": False, "error": str(e)}
|
||
elif cmd == "reset_migan":
|
||
# MIGAN 재구성(토글상 migan_use_cuda 등 변경 반영)
|
||
try:
|
||
from modules.migan_module import build_migan_from_toggle
|
||
enhanced_toggle_states = processor.toggle_states.copy()
|
||
# 가속 사용 플래그 명칭 정리(호환): migan_use_cuda -> migan_use_accel
|
||
if 'migan_use_cuda' in enhanced_toggle_states and 'migan_use_accel' not in enhanced_toggle_states:
|
||
enhanced_toggle_states['migan_use_accel'] = enhanced_toggle_states['migan_use_cuda']
|
||
# provider override가 들어왔으면 반영(auto|dml|cpu)
|
||
prov = kwargs.get('provider')
|
||
if prov:
|
||
enhanced_toggle_states['migan_provider_override'] = prov
|
||
# gpu_manager 상태와 무관하게 토글을 그대로 전달(직접 폴백은 모듈 내부)
|
||
processor.migan = build_migan_from_toggle(enhanced_toggle_states, logger=logger, gpu_manager=getattr(processor, 'gpu_manager', None))
|
||
data = {"ok": bool(processor.migan is not None)}
|
||
except Exception as mm_err:
|
||
logger.error(f"MIGAN 재설정 실패: {mm_err}")
|
||
data = {"ok": False, "error": str(mm_err)}
|
||
elif cmd == "__PING__":
|
||
# 하트비트 응답
|
||
data = "__PONG__"
|
||
elif cmd == "update_toggle_states":
|
||
# 토글 업데이트는 위쪽 공통 로직에서 이미 수행됨
|
||
data = {"ok": True}
|
||
else:
|
||
raise ValueError(f"unknown cmd: {cmd}")
|
||
|
||
logger.debug(f"작업 결과 반환 중: uid={uid}")
|
||
result_q.put({"id": uid, "data": data})
|
||
logger.debug(f"작업 결과 반환 완료: uid={uid}")
|
||
except Exception:
|
||
logger.error(f"작업 처리 중 오류: cmd={cmd}, uid={uid}")
|
||
logger.error("작업 처리 중 오류", exc_info=True)
|
||
result_q.put({"id": uid, "error": traceback.format_exc()})
|