### 基于 PyTorch 的叶片图像分割方法及代码示例 #### UNet 模型简介 UNet 是一种经典的卷积神经网络架构，广泛应用于医学影像分析和自然图像分割领域。它通过编码器-解码器结构实现特征提取与上采样重建，适用于多种类型的图像分割任务。 在当前场景下，UNet 已被成功用于西瓜叶片图像分割任务[^1]。以下是具体实现细节： --- #### 代码实现 ##### 1. UNet 模型构建 (`Unet_model.py`) 以下是一个标准的 UNet 构建代码片段，支持自定义输入和输出通道数： ```python import torch.nn as nn import torch class DoubleConv(nn.Module): def __init__(self, in_channels, out_channels): super(DoubleConv, self).__init__() self.conv = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, input_channel=1, output_channel=1): super(UNet, self).__init__() filters = [64, 128, 256, 512] # 下采样路径 (Encoder) self.down_conv1 = DoubleConv(input_channel, filters[0]) self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) self.down_conv2 = DoubleConv(filters[0], filters[1]) self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) self.down_conv3 = DoubleConv(filters[1], filters[2]) self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) self.bottle_neck = DoubleConv(filters[2], filters[3]) # 上采样路径 (Decoder) self.up_transpose1 = nn.ConvTranspose2d(filters[3], filters[2], kernel_size=2, stride=2) self.up_conv1 = DoubleConv(filters[3], filters[2]) self.up_transpose2 = nn.ConvTranspose2d(filters[2], filters[1], kernel_size=2, stride=2) self.up_conv2 = DoubleConv(filters[2], filters[1]) self.up_transpose3 = nn.ConvTranspose2d(filters[1], filters[0], kernel_size=2, stride=2) self.up_conv3 = DoubleConv(filters[1], filters[0]) self.final_layer = nn.Conv2d(filters[0], output_channel, kernel_size=1) def forward(self, x): conv1_out = self.down_conv1(x) pool1_out = self.pool1(conv1_out) conv2_out = self.down_conv2(pool1_out) pool2_out = self.pool2(conv2_out) conv3_out = self.down_conv3(pool2_out) pool3_out = self.pool3(conv3_out) bottle_neck_out = self.bottle_neck(pool3_out) up_sample1_out = self.up_transpose1(bottle_neck_out) concat1_out = torch.cat([up_sample1_out, conv3_out], dim=1) up_conv1_out = self.up_conv1(concat1_out) up_sample2_out = self.up_transpose2(up_conv1_out) concat2_out = torch.cat([up_sample2_out, conv2_out], dim=1) up_conv2_out = self.up_conv2(concat2_out) up_sample3_out = self.up_transpose3(up_conv2_out) concat3_out = torch.cat([up_sample3_out, conv1_out], dim=1) up_conv3_out = self.up_conv3(concat3_out) final_output = self.final_layer(up_conv3_out) return final_output ``` --- ##### 2. 数据加载模块 (`loaddata.py`) 利用 `torch.utils.data.Dataset` 和 `DataLoader` 加载数据集。注意调整图片尺寸至 `(512, 512)` 并转换为灰度图。 ```python from torchvision import transforms from torch.utils.data import Dataset, DataLoader from PIL import Image import os class DealDataset(Dataset): def __init__(self, image_dir, mask_dir, transform=None): self.image_dir = image_dir self.mask_dir = mask_dir self.transform = transform self.images = sorted(os.listdir(image_dir)) self.masks = sorted(os.listdir(mask_dir)) def __len__(self): return len(self.images) def __getitem__(self, idx): img_path = os.path.join(self.image_dir, self.images[idx]) mask_path = os.path.join(self.mask_dir, self.masks[idx]) 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, mask transform = transforms.Compose([ transforms.Resize((512, 512)), transforms.ToTensor() ]) dataset = DealDataset(image_dir='path_to_images', mask_dir='path_to_masks', transform=transform) train_loader = DataLoader(dataset, batch_size=4, shuffle=True) ``` --- ##### 3. 训练过程 (`train.py`) 采用 BCEWithLogitsLoss 作为损失函数，并使用 RMSprop 进行优化。 ```python import torch.optim as optim from torch.nn import BCEWithLogitsLoss def train(model, device, train_loader, optimizer, criterion): model.train() running_loss = 0.0 for images, masks in train_loader: images, masks = images.to(device), masks.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, masks.unsqueeze(1)) # 添加维度匹配 loss.backward() optimizer.step() running_loss += loss.item() epoch_loss = running_loss / len(train_loader) return epoch_loss device = 'cuda' if torch.cuda.is_available() else 'cpu' model = UNet().to(device) optimizer = optim.RMSprop(model.parameters(), lr=0.001) criterion = BCEWithLogitsLoss() for epoch in range(10): # 设置适当轮次 loss = train(model, device, train_loader, optimizer, criterion) print(f"Epoch {epoch+1}, Loss: {loss}") ``` --- ##### 4. 准确率评估 (`accuracy.py`) 使用 MIOU（Mean Intersection over Union）衡量模型性能。 ```python def calculate_iou(pred_mask, true_mask): intersection = torch.sum(torch.logical_and(pred_mask, true_mask)).item() union = torch.sum(torch.logical_or(pred_mask, true_mask)).item() iou = intersection / union if union != 0 else 0 return iou def evaluate_miou(model, dataloader, device): model.eval() total_iou = 0.0 with torch.no_grad(): for images, masks in dataloader: images, masks = images.to(device), masks.to(device) preds = model(images) pred_masks = (preds > 0.5).float() # 阈值处理 miou = calculate_iou(pred_masks.squeeze(1), masks) total_iou += miou avg_miou = total_iou / len(dataloader) return avg_miou miou_score = evaluate_miou(model, test_loader, device) print(f"MIOU Score: {miou_score:.4f}") ``` --- ### 总结 上述代码展示了如何基于 PyTorch 使用 UNet 实现叶片图像分割的任务。完整的流程包括模型设计、数据预处理、训练以及准确率评估[^1][^2]。