户型图识别（实践篇）

前言：嘿，朋友们！今天咱们来唠唠怎么用 AI 把户型图里的东西（像墙、窗户、门啥的）给认出来。听起来是不是有点高大上？别担心，其实挺简单的！

文章重点：

数据标注：用工具给户型图里的东西（墙、窗户、门等）打上标签，让模型知道它们是什么。
模型训练：把标注好的数据喂给模型，让它学会识别这些东西。
模型导出 ：把训练好的模型转换成通用格式（比如 .onnx），方便在其他地方用。
本地运行：在本地电脑上测试模型，看看它能不能正确识别户型图里的内容。
部署模型：把模型放到服务器上，通过 API 调用，让其他人也能用它来识别户型图。
实用技巧：分享一些训练和部署过程中可能会遇到的问题，以及怎么解决它们。

数据标注

数据标注工具选型对比

工具名称	类型	主要功能	支持标注类型	输出格式	适用场景	优点	缺点
LabelImg	桌面工具	简单的图像标注	矩形框	XML（Pascal VOC）、TXT（YOLO）、CSV等	小型项目、简单标注任务	界面直观、轻量级、跨平台	功能有限，不支持复杂形状标注，无团队协作功能
LabelMe	桌面/在线工具	多类型标注（实例分割、语义分割等）	矩形、多边形、点、线等	JSON、VOC等	中小型项目、复杂标注任务	功能强大、支持自定义属性、开源可定制	入门版功能有限，Pro版需付费
Make Sense	在线工具	轻量化标注，支持AI辅助	矩形、多边形、点、线等	YOLO、VOCXML、JSON、CSV等	中小型项目、AI辅助标注	无需部署、轻量化、支持多种格式	社区响应慢，近两年代码未更新
T-Rex Label	在线工具	视觉提示的Bbox自动标注	矩形框	COCO、YOLO等	难以用文本描述的物体标注	交互流畅、社区反馈快	功能相对单一
VoTT	桌面工具	图像和视频标注	矩形、多边形	CSV、CNTK、TensorFlow、自定义JSON	中大型项目、视频标注	支持多种模型集成、自定义导入导出	需要一定系统资源
CVAT	在线/服务器工具	复杂标注任务，支持团队协作	矩形、多边形、点、线等	COCO、YOLO等	大型项目、多用户协作	功能强大、支持复杂标注、团队协作	部署复杂，对服务器要求高
LabelU	桌面工具	综合性标注，支持多种数据类型	2D框、3D框、多边形、点、线等	JSON、COCO、MASK等	复杂标注任务、多数据类型	工具类型全面、支持预标注载入	缺乏智能标注功能
Label Studio	在线工具	多类型标注，支持多种数据格式	矩形、多边形、文本标注等	JSON、CSV等	多样化标注任务	功能丰富、支持多种数据类型	部署复杂，需一定技术背景

我用的T-Rex Label

选择建议

简单标注（单人、小项目）：
- LabelImg：简单好用，上手快，适合标注矩形框，输出格式多，但功能有限。
- Make Sense：在线工具，轻量化，支持AI辅助标注，适合快速标注。
复杂标注（多边形、分割等）：
- LabelMe：功能强大，支持多种标注类型，适合复杂任务，但免费版功能有限。
- LabelU：工具全面，支持多种标注类型，适合复杂任务，但缺乏智能辅助。
团队协作（多人标注）：
- CVAT：功能最强，支持团队协作，适合大规模项目，但部署复杂，需要服务器支持。
- Label Studio：支持多种数据类型，适合多样化标注任务，部署稍复杂。
视频标注：
- VoTT：支持图像和视频标注，能和多种模型集成，适合视频标注任务。
AI辅助标注：
- Make Sense：轻量化，支持AI辅助，适合快速标注。
- CVAT：有AI辅助功能，适合复杂任务。
- T-Rex Label：轻量化，支持AI辅助，适合快速标注。

如果预算有限，优先考虑开源免费的工具，比如LabelImg、LabelMe、CVAT；如果任务复杂或需要团队协作，可以考虑CVAT或Label Studio。

数据标注格式

数据标注格式就像是给机器"喂食"的方式，不同的格式适合不同的"口味"（模型）。以下是几种常见标注格式的对比：

