增加多模态模型识别功能。

2025-11-03 09:29:57 +08:00
parent 9d8eccde84
commit 9e12a5fbc3
7 changed files with 344 additions and 0 deletions
@@ -0,0 +1,64 @@
+import json
+from typing import Dict, List, Optional
+
+from .ollama_client import chat_vision
+from .prompts import build_closedset_prompt
+
+
+def classify_image(
+    image_path: str,
+    labels: List[str],
+    fewshot_hints: Optional[Dict[str, str]] = None,
+    ingredient_only: bool = False,
+) -> Dict:
+    if not labels:
+        print("[VLM] Empty labels provided, return Unknown")
+        return {"label": "Unknown", "confidence": 0.0}
+
+    print(
+        f"[VLM] Start classify image={image_path} labels={len(labels)} "
+        f"ingredient_only={ingredient_only} fewshot={bool(fewshot_hints)}"
+    )
+    prompt = build_closedset_prompt(labels, fewshot_hints, ingredient_only)
+    _preview_prompt = (prompt[:120]).replace("\n", " ")
+    print(f"[VLM] Prompt length={len(prompt)} preview={_preview_prompt}...")
+
+    print("[VLM] Calling chat_vision...")
+    text = chat_vision(prompt, [image_path], temperature=0.1)
+    _preview_text = (text[:120]).replace("\n", " ")
+    print(f"[VLM] Received text length={len(text)} preview={_preview_text}...")
+
+    # 主路径：严格 JSON 解析
+    try:
+        obj = json.loads(text)
+        if isinstance(obj, dict) and "label" in obj:
+            label = obj.get("label", "")
+            # 纠偏：若输出不在封闭集，尝试包含匹配；否则 Unknown
+            if label not in labels:
+                for lb in labels:
+                    if lb in str(label):
+                        label = lb
+                        break
+                else:
+                    label = "Unknown"
+            conf = obj.get("confidence", 0.0)
+            try:
+                conf = float(conf)
+            except Exception:
+                conf = 0.0
+            result = {"label": label, "confidence": conf}
+            print(f"[VLM] Parsed JSON result={result}")
+            return result
+    except Exception as e:
+        print(f"[VLM] JSON parse failed: {e}")
+
+    # 回退：字符串包含匹配
+    for lb in labels:
+        if lb in text:
+            print(f"[VLM] Fallback match label={lb}")
+            return {"label": lb, "confidence": 0.5}
+
+    print("[VLM] No match, return Unknown")
+    return {"label": "Unknown", "confidence": 0.0}
+
+