107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from typing import Dict
|
|
|
|
# 兼容：支持直接运行脚本或用 -m 模块方式运行
|
|
if __name__ == "__main__" and __package__ is None:
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
# 统一使用绝对导入，避免相对导入在脚本直跑时失败
|
|
from exp_multimodal.labels import build_labels
|
|
from exp_multimodal.vlm_classifier import classify_image, classify_image_openset
|
|
from exp_multimodal.ollama_client import OLLAMA_URL, DEFAULT_MODEL
|
|
from exp_multimodal.text_embedder import OllamaEmbedder
|
|
from exp_multimodal.vector_matcher import DishNameMatcher
|
|
from exp_multimodal.vlm_providers.ollama_provider import OllamaVLMProvider
|
|
|
|
|
|
FEWSHOT_HINTS: Dict[str, str] = {
|
|
# 可选：仅对菜品模式提供少量文字提示，帮助区分相似菜
|
|
# "西红柿鸡蛋": "红黄对比明显，蛋块与番茄块同炒，汤汁偏红",
|
|
# "麻婆豆腐": "豆腐块+红油辣椒+花椒，肉末点缀",
|
|
}
|
|
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--mode", choices=["dish", "whole", "processed", "openset_dish"], default="openset_dish")
|
|
ap.add_argument("--image", default=r"D:\MyProjects\PythonProjects\FoodClassifier\dataset\DishClassification\test\红烧肉\img04.png")
|
|
ap.add_argument("--alias_map", default=None)
|
|
|
|
# 开放式识别专用参数
|
|
ap.add_argument("--openset-top-k", type=int, default=3, help="开放式识别：向量检索Top-K候选数")
|
|
ap.add_argument("--openset-min-score", type=float, default=0.5, help="开放式识别：最低匹配分数阈值")
|
|
ap.add_argument("--openset-index-path", default="exp_multimodal/dish_name_index", help="开放式识别：向量索引路径")
|
|
|
|
args = ap.parse_args()
|
|
|
|
print(
|
|
f"[Main] mode={args.mode} image={args.image} alias_map={args.alias_map} "
|
|
f"env_OLLAMA_URL={os.getenv('OLLAMA_URL')} env_VLM_MODEL={os.getenv('VLM_MODEL')} "
|
|
f"defaults url={OLLAMA_URL} model={DEFAULT_MODEL}"
|
|
)
|
|
print(f"[Main] Image exists={os.path.exists(args.image)} size={os.path.getsize(args.image) if os.path.exists(args.image) else 'N/A'}")
|
|
|
|
# 开放式识别分支
|
|
if args.mode == "openset_dish":
|
|
print(f"[Main] Openset mode: top_k={args.openset_top_k} min_score={args.openset_min_score} index_path={args.openset_index_path}")
|
|
|
|
# 初始化Provider、Embedder和Matcher
|
|
provider = OllamaVLMProvider(base_url=OLLAMA_URL, model_name=DEFAULT_MODEL)
|
|
embedder = OllamaEmbedder(base_url=OLLAMA_URL, model_name="bge-large-zh-v1.5")
|
|
matcher = DishNameMatcher(index_dir=args.openset_index_path)
|
|
|
|
print("[Main] Loading openset components...")
|
|
matcher.load()
|
|
print(f"[Main] Loaded {len(matcher.dish_names)} dish names from index")
|
|
|
|
res = classify_image_openset(
|
|
image_path=args.image,
|
|
provider=provider,
|
|
embedder=embedder,
|
|
matcher=matcher,
|
|
top_k=args.openset_top_k,
|
|
min_match_score=args.openset_min_score,
|
|
)
|
|
|
|
print(
|
|
json.dumps(
|
|
{"mode": args.mode, "result": res},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
)
|
|
|
|
# 封闭式识别分支（原有逻辑）
|
|
else:
|
|
labels = build_labels(args.mode, args.alias_map)
|
|
ingredient_only = args.mode in {"whole", "processed"}
|
|
fewshot = FEWSHOT_HINTS if args.mode == "dish" else None
|
|
print(f"[Main] Built labels count={len(labels)} ingredient_only={ingredient_only} fewshot={bool(fewshot)}")
|
|
|
|
provider = OllamaVLMProvider(base_url=OLLAMA_URL, model_name=DEFAULT_MODEL)
|
|
|
|
res = classify_image(
|
|
args.image,
|
|
labels,
|
|
provider=provider,
|
|
fewshot_hints=fewshot,
|
|
ingredient_only=ingredient_only,
|
|
)
|
|
|
|
print(
|
|
json.dumps(
|
|
{"mode": args.mode, "result": res, "num_labels": len(labels)},
|
|
ensure_ascii=False,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
|
|