使用 Mask R-CNN 进行血细胞分割

介绍

血细胞分析是诊断各种医学疾病的重要步骤，从感染和贫血到更严重的疾病如白血病。传统上，这一过程是通过老方法进行的——实验室技术员通过显微镜查看血涂片玻片，花费几个小时。这一过程不仅令人乏味，还容易出现人为错误，尤其是在处理大量样本或复杂病例时。

难怪医疗专业人员一直渴望自动化这一重要分析。借助计算机视觉和深度学习算法的力量，我们可以以更高的准确性和效率处理血细胞检查。一项改变这一应用的技术是图像分割——本质上是从图像的周围区域中挑选出并分离单个细胞。

图像分割与 Mask R-CNN

图像分割涉及将图像分成几个片段或区域，每个片段或区域代表图像中单独的对象或对象的一部分。此过程对于获取有价值的数据和理解图像的内容至关重要。语义分割和实例分割是分割的两个基本类别。

语义分割：语义分割为图像中的每个像素分配一个类标签，而不区分同一类的不同实例。
实例分割：实例分割为像素分配类标签，这有助于区分同一类的许多实例。

图像分割的应用多种多样，包括医学成像（如肿瘤检测和器官描绘）到自动驾驶（识别和跟踪行人和车辆等物体）、卫星图像（土地覆盖分类）和增强现实。

Mask R-CNN 简介及其在实例分割中的作用

现代深度学习模型（如 Mask R-CNN（基于掩码区域的卷积神经网络））用于处理实例分割。它在每个感兴趣区域（RoI）上增加了分割掩码预测分支，扩展了用于对象检测的Faster R-CNN模型。通过这一新增强，Mask R-CNN现在可以通过检测图像中的对象并为每个对象生成像素级的掩码来实现实例分割。

Mask R-CNN是一种在需要精确对象边界的应用中非常成功的方法，例如在医学影像中分割血液样本中的不同类型细胞。它在正确识别和勾画图像中的特定对象方面表现出色。

Mask R-CNN 架构和关键组件概述

Mask R-CNN 架构建立在 Faster R-CNN 框架之上，并包含几个关键组件：

主干网络：通常为深度卷积神经网络（例如 ResNet 或 ResNeXt），充当特征提取器。该网络处理输入图像并生成特征图。
区域提议网络 (RPN)：此组件生成区域提议，即特征图中可能包含对象的潜在区域。RPN 是一个轻量级神经网络，可输出这些区域的边界框和对象性分数。
RoI Align：对 RoI Pooling 的改进，RoI Align 通过避免量化问题，准确地从建议的感兴趣区域中提取特征，确保特征的精确对齐。
边界框头：一个完全连接的网络，采用 RoI 特征并执行对象分类和边界框回归以细化初始区域提议。
Mask Head：一个小型卷积网络，采用 RoI 特征并预测每个对象的二进制掩码，在像素级别对对象进行分割。

这些组件的集成使 Mask R-CNN 能够有效检测物体并生成高质量的分割蒙版，使其成为执行详细而准确的实例分割任务的强大工具。血细胞分割等医疗应用尤其受益于此架构，其中精确的物体边界对于准确分析和诊断至关重要。

使用 Mask R-CNN 实现血细胞分割

现在让我们实现 Mask RCNN 进行血细胞分割。

步骤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