标注格式	文件类型	主要特点	适用场景	优点	缺点
Pascal VOC	XML	结构化标注，每个目标有单独标签	目标检测（矩形框标注）	结构清晰，通用性强	文件较大，处理稍慢
YOLO	TXT	简单文本格式，每行一个标注	目标检测（矩形框标注）	文件小，读取快	信息简略，功能有限
COCO	JSON	支持复杂标注（多边形、分割等）	复杂目标检测、实例分割	功能强大，支持多种标注	文件大，格式复杂
LabelMe	JSON	支持多种形状（矩形、多边形等）	复杂标注任务（不规则物体）	灵活，支持自定义属性	文件结构复杂
CSV	CSV	表格形式，简单直观	简单标注任务（分类）	简单易用，容易处理	功能有限，不适合复杂标注

总结：

简单目标检测：Pascal VOC 或 YOLO
复杂标注任务：COCO 或 LabelMe
分类任务：CSV

根据任务需求和标注复杂度选择合适的格式就好啦！

如果你需要将数据从一种格式转换为另一种格式，可以使用一些开源工具（如 labelImg、cvat 或自定义脚本）来完成转换。

模型选型好的！进行模型训练是一个关键步骤，以下是基于你提供的背景信息（户型图数据标注、预处理、选择模型架构）的详细指导，帮助你顺利开始模型训练。

模型训练

1. 数据准备

在开始训练之前，确保你的数据已经准备好，包括标注好的图片和对应的 JSON 文件（YOLO 格式）。

数据划分

将数据划分为训练集和验证集，通常比例为 80%（训练集）和 20%（验证集）。如果数据很多，还可以再分一个测试集，用来最后再考一次，看看模型的最终表现。

可以使用以下 Python 脚本进行划分：

python 复制代码

// kimi自动生成，没有经过验证
import os
import random
import shutil

# 数据目录
data_dir = "path/to/your/data"
train_dir = "path/to/train"
val_dir = "path/to/val"

# 创建目录
os.makedirs(os.path.join(train_dir, "images"), exist_ok=True)
os.makedirs(os.path.join(train_dir, "annotations"), exist_ok=True)
os.makedirs(os.path.join(val_dir, "images"), exist_ok=True)
os.makedirs(os.path.join(val_dir, "annotations"), exist_ok=True)

# 获取所有图片和标注文件
images = [f for f in os.listdir(data_dir) if f.endswith(".jpg")]
annotations = [f for f in os.listdir(data_dir) if f.endswith(".json")]

# 随机划分
random.seed(42)
random.shuffle(images)
split = int(0.8 * len(images))

train_images = images[:split]
val_images = images[split:]

# 复制文件到训练集和验证集目录
for img in train_images:
    shutil.copy(os.path.join(data_dir, img), os.path.join(train_dir, "images"))
    shutil.copy(os.path.join(data_dir, img.replace(".jpg", ".json")), os.path.join(train_dir, "annotations"))

for img in val_images:
    shutil.copy(os.path.join(data_dir, img), os.path.join(val_dir, "images"))
    shutil.copy(os.path.join(data_dir, img.replace(".jpg", ".json")), os.path.join(val_dir, "annotations"))

print("Data split complete!")

PS：

训练集：就是让模型学习的数据，就像课本上的题。
验证集：就是用来"考试"的数据，看看模型是不是真的学会了，会不会只背答案。

简单来说，训练集是让模型学习的，验证集是让模型考试的，这样才能保证模型不是只会背答案，而是真的学会了。

2. 模型选择

根据您的需求，以下是关于常见模型的训练对比表格，同时补充了是否可用于模型训练的信息：

模型名称	类型	优点	缺点	适用场景	是否可用于模型训练
ResNet	图像分类/通用	高精度，适合复杂任务	模型较大，推理速度慢	高精度图像分类、医学图像分析	是
MobileNet	图像分类/通用	轻量级，速度快	精度略低	移动端应用、实时图像分类	是
YOLOv5/v7	目标检测	轻量级，推理速度快	精度略低于高精度模型	实时目标检测（安防监控、自动驾驶）	是
Faster R-CNN	目标检测	精度高，适合复杂场景	速度慢，不适合实时应用	高精度目标检测（医学图像、复杂场景）	是
SSD	目标检测	速度和精度平衡	性能介于 YOLO 和 Faster R-CNN 之间	通用目标检测（安防监控）	是
EfficientDet	目标检测	速度和精度平衡，适合多种场景	性能介于 YOLO 和 Faster R-CNN 之间	通用目标检测（安防监控）	是
DETR	目标检测	基于 Transformer，精度高	速度慢，不适合实时应用	高精度目标检测（复杂图像场景）	是

