64 lines
2.1 KiB
Python
64 lines
2.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
|
|
from exp_multimodal.ollama_client import OLLAMA_URL, DEFAULT_MODEL
|
|
|
|
|
|
FEWSHOT_HINTS: Dict[str, str] = {
|
|
# 可选：仅对菜品模式提供少量文字提示，帮助区分相似菜
|
|
# "西红柿鸡蛋": "红黄对比明显，蛋块与番茄块同炒，汤汁偏红",
|
|
# "麻婆豆腐": "豆腐块+红油辣椒+花椒，肉末点缀",
|
|
}
|
|
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--mode", choices=["dish", "whole", "processed"], default="dish")
|
|
ap.add_argument("--image", default=r"D:\MyProjects\PythonProjects\FoodClassifier\dataset\DishClassification\test\红烧肉\img04.png")
|
|
ap.add_argument("--alias_map", default=None)
|
|
|
|
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'}")
|
|
|
|
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)}")
|
|
|
|
res = classify_image(
|
|
args.image,
|
|
labels,
|
|
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()
|
|
|
|
|