284 lines
8.0 KiB
Python
284 lines
8.0 KiB
Python
"""
|
|
OCR 结果缓存模块
|
|
使用文件系统缓存 OCR 解析结果，避免重复解析
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import hashlib
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Dict, Any, Optional
|
|
from datetime import datetime, timedelta
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class OCRResultCache:
|
|
"""
|
|
OCR 结果缓存管理器
|
|
|
|
功能：
|
|
- 基于文件哈希的缓存键生成
|
|
- 支持缓存过期时间
|
|
- 自动清理过期缓存
|
|
- 统计缓存命中率
|
|
"""
|
|
|
|
def __init__(self, cache_dir: str = "cache/ocr_results", expire_days: int = 30):
|
|
"""
|
|
初始化缓存管理器
|
|
|
|
Args:
|
|
cache_dir: 缓存目录
|
|
expire_days: 缓存过期天数
|
|
"""
|
|
self.cache_dir = Path(cache_dir)
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
self.expire_days = expire_days
|
|
|
|
# 统计信息
|
|
self.stats = {
|
|
"hits": 0,
|
|
"misses": 0,
|
|
"saves": 0
|
|
}
|
|
|
|
logger.info(f"OCR 缓存初始化: {self.cache_dir}, 过期: {expire_days} 天")
|
|
|
|
def _get_file_hash(self, file_path: str) -> str:
|
|
"""
|
|
计算文件的 MD5 哈希值
|
|
|
|
Args:
|
|
file_path: 文件路径
|
|
|
|
Returns:
|
|
文件的 MD5 哈希值
|
|
"""
|
|
hasher = hashlib.md5()
|
|
|
|
with open(file_path, 'rb') as f:
|
|
# 分块读取，避免大文件内存问题
|
|
for chunk in iter(lambda: f.read(4096), b''):
|
|
hasher.update(chunk)
|
|
|
|
return hasher.hexdigest()
|
|
|
|
def _get_cache_key(
|
|
self,
|
|
file_path: str,
|
|
ocr_engine: str,
|
|
additional_params: Optional[Dict] = None
|
|
) -> str:
|
|
"""
|
|
生成缓存键
|
|
|
|
Args:
|
|
file_path: 文件路径
|
|
ocr_engine: OCR 引擎
|
|
additional_params: 额外参数（如 use_gpu, dpi 等）
|
|
|
|
Returns:
|
|
缓存键
|
|
"""
|
|
# 计算文件哈希
|
|
file_hash = self._get_file_hash(file_path)
|
|
|
|
# 组合参数生成唯一键
|
|
key_parts = [file_hash, ocr_engine]
|
|
|
|
if additional_params:
|
|
# 将参数排序后加入键
|
|
sorted_params = sorted(additional_params.items())
|
|
params_str = "_".join(f"{k}={v}" for k, v in sorted_params)
|
|
key_parts.append(params_str)
|
|
|
|
cache_key = "_".join(key_parts)
|
|
return cache_key
|
|
|
|
def _get_cache_file_path(self, cache_key: str) -> Path:
|
|
"""获取缓存文件路径"""
|
|
return self.cache_dir / f"{cache_key}.json"
|
|
|
|
def get(
|
|
self,
|
|
file_path: str,
|
|
ocr_engine: str,
|
|
additional_params: Optional[Dict] = None
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
从缓存获取结果
|
|
|
|
Args:
|
|
file_path: 文件路径
|
|
ocr_engine: OCR 引擎
|
|
additional_params: 额外参数
|
|
|
|
Returns:
|
|
缓存的结果，如果不存在或已过期则返回 None
|
|
"""
|
|
try:
|
|
cache_key = self._get_cache_key(file_path, ocr_engine, additional_params)
|
|
cache_file = self._get_cache_file_path(cache_key)
|
|
|
|
if not cache_file.exists():
|
|
self.stats["misses"] += 1
|
|
logger.debug(f"缓存未命中: {cache_key}")
|
|
return None
|
|
|
|
# 检查缓存是否过期
|
|
file_mtime = datetime.fromtimestamp(cache_file.stat().st_mtime)
|
|
expire_time = datetime.now() - timedelta(days=self.expire_days)
|
|
|
|
if file_mtime < expire_time:
|
|
logger.info(f"缓存已过期: {cache_key}")
|
|
cache_file.unlink() # 删除过期缓存
|
|
self.stats["misses"] += 1
|
|
return None
|
|
|
|
# 读取缓存
|
|
with open(cache_file, 'r', encoding='utf-8') as f:
|
|
cached_data = json.load(f)
|
|
|
|
self.stats["hits"] += 1
|
|
logger.info(f"缓存命中: {cache_key}")
|
|
return cached_data
|
|
|
|
except Exception as e:
|
|
logger.error(f"读取缓存失败: {str(e)}", exc_info=True)
|
|
self.stats["misses"] += 1
|
|
return None
|
|
|
|
def set(
|
|
self,
|
|
file_path: str,
|
|
ocr_engine: str,
|
|
result: Dict[str, Any],
|
|
additional_params: Optional[Dict] = None
|
|
):
|
|
"""
|
|
保存结果到缓存
|
|
|
|
Args:
|
|
file_path: 文件路径
|
|
ocr_engine: OCR 引擎
|
|
result: OCR 结果
|
|
additional_params: 额外参数
|
|
"""
|
|
try:
|
|
cache_key = self._get_cache_key(file_path, ocr_engine, additional_params)
|
|
cache_file = self._get_cache_file_path(cache_key)
|
|
|
|
# 添加元数据
|
|
cache_data = {
|
|
"cached_at": datetime.now().isoformat(),
|
|
"file_path": file_path,
|
|
"ocr_engine": ocr_engine,
|
|
"additional_params": additional_params or {},
|
|
"result": result
|
|
}
|
|
|
|
# 保存到文件
|
|
with open(cache_file, 'w', encoding='utf-8') as f:
|
|
json.dump(cache_data, f, ensure_ascii=False, indent=2)
|
|
|
|
self.stats["saves"] += 1
|
|
logger.info(f"结果已缓存: {cache_key}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"保存缓存失败: {str(e)}", exc_info=True)
|
|
|
|
def clear_expired(self):
|
|
"""清理过期缓存"""
|
|
try:
|
|
expire_time = datetime.now() - timedelta(days=self.expire_days)
|
|
cleared_count = 0
|
|
|
|
for cache_file in self.cache_dir.glob("*.json"):
|
|
file_mtime = datetime.fromtimestamp(cache_file.stat().st_mtime)
|
|
if file_mtime < expire_time:
|
|
cache_file.unlink()
|
|
cleared_count += 1
|
|
|
|
logger.info(f"清理过期缓存: 删除 {cleared_count} 个文件")
|
|
return cleared_count
|
|
|
|
except Exception as e:
|
|
logger.error(f"清理缓存失败: {str(e)}", exc_info=True)
|
|
return 0
|
|
|
|
def clear_all(self):
|
|
"""清空所有缓存"""
|
|
try:
|
|
cleared_count = 0
|
|
for cache_file in self.cache_dir.glob("*.json"):
|
|
cache_file.unlink()
|
|
cleared_count += 1
|
|
|
|
logger.info(f"清空所有缓存: 删除 {cleared_count} 个文件")
|
|
return cleared_count
|
|
|
|
except Exception as e:
|
|
logger.error(f"清空缓存失败: {str(e)}", exc_info=True)
|
|
return 0
|
|
|
|
def get_cache_size(self) -> Dict[str, Any]:
|
|
"""获取缓存大小统计"""
|
|
try:
|
|
total_size = 0
|
|
file_count = 0
|
|
|
|
for cache_file in self.cache_dir.glob("*.json"):
|
|
total_size += cache_file.stat().st_size
|
|
file_count += 1
|
|
|
|
return {
|
|
"file_count": file_count,
|
|
"total_size_bytes": total_size,
|
|
"total_size_mb": round(total_size / 1024 / 1024, 2),
|
|
"cache_dir": str(self.cache_dir)
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"获取缓存大小失败: {str(e)}", exc_info=True)
|
|
return {}
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
"""获取缓存统计信息"""
|
|
total_requests = self.stats["hits"] + self.stats["misses"]
|
|
hit_rate = (self.stats["hits"] / total_requests * 100) if total_requests > 0 else 0
|
|
|
|
cache_size = self.get_cache_size()
|
|
|
|
return {
|
|
**self.stats,
|
|
"total_requests": total_requests,
|
|
"hit_rate": f"{hit_rate:.2f}%",
|
|
**cache_size
|
|
}
|
|
|
|
|
|
# 全局缓存实例
|
|
_global_cache = None
|
|
|
|
|
|
def get_cache() -> OCRResultCache:
|
|
"""获取全局缓存实例"""
|
|
global _global_cache
|
|
if _global_cache is None:
|
|
_global_cache = OCRResultCache()
|
|
return _global_cache
|
|
|
|
|
|
def clear_cache():
|
|
"""清空全局缓存"""
|
|
cache = get_cache()
|
|
return cache.clear_all()
|
|
|
|
|
|
def get_cache_stats():
|
|
"""获取缓存统计"""
|
|
cache = get_cache()
|
|
return cache.get_stats()
|