这里作者有一个疑惑点："deepSeek都出来了"，我们还选择个屁的别的模型呀？

打脸解惑：

如果您的任务是图像分类或目标检测，可以选择 ResNet、MobileNet、YOLOv5/v7、Faster R-CNN、SSD、EfficientDet 或 DETR 进行模型训练。（针对性训练的模型，从娃娃培养起来的，比如这篇专门培养一个户型图识别的小模型）
如果您的任务是复杂推理、语音识别或语言生成，建议选择预训练模型（如 DeepSeek-R1、Whisper-large-v3、Mistral-Small-24B）进行微调或强化学习优化。（就是已经被别人有针对性训练好的模型）

简单来说，预训练模型就是别人已经训练好的模型，你可以直接拿来用，或者在它的基础上再训练一下，让它更适合你的任务。

总结

如果您的任务是图像分类或目标检测，可以选择 ResNet、MobileNet、YOLOv5/v7、Faster R-CNN、SSD、EfficientDet 或 DETR 进行模型训练。
如果您的任务是复杂推理、语音识别或语言生成，建议选择预训练模型（如 DeepSeek-R1、Whisper-large-v3、Mistral-Small-24B）进行微调或强化学习优化。

3. 模型训练

安装 YOLO 框架

bash 复制代码

git clone https://github.com/ultralytics/yolov5

cd yolov5

// 安装必要的Python依赖
pip install -r requirements.txt
// PS：安装时终端可能会报环境不安全（这个错误是由于 **PEP 668** 的引入导致的）
pip install -r requirements.txt error: externally-managed-environment 
× This environment is externally managed 
╰─> To install Python packages system-wide, try brew install

解决：
# 创建虚拟环境
python3 -m venv path/to/venv
# 激活虚拟环境
source path/to/venv/bin/activate
# 关闭虚拟环境
deactivate
# 在虚拟环境中重新安装依赖
pip install -r requirements.txt

安装完终端名称：(venv) shuaiBi@shuaiBideMacBook-Pro yolov5 %
优点：不会破坏系统级 Python 环境，适合开发和测试。
缺点：需要额外的步骤来创建和激活虚拟环境。

开始训练

yolov5项目拉取下来之后，把你生成好的dataset（标注导出数据）放到yolov5中

主目录新增dataset.yaml

bash 复制代码

dataset.yaml
train: dataset/images
val: dataset/images

nc: 8
names: ['承重墙','非承重墙','一字型窗','一字型飘窗','落地窗','单开门','双开门','移门']

终端使用以下命令开始训练：

bash 复制代码

python train.py --img 640 --batch 16 --epochs 150 --data dataset.yaml
--weights yolov5s.pt --name my_model

参数说明：

css 复制代码

--img 640`：训练时的图片尺寸。
--batch 16`：每批次的图片数量。
--epochs 100`：训练的总轮数。
--data dataset.yaml`：数据集配置文件。
--weights yolov5s.pt`：使用预训练的 YOLOv5s 模型权重。
--name house_model`：保存训练结果的文件夹名称。

训练完成

less 复制代码

Model summary: 157 layers, 7031701 parameters, 0 gradients, 15.8 GFLOPs 
Class Images Instances P R mAP50 all 2 79 0.00161 0.0222 0.000966 0.00029 
承重墙 2 28 0 0 0 0 
非承重墙 2 32 0 0 0 0 
一字型窗 2 8 0 0 0 0 
单开门 2 9 0.00806 0.111 0.00483 0.00145 
移门 2 2 0 0 0 0 
Results saved to runs/train/my_model3

markdown 复制代码

-  **训练周期**：50 epochs 完成，耗时 0.025 小时（约 1.5 分钟）。

-  **模型权重文件**：
    -   `runs/train/my_model3/weights/last.pt`：最后一个 epoch 的权重文件，大小为 14.5MB。
    -   `runs/train/my_model3/weights/best.pt`：最佳 epoch 的权重文件，大小为 14.5MB。

-   **验证结果**：
    -   **总类别**：7 个类别（承重墙、非承重墙、一字型窗、单开门、移门等）。
    -   **验证指标**：
        -   **P（Precision）** ：精确度，表示预测为正的样本中实际为正的比例。
        -   **R（Recall）** ：召回率，表示实际为正的样本中预测为正的比例。
        -   **mAP50**：平均精度（IoU=0.5），表示模型在不同类别上的平均性能。
        
