diff --git a/exp_multimodal/vlm_classifier.py b/exp_multimodal/vlm_classifier.py index 12de74b..bd88149 100644 --- a/exp_multimodal/vlm_classifier.py +++ b/exp_multimodal/vlm_classifier.py @@ -1,10 +1,39 @@ 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], @@ -25,38 +54,37 @@ def classify_image( 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}...") + _preview_text = (str(text)[:200]).replace("\n", " ") + print(f"[VLM] Received text length={len(str(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}") + # 解析 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 text: + if lb in str(text): print(f"[VLM] Fallback match label={lb}") - return {"label": lb, "confidence": 0.5} + return {"label": lb, "confidence": 0.0} print("[VLM] No match, return Unknown") return {"label": "Unknown", "confidence": 0.0}