使用 Mask R-CNN 实现分割
步骤1. 导入依赖项
import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import Compose, ToTensor, Resize
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
步骤2. 设置种子
设置种子将确保我们每次运行代码时都会获得相同的随机生成。
seed = 42
np.random.seed(seed)
torch.manual_seed(seed)
步骤3. 定义文件路径
初始化图像路径和用于检索图像的目标（掩码）。
images_dir = '/content/images_BloodCellSegmentation'
targets_dir = '/content/targets_BloodCellSegmentation'
步骤4. 定义自定义数据集类
BloodCellSegDataset：创建一个自定义数据集类，用于加载和预处理血细胞图像及其掩模。
init：此构造函数通过列出所有图像文件名，并为图像和蒙版构建完整路径来初始化数据集。
getitem：该函数加载
-
图像及其掩码，对掩码进行预处理以创建二进制掩码
-
计算边界框
-
调整图像和蒙版的大小
-
应用转换。
len：此函数返回数据集中的图像总数。
class BloodCellSegDataset(Dataset):
def __init__(self, images_dir, masks_dir):
self.image_names = os.listdir(images_dir)
self.images_paths = [os.path.join(images_dir, image_name) for image_name in self.image_names]
self.masks_paths = [os.path.join(masks_dir, image_name.split('.')[0] + '.png') for image_name in self.image_names]
def __getitem__(self, idx):
image = Image.open(self.images_paths[idx])
mask = Image.open(self.masks_paths[idx])
mask = np.array(mask)
mask = ((mask == 128) | (mask == 255))
get_x = (mask.sum(axis=0) > 0).astype(int)
get_y = (mask.sum(axis=1) > 0).astype(int)
x1, x2 = get_x.argmax(), get_x.shape[0] - get_x[::-1].argmax()
y1, y2 = get_y.argmax(), get_y.shape[0] - get_y[::-1].argmax()
boxes = torch.as_tensor([[x1, y1, x2, y2]], dtype=torch.float32)
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
mask = Image.fromarray(mask)
label = torch.ones((1,), dtype=torch.int64)
image_id = torch.tensor([idx])
iscrowd = torch.zeros((1,), dtype=torch.int64)
transform = Compose([Resize(224), ToTensor()])
boxes *= (224 / image.size[0])
image = transform(image)
mask = transform(mask)
target = {'masks': mask, 'labels': label, 'boxes': boxes, "image_id": image_id, "area": area, "iscrowd": iscrowd}
return image, target
def __len__(self):
return len(self.image_names)
步骤5. 创建DataLoader
collate_fn：此函数用于处理批量数据，确保格式正确。
DataLoader：用于创建 pytorch 数据加载器
-
处理批处理
-
改组
-
并行加载数据。
def collate_fn(batch):
return tuple(zip(*batch))
dataset = BloodCellSegDataset(images_dir, targets_dir)
data_loader = DataLoader(dataset, batch_size=8, num_workers=2, shuffle=True, collate_fn=collate_fn)
步骤6. 定义和修改模型
maskrcnn_resnet50_fpn：这将加载一个预先训练的 Mask R-CNN 模型，该模型具有 ResNet-50 主干和特征金字塔网络 (FPN)。
num_classes：这设置了我们的数据集中的类别数量。
FastRCNNPredictor：这取代了适合自定义类别数量的分类头。
MaskRCNNPredictor：这取代了适合自定义类别数量的掩码预测头。
model = maskrcnn_resnet50_fpn(pretrained=True)
num_classes = 2
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
num_filters = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, num_filters, num_classes)
步骤7. 训练模型
model.to（“cuda”）：这将我们的模型转移到 GPU 以加速训练。
torch.optim.Adam：这定义了我们的优化器，用于更新模型参数。
model.train()：将模型设置为训练模式并使其能够改变权重。
训练循环：
-
我们经历了多个时期的迭代。
-
一批批图像和目标被传输到 GPU。
-
该模型清除下一个时期的梯度，并传递图像来计算损失。
-
损失反向传播，模型参数更新。
-
计算并打印每个时期的平均损失。
model = model.to("cuda")
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
model.train()
for epoch in range(10):
epoch_loss = cnt = 0
for batch_x, batch_y in tqdm(data_loader):
batch_x = list(image.to("cuda") for image in batch_x)
batch_y = [{k: v.to("cuda") for k, v in t.items()} for t in batch_y]
optimizer.zero_grad()
loss_dict = model(batch_x, batch_y)
losses = sum(loss for loss in loss_dict.values())
losses.backward()
optimizer.step()
epoch_loss += loss_dict['loss_mask'].item()
cnt += 1
epoch_loss /= cnt
print("Training loss for epoch {} is {} ".format(epoch + 1, epoch_loss))
步骤8. 评估模型
-
我们加载一个示例图像及其原始蒙版。
-
我们对图像和蒙版应用变换。
-
我们将模型设置为评估模式，这样模型就不会计算梯度。
-
我们将图像传入模型以获得预测的掩码。
-
最后，我们使用 Matplotlib 将原始和预测的蒙版可视化。
image = Image.open('/content/images_BloodCellSegmentation/002.bmp')
gt_mask = Image.open('/content/targets_BloodCellSegmentation/002.png')
gt_mask = np.array(gt_mask)
gt_mask = ((gt_mask == 128) | (gt_mask == 255))
gt_mask = Image.fromarray(gt_mask)
transform = Compose([Resize(224), ToTensor()])
image = transform(image)
gt_mask = transform(gt_mask)
model.eval()
output = model(image.unsqueeze(dim=0).to('cuda'))
output = output[0]['masks'][0].cpu().detach().numpy()
plt.imshow(gt_mask.squeeze(), cmap='gray')
plt.imshow((output.squeeze() > 0.5).astype(int), cmap='gray')
步骤9. 计算交并比（IoU）
IoU计算：
-
在这里，我们将预测的和原始的蒙版压平。
-
然后我们计算预测掩码和原始掩码的交集和并集。
-
现在我们计算 IoU 分数，这是评估分割性能的指标。
mask = (output.squeeze() > 0.5).astype(int)
pred = mask.ravel().copy()
gt_mask = gt_mask.numpy()
target = gt_mask.ravel().copy().astype(int)
pred_inds = pred == 1
target_inds = target == 1
intersection = pred_inds[target_inds].sum()
union = pred_inds.sum() + target_inds.sum() - intersection
iou = (float(intersection) / float(max(union, 1)))
iou