【源码】基于PyTorch的图像分割工具箱

项目简介

本项目是基于PyTorch框架构建的图像分割工具箱，主要服务于生物医学图像分割任务。项目参考多个开源仓库，实现了U-Net模型，具备训练、验证、多GPU训练以及预测等功能。

项目的主要特性和功能

• U-Net模型实现：基于PyTorch实现U-Net模型，用于生物医学图像分割。
• 多GPU训练支持：提供单GPU和多GPU训练脚本，支持分布式训练。
• 数据集处理：自定义数据集读取模块，支持DRIVE数据集（视网膜血管分割）。
• 预测功能：提供简易预测脚本，可使用训练好的权重进行预测测试。
• 数据统计：提供脚本统计数据集各通道均值和标准差。

安装使用步骤

环境配置

1. Python版本：Python 3.6/3.7/3.8
2. PyTorch版本：PyTorch 1.10
3. 操作系统：Ubuntu或CentOS（Windows暂不支持多GPU训练）
4. GPU支持：建议使用GPU进行训练
5. 依赖安装：根据requirements.txt文件安装所有依赖。

数据准备

1. 下载DRIVE数据集：
2. 官网地址：https://drive.grand-challenge.org/
3. 百度云链接：https://pan.baidu.com/s/1Tjkrx2B9FgoJk0KviA-rDw 密码: 8no8

训练步骤

1. 单GPU训练：使用train.py脚本进行训练，确保设置--data-path参数为DRIVE数据集的根目录。
2. 多GPU训练：使用torchrun --nproc_per_node=8 train_multi_GPU.py指令进行多GPU训练，可通过CUDA_VISIBLE_DEVICES=0,3指定使用的GPU设备。

预测步骤

使用predict.py脚本进行预测，设置weights_path参数为训练好的权重路径。

其他功能

使用compute_mean_std.py脚本统计数据集各通道的均值和标准差。

注意事项

• 使用训练脚本时，注意设置--data-path参数。
• 使用预测脚本时，确保设置weights_path参数。
• 使用验证文件时，确保验证集或测试集中包含每个类别的目标。

下载地址

点击下载 【提取码: 4003】【解压密码: www.makuang.net】