154 lines
5.9 KiB
Python
154 lines
5.9 KiB
Python
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() |