#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ ChildHood 户型图识别(生产版) 固定用 qwen-vl-max(经 API易)把户型图解析成结构化 JSON, 并额外生成「路径 B 文字钉死」Markdown 片段,可直接喂给居民楼立面推导 prompt。 与 floorplan_eval.py(评测版)共享解析逻辑;本脚本固定模型、聚焦产出可用文本。 评测结论见:生成结果/户型图识别评测/评测对比.md(qwen-vl-max 领先)。 Usage: python floorplan_parse.py --image 户型图.png python floorplan_parse.py --image 户型图.png --out 结果.json python floorplan_parse.py --image 户型图.png --no-pin # 不生成文字钉死片段 Requirements: pip install openai python-dotenv """ import argparse import io import json 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/ OUTPUT_BASE = RESEARCH_DIR / "生成结果" / "户型图识别" # 让同目录的 floorplan_eval 可被 import sys.path.insert(0, str(SCRIPT_DIR)) # --------------------------------------------------------------------------- # 生产配置(评测选出的最优组合,固定不变) # --------------------------------------------------------------------------- PROVIDER = "apiyi" MODEL = "qwen-vl-max" from floorplan_eval import ( # noqa: E402 (import 在 sys.path 调整之后) EXTRACT_PROMPT, get_client, encode_image, parse_floorplan, extract_json, summarize, ) # --------------------------------------------------------------------------- # 增值:生成「路径 B 文字钉死」片段 # --------------------------------------------------------------------------- _FACES = ("南", "北", "东", "西") def _face_key(orientation: str) -> str: """把朝向归一到 南/北/东/西/未知""" s = str(orientation or "") for k in _FACES: if k in s: return k return "未知" def _window_desc(room: dict) -> str: wins = room.get("窗户", []) if not isinstance(wins, list) or not wins: return "无窗" parts = [] for w in wins: if isinstance(w, dict): parts.append(f"{w.get('类型', '?')}×{w.get('数量', '?')}({w.get('推测大小', '?')})") return "、".join(parts) if parts else "有窗" def generate_pin_text(parsed: dict) -> str: """把解析结果转成「路径 B 文字钉死」Markdown 片段,按立面组织窗户形态""" if not isinstance(parsed, dict): return "(解析结果非结构化,无法生成钉死片段)" lines = [] lines.append("# 户型图解析 → 路径 B「文字钉死」片段") lines.append("") lines.append(f"> qwen-vl-max 解析 · 户型 {parsed.get('户型结构', '未知')} · 整体朝向 {parsed.get('整体朝向', '未知')}") lines.append("> ⚠️ 窗户类型/朝向/大小均为模型推测,**须对照原图核对后再用于生图**。") lines.append("") rooms = parsed.get("房间", []) rooms = [r for r in rooms if isinstance(r, dict)] if isinstance(rooms, list) else [] # 按立面分组 by_face = {k: [] for k in _FACES} by_face["未知"] = [] review_notes = [] for r in rooms: face = _face_key(r.get("朝向")) area = r.get("面积估算", "未知") size = str(r.get("长宽尺寸", "") or "") size_str = f",{size}" if size and size != "未知" else "" desc = (f"- {r.get('名称', '?')}({area}{size_str}):{_window_desc(r)};" f"阳台={r.get('是否带阳台', '?')}") by_face[face].append(desc) # 收集脑补重灾区(卫生间暗卫 / 疑似暗房) name = str(r.get("名称", "")) rmk = str(r.get("备注", "") or "") if "卫" in name: review_notes.append(f"卫生间「{name}」朝向={face},是否暗卫?备注={rmk}") elif "暗" in rmk: review_notes.append(f"「{name}」疑似暗房(无直接采光),核对是否有窗:备注={rmk}") lines.append("## 各立面窗户形态(直接钉死立面图的窗户)") for face in _FACES: items = by_face.get(face, []) if items: lines.append(f"\n### {face}立面") lines.extend(items) if by_face.get("未知"): lines.append("\n### 朝向未知 / 内部空间") lines.extend(by_face["未知"]) # 模型给出的立面对应表 lmap = parsed.get("立面对应表", {}) if isinstance(lmap, dict) and any(k != "说明" for k in lmap): lines.append("") lines.append("## 模型立面对应表(参考)") for k, v in lmap.items(): if k == "说明": continue lines.append(f"- {k}: {v}") # 外墙尺寸标尺(用于推导立面开间比例) dim = parsed.get("尺寸标注", {}) if isinstance(dim, dict) and any(k != "说明" for k in dim): lines.append("") lines.append("## 外墙尺寸标尺(推导立面开间比例用)") for k, v in dim.items(): if k == "说明": continue lines.append(f"- {k}: {v}") # 需核对项 lines.append("") lines.append("## ⚠️ 人工核对重点") if review_notes: for n in review_notes: lines.append(f"- {n}") lines.append("- 窗户「类型 / 大小」均为推测,对照原图逐间确认") total = parsed.get("总面积", "") if not total or total == "未知": lines.append("- 总面积模型未识别,精确值看原图标注") note = parsed.get("识别备注", "") if note: lines.append(f"- 模型自述:{note}") return "\n".join(lines) # --------------------------------------------------------------------------- # CLI 入口 # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser( description="ChildHood 户型图识别(生产版)— qwen-vl-max 解析户型图 + 生成路径 B 文字钉死片段", formatter_class=argparse.RawDescriptionHelpFormatter, epilog="""\ 示例: python floorplan_parse.py --image 参考图/真实场景/户型图/xxx.png python floorplan_parse.py --image 户型.png --out 我的结果.json python floorplan_parse.py --image 户型.png --no-pin # 只要 JSON,不要钉死片段 环境变量 (ai-art-research/.env): APIYI_API_KEY API易 令牌(与 image_gen.py 共用) """, ) parser.add_argument("-i", "--image", type=str, required=True, help="户型图路径") parser.add_argument("--out", type=str, help="自定义输出 JSON 路径(默认存 生成结果/户型图识别/)") parser.add_argument("--no-pin", action="store_true", help="不生成「路径 B 文字钉死」片段") 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) client = get_client(PROVIDER) print(f"🏠 户型图识别(生产版)") print(f" 图片: {image_path.name}") print(f" Provider: {PROVIDER} 模型: {MODEL}") print() raw = parse_floorplan(client, MODEL, image_path, EXTRACT_PROMPT) parsed = extract_json(raw) metadata = { "provider": PROVIDER, "model": MODEL, "image": str(image_path), "parse_success": parsed is not None, "timestamp": datetime.now().isoformat(), } if parsed is None: print("⚠️ 未能从模型输出中解析出 JSON。原始输出前 500 字:") print(" " + raw[:500].replace("\n", "\n ")) # 保存 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}") base_path = out_path else: ts = datetime.now().strftime("%y%m%d-%H%M%S") model_slug = MODEL.replace("/", "-").replace(".", "-") base = f"{ts}_{PROVIDER}_{model_slug}" out_dir = OUTPUT_BASE / args.subdir if args.subdir else OUTPUT_BASE out_dir.mkdir(parents=True, exist_ok=True) payload = parsed if parsed is not None else {"_raw": raw} json_path = out_dir / f"{base}.json" json_path.write_text(json.dumps(payload, 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}") base_path = json_path # 生成「路径 B 文字钉死」片段 if parsed is not None and not args.no_pin: pin_text = generate_pin_text(parsed) pin_path = base_path.with_suffix(".钉死.md") pin_path.write_text(pin_text, encoding="utf-8") print(f"📝 文字钉死: {pin_path.name}") print() if parsed is not None: print("📋 解析预览:") print(summarize(parsed)) print() print(f"📍 结果: {base_path}") if parsed is not None and not args.no_pin: print(f"📍 钉死片段: {base_path.with_suffix('.钉死.md')}") print() print("🔍 下一步:对照原图核对 ⚠️ 项(卫生间朝向/暗卫、阳台朝向、窗户大小),") print(" 再把钉死片段贴进 prompts/居民楼-*-户型图推导-v*.txt 走路径 B 生图。") if __name__ == "__main__": main()