787 lines
29 KiB
Python
787 lines
29 KiB
Python
"""
|
||
图片OCR热量识别模块 - 基于基座架构
|
||
支持多种OCR技术识别食物热量信息，包含智能验证和修正机制
|
||
"""
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import re
|
||
import json
|
||
import logging
|
||
from typing import Dict, List, Optional, Any, Tuple
|
||
from dataclasses import dataclass
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
import requests
|
||
import base64
|
||
from PIL import Image, ImageEnhance, ImageFilter
|
||
import pytesseract
|
||
from core.base import BaseModule, ModuleType, UserData, AnalysisResult, BaseConfig
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class OCRResult:
|
||
"""OCR识别结果"""
|
||
text: str
|
||
confidence: float
|
||
bounding_boxes: List[Dict[str, Any]]
|
||
processing_time: float
|
||
method: str
|
||
|
||
|
||
@dataclass
|
||
class CalorieInfo:
|
||
"""热量信息"""
|
||
food_name: str
|
||
calories: Optional[float]
|
||
serving_size: Optional[str]
|
||
confidence: float
|
||
source: str # 'ocr', 'database', 'user_confirmed'
|
||
raw_text: str
|
||
validation_status: str # 'pending', 'validated', 'corrected'
|
||
|
||
|
||
@dataclass
|
||
class FoodRecognitionResult:
|
||
"""食物识别结果"""
|
||
image_path: str
|
||
ocr_results: List[OCRResult]
|
||
calorie_infos: List[CalorieInfo]
|
||
overall_confidence: float
|
||
processing_time: float
|
||
suggestions: List[str]
|
||
|
||
|
||
class OCRCalorieRecognitionModule(BaseModule):
|
||
"""OCR热量识别模块"""
|
||
|
||
def __init__(self, config: BaseConfig):
|
||
super().__init__(config, ModuleType.DATA_COLLECTION)
|
||
|
||
# OCR配置
|
||
self.ocr_methods = ['tesseract', 'paddleocr', 'easyocr']
|
||
self.min_confidence = 0.6
|
||
self.max_processing_time = 30.0
|
||
|
||
# 热量识别模式
|
||
self.calorie_patterns = [
|
||
r'(\d+(?:\.\d+)?)\s*[kK]?[cC][aA][lL](?:ories?)?',
|
||
r'(\d+(?:\.\d+)?)\s*[kK][cC][aA][lL]',
|
||
r'(\d+(?:\.\d+)?)\s*卡路里',
|
||
r'(\d+(?:\.\d+)?)\s*千卡',
|
||
r'(\d+(?:\.\d+)?)\s*大卡',
|
||
r'(\d+(?:\.\d+)?)\s*[kK][jJ]', # 千焦
|
||
]
|
||
|
||
# 食物名称模式
|
||
self.food_patterns = [
|
||
r'([a-zA-Z\u4e00-\u9fff]+)\s*(?:\d+(?:\.\d+)?)',
|
||
r'(\d+(?:\.\d+)?)\s*([a-zA-Z\u4e00-\u9fff]+)',
|
||
]
|
||
|
||
# 食物数据库
|
||
self.food_database = self._load_food_database()
|
||
|
||
# 用户学习数据
|
||
self.user_corrections = {}
|
||
|
||
# 初始化OCR引擎
|
||
self.ocr_engines = {}
|
||
self._initialize_ocr_engines()
|
||
|
||
def initialize(self) -> bool:
|
||
"""初始化模块"""
|
||
try:
|
||
self.logger.info("OCR热量识别模块初始化中...")
|
||
|
||
# 创建必要的目录
|
||
self._create_directories()
|
||
|
||
# 加载用户学习数据
|
||
self._load_user_corrections()
|
||
|
||
self.is_initialized = True
|
||
self.logger.info("OCR热量识别模块初始化完成")
|
||
return True
|
||
except Exception as e:
|
||
self.logger.error(f"OCR热量识别模块初始化失败: {e}")
|
||
return False
|
||
|
||
def process(self, input_data: Any, user_data: UserData) -> AnalysisResult:
|
||
"""处理OCR识别请求"""
|
||
try:
|
||
request_type = input_data.get('type', 'unknown')
|
||
|
||
if request_type == 'recognize_image':
|
||
result = self._recognize_image_calories(input_data, user_data)
|
||
elif request_type == 'validate_result':
|
||
result = self._validate_recognition_result(input_data, user_data)
|
||
elif request_type == 'learn_correction':
|
||
result = self._learn_from_correction(input_data, user_data)
|
||
else:
|
||
result = self._create_error_result("未知的请求类型")
|
||
|
||
return AnalysisResult(
|
||
module_type=self.module_type,
|
||
user_id=user_data.user_id,
|
||
input_data=input_data,
|
||
result=result,
|
||
confidence=result.get('confidence', 0.5)
|
||
)
|
||
except Exception as e:
|
||
self.logger.error(f"OCR识别处理失败: {e}")
|
||
return self._create_error_result(f"处理失败: {str(e)}")
|
||
|
||
def _recognize_image_calories(self, input_data: Dict[str, Any], user_data: UserData) -> Dict[str, Any]:
|
||
"""识别图片中的热量信息"""
|
||
start_time = datetime.now()
|
||
|
||
try:
|
||
image_path = input_data.get('image_path')
|
||
if not image_path or not Path(image_path).exists():
|
||
return self._create_error_result("图片文件不存在")
|
||
|
||
# 预处理图片
|
||
processed_image = self._preprocess_image(image_path)
|
||
|
||
# 多OCR引擎识别
|
||
ocr_results = []
|
||
for method in self.ocr_methods:
|
||
try:
|
||
result = self._ocr_recognize(processed_image, method)
|
||
if result:
|
||
ocr_results.append(result)
|
||
except Exception as e:
|
||
self.logger.warning(f"OCR方法 {method} 失败: {e}")
|
||
|
||
# 合并和去重OCR结果
|
||
merged_text = self._merge_ocr_results(ocr_results)
|
||
|
||
# 提取热量信息
|
||
calorie_infos = self._extract_calorie_info(merged_text, user_data)
|
||
|
||
# 数据库匹配和验证
|
||
validated_infos = self._validate_with_database(calorie_infos, user_data)
|
||
|
||
# 生成建议
|
||
suggestions = self._generate_suggestions(validated_infos, user_data)
|
||
|
||
processing_time = (datetime.now() - start_time).total_seconds()
|
||
|
||
result = FoodRecognitionResult(
|
||
image_path=image_path,
|
||
ocr_results=ocr_results,
|
||
calorie_infos=validated_infos,
|
||
overall_confidence=self._calculate_overall_confidence(ocr_results, validated_infos),
|
||
processing_time=processing_time,
|
||
suggestions=suggestions
|
||
)
|
||
|
||
return {
|
||
'success': True,
|
||
'result': result,
|
||
'confidence': result.overall_confidence,
|
||
'message': f"识别完成，处理时间: {processing_time:.2f}秒"
|
||
}
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"图片热量识别失败: {e}")
|
||
return self._create_error_result(f"识别失败: {str(e)}")
|
||
|
||
def _preprocess_image(self, image_path: str) -> np.ndarray:
|
||
"""预处理图片以提高OCR准确性"""
|
||
try:
|
||
# 读取图片
|
||
image = cv2.imread(image_path)
|
||
if image is None:
|
||
raise ValueError("无法读取图片")
|
||
|
||
# 转换为灰度图
|
||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
|
||
# 降噪
|
||
denoised = cv2.medianBlur(gray, 3)
|
||
|
||
# 增强对比度
|
||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||
enhanced = clahe.apply(denoised)
|
||
|
||
# 二值化
|
||
_, binary = cv2.threshold(enhanced, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
||
|
||
# 形态学操作
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
|
||
cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
|
||
|
||
return cleaned
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"图片预处理失败: {e}")
|
||
return cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
|
||
|
||
def _ocr_recognize(self, image: np.ndarray, method: str) -> Optional[OCRResult]:
|
||
"""使用指定方法进行OCR识别"""
|
||
start_time = datetime.now()
|
||
|
||
try:
|
||
if method == 'tesseract':
|
||
return self._tesseract_ocr(image)
|
||
elif method == 'paddleocr':
|
||
return self._paddleocr_recognize(image)
|
||
elif method == 'easyocr':
|
||
return self._easyocr_recognize(image)
|
||
else:
|
||
return None
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"OCR方法 {method} 失败: {e}")
|
||
return None
|
||
|
||
def _tesseract_ocr(self, image: np.ndarray) -> OCRResult:
|
||
"""使用Tesseract进行OCR识别"""
|
||
try:
|
||
# 配置Tesseract
|
||
config = '--oem 3 --psm 6 -c tessedit_char_whitelist=0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\u4e00-\u9fff'
|
||
|
||
# OCR识别
|
||
text = pytesseract.image_to_string(image, config=config, lang='chi_sim+eng')
|
||
|
||
# 获取置信度
|
||
data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT, config=config, lang='chi_sim+eng')
|
||
confidences = [int(conf) for conf in data['conf'] if int(conf) > 0]
|
||
avg_confidence = sum(confidences) / len(confidences) / 100.0 if confidences else 0.0
|
||
|
||
# 获取边界框
|
||
bounding_boxes = []
|
||
for i in range(len(data['text'])):
|
||
if int(data['conf'][i]) > 0:
|
||
bounding_boxes.append({
|
||
'text': data['text'][i],
|
||
'confidence': int(data['conf'][i]) / 100.0,
|
||
'bbox': [data['left'][i], data['top'][i], data['width'][i], data['height'][i]]
|
||
})
|
||
|
||
processing_time = (datetime.now() - datetime.now()).total_seconds()
|
||
|
||
return OCRResult(
|
||
text=text.strip(),
|
||
confidence=avg_confidence,
|
||
bounding_boxes=bounding_boxes,
|
||
processing_time=processing_time,
|
||
method='tesseract'
|
||
)
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"Tesseract OCR失败: {e}")
|
||
return None
|
||
|
||
def _paddleocr_recognize(self, image: np.ndarray) -> Optional[OCRResult]:
|
||
"""使用PaddleOCR进行识别"""
|
||
try:
|
||
# 这里需要安装paddleocr: pip install paddleocr
|
||
from paddleocr import PaddleOCR
|
||
|
||
if 'paddleocr' not in self.ocr_engines:
|
||
self.ocr_engines['paddleocr'] = PaddleOCR(use_angle_cls=True, lang='ch')
|
||
|
||
ocr = self.ocr_engines['paddleocr']
|
||
result = ocr.ocr(image, cls=True)
|
||
|
||
if not result or not result[0]:
|
||
return None
|
||
|
||
# 提取文本和置信度
|
||
texts = []
|
||
confidences = []
|
||
bounding_boxes = []
|
||
|
||
for line in result[0]:
|
||
if line:
|
||
bbox, (text, confidence) = line
|
||
texts.append(text)
|
||
confidences.append(confidence)
|
||
bounding_boxes.append({
|
||
'text': text,
|
||
'confidence': confidence,
|
||
'bbox': bbox
|
||
})
|
||
|
||
merged_text = ' '.join(texts)
|
||
avg_confidence = sum(confidences) / len(confidences) if confidences else 0.0
|
||
|
||
return OCRResult(
|
||
text=merged_text,
|
||
confidence=avg_confidence,
|
||
bounding_boxes=bounding_boxes,
|
||
processing_time=0.0,
|
||
method='paddleocr'
|
||
)
|
||
|
||
except ImportError:
|
||
self.logger.warning("PaddleOCR未安装，跳过此方法")
|
||
return None
|
||
except Exception as e:
|
||
self.logger.error(f"PaddleOCR识别失败: {e}")
|
||
return None
|
||
|
||
def _easyocr_recognize(self, image: np.ndarray) -> Optional[OCRResult]:
|
||
"""使用EasyOCR进行识别"""
|
||
try:
|
||
# 这里需要安装easyocr: pip install easyocr
|
||
import easyocr
|
||
|
||
if 'easyocr' not in self.ocr_engines:
|
||
self.ocr_engines['easyocr'] = easyocr.Reader(['ch_sim', 'en'])
|
||
|
||
reader = self.ocr_engines['easyocr']
|
||
result = reader.readtext(image)
|
||
|
||
if not result:
|
||
return None
|
||
|
||
# 提取文本和置信度
|
||
texts = []
|
||
confidences = []
|
||
bounding_boxes = []
|
||
|
||
for bbox, text, confidence in result:
|
||
texts.append(text)
|
||
confidences.append(confidence)
|
||
bounding_boxes.append({
|
||
'text': text,
|
||
'confidence': confidence,
|
||
'bbox': bbox
|
||
})
|
||
|
||
merged_text = ' '.join(texts)
|
||
avg_confidence = sum(confidences) / len(confidences) if confidences else 0.0
|
||
|
||
return OCRResult(
|
||
text=merged_text,
|
||
confidence=avg_confidence,
|
||
bounding_boxes=bounding_boxes,
|
||
processing_time=0.0,
|
||
method='easyocr'
|
||
)
|
||
|
||
except ImportError:
|
||
self.logger.warning("EasyOCR未安装，跳过此方法")
|
||
return None
|
||
except Exception as e:
|
||
self.logger.error(f"EasyOCR识别失败: {e}")
|
||
return None
|
||
|
||
def _merge_ocr_results(self, ocr_results: List[OCRResult]) -> str:
|
||
"""合并多个OCR结果"""
|
||
if not ocr_results:
|
||
return ""
|
||
|
||
# 按置信度排序
|
||
sorted_results = sorted(ocr_results, key=lambda x: x.confidence, reverse=True)
|
||
|
||
# 使用最高置信度的结果作为主要结果
|
||
primary_result = sorted_results[0]
|
||
|
||
# 如果有多个高置信度结果，尝试合并
|
||
if len(sorted_results) > 1 and sorted_results[1].confidence > 0.7:
|
||
# 简单的文本合并策略
|
||
merged_text = self._smart_text_merge([r.text for r in sorted_results[:3]])
|
||
return merged_text
|
||
|
||
return primary_result.text
|
||
|
||
def _smart_text_merge(self, texts: List[str]) -> str:
|
||
"""智能文本合并"""
|
||
if not texts:
|
||
return ""
|
||
|
||
if len(texts) == 1:
|
||
return texts[0]
|
||
|
||
# 简单的合并策略：选择最长的文本
|
||
return max(texts, key=len)
|
||
|
||
def _extract_calorie_info(self, text: str, user_data: UserData) -> List[CalorieInfo]:
|
||
"""从文本中提取热量信息"""
|
||
calorie_infos = []
|
||
|
||
try:
|
||
# 查找热量数值
|
||
for pattern in self.calorie_patterns:
|
||
matches = re.finditer(pattern, text, re.IGNORECASE)
|
||
for match in matches:
|
||
calories = float(match.group(1))
|
||
|
||
# 查找对应的食物名称
|
||
food_name = self._extract_food_name(text, match.start())
|
||
|
||
calorie_info = CalorieInfo(
|
||
food_name=food_name,
|
||
calories=calories,
|
||
serving_size=None,
|
||
confidence=0.8, # OCR基础置信度
|
||
source='ocr',
|
||
raw_text=match.group(0),
|
||
validation_status='pending'
|
||
)
|
||
|
||
calorie_infos.append(calorie_info)
|
||
|
||
# 如果没有找到热量信息，尝试查找食物名称
|
||
if not calorie_infos:
|
||
food_names = self._extract_all_food_names(text)
|
||
for food_name in food_names:
|
||
calorie_info = CalorieInfo(
|
||
food_name=food_name,
|
||
calories=None,
|
||
serving_size=None,
|
||
confidence=0.6,
|
||
source='ocr',
|
||
raw_text=food_name,
|
||
validation_status='pending'
|
||
)
|
||
calorie_infos.append(calorie_info)
|
||
|
||
return calorie_infos
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"热量信息提取失败: {e}")
|
||
return []
|
||
|
||
def _extract_food_name(self, text: str, calorie_position: int) -> str:
|
||
"""提取食物名称"""
|
||
try:
|
||
# 在热量数值前后查找食物名称
|
||
context_start = max(0, calorie_position - 50)
|
||
context_end = min(len(text), calorie_position + 50)
|
||
context = text[context_start:context_end]
|
||
|
||
# 查找中文和英文食物名称
|
||
food_pattern = r'([a-zA-Z\u4e00-\u9fff]{2,20})'
|
||
matches = re.findall(food_pattern, context)
|
||
|
||
if matches:
|
||
# 选择最可能的食物名称
|
||
return matches[0]
|
||
|
||
return "未知食物"
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"食物名称提取失败: {e}")
|
||
return "未知食物"
|
||
|
||
def _extract_all_food_names(self, text: str) -> List[str]:
|
||
"""提取所有可能的食物名称"""
|
||
try:
|
||
food_pattern = r'([a-zA-Z\u4e00-\u9fff]{2,20})'
|
||
matches = re.findall(food_pattern, text)
|
||
|
||
# 去重并过滤
|
||
unique_foods = list(set(matches))
|
||
return unique_foods[:5] # 最多返回5个
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"食物名称提取失败: {e}")
|
||
return []
|
||
|
||
def _validate_with_database(self, calorie_infos: List[CalorieInfo], user_data: UserData) -> List[CalorieInfo]:
|
||
"""使用数据库验证热量信息"""
|
||
validated_infos = []
|
||
|
||
for info in calorie_infos:
|
||
try:
|
||
# 在食物数据库中查找匹配
|
||
db_match = self._find_database_match(info.food_name)
|
||
|
||
if db_match:
|
||
# 使用数据库信息更新
|
||
info.calories = db_match.get('calories', info.calories)
|
||
info.serving_size = db_match.get('serving_size', info.serving_size)
|
||
info.confidence = max(info.confidence, 0.9)
|
||
info.source = 'database'
|
||
|
||
# 应用用户学习数据
|
||
user_correction = self._get_user_correction(user_data.user_id, info.food_name)
|
||
if user_correction:
|
||
info.calories = user_correction.get('calories', info.calories)
|
||
info.confidence = max(info.confidence, 0.95)
|
||
info.source = 'user_confirmed'
|
||
|
||
validated_infos.append(info)
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"数据库验证失败: {e}")
|
||
validated_infos.append(info)
|
||
|
||
return validated_infos
|
||
|
||
def _find_database_match(self, food_name: str) -> Optional[Dict[str, Any]]:
|
||
"""在数据库中查找食物匹配"""
|
||
try:
|
||
# 精确匹配
|
||
if food_name in self.food_database:
|
||
return self.food_database[food_name]
|
||
|
||
# 模糊匹配
|
||
for db_food, info in self.food_database.items():
|
||
if food_name in db_food or db_food in food_name:
|
||
return info
|
||
|
||
return None
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"数据库匹配失败: {e}")
|
||
return None
|
||
|
||
def _get_user_correction(self, user_id: str, food_name: str) -> Optional[Dict[str, Any]]:
|
||
"""获取用户修正数据"""
|
||
try:
|
||
user_data = self.user_corrections.get(user_id, {})
|
||
return user_data.get(food_name)
|
||
except Exception as e:
|
||
self.logger.error(f"获取用户修正数据失败: {e}")
|
||
return None
|
||
|
||
def _generate_suggestions(self, calorie_infos: List[CalorieInfo], user_data: UserData) -> List[str]:
|
||
"""生成建议"""
|
||
suggestions = []
|
||
|
||
try:
|
||
for info in calorie_infos:
|
||
if info.calories is None:
|
||
suggestions.append(f"未识别到 {info.food_name} 的热量信息，请手动输入")
|
||
elif info.confidence < 0.8:
|
||
suggestions.append(f"{info.food_name} 的热量 {info.calories} 可能不准确，请确认")
|
||
else:
|
||
suggestions.append(f"{info.food_name}: {info.calories} 卡路里")
|
||
|
||
if not calorie_infos:
|
||
suggestions.append("未识别到任何食物信息，请检查图片质量或手动输入")
|
||
|
||
return suggestions
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"生成建议失败: {e}")
|
||
return ["识别过程中出现错误"]
|
||
|
||
def _calculate_overall_confidence(self, ocr_results: List[OCRResult], calorie_infos: List[CalorieInfo]) -> float:
|
||
"""计算整体置信度"""
|
||
try:
|
||
if not ocr_results and not calorie_infos:
|
||
return 0.0
|
||
|
||
# OCR置信度
|
||
ocr_confidence = sum(r.confidence for r in ocr_results) / len(ocr_results) if ocr_results else 0.0
|
||
|
||
# 热量信息置信度
|
||
calorie_confidence = sum(info.confidence for info in calorie_infos) / len(calorie_infos) if calorie_infos else 0.0
|
||
|
||
# 综合置信度
|
||
overall_confidence = (ocr_confidence * 0.4 + calorie_confidence * 0.6)
|
||
|
||
return min(overall_confidence, 1.0)
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"计算置信度失败: {e}")
|
||
return 0.0
|
||
|
||
def _validate_recognition_result(self, input_data: Dict[str, Any], user_data: UserData) -> Dict[str, Any]:
|
||
"""验证识别结果"""
|
||
try:
|
||
food_name = input_data.get('food_name')
|
||
calories = input_data.get('calories')
|
||
is_correct = input_data.get('is_correct', True)
|
||
|
||
if not is_correct:
|
||
# 用户修正
|
||
corrected_calories = input_data.get('corrected_calories')
|
||
self._save_user_correction(user_data.user_id, food_name, corrected_calories)
|
||
|
||
return {
|
||
'success': True,
|
||
'message': '验证结果已保存',
|
||
'confidence': 1.0
|
||
}
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"验证识别结果失败: {e}")
|
||
return self._create_error_result(f"验证失败: {str(e)}")
|
||
|
||
def _learn_from_correction(self, input_data: Dict[str, Any], user_data: UserData) -> Dict[str, Any]:
|
||
"""从用户修正中学习"""
|
||
try:
|
||
food_name = input_data.get('food_name')
|
||
corrected_calories = input_data.get('corrected_calories')
|
||
|
||
self._save_user_correction(user_data.user_id, food_name, corrected_calories)
|
||
|
||
return {
|
||
'success': True,
|
||
'message': '学习数据已保存',
|
||
'confidence': 1.0
|
||
}
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"学习修正数据失败: {e}")
|
||
return self._create_error_result(f"学习失败: {str(e)}")
|
||
|
||
def _save_user_correction(self, user_id: str, food_name: str, calories: float):
|
||
"""保存用户修正数据"""
|
||
try:
|
||
if user_id not in self.user_corrections:
|
||
self.user_corrections[user_id] = {}
|
||
|
||
self.user_corrections[user_id][food_name] = {
|
||
'calories': calories,
|
||
'timestamp': datetime.now().isoformat(),
|
||
'correction_count': self.user_corrections[user_id].get(food_name, {}).get('correction_count', 0) + 1
|
||
}
|
||
|
||
# 保存到文件
|
||
self._save_user_corrections_to_file()
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"保存用户修正数据失败: {e}")
|
||
|
||
def _load_food_database(self) -> Dict[str, Dict[str, Any]]:
|
||
"""加载食物数据库"""
|
||
try:
|
||
# 基础食物数据库
|
||
food_db = {
|
||
"米饭": {"calories": 130, "serving_size": "100g"},
|
||
"面条": {"calories": 110, "serving_size": "100g"},
|
||
"馒头": {"calories": 221, "serving_size": "100g"},
|
||
"包子": {"calories": 250, "serving_size": "100g"},
|
||
"饺子": {"calories": 250, "serving_size": "100g"},
|
||
"鸡蛋": {"calories": 155, "serving_size": "100g"},
|
||
"豆腐": {"calories": 76, "serving_size": "100g"},
|
||
"鱼肉": {"calories": 206, "serving_size": "100g"},
|
||
"鸡肉": {"calories": 165, "serving_size": "100g"},
|
||
"瘦肉": {"calories": 250, "serving_size": "100g"},
|
||
"青菜": {"calories": 25, "serving_size": "100g"},
|
||
"西红柿": {"calories": 18, "serving_size": "100g"},
|
||
"胡萝卜": {"calories": 41, "serving_size": "100g"},
|
||
"土豆": {"calories": 77, "serving_size": "100g"},
|
||
"西兰花": {"calories": 34, "serving_size": "100g"},
|
||
"苹果": {"calories": 52, "serving_size": "100g"},
|
||
"香蕉": {"calories": 89, "serving_size": "100g"},
|
||
"橙子": {"calories": 47, "serving_size": "100g"},
|
||
"葡萄": {"calories": 67, "serving_size": "100g"},
|
||
"草莓": {"calories": 32, "serving_size": "100g"},
|
||
"牛奶": {"calories": 42, "serving_size": "100ml"},
|
||
"酸奶": {"calories": 59, "serving_size": "100g"},
|
||
"豆浆": {"calories": 31, "serving_size": "100ml"},
|
||
"坚果": {"calories": 607, "serving_size": "100g"},
|
||
"红枣": {"calories": 264, "serving_size": "100g"},
|
||
}
|
||
|
||
# 尝试从文件加载扩展数据库
|
||
db_file = Path("data/food_database.json")
|
||
if db_file.exists():
|
||
with open(db_file, 'r', encoding='utf-8') as f:
|
||
extended_db = json.load(f)
|
||
food_db.update(extended_db)
|
||
|
||
return food_db
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"加载食物数据库失败: {e}")
|
||
return {}
|
||
|
||
def _create_directories(self):
|
||
"""创建必要的目录"""
|
||
directories = [
|
||
'data/ocr_cache',
|
||
'data/user_corrections',
|
||
'data/food_images'
|
||
]
|
||
|
||
for directory in directories:
|
||
Path(directory).mkdir(parents=True, exist_ok=True)
|
||
|
||
def _load_user_corrections(self):
|
||
"""加载用户修正数据"""
|
||
try:
|
||
corrections_file = Path("data/user_corrections.json")
|
||
if corrections_file.exists():
|
||
with open(corrections_file, 'r', encoding='utf-8') as f:
|
||
self.user_corrections = json.load(f)
|
||
else:
|
||
self.user_corrections = {}
|
||
except Exception as e:
|
||
self.logger.error(f"加载用户修正数据失败: {e}")
|
||
self.user_corrections = {}
|
||
|
||
def _save_user_corrections_to_file(self):
|
||
"""保存用户修正数据到文件"""
|
||
try:
|
||
corrections_file = Path("data/user_corrections.json")
|
||
with open(corrections_file, 'w', encoding='utf-8') as f:
|
||
json.dump(self.user_corrections, f, ensure_ascii=False, indent=2)
|
||
except Exception as e:
|
||
self.logger.error(f"保存用户修正数据失败: {e}")
|
||
|
||
def _initialize_ocr_engines(self):
|
||
"""初始化OCR引擎"""
|
||
try:
|
||
# 检查Tesseract是否可用
|
||
try:
|
||
pytesseract.get_tesseract_version()
|
||
self.logger.info("Tesseract OCR引擎可用")
|
||
except Exception:
|
||
self.logger.warning("Tesseract OCR引擎不可用")
|
||
|
||
# 其他OCR引擎将在需要时初始化
|
||
self.logger.info("OCR引擎初始化完成")
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"OCR引擎初始化失败: {e}")
|
||
|
||
def _create_error_result(self, message: str) -> Dict[str, Any]:
|
||
"""创建错误结果"""
|
||
return {
|
||
'success': False,
|
||
'error': message,
|
||
'confidence': 0.0
|
||
}
|
||
|
||
def cleanup(self) -> bool:
|
||
"""清理资源"""
|
||
try:
|
||
# 保存用户修正数据
|
||
self._save_user_corrections_to_file()
|
||
|
||
# 清理OCR引擎
|
||
self.ocr_engines.clear()
|
||
|
||
self.is_initialized = False
|
||
self.logger.info("OCR热量识别模块清理完成")
|
||
return True
|
||
except Exception as e:
|
||
self.logger.error(f"OCR热量识别模块清理失败: {e}")
|
||
return False
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 测试OCR模块
|
||
from core.base import BaseConfig
|
||
|
||
config = BaseConfig()
|
||
ocr_module = OCRCalorieRecognitionModule(config)
|
||
|
||
if ocr_module.initialize():
|
||
print("OCR模块初始化成功")
|
||
|
||
# 测试图片识别
|
||
test_data = {
|
||
'type': 'recognize_image',
|
||
'image_path': 'test_image.jpg' # 需要提供测试图片
|
||
}
|
||
|
||
# 这里需要用户数据，暂时跳过实际测试
|
||
print("OCR模块测试完成")
|
||
else:
|
||
print("OCR模块初始化失败")
|