262 lines
10 KiB
Python
262 lines
10 KiB
Python
"""未知菜品待学习：持久化存储与管理。"""
|
||
import os
|
||
import json
|
||
import shutil
|
||
import threading
|
||
import datetime
|
||
import numpy as np
|
||
|
||
|
||
# 新未知图与库内已有 pending 特征相似度 ≥ 此值则不入待学习库
|
||
PENDING_DEDUP_SIMILARITY = 0.90
|
||
|
||
|
||
class PendingLearningStore:
|
||
STATUS_PENDING = "pending"
|
||
STATUS_LEARNED = "learned"
|
||
STATUS_IGNORED = "ignored"
|
||
|
||
def __init__(self, base_dir):
|
||
self.base_dir = base_dir
|
||
self.pending_dir = os.path.join(base_dir, "pending_images")
|
||
self.data_file = os.path.join(base_dir, "pending_learning.json")
|
||
self.lock = threading.Lock()
|
||
self._items = []
|
||
self._next_id = 1
|
||
os.makedirs(self.pending_dir, exist_ok=True)
|
||
self._load()
|
||
|
||
def _load(self):
|
||
if not os.path.exists(self.data_file):
|
||
return
|
||
try:
|
||
with open(self.data_file, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
self._items = data.get("items", [])
|
||
self._next_id = data.get("next_id", 1)
|
||
for item in self._items:
|
||
if item.get("id") is None:
|
||
continue
|
||
self._next_id = max(self._next_id, item["id"] + 1)
|
||
except Exception as e:
|
||
print(f"加载待学习库失败: {e}")
|
||
self._items = []
|
||
self._next_id = 1
|
||
|
||
def _save(self):
|
||
try:
|
||
with open(self.data_file, "w", encoding="utf-8") as f:
|
||
json.dump({"next_id": self._next_id, "items": self._items}, f, ensure_ascii=False, indent=2)
|
||
except Exception as e:
|
||
print(f"保存待学习库失败: {e}")
|
||
|
||
def _find_pending_by_path(self, image_path):
|
||
norm = os.path.normcase(os.path.abspath(image_path))
|
||
for item in self._items:
|
||
if item.get("status") != self.STATUS_PENDING:
|
||
continue
|
||
stored = item.get("source_path") or item.get("image_path")
|
||
if stored and os.path.normcase(os.path.abspath(stored)) == norm:
|
||
return item
|
||
return None
|
||
|
||
def _is_similar_pending(self, feature, threshold=PENDING_DEDUP_SIMILARITY):
|
||
"""若与任一 pending 特征相似度 ≥ threshold，返回 (item, sim)；否则 (None, 0)。"""
|
||
feat = np.asarray(feature, dtype=np.float32).reshape(-1)
|
||
best_item = None
|
||
best_sim = 0.0
|
||
for item in self._items:
|
||
if item.get("status") != self.STATUS_PENDING:
|
||
continue
|
||
stored = item.get("feature")
|
||
if stored is None:
|
||
continue
|
||
other = np.asarray(stored, dtype=np.float32).reshape(-1)
|
||
if other.shape != feat.shape:
|
||
continue
|
||
sim = float(np.dot(feat, other))
|
||
if sim >= threshold and sim >= best_sim:
|
||
best_sim = sim
|
||
best_item = item
|
||
return best_item, best_sim
|
||
|
||
def add(self, image_path, feature, suggested_category=None, score=None, reason=None):
|
||
if not image_path or not os.path.exists(image_path):
|
||
return None
|
||
feature = np.asarray(feature, dtype=np.float32)
|
||
new_sug = (suggested_category or "").strip()
|
||
|
||
with self.lock:
|
||
existing = self._find_pending_by_path(image_path)
|
||
if existing:
|
||
old_sug = (existing.get("suggested_category") or "").strip()
|
||
# 同一路径但建议菜名已变：视为新的未知菜，不覆盖旧记录
|
||
if old_sug and new_sug and old_sug != new_sug:
|
||
existing = None
|
||
else:
|
||
existing["suggested_category"] = suggested_category or existing.get("suggested_category")
|
||
existing["score"] = score if score is not None else existing.get("score")
|
||
existing["feature"] = feature.tolist()
|
||
existing["capture_time"] = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
if reason:
|
||
existing["reason"] = reason
|
||
# 刷新图片副本，避免源文件被覆盖后内容过期
|
||
try:
|
||
dest = existing.get("image_path")
|
||
if dest and dest.startswith(self.pending_dir):
|
||
shutil.copy2(image_path, dest)
|
||
except Exception:
|
||
pass
|
||
self._save()
|
||
return existing["id"]
|
||
|
||
# 与库内已有图相似度 ≥ 0.90 → 不加入待学习库
|
||
similar, sim = self._is_similar_pending(feature, PENDING_DEDUP_SIMILARITY)
|
||
if similar is not None:
|
||
print(
|
||
f"[pending] skip add (sim={sim:.3f}≥{PENDING_DEDUP_SIMILARITY:.2f} "
|
||
f"vs #{similar.get('id')})"
|
||
)
|
||
return None
|
||
|
||
pending_id = self._next_id
|
||
self._next_id += 1
|
||
stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
ext = os.path.splitext(image_path)[1] or ".jpeg"
|
||
dest_name = f"pending_{stamp}_{pending_id}{ext}"
|
||
dest_path = os.path.join(self.pending_dir, dest_name)
|
||
try:
|
||
shutil.copy2(image_path, dest_path)
|
||
except Exception as e:
|
||
print(f"复制待学习图片失败: {e}")
|
||
dest_path = image_path
|
||
|
||
item = {
|
||
"id": pending_id,
|
||
"image_path": dest_path,
|
||
"source_path": image_path,
|
||
"feature": feature.tolist(),
|
||
"suggested_category": suggested_category,
|
||
"score": float(score) if score is not None else None,
|
||
"status": self.STATUS_PENDING,
|
||
"capture_time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||
"learned_category": None,
|
||
"reason": reason,
|
||
}
|
||
self._items.append(item)
|
||
self._save()
|
||
return pending_id
|
||
|
||
def get(self, pending_id):
|
||
for item in self._items:
|
||
if item.get("id") == pending_id:
|
||
return item.copy()
|
||
return None
|
||
|
||
def get_list(self, status=STATUS_PENDING):
|
||
rows = []
|
||
for item in self._items:
|
||
if item.get("status") != status:
|
||
continue
|
||
row = item.copy()
|
||
feat = row.get("feature")
|
||
if feat is not None:
|
||
row["feature"] = np.array(feat, dtype=np.float32)
|
||
rows.append(row)
|
||
rows.sort(key=lambda x: x.get("capture_time", ""), reverse=True)
|
||
return rows
|
||
|
||
def count(self, status=STATUS_PENDING):
|
||
return sum(1 for item in self._items if item.get("status") == status)
|
||
|
||
def mark_learned(self, pending_id, category):
|
||
with self.lock:
|
||
for item in self._items:
|
||
if item.get("id") == pending_id:
|
||
item["status"] = self.STATUS_LEARNED
|
||
item["learned_category"] = category
|
||
item["learn_time"] = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
self._save()
|
||
return True
|
||
return False
|
||
|
||
def update_suggestion(self, pending_id, suggested_category=None, score=None):
|
||
"""刷新仍待处理项的建议菜名 / 相似度（不改 status）。"""
|
||
with self.lock:
|
||
for item in self._items:
|
||
if item.get("id") != pending_id:
|
||
continue
|
||
if item.get("status") != self.STATUS_PENDING:
|
||
return False
|
||
if suggested_category is not None:
|
||
item["suggested_category"] = suggested_category
|
||
if score is not None:
|
||
item["score"] = float(score)
|
||
self._save()
|
||
return True
|
||
return False
|
||
|
||
def _unlink_image(self, item):
|
||
"""删除 pending_images 下的本地图（仅限本库目录，防误删）。"""
|
||
path = item.get("image_path")
|
||
if not path:
|
||
return
|
||
try:
|
||
abs_path = os.path.abspath(path)
|
||
if not abs_path.startswith(os.path.abspath(self.pending_dir)):
|
||
return
|
||
if os.path.exists(abs_path):
|
||
os.remove(abs_path)
|
||
except Exception as e:
|
||
print(f"删除待学习图片失败: {e}")
|
||
|
||
def mark_ignored(self, pending_id, remove_image=True):
|
||
"""忽略一条：删除图片并从库中移除（不再只改 status 留尸）。"""
|
||
with self.lock:
|
||
for idx, item in enumerate(self._items):
|
||
if item.get("id") != pending_id:
|
||
continue
|
||
if remove_image:
|
||
self._unlink_image(item)
|
||
del self._items[idx]
|
||
self._save()
|
||
return True
|
||
return False
|
||
|
||
def ignore_all_pending(self, remove_image=True):
|
||
"""忽略全部 pending：删图 + 从库移除，只写盘一次。返回删除条数。"""
|
||
with self.lock:
|
||
kept = []
|
||
removed = 0
|
||
for item in self._items:
|
||
if item.get("status") == self.STATUS_PENDING:
|
||
if remove_image:
|
||
self._unlink_image(item)
|
||
removed += 1
|
||
else:
|
||
kept.append(item)
|
||
if removed:
|
||
self._items = kept
|
||
self._save()
|
||
return removed
|
||
|
||
def delete(self, pending_id, remove_image=True):
|
||
with self.lock:
|
||
for idx, item in enumerate(self._items):
|
||
if item.get("id") != pending_id:
|
||
continue
|
||
if remove_image:
|
||
self._unlink_image(item)
|
||
del self._items[idx]
|
||
self._save()
|
||
return True
|
||
return False
|
||
|
||
def clear_all(self):
|
||
with self.lock:
|
||
for item in self._items:
|
||
self._unlink_image(item)
|
||
self._items = []
|
||
self._next_id = 1
|
||
self._save()
|