Dish_Recognition/service/local_cache.py

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()