95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
import json
|
|
import re
|
|
from typing import Dict, List, Optional
|
|
|
|
from .ollama_client import chat_vision
|
|
from .prompts import build_closedset_prompt
|
|
|
|
|
|
def _extract_json_obj(text: str):
|
|
"""尽量从返回文本中提取出合法 JSON（支持代码块/前后说明）。"""
|
|
if not isinstance(text, str):
|
|
return None
|
|
s = text.strip()
|
|
# 去掉 Markdown 代码块包裹
|
|
if s.startswith("```"):
|
|
s = re.sub(r"^```[a-zA-Z]*\n|\n```$", "", s).strip()
|
|
# 直接解析
|
|
try:
|
|
obj = json.loads(s)
|
|
if isinstance(obj, dict):
|
|
return obj
|
|
except Exception:
|
|
pass
|
|
# 提取第一个花括号对象
|
|
m = re.search(r"\{[\s\S]*?\}", s)
|
|
if m:
|
|
frag = m.group(0)
|
|
try:
|
|
obj = json.loads(frag)
|
|
if isinstance(obj, dict):
|
|
return obj
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
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", " ")
|
|
final_prompt = prompt.replace("\n", " ")
|
|
# print(f"[VLM] Prompt length={len(prompt)} preview={_preview_prompt}...")
|
|
print(f"[VLM] Prompt length={len(prompt)} preview={final_prompt}")
|
|
|
|
print("[VLM] Calling chat_vision...")
|
|
text = chat_vision(prompt, [image_path], temperature=0.1)
|
|
_preview_text = (str(text)[:200]).replace("\n", " ")
|
|
print(f"[VLM] Received text length={len(str(text))} preview={_preview_text}...")
|
|
|
|
# 解析 JSON（更健壮）
|
|
obj = _extract_json_obj(str(text))
|
|
if obj 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
|
|
|
|
print("[VLM] JSON parse failed or missing 'label', enter fallback")
|
|
|
|
# 回退：字符串包含匹配（无法确定置信度时返回 0.0）
|
|
for lb in labels:
|
|
if lb in str(text):
|
|
print(f"[VLM] Fallback match label={lb}")
|
|
return {"label": lb, "confidence": 0.0}
|
|
|
|
print("[VLM] No match, return Unknown")
|
|
return {"label": "Unknown", "confidence": 0.0}
|
|
|
|
|