初始化 dish-image:菜品图爬虫 + 图片名称检索平台
- crawler/:豆果(主)/下厨房/Bing 爬虫 + dishes.jsonl(3056 菜→9167 图映射)+ verify_repair 按 URL 重下 - retrieval/:三路检索(BM25 + 本地 BGE-M3 向量 + RRF 融合),FastAPI + 前端;写死图片目录(默认 crawler/images,可 IMAGE_DIR 覆盖)、绑 0.0.0.0 局域网访问、启动自动建索引、服务器 serve 图片 - 图片(1.7G)与向量模型(2.3G)不进 git Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
陈世睿
@@ -0,0 +1,262 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""图片名称检索测试台 —— 后端(FastAPI)。
|
||||
|
||||
三路检索,全部基于「图片名称」:
|
||||
1) BM25 —— 关键词匹配(jieba 分词)
|
||||
2) 向量 —— 本地 BAAI/bge-m3 余弦相似度
|
||||
3) 融合 —— RRF 合并上面两路排名
|
||||
|
||||
【写死目录 + 局域网模式】图片目录写死在 IMAGE_DIR,服务器启动即扫描并自动建索引、
|
||||
自己通过 /images 把图片发出去。绑 0.0.0.0,mentor 用局域网 IP 打开即可直接检索,
|
||||
无需选文件夹、无需建索引。
|
||||
|
||||
跑法:python server.py → 控制台会打印可发给 mentor 的局域网地址。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import sys
|
||||
import threading
|
||||
|
||||
import numpy as np
|
||||
from fastapi import Body, FastAPI, Query
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from bm25 import BM25Index
|
||||
from embed_local import LocalEmbedder
|
||||
from fuse import rrf
|
||||
|
||||
for _s in (sys.stdout, sys.stderr):
|
||||
try:
|
||||
_s.reconfigure(encoding="utf-8", errors="replace")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
# 图片目录:默认 = 本仓库的 crawler/images(爬虫默认就把图下到那里)。
|
||||
# 想指向别处已爬好的图,启动前设环境变量 IMAGE_DIR=<绝对路径> 覆盖即可,无需改代码。
|
||||
IMAGE_DIR = os.environ.get("IMAGE_DIR") or os.path.abspath(
|
||||
os.path.join(HERE, os.pardir, "crawler", "images"))
|
||||
IMG_EXT = (".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp")
|
||||
PORT = 8799
|
||||
|
||||
app = FastAPI(title="图片名称检索测试台")
|
||||
|
||||
|
||||
def scan_images(img_dir: str) -> dict[str, list[str]]:
|
||||
"""扫描图片目录 → {菜名: [文件名,...]}。去掉尾部 _NN 把同名多图归并为一条。"""
|
||||
name2files: dict[str, list[str]] = {}
|
||||
if not os.path.isdir(img_dir):
|
||||
return name2files
|
||||
for fn in os.listdir(img_dir):
|
||||
if not fn.lower().endswith(IMG_EXT):
|
||||
continue
|
||||
stem = os.path.splitext(fn)[0]
|
||||
name = re.sub(r"_\d+$", "", stem) or stem
|
||||
name2files.setdefault(name, []).append(fn)
|
||||
for files in name2files.values():
|
||||
files.sort()
|
||||
return name2files
|
||||
|
||||
|
||||
class State:
|
||||
"""全局索引状态(单实例工具,够用)。"""
|
||||
|
||||
def __init__(self):
|
||||
self.lock = threading.Lock()
|
||||
self.names: list[str] = []
|
||||
self.name2files: dict[str, list[str]] = {}
|
||||
self.bm25: BM25Index | None = None
|
||||
self.embedder: LocalEmbedder | None = None
|
||||
self.model = "BAAI/bge-m3"
|
||||
self.vec_names: list[str] = []
|
||||
self.vec_matrix: np.ndarray | None = None
|
||||
self.state = "idle" # idle | building | ready | error
|
||||
self.phase = "" # bm25 | loadmodel | embedding | done
|
||||
self.emb_done = 0
|
||||
self.emb_total = 0
|
||||
self.failed = 0
|
||||
self.message = ""
|
||||
self.device = ""
|
||||
|
||||
|
||||
ST = State()
|
||||
|
||||
|
||||
def build_index(names: list[str], model: str) -> None:
|
||||
try:
|
||||
with ST.lock:
|
||||
ST.state, ST.phase, ST.message = "building", "bm25", ""
|
||||
ST.names, ST.model = names, model
|
||||
ST.bm25 = None
|
||||
ST.vec_matrix, ST.vec_names = None, []
|
||||
ST.emb_done = ST.emb_total = ST.failed = 0
|
||||
|
||||
# 1) BM25(秒级)
|
||||
bm = BM25Index(names)
|
||||
with ST.lock:
|
||||
ST.bm25, ST.phase = bm, "loadmodel"
|
||||
|
||||
# 2) 本地 BGE-M3:先加载模型(首次会下载约 2.3G),再编码(有 CUDA 走 GPU)
|
||||
emb = LocalEmbedder(model=model, cache_dir=os.path.join(HERE, "cache"))
|
||||
dev = emb.load_model()
|
||||
with ST.lock:
|
||||
ST.device, ST.phase = dev, "embedding"
|
||||
|
||||
def prog(done: int, total: int) -> None:
|
||||
with ST.lock:
|
||||
ST.emb_done, ST.emb_total = done, total
|
||||
|
||||
emb.embed_corpus(names, progress=prog)
|
||||
vnames, mat = emb.matrix(names)
|
||||
with ST.lock:
|
||||
ST.embedder, ST.vec_names, ST.vec_matrix = emb, vnames, mat
|
||||
ST.failed = emb.last_failed
|
||||
ST.message = "" if vnames else ("向量未生成:" + emb.last_error if emb.last_error else "")
|
||||
ST.state, ST.phase = "ready", "done"
|
||||
except Exception as e: # noqa: BLE001
|
||||
with ST.lock:
|
||||
ST.state, ST.message = "error", f"{type(e).__name__}: {e}"
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def _startup_autobuild() -> None:
|
||||
"""启动即扫描写死目录并后台建索引,mentor 打开就能搜。"""
|
||||
n2f = scan_images(IMAGE_DIR)
|
||||
with ST.lock:
|
||||
ST.name2files = n2f
|
||||
names = sorted(n2f.keys())
|
||||
if not names:
|
||||
with ST.lock:
|
||||
ST.state = "error"
|
||||
ST.message = f"图片目录为空或不存在:{IMAGE_DIR}"
|
||||
return
|
||||
threading.Thread(target=build_index, args=(names, "BAAI/bge-m3"), daemon=True).start()
|
||||
|
||||
|
||||
@app.post("/api/index")
|
||||
def api_index(payload: dict = Body(...)):
|
||||
"""(保留)手动指定名称重建索引;写死目录模式下一般用不到。"""
|
||||
names = [str(n).strip() for n in payload.get("names", []) if str(n).strip()]
|
||||
names = list(dict.fromkeys(names))
|
||||
model = (payload.get("model") or "BAAI/bge-m3").strip()
|
||||
if not names:
|
||||
return JSONResponse({"ok": False, "error": "没有菜名"}, status_code=400)
|
||||
if ST.state == "building":
|
||||
return JSONResponse({"ok": False, "error": "正在建索引,请稍候"}, status_code=409)
|
||||
threading.Thread(target=build_index, args=(names, model), daemon=True).start()
|
||||
return {"ok": True, "total": len(names), "model": model}
|
||||
|
||||
|
||||
@app.get("/api/status")
|
||||
def api_status():
|
||||
with ST.lock:
|
||||
return {
|
||||
"state": ST.state,
|
||||
"phase": ST.phase,
|
||||
"bm25_ready": ST.bm25 is not None,
|
||||
"vec_ready": ST.vec_matrix is not None and len(ST.vec_names) > 0,
|
||||
"total": len(ST.names),
|
||||
"n_images": sum(len(v) for v in ST.name2files.values()),
|
||||
"emb_done": ST.emb_done,
|
||||
"emb_total": ST.emb_total,
|
||||
"vec_count": len(ST.vec_names),
|
||||
"failed": ST.failed,
|
||||
"model": ST.model,
|
||||
"device": ST.device,
|
||||
"message": ST.message,
|
||||
}
|
||||
|
||||
|
||||
def _bm25_search(q: str, topk: int) -> list[tuple[str, float]]:
|
||||
return ST.bm25.search(q, topk=topk) if ST.bm25 else []
|
||||
|
||||
|
||||
def _vec_search(q: str, topk: int) -> list[tuple[str, float]]:
|
||||
mat, names, emb = ST.vec_matrix, ST.vec_names, ST.embedder
|
||||
if mat is None or len(names) == 0 or emb is None:
|
||||
return []
|
||||
qv = emb.embed_query(q)
|
||||
sims = mat @ qv
|
||||
order = np.argsort(sims)[::-1][:topk]
|
||||
return [(names[i], float(sims[i])) for i in order]
|
||||
|
||||
|
||||
@app.get("/api/search")
|
||||
def api_search(q: str = Query(...), topk: int = 30, rrf_k: int = 60,
|
||||
w_bm25: float = 1.0, w_vec: float = 1.0):
|
||||
q = q.strip()
|
||||
if not q:
|
||||
return {"q": q, "bm25": [], "vector": [], "fusion": []}
|
||||
|
||||
bm = _bm25_search(q, topk)
|
||||
ve = _vec_search(q, topk)
|
||||
bm_rank = [n for n, _ in bm]
|
||||
ve_rank = [n for n, _ in ve]
|
||||
fused = rrf({"bm25": bm_rank, "vector": ve_rank},
|
||||
{"bm25": w_bm25, "vector": w_vec}, k=rrf_k)[:topk]
|
||||
|
||||
bm_s, ve_s = dict(bm), dict(ve)
|
||||
bm_pos = {n: i + 1 for i, n in enumerate(bm_rank)}
|
||||
ve_pos = {n: i + 1 for i, n in enumerate(ve_rank)}
|
||||
|
||||
def card(name: str, score: float) -> dict:
|
||||
return {
|
||||
"name": name,
|
||||
"score": round(score, 4),
|
||||
"files": ST.name2files.get(name, []), # 服务器 serve 的图片文件名
|
||||
"bm25": round(bm_s[name], 3) if name in bm_s else None,
|
||||
"bm25_rank": bm_pos.get(name),
|
||||
"vec": round(ve_s[name], 4) if name in ve_s else None,
|
||||
"vec_rank": ve_pos.get(name),
|
||||
}
|
||||
|
||||
return {
|
||||
"q": q,
|
||||
"bm25": [card(n, s) for n, s in bm],
|
||||
"vector": [card(n, s) for n, s in ve],
|
||||
"fusion": [card(n, s) for n, s in fused],
|
||||
}
|
||||
|
||||
|
||||
# 图片由服务器 serve(局域网里 mentor 的浏览器据此显示);必须在 "/" 之前挂载。
|
||||
if os.path.isdir(IMAGE_DIR):
|
||||
app.mount("/images", StaticFiles(directory=IMAGE_DIR), name="images")
|
||||
# 静态前端(放在所有 /api 与 /images 之后挂载)
|
||||
app.mount("/", StaticFiles(directory=os.path.join(HERE, "static"), html=True), name="static")
|
||||
|
||||
|
||||
def _lan_ips() -> list[str]:
|
||||
ips: list[str] = []
|
||||
try:
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
s.connect(("8.8.8.8", 80))
|
||||
ips.append(s.getsockname()[0])
|
||||
s.close()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
for info in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET):
|
||||
ip = info[4][0]
|
||||
if ip not in ips:
|
||||
ips.append(ip)
|
||||
except Exception:
|
||||
pass
|
||||
return [i for i in ips if not i.startswith("127.")]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
print("=" * 56, flush=True)
|
||||
print(f"图片目录:{IMAGE_DIR}", flush=True)
|
||||
print("发给 mentor 的局域网地址:", flush=True)
|
||||
for ip in _lan_ips() or ["<查不到,用 ipconfig 看本机 IPv4>"]:
|
||||
print(f" http://{ip}:{PORT}", flush=True)
|
||||
print(f"本机自测:http://127.0.0.1:{PORT}", flush=True)
|
||||
print("(若 mentor 连不上:Windows 防火墙放行该端口入站)", flush=True)
|
||||
print("=" * 56, flush=True)
|
||||
uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info")
|
||||
Reference in New Issue
Block a user