# 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()})