# medical_segmentation **Repository Path**: wangwei83ky/medical_segmentation ## Basic Information - **Project Name**: medical_segmentation - **Description**: 陈博代码 - **Primary Language**: Unknown - **License**: Not specified - **Default Branch**: master - **Homepage**: None - **GVP Project**: No ## Statistics - **Stars**: 0 - **Forks**: 0 - **Created**: 2025-08-06 - **Last Updated**: 2025-08-06 ## Categories & Tags **Categories**: Uncategorized **Tags**: None ## README # 医学图像分割项目 这是一个完整的医学图像分割项目，支持多种网络架构（U-Net、VGG、ResNet），提供训练、评估和推理功能。项目采用模块化设计，代码清晰易懂，可直接在VSCode中运行。 ## 项目特色 - **模块化架构**: 职责单一，易于扩展和维护 - **配置驱动**: 所有参数集中在config.yaml，无需修改代码 - **多网络支持**: U-Net、VGG、ResNet等主流分割网络 - **组合损失函数**: 交叉熵与Dice损失的组合，提升分割效果 - **完整工作流**: 训练、评估、推理一体化 - **可视化丰富**: 训练曲线、混淆矩阵、预测结果可视化 - **灵活的推理模式**: 支持单张或批量图像推理，可指定模型和输出颜色 - **自动类别管理**: 根据name_classes自动计算num_classes - **实时训练监控**: 支持设置绘图间隔，实时查看训练和验证集曲线 ## 项目结构 ``` medical_segmentation_project/ ├── network/ # 网络架构实现 │ ├── u_net.py # U-Net 网络实现 │ ├── vgg.py # VGG 作为 backbone 的分割网络 │ ├── resnet.py # ResNet 作为 backbone 的分割网络 │ └── __init__.py # 网络模块初始化 ├── data/ # 数据处理模块 │ ├── data_loader.py # 数据集加载和DataLoader创建 │ ├── augment_utils.py # 数据增强工具 │ └── __init__.py # 数据模块初始化 ├── utils/ # 工具函数模块 │ ├── metrics_utils.py # 评估指标计算 │ ├── plot_utils.py # 可视化工具 │ ├── model_utils.py # 模型工具（保存、加载、设备管理等） │ └── __init__.py # 工具模块初始化 ├── logs/ # 训练日志目录 ├── saved_models/ # 模型权重保存目录 ├── results/ # 评估结果目录 │ ├── metrics/ # 指标结果文件 │ └── plots/ # 可视化图表 ├── pred_results/ # 推理结果目录 ├── docs/ # 文档目录 ├── config.yaml # 全局配置文件 ├── requirements.txt # 依赖库列表 ├── train.py # 训练入口脚本 ├── trainer.py # 训练核心逻辑 ├── evaluation.py # 评估脚本 ├── inference.py # 推理脚本 └── README.md # 项目说明文档 ``` ## 环境配置 ### 1. 安装依赖 ```bash pip install -r requirements.txt ``` ### 2. 数据准备 数据集应按以下结构组织： ``` data/ ├── images/ # 原始图像 │ ├── image_001.png │ ├── image_002.png │ └── ... ├── labels/ # 分割标签 │ ├── image_001.png │ ├── image_002.png │ └── ... ├── train.txt # 训练集文件列表（文件名不带后缀） ├── val.txt # 验证集文件列表（文件名不带后缀） └── test.txt # 测试集文件列表（文件名不带后缀） ``` **重要**: 文件列表格式（每行一个文件名，不带后缀）： ``` image_001 image_002 image_003 ``` ## 配置文件说明 `config.yaml` 包含所有可配置参数： ### 数据配置 - `image_dir`: 图像目录路径 - `label_dir`: 标签目录路径 - `train_list/val_list/test_list`: 数据集划分文件 - `image_size`: 输入图像尺寸 [height, width] - `name_classes`: **类别名称列表**，例如 `["Background", "Right_Ventricle", "Myocardium", "Left_Ventricle"]`。`num_classes` 将根据此列表的长度自动计算。 ### 数据增强配置 - `enabled`: 是否启用数据增强 - `augment_stop_epoch`: 数据增强停止的epoch - `flip/rotate/scale`: 具体增强方式及参数 ### 模型配置 - `name`: 模型名称 (unet/vgg/resnet) - `pretrained`: 是否使用预训练权重 - `pretrained_path`: 预训练权重路径 ### 训练配置 - `epochs`: 训练轮数 - `batch_size`: 批次大小 - `learning_rate`: 学习率 - `optimizer`: 优化器 (Adam/SGD/AdamW) - `loss_function`: 损失函数 (CrossEntropyDiceLoss为交叉熵+Dice组合) - `log_dir`: 日志保存目录 - `save_dir`: 模型保存目录 - `checkpoint_interval`: 每隔多少个epoch保存一次断点续训模型 - **`plot_interval`**: **每隔多少个epoch绘制一次训练曲线**（包括验证集损失和Dice） ### 评估配置 - `metric_save_path`: 指标保存路径 - `plot_confusion_matrix`: 是否绘制混淆矩阵 - **`model_path`**: **用于评估的模型权重路径**，如果为空则使用训练后的最新模型 ### 推理配置 - **`mode`**: **推理模式**：`single` (单张图像) 或 `batch` (批量图像) - **`input_path`**: **输入路径**。如果是 `single` 模式，为图像文件路径；如果是 `batch` 模式，可以是一个包含图像的目录路径，也可以是一个包含图像文件名的 `.txt` 文件路径（例如 `test.txt`）。 - `output_dir`: 推理结果输出目录 - **`model_path`**: **推理使用的模型权重路径**，如果为空则使用训练后的最新模型 - `visualize`: 是否保存可视化结果（原图 + 预测叠加） - **`colors`**: **分割类别的颜色配置**，与 `name_classes` 顺序对应，例如 `[[0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255]]`。 ## 使用方法 ### 1. 训练模型 ```bash # 使用默认配置训练 python train.py # 使用自定义配置文件 python train.py --config custom_config.yaml # 从检查点恢复训练 python train.py --resume # 指定GPU设备 python train.py --gpu 0 ``` **训练过程中的新特性**： - 每隔 `plot_interval` 个epoch自动绘制并保存训练集和验证集的损失和Dice曲线 - 训练曲线保存在 `logs/training_curves.png` ### 2. 评估模型 ```bash # 评估测试集 (使用config.yaml中evaluation.model_path指定的模型) python evaluation.py # 评估验证集 python evaluation.py --dataset val # 使用指定模型进行评估 python evaluation.py --model_path ./saved_models/best_model.pth # 保存预测结果 python evaluation.py --save_predictions ``` ### 3. 推理预测 ```bash # 根据config.yaml中的设置进行推理 (例如，如果mode为batch，input_path为./data/inference_images) python inference.py # 覆盖config.yaml中的设置进行单张图像推理 python inference.py --mode single --input_path /path/to/your/image.png # 覆盖config.yaml中的设置进行批量图像推理 (指定输入目录) python inference.py --mode batch --input_path /path/to/images/ # 覆盖config.yaml中的设置进行批量图像推理 (指定包含文件名的txt文件，例如test.txt) python inference.py --mode batch --input_path ./data/test.txt # 不保存可视化结果 python inference.py --no_visualize ``` ## 支持的网络架构 ### 1. U-Net - 经典的医学图像分割网络 - 编码器-解码器结构，带跳跃连接 - 适合各种医学图像分割任务 ### 2. VGG-based - 使用VGG16作为编码器 - 支持预训练权重 - 适合需要强特征提取能力的任务 ### 3. ResNet-based - 使用ResNet作为编码器 - 支持ResNet18/34/50等变体 - 适合深层网络训练 ## 评估指标 项目支持多种评估指标： - **Dice系数**: 衡量分割重叠度 - **IoU (Jaccard Index)**: 交并比 - **像素准确率**: 整体像素分类准确率 - **精确率/召回率/F1分数**: 各类别的分类性能 - **混淆矩阵**: 详细的分类结果分析 ## 可视化功能 - **训练曲线**（训练集和验证集的损失和Dice分数）- 支持设置绘图间隔 - 混淆矩阵热力图 - 各类别指标对比图 - **预测结果可视化**（原图+预测叠加）- 支持自定义颜色 ## 高级功能 ### 1. 组合损失函数 - 交叉熵损失：处理类别不平衡 - Dice损失：优化分割重叠度 - 组合使用：提升整体分割效果 ### 2. 数据增强控制 - 支持按epoch控制数据增强的启用/停用 - 多种增强方式：翻转、旋转、缩放等 - 验证集和测试集不使用数据增强 ### 3. 断点续训 - 自动保存训练检查点 - 支持从任意检查点恢复训练 - 保留训练历史和最佳模型 ### 4. 模型管理 - 自动保存最佳模型（基于验证损失） - 定期保存检查点 - 支持预训练权重加载 ### 5. 灵活的推理 - 支持单张图像和批量推理，输入路径可以是目录或文件列表 - 自动处理图像尺寸调整 - 可选的可视化结果保存，并可自定义分割颜色 ### 6. 自动类别管理 - 通过 `name_classes` 自动计算 `num_classes` - 无需手动设置类别数量，避免不一致错误 ### 7. 实时训练监控 - 通过 `plot_interval` 设置绘图频率 - 实时查看训练和验证集的损失和Dice曲线变化 ## 快速开始 1. **准备数据**: 按照上述格式组织数据集 2. **修改配置**: 编辑config.yaml中的数据路径和参数，特别是 `name_classes` 和 `colors` 3. **开始训练**: `python train.py` 4. **评估模型**: `python evaluation.py` 5. **推理预测**: `python inference.py` (根据config.yaml中的设置进行推理) ## 常见问题 ### Q: 如何添加新的网络架构？ A: 在 `network/` 目录下创建新的网络文件，实现网络类，然后在 `network/__init__.py` 中注册。 ### Q: 如何修改损失函数？ A: 在 `utils/model_utils.py` 的 `get_loss_function()` 函数中添加新的损失函数。 ### Q: 如何修改数据增强策略？ A: 在 `data/augment_utils.py` 中添加新的增强类，然后在相应的获取函数中使用。 ### Q: 训练过程中出现内存不足怎么办？ A: 减小 `batch_size`，或者减小 `image_size`。 ### Q: 为什么文件列表不带后缀？ A: 这样设计是为了更好的灵活性，代码会自动添加.png后缀。如需支持其他格式，可修改data_loader.py。 ### Q: 如何设置自定义的分割颜色？ A: 在 `config.yaml` 的 `inference.colors` 中设置每个类别对应的RGB颜色值，顺序与 `name_classes` 一致。 ### Q: 如何控制训练曲线的绘制频率？ A: 在 `config.yaml` 的 `training.plot_interval` 中设置间隔的epoch数，例如设置为5表示每5个epoch绘制一次。 ## 许可证 本项目采用 MIT 许可证。 ## 贡献 欢迎提交 Issue 和 Pull Request 来改进这个项目。 ## 联系方式 如有问题，请通过 GitHub Issues 联系。