267 lines
10 KiB
Python
267 lines
10 KiB
Python
#!/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()
|