使用深度学习框架（例如PyTorch）进行胸腹部骨骼分割任务，并提供详细的训练代码和数据集准备步骤 ——胸腹部骨骼分割CT数据集_
胸腹部骨骼分割CT数据集:
（脊椎和肋骨的截面切片）
包含原图和标签各2832张
如何使用深度学习框架（例如PyTorch）进行胸腹部骨骼分割任务，并提供详细的训练代码和数据集准备步骤。假设你已经有一个包含2832张CT图像及其对应标签的数据集。
项目结构
深色版本
skeletal_segmentation/
├── dataset/
│ ├── images/
│ │ └── *.png
│ ├── masks/
│ │ └── *.png
├── models/
│ └── unet.py
├── src/
│ ├── train.py
│ ├── predict.py
│ ├── utils.py
├── weights/
│ └── best_model.pth
├── requirements.txt
└── README.md
- 安装依赖
首先，确保你已经安装了必要的库。创建一个requirements.txt文件，内容如下：
深色版本
torch
torchvision
numpy
pandas
matplotlib
scikit-image
tqdm
然后，使用以下命令安装依赖：
bash
深色版本
pip install -r requirements.txt
2. 数据集准备
确保你的数据集已经按照以下结构组织：
深色版本
dataset/
├── images/
│ └── *.png
└── masks/
└── *.png
每个文件夹中包含对应的图像文件和标签文件。确保所有图像文件和标签文件的名字一致，例如image_0001.png和mask_0001.png。
- 数据集配置
创建一个数据集类，用于加载和预处理数据。
3.1 src/utils.py
python
深色版本
import os
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import numpy as np
class SkeletalSegmentationDataset(Dataset):
definit(self, image_dir, mask_dir, transform=None):
self.image_dir = image_dir
self.mask_dir = mask_dir
self.transform = transform
self.images = os.listdir(image_dir)
def __len__(self): return len(self.images) def __getitem__(self, index): img_path = os.path.join(self.image_dir, self.images[index]) mask_path = os.path.join(self.mask_dir, self.images[index]) image = Image.open(img_path).convert("L") mask = Image.open(mask_path).convert("L") if self.transform: image = self.transform(image) mask = self.transform(mask) return image, maskdef get_data_loaders(image_dir, mask_dir, batch_size=16, num_workers=4):
transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
])
dataset = SkeletalSegmentationDataset(image_dir, mask_dir, transform=transform) train_size = int(0.8 * len(dataset)) val_size = int(0.1 * len(dataset)) test_size = len(dataset) - train_size - val_size train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, val_size, test_size]) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers) return train_loader, val_loader, test_loader- 模型定义
定义一个U-Net模型用于分割任务。
4.1 models/unet.py
python
深色版本
import torch
import torch.nn as nn
import torch.nn.functional as F
class UNet(nn.Module):
definit(self, in_channels=1, out_channels=1, init_features=32):
super(UNet, self).init()
features = init_features self.encoder1 = self._block(in_channels, features, name="enc1") self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) self.encoder2 = self._block(features, features * 2, name="enc2") self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) self.encoder3 = self._block(features * 2, features * 4, name="enc3") self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) self.encoder4 = self._block(features * 4, features * 8, name="enc4") self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2) self.bottleneck = self._block(features * 8, features * 16, name="bottleneck") self.upconv4 = nn.ConvTranspose2d(features * 16, features * 8, kernel_size=2, stride=2) self.decoder4 = self._block((features * 8) * 2, features * 8, name="dec4") self.upconv3 = nn.ConvTranspose2d(features * 8, features * 4, kernel_size=2, stride=2) self.decoder3 = self._block((features * 4) * 2, features * 4, name="dec3") self.upconv2 = nn.ConvTranspose2d(features * 4, features * 2, kernel_size=2, stride=2) self.decoder2 = self._block((features * 2) * 2, features * 2, name="dec2") self.upconv1 = nn.ConvTranspose2d(features * 2, features, kernel_size=2, stride=2) self.decoder1 = self._block(features * 2, features, name="dec1") self.conv = nn.Conv2d(in_channels=features, out_channels=out_channels, kernel_size=1) def forward(self, x): enc1 = self.encoder1(x) enc2 = self.encoder2(self.pool1(enc1)) enc3 = self.encoder3(self.pool2(enc2)) enc4 = self.encoder4(self.pool3(enc3)) bottleneck = self.bottleneck(self.pool4(enc4)) dec4 = self.upconv4(bottleneck) dec4 = torch.cat((dec4, enc4), dim=1) dec4 = self.decoder4(dec4) dec3 = self.upconv3(dec4) dec3 = torch.cat((dec3, enc3), dim=1) dec3 = self.decoder3(dec3) dec2 = self.upconv2(dec3) dec2 = torch.cat((dec2, enc2), dim=1) dec2 = self.decoder2(dec2) dec1 = self.upconv1(dec2) dec1 = torch.cat((dec1, enc1), dim=1) dec1 = self.decoder1(dec1) return torch.sigmoid(self.conv(dec1)) def _block(self, in_channels, features, name): return nn.Sequential( nn.Conv2d(in_channels, features, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(features), nn.ReLU(inplace=True), nn.Conv2d(features, features, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(features), nn.ReLU(inplace=True), )- 训练代码
5.1 src/train.py
python
深色版本
import torch
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from models.unet import UNet
from src.utils import get_data_loaders
def train_model(image_dir, mask_dir, epochs=100, batch_size=16, learning_rate=1e-4):
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
model = UNet().to(device) criterion = nn.BCELoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) train_loader, val_loader, _ = get_data_loaders(image_dir, mask_dir, batch_size=batch_size) writer = SummaryWriter() for epoch in range(epochs): model.train() running_loss = 0.0 for images, masks in tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs}"): images = images.to(device) masks = masks.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, masks) loss.backward() optimizer.step() running_loss += loss.item() train_loss = running_loss / len(train_loader) writer.add_scalar('Training Loss', train_loss, epoch) model.eval() running_val_loss = 0.0 with torch.no_grad(): for images, masks in val_loader: images = images.to(device) masks = masks.to(device) outputs = model(images) loss = criterion(outputs, masks) running_val_loss += loss.item() val_loss = running_val_loss / len(val_loader) writer.add_scalar('Validation Loss', val_loss, epoch) print(f"Epoch {epoch + 1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}") torch.save(model.state_dict(), "weights/best_model.pth") writer.close()ifname== “main”:
image_dir = “dataset/images”
mask_dir = “dataset/masks”
train_model(image_dir, mask_dir)
6. 模型评估
训练完成后，可以通过验证集和测试集来评估模型的性能。示例如下：
6.1 src/predict.py
python
深色版本
import torch
import matplotlib.pyplot as plt
from models.unet import UNet
from src.utils import get_data_loaders
def predict_and_plot(image_dir, mask_dir, model_path, num_samples=5):
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
model = UNet().to(device) model.load_state_dict(torch.load(model_path)) model.eval() _, _, test_loader = get_data_loaders(image_dir, mask_dir) fig, axes = plt.subplots(num_samples, 3, figsize=(15, 5 * num_samples)) with torch.no_grad(): for i, (images, masks) in enumerate(test_loader): if i >= num_samples: break images = images.to(device) masks = masks.to(device) outputs = model(images) outputs = (outputs > 0.5).float() images = images.cpu().numpy().squeeze() masks = masks.cpu().numpy().squeeze() outputs = outputs.cpu().numpy().squeeze() axes[i, 0].imshow(images, cmap='gray') axes[i, 0].set_title("Input Image") axes[i, 0].axis('off') axes[i, 1].imshow(masks, cmap='gray') axes[i, 1].set_title("Ground Truth Mask") axes[i, 1].axis('off') axes[i, 2].imshow(outputs, cmap='gray') axes[i, 2].set_title("Predicted Mask") axes[i, 2].axis('off') plt.tight_layout() plt.show()ifname== “main”:
image_dir = “dataset/images”
mask_dir = “dataset/masks”
model_path = “weights/best_model.pth”
predict_and_plot(image_dir, mask_dir, model_path)
7. 运行项目
确保你的数据集已经放在相应的文件夹中。
在项目根目录下运行以下命令启动训练：
bash
深色版本
python src/train.py
训练完成后，运行以下命令进行评估和可视化：
bash
深色版本
python src/predict.py
8. 功能说明
训练模型：train.py脚本用于训练U-Net模型，使用BCE损失函数和Adam优化器。
评估模型：predict.py脚本用于评估模型性能，并可视化输入图像、真实标签和预测结果。
9. 详细注释
utils.py
数据集类：定义了一个SkeletalSegmentationDataset类，用于加载和预处理数据。
数据加载器：定义了一个get_data_loaders函数，用于创建训练、验证和测试数据加载器。
unet.py
U-Net模型：定义了一个U-Net模型，用于分割任务。
train.py
训练函数：定义了一个train_model函数，用于训练U-Net模型。
训练过程：在每个epoch中，模型在训练集上进行前向传播和反向传播，并在验证集上进行评估。
predict.py
预测和可视化：定义了一个predict_and_plot函数，用于在测试集上进行预测，并可视化输入图像、真实标签和预测结果。
