65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
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}
|
|
|
|
|