123 lines
4.3 KiB
Python
123 lines
4.3 KiB
Python
import json
|
|
import random
|
|
from pathlib import Path
|
|
|
|
from loguru import logger
|
|
|
|
import config
|
|
|
|
|
|
# ── OCR 캐시 (파일별 json 분리) ───────────────────
|
|
|
|
def _cache_path(image_path: str) -> Path:
|
|
"""
|
|
이미지 경로 → json 저장 경로
|
|
files/data/TSMC_A/1000114581_INV/page_001.png
|
|
→ files/json/TSMC_A/1000114581_INV/page_001.json
|
|
"""
|
|
img = Path(image_path)
|
|
try:
|
|
rel = img.relative_to(config.DATA_DIR)
|
|
except ValueError:
|
|
rel = Path(img.parent.name) / img.name
|
|
return config.OCR_CACHE_DIR / rel.with_suffix(".json")
|
|
|
|
|
|
def save_cache_by_file(ocr_result: dict):
|
|
"""이미지 1장의 OCR 결과를 개별 json으로 저장"""
|
|
path = _cache_path(ocr_result["image_path"])
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(ocr_result, f, ensure_ascii=False, indent=2)
|
|
logger.debug(f"캐시 저장: {path}")
|
|
|
|
|
|
def load_cache() -> list:
|
|
"""전체 캐시 통합 로드 (학습/평가 시 사용)"""
|
|
samples = []
|
|
for p in config.OCR_CACHE_DIR.rglob("*.json"):
|
|
try:
|
|
with open(p, encoding="utf-8") as f:
|
|
samples.append(json.load(f))
|
|
except (json.JSONDecodeError, UnicodeDecodeError) as e:
|
|
logger.warning(f"캐시 로드 실패 (건너뜀): {p} | {e}")
|
|
return samples
|
|
|
|
|
|
def get_cached_paths() -> set:
|
|
"""캐시된 이미지 경로 전체 반환 (중복 방지용)"""
|
|
cached = []
|
|
for p in config.OCR_CACHE_DIR.rglob("*.json"):
|
|
try:
|
|
with open(p, encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
cached.append(data.get("image_path", ""))
|
|
# image_path = data.get("image_path", "")
|
|
# if image_path:
|
|
# cached.append(str(Path(image_path))) # 정규화
|
|
# print(str(Path(image_path)))
|
|
except (json.JSONDecodeError, UnicodeDecodeError):
|
|
pass
|
|
return set(cached)
|
|
|
|
|
|
def cache_status() -> dict:
|
|
"""양식별 캐시 현황 (폴더별 json 파일 수 집계)"""
|
|
if not config.OCR_CACHE_DIR.exists():
|
|
return {}
|
|
status = {}
|
|
for cls_dir in sorted(config.OCR_CACHE_DIR.iterdir()):
|
|
if cls_dir.is_dir():
|
|
status[cls_dir.name] = len(list(cls_dir.rglob("*.json")))
|
|
return status
|
|
|
|
|
|
# ── label2id ──────────────────────────────────────
|
|
|
|
def load_label2id() -> dict:
|
|
path = config.MODEL_SAVE_PATH / "label2id.json"
|
|
if path.exists():
|
|
with open(path, encoding="utf-8") as f:
|
|
return json.load(f)
|
|
if config.DATA_DIR.exists():
|
|
classes = sorted([d.name for d in config.DATA_DIR.iterdir() if d.is_dir()])
|
|
return {c: i for i, c in enumerate(classes)}
|
|
return {}
|
|
|
|
|
|
def save_label2id(label2id: dict):
|
|
config.MODEL_SAVE_PATH.mkdir(exist_ok=True)
|
|
with open(config.MODEL_SAVE_PATH / "label2id.json", "w", encoding="utf-8") as f:
|
|
json.dump(label2id, f, ensure_ascii=False, indent=2)
|
|
logger.info(f"label2id 저장: {label2id}")
|
|
|
|
|
|
# ── 리플레이 버퍼 ──────────────────────────────────
|
|
|
|
def load_replay_buffer() -> dict:
|
|
if not config.REPLAY_BUFFER_FILE.exists():
|
|
return {}
|
|
try:
|
|
with open(config.REPLAY_BUFFER_FILE, encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, UnicodeDecodeError) as e:
|
|
backup = config.REPLAY_BUFFER_FILE.with_suffix(".json.bak")
|
|
config.REPLAY_BUFFER_FILE.rename(backup)
|
|
logger.warning(f"replay_buffer.json 손상 → 백업 후 초기화: {e}")
|
|
return {}
|
|
|
|
|
|
def update_replay_buffer(new_samples: list, id2label: dict, per_class: int = 10):
|
|
buf = load_replay_buffer()
|
|
by_label = {}
|
|
for s in new_samples:
|
|
name = id2label[s["label"]]
|
|
by_label.setdefault(name, []).append(s)
|
|
|
|
for name, samples in by_label.items():
|
|
combined = buf.get(name, []) + samples
|
|
buf[name] = random.sample(combined, min(per_class, len(combined)))
|
|
|
|
with open(config.REPLAY_BUFFER_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(buf, f, ensure_ascii=False, indent=2)
|
|
logger.info(f"리플레이 버퍼 갱신: { {k: len(v) for k, v in buf.items()} }") |