import sqlite3 import numpy as np import threading from datetime import datetime class LocalFeatureCache: def __init__(self, db_path="dish_cache.db"): self.conn = sqlite3.connect(db_path, check_same_thread=False) self.conn.row_factory = sqlite3.Row self.lock = threading.Lock() self._init_tables() self.threshold = 0.7 def _init_tables(self): # 模板缓存表(与MySQL结构对应) self.conn.execute(''' CREATE TABLE IF NOT EXISTS template_cache ( id INTEGER PRIMARY KEY, category TEXT NOT NULL, feature_blob BLOB NOT NULL, source_base TEXT, server_version DATETIME, is_active INTEGER DEFAULT 1 ) ''') # 待学习表 self.conn.execute(''' CREATE TABLE IF NOT EXISTS pending_learning ( id INTEGER PRIMARY KEY AUTOINCREMENT, image_path TEXT NOT NULL, feature_blob BLOB NOT NULL, capture_time DATETIME DEFAULT CURRENT_TIMESTAMP, status TEXT DEFAULT 'pending', category TEXT ) ''') # 上传队列表 self.conn.execute(''' CREATE TABLE IF NOT EXISTS upload_queue ( id INTEGER PRIMARY KEY AUTOINCREMENT, category TEXT NOT NULL, feature_blob BLOB NOT NULL, image_path TEXT, create_time DATETIME DEFAULT CURRENT_TIMESTAMP, status TEXT DEFAULT 'pending', retry_count INTEGER DEFAULT 0 ) ''') # 同步状态表 self.conn.execute(''' CREATE TABLE IF NOT EXISTS sync_status ( key TEXT PRIMARY KEY, value TEXT ) ''') self.conn.execute("INSERT OR IGNORE INTO sync_status (key, value) VALUES ('last_sync_time', '1970-01-01 00:00:00')") self.conn.commit() # ---------- 核心识别方法 ---------- def predict(self, feature_vec): """输入特征向量,返回(类别, 相似度)或(None, 最高相似度)""" cursor = self.conn.execute("SELECT category, feature_blob FROM template_cache WHERE is_active=1") max_sim = -1.0 best_cat = None for row in cursor: db_feat = np.frombuffer(row[1], dtype=np.float32) sim = np.dot(feature_vec, db_feat) if sim > max_sim: max_sim = sim best_cat = row[0] if max_sim >= self.threshold: return best_cat, max_sim else: return None, max_sim # ---------- 添加模板(本地立即生效,同时入队待上传) ---------- def add_template(self, category, feature_vec, image_path=None, source_base=""): blob = feature_vec.tobytes() with self.lock: # 插入缓存表 cur = self.conn.execute( "INSERT INTO template_cache (category, feature_blob, source_base) VALUES (?, ?, ?)", (category, blob, source_base) ) # 插入上传队列 self.conn.execute( "INSERT INTO upload_queue (category, feature_blob, image_path) VALUES (?, ?, ?)", (category, blob, image_path) ) self.conn.commit() return cur.lastrowid # ---------- 待学习管理 ---------- def add_pending(self, image_path, feature_vec, category=None): blob = feature_vec.tobytes() self.conn.execute( "INSERT INTO pending_learning (image_path, feature_blob, category) VALUES (?, ?, ?)", (image_path, blob, category) ) self.conn.commit() def get_pending_list(self, status='pending'): cursor = self.conn.execute( "SELECT id, image_path, feature_blob, category, capture_time FROM pending_learning WHERE status=?", (status,) ) rows = cursor.fetchall() result = [] for row in rows: result.append({ 'id': row[0], 'image_path': row[1], 'feature': np.frombuffer(row[2], dtype=np.float32), 'category': row[3], 'capture_time': row[4] }) return result def mark_pending_learned(self, pending_id, category): # 学习后,将待学习记录标记为已学习,并同时加入模板库(但需后续流程) self.conn.execute( "UPDATE pending_learning SET status='learned', category=? WHERE id=?", (category, pending_id) ) self.conn.commit() # ---------- 同步状态 ---------- def get_last_sync_time(self): row = self.conn.execute("SELECT value FROM sync_status WHERE key='last_sync_time'").fetchone() return row[0] if row else '1970-01-01 00:00:00' def update_last_sync_time(self, time_str): self.conn.execute("UPDATE sync_status SET value=? WHERE key='last_sync_time'", (time_str,)) self.conn.commit() # ---------- 获取所有待上传模板 ---------- def get_pending_uploads(self): cursor = self.conn.execute( "SELECT id, category, feature_blob, image_path FROM upload_queue WHERE status='pending'" ) rows = cursor.fetchall() return [{'id': r[0], 'category': r[1], 'feature': np.frombuffer(r[2], dtype=np.float32), 'image_path': r[3]} for r in rows] def mark_uploaded(self, upload_id): self.conn.execute("UPDATE upload_queue SET status='uploaded' WHERE id=?", (upload_id,)) self.conn.commit() # ---------- 清空/重置(调试用) ---------- def clear_all(self): with self.lock: self.conn.execute("DELETE FROM template_cache") self.conn.execute("DELETE FROM upload_queue") self.conn.execute("DELETE FROM pending_learning") self.conn.commit()