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