-   **下一步建议**
    -   **检查数据集**：
        -   确保数据集标注准确无误。
        -   如果数据量较少，尝试增加更多标注数据。
    -   **延长训练时间**：
        -   增加训练周期（epochs），例如从 50 增加到 100 或更多。
    -   **调整超参数**：
        -   学习率（`lr0`）、优化器（`optimizer`）、批量大小（`batch_size`）等超参数可能需要调整。例如，尝试使用 Adam 优化器或调整学习率。
    -   **数据增强**：
        -   如果数据量有限，可以使用数据增强技术（如随机裁剪、翻转、颜色抖动等）来增加数据多样性。
    -   **模型可视化**：
        -   使用 TensorBoard 或其他工具可视化训练过程，观察损失函数和指标的变化。
    **验证和测试**：
        -   在验证集上评估模型性能，确保模型在未见过的数据上表现良好。

4.模型导出

YOLOv5 支持多种格式，如 ONNX、TorchScript 等，当我们跑完数据的时候，会生成（my_model）文件。

以下是导出模型的命令：

bash 复制代码

python export.py --weights runs/train/my_model4/weights/best.pt --include onnx

这将生成一个 .onnx 文件，适合在多种平台上部署。

5.本地运行

安装依赖

我的需求是本地运行我的模型，实现给一张图片并返回检测结果的功能：

bash 复制代码

pip install onnxruntime pillow numpy

onnxruntime用于加载和运行 ONNX 格式的模型
Pillow用于图像处理
NumPy用于数值计算

编写推理脚本

Python 复制代码

import onnxruntime as ort
import numpy as np
from PIL import Image
import json

def load_image(image_path):
    """加载并预处理图片"""
    image = Image.open(image_path)
    if image.mode != 'RGB':
        image = image.convert('RGB')  # 确保图片是 RGB 格式
    image = image.resize((640, 640))  # 调整图片大小
    image = np.array(image) / 255.0  # 归一化
    image = np.transpose(image, (2, 0, 1))  # 转换为 CHW 格式
    image = np.expand_dims(image, axis=0).astype(np.float32)  # 添加批次维度
    return image

def run_inference(session, image):
    """运行推理"""
    inputs = {session.get_inputs()[0].name: image}
    outputs = session.run(None, inputs)
    return outputs

def format_results(outputs, class_names):
    """格式化结果"""
    results = []
    for detection in outputs[0][0]:  # 遍历每个检测结果
        class_id = int(detection[5])  # 获取类别 ID
        class_name = class_names[class_id]  # 获取类别名称
        bbox = {
            "x": {"start": int(detection[0]), "end": int(detection[2])},
            "y": {"start": int(detection[1]), "end": int(detection[3])}
        }
        results.append({"class": class_name, "bbox": bbox})
    return results

def main(image_path, model_path, class_names):
    """主函数"""
    # 加载模型
    session = ort.InferenceSession(model_path)
    
    # 加载图片
    image = load_image(image_path)
    
    # 运行推理
    outputs = run_inference(session, image)
    
    # 格式化结果
    results = format_results(outputs, class_names)
    
    # 打印结果
    print(json.dumps(results, indent=4))

if __name__ == "__main__":
    # 图片路径
    image_path = "/Users/jinchunye/Documents/靳春野/项目/代码片段练习/练习/模型训练/test.png"
    
    # 模型路径
    model_path = "runs/train/my_model4/weights/best.onnx"
    
    # 类别名称
    class_names = ['承重墙','非承重墙','一字型窗','一字型飘窗','落地窗','单开门','双开门','移门']
    
    # 运行主函数
    main(image_path, model_path, class_names)

执行此脚本：

bash 复制代码

python /Users/jinchunye/Documents/shuaibi/项目/代码片段练习/练习/模型训练/infer.py

6.部署模型

简单说一下

上传 .onnx 模型到 Ollama：通过 Ollama 管理和加载模型。

在dify应用流程编排中直接引用

四句打油诗

数据标注要准确，不然模型学不好。
训练切记要微调，一给我哩giaogiao。

最后

水平有限，还不能写到尽善尽美，希望大家多多交流，跟春野一同进步！！！

往期相关精彩文章链接

户型图识别（胡思乱想篇）