#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ ChildHood 户型图识别评测工具 通过 vision 大模型把户型图解析成结构化 JSON（房间 / 朝向 / 窗户形态）， 用于「路径 B — 户型图推导」的"文字钉死"环节，替代人工逐项校验。 配套：与 image_gen.py 同一工具链（共用 ai-art-research/.env、OpenAI 兼容客户端、 生成结果/ 输出目录、.meta.json 元数据约定）。 Usage: python floorplan_eval.py --image 户型图.png # 默认 provider + 默认模型 python floorplan_eval.py --image 户型图.png -m qwen-vl-max # 指定模型 python floorplan_eval.py --image 户型图.png --provider dashscope # 走阿里云百炼直连 python floorplan_eval.py --image 户型图.png --prompt-file 自定义.txt # 自定义提取 prompt Requirements: pip install openai python-dotenv """ import argparse import base64 import io import json import os import re import sys from datetime import datetime from pathlib import Path # Windows 控制台 UTF-8 支持（用 reconfigure 幂等重配，避免被其他脚本 import 时 # 产生双层 TextIOWrapper 包装、损坏输出流） if sys.platform == "win32": for _stream in (sys.stdout, sys.stderr): try: _stream.reconfigure(encoding="utf-8", errors="replace") except Exception: pass # --------------------------------------------------------------------------- # 路径配置 # --------------------------------------------------------------------------- SCRIPT_DIR = Path(__file__).resolve().parent # ai-art-research/tools/ RESEARCH_DIR = SCRIPT_DIR.parent # ai-art-research/ DOTENV_PATH = RESEARCH_DIR / ".env" OUTPUT_BASE = RESEARCH_DIR / "生成结果" / "户型图识别评测" # --------------------------------------------------------------------------- # 加载 .env # --------------------------------------------------------------------------- if DOTENV_PATH.exists(): from dotenv import load_dotenv load_dotenv(DOTENV_PATH) # provider 配置 APIYI_API_KEY = os.getenv("APIYI_API_KEY", "") APIYI_BASE_URL = os.getenv("APIYI_BASE_URL", "https://api.apiyi.com/v1") DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY", "") DASHSCOPE_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1" # 各 provider 的默认模型 DEFAULT_MODEL = { "apiyi": "qwen-vl-max", "dashscope": "qwen-vl-max-latest", } # --------------------------------------------------------------------------- # 统一提取 prompt（核心：直接对接路径 B 的"文字钉死"） # --------------------------------------------------------------------------- EXTRACT_PROMPT = """\ 你是一个专业的户型图解析助手。请仔细分析这张户型图，提取结构化信息。 【户型图符号约定】 - 墙体：粗实线或双线 - 门：弧线（表示开门方向） - 窗：细的双平行线 - 阳台：半开放区域，常带栏杆符号 - 房间名：通常文字标注在房间内部（客厅/卧室/厨房/卫生间/阳台/餐厅/玄关等） - 面积：常以 m² 数字标注在房间内或户型角落 - 尺寸标尺：墙体旁的尺寸线（细线 + 两端短斜线/箭头）+ 数字，标注墙段长度 / 房间长宽，单位多为米(m)或毫米(mm) 【严格规则】 1. 只描述图上实际可见、或可明确推断的结构。禁止用"住宅常识"脑补—— 不要默认卧室有大窗、阳台是落地窗、北面是厨房。 2. 看不清或不确定的字段，一律填 "未知"，绝不瞎猜。 3. **重点提取尺寸标尺**：户型图外墙周边常有尺寸线（细线 + 数字）标注各开间/进深分段。请尽量读出南向开间分段、北向开间分段、进深分段的数字（多为米），填入"尺寸标注"——这是推导立面窗户水平比例的关键；看不清的分段填"未知"。 4. 客餐厅若为一体空间（无明显实体隔断），合并为一个"客餐厅"，不要拆成"客厅"+"餐厅"两个。 5. 多个阳台若朝向不同，分别标注各自朝向（如一个朝南、一个朝北）。 6. 输出必须是纯 JSON，不要任何解释文字、不要用 ```json 代码块包裹、不要前后缀。 【输出 JSON 结构】 { "整体朝向": "主采光面朝向，如 南/北/东南；看不清填 未知", "户型结构": "几室几厅几卫，如 三室两厅一卫", "总面积": "图上标注的建筑面积，如 约 95 m²；无标注填 未知", "房间": [ { "名称": "客厅/主卧/次卧/厨房/卫生间/阳台/餐厅/玄关...", "朝向": "该房间主要采光朝向：南/北/东/西", "面积估算": "图上有标注则填，否则 未知", "长宽尺寸": "图上标尺标注的房间长与宽，如 3.6m×4.2m；无标尺填 未知", "窗户": [ {"所在立面": "南/北/东/西", "类型": "落地窗/普通窗/飘窗/无", "数量": 1, "推测大小": "大/中/小/未知"} ], "是否带阳台": "是/否/未知", "备注": "其他可见特征，如 是否飘窗、是否暗卫（无窗）等" } ], "尺寸标注": { "说明": "图上外墙周边标尺的数字，用于推导立面窗户的水平开间比例", "南向开间分段": ["如 3.6m", "4.2m"], "北向开间分段": ["如 3.0m", "4.8m"], "进深分段": ["如 5.1m", "3.2m"], "总开间": "南立面总宽（有则填，无则未知）", "总进深": "进深总长（有则填，无则未知）" }, "立面对应表": { "说明": "哪些房间的窗户落在哪个外立面，用于推导正投影立面", "南立面": ["客厅", "主卧"], "北立面": ["厨房", "卫生间"] }, "识别备注": "图纸清晰度、符号可能误读之处、需要人工重点核对的项" } 现在请解析这张户型图，只输出 JSON。""" # --------------------------------------------------------------------------- # 工具函数 # --------------------------------------------------------------------------- def get_client(provider: str): """根据 provider 创建 OpenAI 兼容客户端，返回 (client, default_model)""" from openai import OpenAI if provider == "apiyi": if not APIYI_API_KEY: print("❌ 未找到 APIYI_API_KEY") print(f" 请在 {DOTENV_PATH} 中设置: APIYI_API_KEY=sk-xxx") sys.exit(1) return OpenAI(api_key=APIYI_API_KEY, base_url=APIYI_BASE_URL) elif provider == "dashscope": if not DASHSCOPE_API_KEY: print("❌ 未找到 DASHSCOPE_API_KEY") print(" 阿里云百炼直连需要单独的 key：") print(" 1. 前往 https://bailian.console.aliyun.com 申请 API Key") print(f" 2. 在 {DOTENV_PATH} 中添加: DASHSCOPE_API_KEY=sk-xxx") sys.exit(1) return OpenAI(api_key=DASHSCOPE_API_KEY, base_url=DASHSCOPE_BASE_URL) else: print(f"❌ 未知 provider: {provider}（支持: apiyi / dashscope）") sys.exit(1) _EXT_MIME = { ".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp", ".gif": "image/gif", } def encode_image(path: Path) -> str: """读图为 base64 data URL""" mime = _EXT_MIME.get(path.suffix.lower(), "image/png") b64 = base64.b64encode(path.read_bytes()).decode() return f"data:{mime};base64,{b64}" def parse_floorplan(client, model: str, image_path: Path, prompt: str) -> str: """调用 vision 模型，返回原始文本输出""" data_url = encode_image(image_path) response = client.chat.completions.create( model=model, messages=[{ "role": "user", "content": [ {"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": data_url}}, ], }], ) return response.choices[0].message.content def extract_json(text: str): """从模型输出中尽量鲁棒地提取 JSON dict；失败返回 None""" text = text.strip() # 1. 直接解析 try: return json.loads(text) except json.JSONDecodeError: pass # 2. 去掉 ```json ... ``` 代码块包裹 fenced = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL) if fenced: try: return json.loads(fenced.group(1)) except json.JSONDecodeError: pass # 3. 截取第一个 { 到最后一个 } start, end = text.find("{"), text.rfind("}") if start != -1 and end != -1 and end > start: candidate = text[start:end + 1] try: return json.loads(candidate) except json.JSONDecodeError: # 4. 容错：清理常见问题（中文逗号、尾随逗号）后再试 cleaned = (candidate .replace("，", ",") .replace("：", ":")) cleaned = re.sub(r",(\s*[}\]])", r"\1", cleaned) # 去尾随逗号 try: return json.loads(cleaned) except json.JSONDecodeError: pass return None def save_result(result, metadata: dict, subdir: str = "") -> Path: """保存解析结果 JSON + .meta.json，返回结果文件路径""" ts = datetime.now().strftime("%y%m%d-%H%M%S") model_slug = metadata.get("model", "unknown").replace("/", "-").replace(".", "-") provider_slug = metadata.get("provider", "unknown") base = f"{ts}_{provider_slug}_{model_slug}" out_dir = OUTPUT_BASE / subdir if subdir else OUTPUT_BASE out_dir.mkdir(parents=True, exist_ok=True) json_path = out_dir / f"{base}.json" json_path.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8") meta_path = out_dir / f"{base}.meta.json" meta_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8") print(f"✅ 已保存: {json_path.name}") return json_path def summarize(parsed: dict) -> str: """从解析结果里摘出要点用于控制台预览""" if not isinstance(parsed, dict): return "(非结构化结果)" lines = [] if parsed.get("户型结构"): lines.append(f" 户型: {parsed['户型结构']}") if parsed.get("整体朝向"): lines.append(f" 朝向: {parsed['整体朝向']}") if parsed.get("总面积"): lines.append(f" 面积: {parsed['总面积']}") rooms = parsed.get("房间", []) if isinstance(rooms, list): lines.append(f" 房间数: {len(rooms)}") for r in rooms: if isinstance(r, dict): win_n = len(r.get("窗户", [])) if isinstance(r.get("窗户"), list) else 0 lines.append(f" - {r.get('名称','?')} 朝向={r.get('朝向','?')} 窗={win_n} 阳台={r.get('是否带阳台','?')}") return "\n".join(lines) # --------------------------------------------------------------------------- # CLI 入口 # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser( description="ChildHood 户型图识别评测工具 — vision 模型解析户型图为结构化 JSON", formatter_class=argparse.RawDescriptionHelpFormatter, epilog="""\ 示例: python floorplan_eval.py --image 参考图/真实场景/户型图/QQ20260618-213247.png python floorplan_eval.py --image 户型.png -m qwen-vl-max python floorplan_eval.py --image 户型.png -m gemini-2.5-flash # baseline 对照 python floorplan_eval.py --image 户型.png --provider dashscope # 百炼直连 环境变量 (ai-art-research/.env): APIYI_API_KEY API易 令牌 (apiyi provider 用，复用 image_gen.py) DASHSCOPE_API_KEY 阿里云百炼 key (dashscope provider 用) """, ) parser.add_argument("-i", "--image", type=str, required=True, help="户型图路径") parser.add_argument("-m", "--model", type=str, help="vision 模型名（默认随 provider）") parser.add_argument("--provider", type=str, default="apiyi", choices=["apiyi", "dashscope"], help="API 来源（默认 apiyi 复用现有 key）") parser.add_argument("--prompt-file", type=str, help="自定义提取 prompt 文件（默认用内置 EXTRACT_PROMPT）") parser.add_argument("--out", type=str, help="自定义输出 JSON 路径（默认存 生成结果/户型图识别评测/）") parser.add_argument("--raw", action="store_true", help="额外保存模型原始输出文本（便于排查解析失败）") parser.add_argument("--subdir", type=str, default="", help="输出子目录名") args = parser.parse_args() image_path = Path(args.image).resolve() if not image_path.exists(): print(f"❌ 户型图不存在: {image_path}") sys.exit(1) if args.prompt_file: pf = Path(args.prompt_file).resolve() if not pf.exists(): print(f"❌ prompt 文件不存在: {pf}") sys.exit(1) prompt = pf.read_text(encoding="utf-8") print(f"📄 从文件读取 prompt: {pf.name}") else: prompt = EXTRACT_PROMPT client = get_client(args.provider) model = args.model or DEFAULT_MODEL[args.provider] print(f"🏠 户型图识别") print(f" 图片: {image_path.name}") print(f" Provider: {args.provider} 模型: {model}") print() raw = parse_floorplan(client, model, image_path, prompt) parsed = extract_json(raw) metadata = { "provider": args.provider, "model": model, "image": str(image_path), "prompt": prompt, "parse_success": parsed is not None, "timestamp": datetime.now().isoformat(), } if parsed is None: print("⚠️ 未能从模型输出中解析出 JSON（可能模型没遵守格式约束）。") print(" 原始输出前 500 字：") print(" " + raw[:500].replace("\n", "\n ")) print() print("💡 建议：换模型重试，或加 --raw 保留原文人工查看。") # 保存 if args.out: out_path = Path(args.out).resolve() out_path.parent.mkdir(parents=True, exist_ok=True) payload = parsed if parsed is not None else {"_raw": raw} out_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") meta_path = out_path.with_suffix(".meta.json") meta_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8") print(f"✅ 已保存: {out_path.name}") result_path = out_path else: payload = parsed if parsed is not None else {"_raw": raw} result_path = save_result(payload, metadata, subdir=args.subdir) if args.raw: raw_path = result_path.with_suffix(".raw.txt") raw_path.write_text(raw, encoding="utf-8") print(f"📝 原始输出: {raw_path.name}") print() if parsed is not None: print("📋 解析预览：") print(summarize(parsed)) print() print(f"📍 完整结果: {result_path}") print() print("🔍 下一步：对照原图人工核对（重点看 房间是否漏识别、窗户数量/类型是否对）。") if __name__ == "__main__": main()