PyTorch 现代计算机视觉(五)

九、图像分割

在前一章中,我们学习了如何检测图像中出现的对象,以及与检测到的对象相对应的类。在这一章中,我们将更进一步,不仅在对象周围画一个边界框,而且识别包含对象的确切像素。除此之外,到本章结束时,我们将能够挑出属于同一个类的实例/对象。

在本章中,我们将通过研究 U-Net 和 Mask R-CNN 架构来了解语义分段和实例分段。具体来说,我们将涵盖以下主题:

  • 探索 U-Net 架构
  • 用 U-Net 实现语义切分
  • 探索掩模 R-CNN 架构
  • 使用掩模 R-CNN 实现实例分割

我们试图通过图像分割(arxiv.org/pdf/1405.0312.pdf)实现的简洁图像如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

我们开始吧!

探索 U-Net 架构

想象一下这样一个场景,给你一张图像,要求你预测哪个像素对应哪个对象。到目前为止,当我们预测对象的类别和对应于该对象的边界框时,我们通过网络传递图像,然后通过主干架构(如 VGG 或 ResNet)传递图像,在某一层平坦化输出,并在预测类别和边界框偏移之前连接附加的密集层。但是,在图像分割的情况下,输出形状与输入图像的形状相同,展平卷积的输出然后重建图像可能会导致信息丢失。此外,在图像分割的情况下,原始图像中存在的轮廓和形状在输出图像中不会变化,因此,当我们执行分割时,我们迄今为止处理的网络(展平最后一层并连接附加的密集层)并不是最佳的。

在本节中,我们将了解如何执行图像分割。

在执行分段时,我们需要记住以下两个方面:

  • 原始图像中对象的形状和结构在分段输出中保持不变。
  • 利用完全卷积架构(而不是我们展平某一层的结构)会有所帮助,因为我们使用一个图像作为输入,另一个图像作为输出。

U-Net 架构帮助我们实现了这一目标。U-Net 的典型表示如下(输入图像是 3×96×128 的形状,而图像中存在的类的数量是 21;这意味着输出包含 21 个通道):

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

前述架构因其“ U ”状的形状而被称为 U 网架构

在上图的左半部分,我们可以看到图像通过卷积层,正如我们在前面章节中看到的那样,图像尺寸不断减小,而通道数量不断增加。然而,在右半部分,我们可以看到我们正在放大缩小的图像,回到原始的高度和宽度,但是通道的数量与类的数量一样多。

此外,在向上扩展的同时,我们还使用跳过连接来利用来自左半部分中相应层的信息,以便我们可以保留原始图像中的结构/对象。

这样,U-Net 架构学习保留原始图像的结构(和对象的形状),同时利用卷积的特征来预测对应于每个像素的类别。

一般来说,输出中的通道数量与我们想要预测的类别数量一样多。

执行升级

在 U-Net 架构中,使用nn.ConvTranspose2d方法执行向上扩展,该方法将输入通道的数量、输出通道的数量、内核大小和步幅作为输入参数。ConvTranspose2d的计算示例如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

在前面的例子中,我们采用了一个形状为 3 x 3 的输入数组(输入数组),应用了一个步长为 2 的输入数组,其中我们分配输入值以适应步长(输入数组针对步长进行了调整),用零填充数组(输入数组针对步长和填充进行了调整),并用一个过滤器(过滤器/内核)对填充后的输入进行卷积以获取输出数组。

通过利用填充和步幅的组合,我们将形状为 3 x 3 的输入升级为形状为 6 x 6 的数组。虽然前面的示例仅用于说明目的,但是最佳滤波器值学习(因为滤波器权重和偏差在模型训练过程中被优化)以尽可能多地重建原始图像。

nn.ConvTranspose2d中的超参数如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

为了理解nn.ConvTranspose2d如何帮助提升一个数组,让我们看一下下面的代码:

  1. 导入相关包:
import torch
import torch.nn as nn
  1. nn.ConvTranspose2d方法初始化网络m:
m = nn.ConvTranspose2d(1, 1, kernel_size=(2,2), 
                       stride=2, padding = 0)

在前面的代码中,我们指定输入通道的值是1,输出通道的值是1,内核的大小是(2,2),步幅是2,填充是0

在内部,填充被计算为膨胀* (kernel_size - 1) -填充。

因此 1*(2-1)-0 = 1,其中我们向输入数组的两个维度添加零填充 1。

  1. 初始化一个输入数组,并通过模型传递它:
input = torch.ones(1, 1, 3, 3)
output = m(input)
output.shape

前面的代码产生了一个形状1x1x6x6,如前面提供的示例图像所示。

既然我们已经了解了 U-Net 架构是如何工作的,以及nn.ConvTranspose2d是如何帮助提升图像的,那么让我们来实现它,这样我们就可以预测道路场景图像中出现的不同对象。

用 U-Net 实现语义切分

在这一节中,我们将利用 U-Net 架构来预测对应于图像中所有像素的类。这种输入-输出组合的一个例子如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

请注意,在前面的图片中,属于同一类别的对象(在左侧图像-输入图像中)具有相同的像素值(在右侧图像-输出图像中),这就是为什么我们要分割彼此语义相似的像素。这也称为语义分割。

现在,让我们学习如何编码语义分段:

以下代码在本书的 GitHub 知识库的Chapter09文件夹中以Semantic_Segmentation_with_U_Net.ipynb的形式提供-【https://tinyurl.com/mcvp-packt 代码包含下载数据的 URL,长度适中。

  1. 让我们首先下载必要的数据集,安装必要的包,然后导入它们。完成后,我们可以定义设备:
import os
if not os.path.exists('dataset1'):
    !wget -q \
     https://www.dropbox.com/s/0pigmmmynbf9xwq/dataset1.zip
    !unzip -q dataset1.zip
    !rm dataset1.zip
    !pip install -q torch_snippets pytorch_model_summary

from torch_snippets import *
from torchvision import transforms
from sklearn.model_selection import train_test_split
device = 'cuda' if torch.cuda.is_available() else 'cpu'
  1. 定义将用于转换图像的函数(tfms):
tfms = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], 
                                 [0.229, 0.224, 0.225]) 
        ])
  1. 定义数据集类(SegData):
  • __init__方法中指定包含图像的文件夹:
class SegData(Dataset):
    def __init__(self, split):
        self.items=stems(f'dataset1/images_prepped_{split}')
        self.split = split
  • 定义__len__方法:
    def __len__(self):
        return len(self.items)
  • 定义__getitem__方法:
    def __getitem__(self, ix):
        image = read(f'dataset1/images_prepped_{self.split}/\
{self.items[ix]}.png', 1)
        image = cv2.resize(image, (224,224))
        mask=read(f'dataset1/annotations_prepped_{self.split}\
/{self.items[ix]}.png')
        mask = cv2.resize(mask, (224,224))
        return image, mask

__getitem__ 方法中,我们调整了输入(image)和输出(mask)图像的大小,使它们具有相同的形状。请注意,屏蔽图像包含范围在[0,11]之间的整数。这表明有 12 个不同的类别。

  • 定义一个用于选择随机图像索引的函数(choose)(主要用于调试目的):
    def choose(self): return self[randint(len(self))]
  • 定义对一批图像进行预处理的collate_fn方法:
    def collate_fn(self, batch):
        ims, masks = list(zip(*batch))
        ims = torch.cat([tfms(im.copy()/255.)[None] \
                         for im in ims]).float().to(device)
        ce_masks = torch.cat([torch.Tensor(mask[None]) for \
                            mask in masks]).long().to(device)
        return ims, ce_masks

在前面的代码中,我们对所有的输入图像进行预处理,以便在我们转换缩放后的图像后,它们有一个通道(以便每个图像可以在以后通过 CNN 传送)。注意ce_masks是一个长整数张量,类似于交叉熵目标。

  1. 定义训练和验证数据集,以及数据加载器:
trn_ds = SegData('train')
val_ds = SegData('test')
trn_dl = DataLoader(trn_ds, batch_size=4, shuffle=True, \
                    collate_fn=trn_ds.collate_fn)
val_dl = DataLoader(val_ds, batch_size=1, shuffle=True, \
                    collate_fn=val_ds.collate_fn)
  1. 定义神经网络模型:
  • 定义卷积块(conv):
def conv(in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels,out_channels,kernel_size=3, \
                    stride=1, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True)
    )

在前面的conv定义中,我们依次执行Conv2d操作、BatchNorm2d操作和ReLU操作。

  • 定义up_conv块:
def up_conv(in_channels, out_channels):
    return nn.Sequential(
        nn.ConvTranspose2d(in_channels, out_channels, \
                           kernel_size=2, stride=2),
        nn.ReLU(inplace=True)
    )

确保我们放大图像。这不同于Conv2d操作,在这里我们缩小图像的尺寸。它将具有in_channels个通道的图像作为输入通道,并产生具有out_channels个输出通道的图像。

  • 定义网络类(UNet):
from torchvision.models import vgg16_bn
class UNet(nn.Module):
    def __init__(self, pretrained=True, out_channels=12):
        super().__init__()

        self.encoder = \
                vgg16_bn(pretrained=pretrained).features
        self.block1 = nn.Sequential(*self.encoder[:6])
        self.block2 = nn.Sequential(*self.encoder[6:13])
        self.block3 = nn.Sequential(*self.encoder[13:20])
        self.block4 = nn.Sequential(*self.encoder[20:27])
        self.block5 = nn.Sequential(*self.encoder[27:34])

        self.bottleneck = nn.Sequential(*self.encoder[34:])
        self.conv_bottleneck = conv(512, 1024)

        self.up_conv6 = up_conv(1024, 512)
        self.conv6 = conv(512 + 512, 512)
        self.up_conv7 = up_conv(512, 256)
        self.conv7 = conv(256 + 512, 256)
        self.up_conv8 = up_conv(256, 128)
        self.conv8 = conv(128 + 256, 128)
        self.up_conv9 = up_conv(128, 64)
        self.conv9 = conv(64 + 128, 64)
        self.up_conv10 = up_conv(64, 32)
        self.conv10 = conv(32 + 64, 32)
        self.conv11 = nn.Conv2d(32, out_channels, \
                                kernel_size=1)

在前面的__init__方法中,我们定义了将在forward方法中使用的所有层。

  • 定义forward方法:
    def forward(self, x):
        block1 = self.block1(x)
        block2 = self.block2(block1)
        block3 = self.block3(block2)
        block4 = self.block4(block3)
        block5 = self.block5(block4)

        bottleneck = self.bottleneck(block5)
        x = self.conv_bottleneck(bottleneck)

        x = self.up_conv6(x)
        x = torch.cat([x, block5], dim=1)
        x = self.conv6(x)

        x = self.up_conv7(x)
        x = torch.cat([x, block4], dim=1)
        x = self.conv7(x)

        x = self.up_conv8(x)
        x = torch.cat([x, block3], dim=1)
        x = self.conv8(x)

        x = self.up_conv9(x)
        x = torch.cat([x, block2], dim=1)
        x = self.conv9(x)

        x = self.up_conv10(x)
        x = torch.cat([x, block1], dim=1)
        x = self.conv10(x)

        x = self.conv11(x)

        return x

在前面的代码中,我们通过在适当的张量对上使用torch.cat,在向下缩放和向上缩放卷积特征之间建立 U 型连接。

  • 定义一个函数(UNetLoss)来计算我们的损失和准确度值:
ce = nn.CrossEntropyLoss()
def UnetLoss(preds, targets):
    ce_loss = ce(preds, targets)
    acc = (torch.max(preds, 1)[1] == targets).float().mean()
    return ce_loss, acc
  • 定义一个函数,该函数将在批处理(train_batch)上进行训练,并在验证数据集(validate_batch)上计算指标:
def train_batch(model, data, optimizer, criterion):
    model.train()
    ims, ce_masks = data
    _masks = model(ims)
    optimizer.zero_grad()
    loss, acc = criterion(_masks, ce_masks)
    loss.backward()
    optimizer.step()
    return loss.item(), acc.item()

@torch.no_grad()
def validate_batch(model, data, criterion):
    model.eval()
    ims, masks = data
    _masks = model(ims)
    loss, acc = criterion(_masks, masks)
    return loss.item(), acc.item()
  • 定义模型、优化器、损失函数和历元数:
model = UNet().to(device)
criterion = UnetLoss
optimizer = optim.Adam(model.parameters(), lr=1e-3)
n_epochs = 20
  1. 在不断增加的时期内训练模型:
log = Report(n_epochs)
for ex in range(n_epochs):
    N = len(trn_dl)
    for bx, data in enumerate(trn_dl):
        loss, acc = train_batch(model, data, optimizer, \
                                criterion)
        log.record(ex+(bx+1)/N,trn_loss=loss,trn_acc=acc, \
                                 end='\r')

    N = len(val_dl)
    for bx, data in enumerate(val_dl):
        loss, acc = validate_batch(model, data, criterion)
        log.record(ex+(bx+1)/N,val_loss=loss,val_acc=acc, \
                                 end='\r')

    log.report_avgs(ex+1)
  1. 绘制递增时期的训练值、验证损失值和准确度值:
log.plot_epochs(['trn_loss','val_loss'])

上述代码生成以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

  1. 计算新图像的预测输出:
  • 获取新图像的模型预测:
im, mask = next(iter(val_dl))
_mask = model(im)
  • 获取概率最高的频道:
_, _mask = torch.max(_mask, dim=1)
  • 显示原始图像和预测图像:
subplots([im[0].permute(1,2,0).detach().cpu()[:,:,0], \
          mask.permute(1,2,0).detach().cpu()[:,:,0], \
          _mask.permute(1,2,0).detach().cpu()[:,:,0]],nc=3, \
          titles=['Original image','Original mask', \
          'Predicted mask'])

上述代码生成以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

从上图中,我们可以看到,我们可以使用 U-Net 架构成功地生成一个分段掩码。但是,同一类的所有实例将具有相同的预测像素值。如果我们想在图像中分离出Person类的实例呢?在下一节中,我们将学习 Mask R-CNN 架构,它有助于生成实例级掩码,以便我们可以区分实例(甚至是同一类的实例)。

探索掩模 R-CNN 架构

Mask R-CNN 架构有助于识别/突出显示图像中给定类别的对象的实例。当图像中存在多个相同类型的对象时,这尤其方便。此外,术语 Mask 表示 Mask R-CNN 在像素级完成的分割。

Mask R-CNN 架构是更快的 R-CNN 网络的扩展,我们在前一章已经了解过。但是,对 Mask R-CNN 架构进行了一些修改,如下所示:

  • RoI Pooling 层已被 RoI Align 层所取代。
  • 除了已经预测最终层中的对象类别和边界框校正的头部之外,还包括了用于预测对象遮罩的遮罩头部。
  • 一个全卷积网络 ( FCN )被用于掩码预测。

在我们了解每个组件如何工作之前,让我们快速浏览一下 Mask R-CNN 中发生的事件(图片来源:arxiv.org/pdf/1703.06870.pdf):

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

在上图中,请注意,我们从一个层获取类和边界框信息,从另一个层获取遮罩信息。

Mask R-CNN 架构的工作细节如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

在我们实现 Mask R-CNN 架构之前,我们需要了解它的组件。我们将从 RoI 对齐开始。

ROI align(ROI 对齐)

通过更快的 R-CNN,我们了解了投资回报率池。RoI 合并的一个缺点是,当我们执行 RoI 合并操作时,我们可能会丢失某些信息。这是因为在汇集之前,我们可能会在图像的所有区域中均匀地呈现内容。

让我们看一下上一章提供的例子:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

在上图中,区域建议的形状是 5 x 7,我们必须将其转换为 2 x 2 的形状。当将其转换为 2 x 2 形状时(一种称为量化的现象),该区域的一部分与该区域的其他部分相比具有较少的表示。这导致信息丢失,因为该区域的某些部分比其他部分具有更大的权重。RoI Align 帮助解决了这种情况。

为了理解 RoI Align 的工作原理,我们来看一个简单的例子。这里,我们尝试将以下区域(用虚线表示)转换为 2 x 2 形状:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

请注意,该区域(虚线中)并非均匀分布在特征地图中的所有像元上。

我们必须执行以下步骤,以在 2 x 2 形状中获得该区域的合理表示:

  1. 首先,将该区域分成相等的 2 x 2 形状:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

  1. 在每个 2 x 2 单元格内定义四个等距点:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

请注意,在上图中,两个连续点之间的距离是 0.75。

  1. 根据每个点到最近已知值的距离计算每个点的加权平均值:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

  1. 对像元中的所有四个点重复上述插值步骤:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

  1. 在一个单元格内的所有四个点上执行平均汇集:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

通过实现上述步骤,我们在执行 RoI 对齐时不会丢失信息;也就是说,当我们把所有的区域放在同一个形状里面。

面具头

使用 RoI Align,我们可以更准确地表示从区域提案网络中获得的区域提案。现在,我们想要获得分割(掩模)输出,给定标准形状的 RoI 对准输出,用于每个区域提议。

通常,在目标检测的情况下,我们将通过展平层传递 RoI Align,以便预测对象的类别和边界框偏移。然而,在图像分割的情况下,我们预测包含对象的边界框内的像素。因此,我们现在有了第三个输出(除了类和边界框偏移之外),这是感兴趣区域内的预测掩膜。

这里,我们预测的是蒙版,它是覆盖在原始图像上的图像。假设我们预测的是一幅图像,我们将把它连接到另一个卷积层以获得另一个类似图像的结构(宽度 x 高度),而不是展平 RoI Align 的输出。让我们通过下图来了解这一现象:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

在上图中,我们使用特征金字塔网络()获得了 7 x 7 x 2048 形状的输出,该网络现在有两个分支:

*** 第一个分支返回对象的类和边界框,后展平 FPN 输出。

  • 第二个分支在 FPN 输出的顶部执行卷积,以获得掩码。

对应于 14×14 输出的基本事实是区域提议的尺寸调整后的图像。如果数据集中有 80 个唯一类,则区域建议的基本事实是 80 x 14 x 14 的形状。80 x 14 x 14 像素中的每一个都是 1 或 0,这表示该像素是否包含对象。因此,我们在预测像素类别的同时执行二进制交叉熵损失最小化。

在模型训练之后,我们能够检测区域,获得类,获得边界框偏移,并获得对应于每个区域的遮罩。当进行推断时,我们首先检测图像中存在的对象,并进行边界框校正。然后,我们将偏移区域传递给掩模头,以预测该区域中不同像素对应的掩模。

现在我们已经了解了 Mask R-CNN 架构的工作原理,让我们对其进行编码,以便我们可以检测图像中的人物实例。

使用掩模 R-CNN 实现实例分割

为了帮助我们理解如何编码屏蔽 R-CNN 用于实例分割,我们将利用屏蔽图像中存在的人的数据集。我们将使用的数据集是从 ADE20K 数据集的子集创建的,该数据集可在groups.csail.mit.edu/vision/datasets/ADE20K/获得。我们将只使用那些有人类面具的图像。

我们将采取的策略如下:

  1. 获取数据集,然后从中创建数据集和数据加载器。
  2. 以 PyTorch 官方实现 Mask R-CNN 所需的格式创建一个地面真相。
  3. 下载预先训练好的更快的 R-CNN 模型,给它贴上一个口罩 R-CNN 头。
  4. 使用 PyTorch 代码片段训练模型,该代码片段已被标准化用于训练 Mask R-CNN。
  5. 通过首先执行非最大值抑制,然后识别与图像中的人相对应的边界框和遮罩,对图像进行推断。

让我们编写前面的策略:

以下代码在本书的 GitHub 知识库的Chapter09文件夹中以Instance_Segmentation.ipynb的形式提供-【https://tinyurl.com/mcvp-packt 代码包含下载数据的 URL,长度适中。我们强烈建议您在 GitHub 中执行 notebook 以重现结果,同时理解执行的步骤和对文本中各种代码组件的解释。

  1. 从 GitHub 导入相关的数据集和训练工具:
!wget --quiet \
 http://sceneparsing.csail.mit.edu/data/ChallengeData2017/images.tar
!wget --quiet \ http://sceneparsing.csail.mit.edu/data/ChallengeData2017/annotations_instance.tar
!tar -xf images.tar
!tar -xf annotations_instance.tar
!rm images.tar annotations_instance.tar
!pip install -qU torch_snippets
!wget --quiet \ https://raw.githubusercontent.com/pytorch/vision/master/references/detection/engine.py
!wget --quiet \ https://raw.githubusercontent.com/pytorch/vision/master/references/detection/utils.py
!wget --quiet \ https://raw.githubusercontent.com/pytorch/vision/master/references/detection/transforms.py
!wget --quiet \ https://raw.githubusercontent.com/pytorch/vision/master/references/detection/coco_eval.py
!wget --quiet \ https://raw.githubusercontent.com/pytorch/vision/master/references/detection/coco_utils.py
!pip install -q -U \
    'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
  1. 导入所有必需的包并定义device:
from torch_snippets import *

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

from engine import train_one_epoch, evaluate
import utils
import transforms as T
device = 'cuda' if torch.cuda.is_available() else 'cpu'
  1. 获取包含人物面具的图像,如下所示:
  • 遍历imagesannotations_instance文件夹以获取文件名:
all_images = Glob('img/training')
all_annots = Glob('annotations_instance/training')
  • 检查原始图像和人物实例的面具表示:
f = 'ADE_train_00014301'

im = read(find(f, all_images), 1)
an = read(find(f, all_annots), 1).transpose(2,0,1)
r,g,b = an
nzs = np.nonzero(r==4) # 4 stands for person
instances = np.unique(g[nzs])
masks = np.zeros((len(instances), *r.shape))
for ix,_id in enumerate(instances):
    masks[ix] = g==_id

subplots([im, *masks], sz=20)

上述代码生成以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

从前面的图像中,我们可以看到已经为每个人生成了一个单独的遮罩。这里有四个Person类的实例。

在这个特定的数据集中,地面实况实例注释以这样一种方式提供,即 RGB 中的红色通道对应于对象的类别,而绿色通道对应于实例编号(如果图像中有多个相同类别的对象,如我们这里的示例)。此外,Person类的编码值为 4。

  • 遍历注释并存储至少包含一个人的文件:
annots = []
for ann in Tqdm(all_annots):
    _ann = read(ann, 1).transpose(2,0,1)
    r,g,b = _ann
    if 4 not in np.unique(r): continue
    annots.append(ann)
  • 将文件分为培训和验证文件:
from sklearn.model_selection import train_test_split
_annots = stems(annots)
trn_items,val_items=train_test_split(_annots,random_state=2)
  1. 定义转换方法:
def get_transform(train):
    transforms = []
    transforms.append(T.ToTensor())
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    return T.Compose(transforms)
  1. 创建数据集类(MasksDataset),如下所示:
  • 定义__init__方法,该方法将图像名称(items)、转换方法(transforms)和要考虑的文件数量(N)作为输入:
class MasksDataset(Dataset):
    def __init__(self, items, transforms, N):
        self.items = items
        self.transforms = transforms
        self.N = N
  • 定义一个方法(get_mask),该方法将获取一些与图像中出现的实例相同的遮罩:
    def get_mask(self, path):
        an = read(path, 1).transpose(2,0,1)
        r,g,b = an
        nzs = np.nonzero(r==4)
        instances = np.unique(g[nzs])
        masks = np.zeros((len(instances), *r.shape))
        for ix,_id in enumerate(instances):
            masks[ix] = g==_id
        return masks
  • 获取要返回的图像和相应的目标值。每个人(实例)被视为不同的对象类;也就是说,每个实例是一个不同的类。注意,类似于训练更快的 R-CNN 模型,目标作为张量的字典返回。让我们定义一下__getitem__方法:
    def __getitem__(self, ix):
        _id = self.items[ix]
        img_path = f'img/training/{_id}.jpg'
        mask_path=f'annotations_instance/training/{_id}.png'
        masks = self.get_mask(mask_path)
        obj_ids = np.arange(1, len(masks)+1)
        img = Image.open(img_path).convert("RGB")
        num_objs = len(obj_ids)
  • 除了遮罩本身,遮罩 R-CNN 还需要边界框信息。但是,这很容易准备,如下面的代码所示:
        boxes = []
        for i in range(num_objs):
            obj_pixels = np.where(masks[i])
            xmin = np.min(obj_pixels[1])
            xmax = np.max(obj_pixels[1])
            ymin = np.min(obj_pixels[0])
            ymax = np.max(obj_pixels[0])
            if (((xmax-xmin)<=10) | (ymax-ymin)<=10):
                xmax = xmin+10
                ymax = ymin+10
            boxes.append([xmin, ymin, xmax, ymax])

在前面的代码中,我们通过在边界框的 x 和 y 坐标的最小值上增加 10 个像素来调整存在可疑事实的场景(Person类的高度或宽度小于 10 个像素)。

  • 将所有目标值转换为张量对象:
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.ones((num_objs,), dtype=torch.int64)
        masks = torch.as_tensor(masks, dtype=torch.uint8)
        area = (boxes[:, 3] - boxes[:, 1]) *\
                    (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
        image_id = torch.tensor([ix])
  • 将目标值存储在字典中:
        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd
  • 指定转换方法并返回图像;也就是,target:
        if self.transforms is not None:
            img, target = self.transforms(img, target)
        return img, target
  • 指定__len__方法:
    def __len__(self):
        return self.N
  • 定义将选择随机图像的函数:
    def choose(self):
        return self[randint(len(self))]
  • 检查输入输出组合:
x = MasksDataset(trn_items, get_transform(train=True), N=100)
im,targ = x[0]
inspect(im,targ)
subplots([im, *targ['masks']], sz=10)

以下是上述代码运行时产生的一些输出示例:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

从前面的输出中,我们可以看到面具的形状是 2 x 512 x 683,表明图像中有两个人。

注意,在__getitem__方法中,我们在一个图像中有和图像中存在的对象(实例)一样多的遮罩和边界框。此外,因为我们只有两个类(Background类和Person类),所以我们将Person类指定为 1。

到这一步结束时,我们在输出字典中有了相当多的信息;即对象类别、边界框、遮罩、遮罩区域以及遮罩是否对应于群组。所有这些信息都可以在target字典中找到。对于我们将要使用的训练函数,将数据标准化为torchvision.models.detection.maskrcnn_resnet50_fpn类要求的格式是很重要的。

  1. 接下来,我们需要定义实例分割模型(get_model_instance_segmentation)。我们将使用预训练的模型,其中仅头部被重新初始化以预测两个类别(背景和人)。首先,我们需要初始化一个预先训练好的模型,替换掉box_predictormask_predictor头,这样就可以从头开始学习了:
def get_model_instance_segmentation(num_classes):
    # load an instance segmentation model pre-trained on 
    # COCO
    model = torchvision.models.detection\
                       .maskrcnn_resnet50_fpn(pretrained=True)

    # get number of input features for the classifier
    in_features = model.roi_heads\
                       .box_predictor.cls_score.in_features
    # replace the pre-trained head with a new one
    model.roi_heads.box_predictor = FastRCNNPredictor(\
                                    in_features,num_classes)
    in_features_mask = model.roi_heads\
                       .mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # and replace the mask predictor with a new one
    model.roi_heads.mask_predictor = MaskRCNNPredictor(\
                                      in_features_mask,\
                                   hidden_layer, num_classes)
    return model

FastRCNNPredictor期望两个输入—in_features(输入通道数)和num_classes(类别数)。根据要预测的类的数量,计算出边界框预测的数量,这是类数量的四倍。

MaskRCNNPredictor期望三个输入—in_features_mask(输入通道数)、hidden_layer(输出通道数)、和num_classes(要预测的类数)。

可以通过指定以下内容来获取已定义模型的详细信息:

model = get_model_instance_segmentation(2).to(device)
model

模型的下半部分(即没有主干)如下所示:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

请注意,更快的 R-CNN 网络(我们在上一章中训练过)和屏蔽 R-CNN 模型之间的主要区别在于roi_heads模块,它本身包含多个子模块。让我们看看他们执行什么任务:

  • roi_heads:对齐从 FPN 网络获取的输入,并创建两个张量。

  • box_predictor:使用我们获得的输出来预测每个 RoI 的类别和边界框偏移。

  • mask_roi_pool : RoI 随后整理来自 FPN 网络的输出。

  • mask_head:将先前获得的校准输出转换成可用于预测掩模的特征图。

  • mask_predictor:获取mask_head的输出,并预测最终掩码。

  1. 获取与训练和验证图像对应的数据集和数据加载器:
dataset = MasksDataset(trn_items, get_transform(train=True), \
                                                    N=3000)
dataset_test = MasksDataset(val_items, \
                           get_transform(train=False), N=800)

# define training and validation data loaders
data_loader=torch.utils.data.DataLoader(dataset,batch_size=2, \
                                shuffle=True, num_workers=0, \
                                 collate_fn=utils.collate_fn)

data_loader_test = torch.utils.data.DataLoader(dataset_test, \
                                batch_size=1, shuffle=False, \
                   num_workers=0,collate_fn=utils.collate_fn)
  1. 定义模型、参数和优化标准:
num_classes = 2
model = get_model_instance_segmentation(\
                        num_classes).to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, \
                            momentum=0.9,weight_decay=0.0005)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, \
                                                step_size=3, \
                                                gamma=0.1)

定义的预训练模型架构将图像和targets字典作为输入,以减少损失。通过运行以下命令,可以看到从模型接收到的输出示例:

# The following code is for illustration purpose only
model.eval()
pred = model(dataset[0][0][None].to(device))
inspect(pred[0])

上述代码会产生以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

在这里,我们可以看到一个具有边界框(BOXES)的字典,对应于边界框的类(LABELS),对应于类预测的置信度得分(SCORES,以及我们的掩码实例的位置(MASKS)。如您所见,该模型被硬编码为返回 100 个预测,这是合理的,因为我们不应该期望在一个典型的图像中有超过 100 个对象。

要获取已检测到的实例数量,我们将使用以下代码:

# The following code is for illustration purpose only
pred[0]['masks'].shape
# torch.Size([100, 1, 536, 559])

前面的代码为一个图像(以及对应于该图像的尺寸)获取最多 100 个遮罩实例(其中这些实例对应于非背景类)。对于这 100 个实例,它还将返回相应的类标签、边界框和该类的 100 个相应的置信度值。

  1. 在不断增加的时期内训练模型:
num_epochs = 5

trn_history = []
for epoch in range(num_epochs):
    # train for one epoch, printing every 10 iterations
    res = train_one_epoch(model, optimizer, data_loader, \
                          device, epoch, print_freq=10)
    trn_history.append(res)
    # update the learning rate
    lr_scheduler.step()
    # evaluate on the test dataset
    res = evaluate(model, data_loader_test, device=device)

通过这样做,我们现在可以将我们的遮罩覆盖在图像中的人物上。我们可以按如下方式记录我们在不断增加的时期内的训练损失变化:

import matplotlib.pyplot as plt
plt.title('Training Loss')
losses =[np.mean(list(trn_history[i].meters['loss'].deque)) \
            for i in range(len(trn_history))]
plt.plot(losses)

上述代码会产生以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

  1. 在测试图像上预测:
model.eval()
im = dataset_test[0][0]
show(im)
with torch.no_grad():
    prediction = model([im.to(device)])
    for i in range(len(prediction[0]['masks'])):
        plt.imshow(Image.fromarray(prediction[0]['masks']\
                      [i, 0].mul(255).byte().cpu().numpy()))
        plt.title('Class: '+str(prediction[0]['labels']\
                   [i].cpu().numpy())+' Score:'+str(\
                  prediction[0]['scores'][i].cpu().numpy()))
        plt.show()

上述代码会产生以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

从前面的图像中,我们可以看到我们可以成功地识别图像中的四个人。此外,该模型预测图像中的多个其他片段(我们在前面的输出中没有显示),尽管这是低置信度的。

既然模型可以很好地检测实例,那么让我们对提供的数据集中不存在的自定义图像运行预测。

  1. 对自己的新形象进行预测:
!wget https://www.dropbox.com/s/e92sui3a4ktvb4j/Hema18.JPG
img = Image.open('Hema18.JPG').convert("RGB")
from torchvision import transforms
pil_to_tensor = transforms.ToTensor()(img).unsqueeze_(0)
Image.fromarray(pil_to_tensor[0].mul(255)\
                        .permute(1, 2, 0).byte().numpy())

输入图像如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

  • 获取输入图像的预测:
model.eval()
with torch.no_grad():
    prediction = model([pil_to_tensor[0].to(device)])
    for i in range(len(prediction[0]['masks'])):
        plt.imshow(Image.fromarray(prediction[0]['masks']\
                        [i, 0].mul(255).byte().cpu().numpy()))
        plt.title('Class: '+str(prediction[0]\
                              ['labels'][i].cpu().numpy())+'\
        Score:'+str(prediction[0]['scores'][i].cpu().numpy()))
        plt.show()

上述代码会产生以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

请注意,在前面的图像中,经过训练的模型不如在测试图像中工作得好。这可能是由于以下原因:

  • 人们在训练时可能不会靠得这么近。
  • 该模型可能没有在感兴趣的类别占据图像的大部分的许多图像上被训练。
  • 我们用来训练模型的数据集中的图像与预测的图像具有不同的数据分布。

但是,即使已经检测到重复的掩码,在这些区域(从第三个掩码开始)中具有较低的类分数是预测中可能存在重复的良好指示。

到目前为止,我们已经学习了如何分割Person类的多个实例。在下一节中,我们将了解我们在这一节中构建的代码中需要调整的内容,以分割图像中多类对象的多个实例。

预测多个类的多个实例

在上一节中,我们学习了如何分割Person类。在这一节中,我们将学习如何通过使用我们在上一节中构建的相同模型,一次性地对 person 和 table 实例进行分段。让我们开始吧:

鉴于大部分代码与前一节中的代码相同,我们将只在这一节中解释额外的代码。在执行代码时,我们鼓励你浏览一下predicting_multiple_instances_of_multiple_classes.ipynb笔记本,它可以在本书的 GitHub 库的Chapter09文件夹中找到

  1. 获取包含感兴趣类别的图像—Person(类别 ID 4)和Table(类别 ID 6):
classes_list = [4,6]
annots = []
for ann in Tqdm(all_annots):
    _ann = read(ann, 1).transpose(2,0,1)
    r,g,b = _ann
    if np.array([num in np.unique(r) for num in \
                classes_list]).sum()==0: continue
    annots.append(ann)
from sklearn.model_selection import train_test_split
_annots = stems(annots)
trn_items, val_items = train_test_split(_annots, \
                                     random_state=2)

在前面的代码中,我们获取了包含至少一个感兴趣的类(classes_list)的图像。

  1. 修改get_mask方法,使其返回两个掩码,以及对应于MasksDataset类中每个掩码的类:
    def get_mask(self,path):
        an = read(path, 1).transpose(2,0,1)
        r,g,b = an
        cls = list(set(np.unique(r)).intersection({4,6}))
        masks = []
        labels = []
        for _cls in cls:
            nzs = np.nonzero(r==_cls)
            instances = np.unique(g[nzs])
            for ix,_id in enumerate(instances):
                masks.append(g==_id)
                labels.append(classes_list.index(_cls)+1)
        return np.array(masks), np.array(labels)

在前面的代码中,我们获取图像中存在的感兴趣的类,并将它们存储在cls中。接下来,我们遍历每个已识别的类(cls),并将红色通道值对应于类(cls)的位置存储在nzs中。接下来,我们获取这些位置的实例 id(instances)。此外,在返回maskslabels的 NumPy 数组之前,我们将instances附加到masks以及对应于labels中实例的类。

  1. 修改__getitem__方法中的labels对象,使其包含从get_mask方法获得的标签,而不是用torch.ones填充。下面代码中的粗体部分是对上一节中的__getitem__方法进行修改的地方:
    def __getitem__(self, ix):
        _id = self.items[ix]
        img_path = f'img/training/{_id}.jpg'
        mask_path = f'annotations_instance/training/{_id}.png'
        masks, labels = self.get_mask(mask_path)
        #print(labels)
        obj_ids = np.arange(1, len(masks)+1)
        img = Image.open(img_path).convert("RGB")
        num_objs = len(obj_ids)
        boxes = []
        for i in range(num_objs):
            obj_pixels = np.where(masks[i])
            xmin = np.min(obj_pixels[1])
            xmax = np.max(obj_pixels[1])
            ymin = np.min(obj_pixels[0])
            ymax = np.max(obj_pixels[0])
            if (((xmax-xmin)<=10) | (ymax-ymin)<=10):
                xmax = xmin+10
                ymax = ymin+10
            boxes.append([xmin, ymin, xmax, ymax])
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.as_tensor(labels, dtype=torch.int64)
        masks = torch.as_tensor(masks, dtype=torch.uint8)
        area = (boxes[:, 3] - boxes[:, 1]) * 
                    (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
        image_id = torch.tensor([ix])
        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd
        if self.transforms is not None:
            img, target = self.transforms(img, target)
        return img, target
    def __len__(self):
        return self.N
    def choose(self):
        return self[randint(len(self))]
  1. 在定义model时,指定您有三个类而不是两个:
num_classes = 3
model=get_model_instance_segmentation(num_classes).to(device)

在训练模型时,正如我们在上一节中所做的,我们将看到训练损失在增加的时期内的变化如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

此外,包含人和桌子的样本图像的预测片段如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

从上图中,我们可以看到,我们能够使用相同的模型预测这两个类别。作为练习,我们鼓励您增加类的数量和纪元的数量,看看您会得到什么结果。

摘要

在本章中,我们学习了如何利用 U-Net 和 Mask R-CNN 在图像上执行分割。我们了解了 U-Net 架构如何使用卷积对图像执行缩小和放大,以保留图像的结构,同时仍然能够预测图像中对象周围的遮罩。然后,我们使用道路场景检测练习巩固了我们对此的理解,在该练习中,我们将图像分割成多个类别。接下来,我们学习了 RoI Align,它有助于确保围绕图像量化的 RoI 汇集问题得到解决。在那之后,我们了解了 Mask R-CNN 是如何工作的,这样我们就可以训练模型来预测图像中的人的实例,以及图像中的人和桌子的实例。

现在,我们已经很好地理解了各种目标检测技术和图像分割技术,在下一章中,我们将学习利用我们到目前为止所学技术的应用程序,以便我们可以扩展我们将要预测的类的数量。此外,我们还将了解 Detectron2 框架,它在我们构建更快的 R-CNN 和 Mask R-CNN 模型时降低了代码复杂性。

问题

  1. 向上扩展对 U-Net 架构有何帮助?
  2. U-Net 中为什么需要全卷积网络?
  3. 在 Mask-RCNN 中,RoI Align 如何改进 RoI pooling?
  4. U-Net 和 Mask-RCNN 在分割方面的主要区别是什么?
  5. 什么是实例分段?**

十、目标检测和分割应用

在之前的章节中,我们学习了各种目标检测技术,例如 R-CNN 系列算法、YOLO、SSD 以及 U-Net 和 Mask R-CNN 图像分割算法。在这一章中,我们将把我们的学习向前推进一步——我们将在更真实的场景中工作,并了解为解决检测和分割问题而更优化的框架/架构。我们将首先利用 Detectron2 框架来训练和检测图像中存在的自定义对象。我们还将使用预先训练的模型来预测图像中出现的人的姿势。此外,我们将学习如何计算图像中人群的人数,然后学习如何利用分割技术来执行图像着色。最后,我们将了解 YOLO 的修改版本,通过使用从激光雷达传感器获得的点云来预测对象周围的 3D 边界框。

本章结束时,您将了解到以下内容:

  • 多对象实例分割
  • 人体姿态检测
  • 人群计数
  • 图像彩色化
  • 基于点云的三维目标检测

多对象实例分割

在前几章中,我们学习了各种目标检测算法。在本节中,我们将了解 Detectron2 平台(ai . Facebook . com/blog/-detectron 2-a-py torch-based-modular-object-detection-library-/),然后使用 Google Open Images 数据集实现它。Detectron2 是由脸书团队构建的平台。Detectron2 包括最先进的目标检测算法的高质量实现,包括 Mask R-CNN 模型系列的 DensePose。最初的 Detectron 框架是用 Caffe2 编写的,而 Detectron2 框架是用 PyTorch 编写的。

Detectron2 支持一系列与目标检测相关的任务。像最初的 Detectron 一样,它支持使用盒子和实例分割遮罩进行目标检测,以及人体姿势预测。除此之外,Detectron2 增加了对语义分段和全景分段的支持(一项结合了语义和实例分段的任务)。通过利用 Detectron2,我们能够在几行代码中构建目标检测、分割和姿态估计。

在本节中,我们将了解以下内容:

  1. open-images存储库中获取数据
  2. 将数据转换成 Detectron2 接受的 COCO 格式
  3. 为实例分割训练模型
  4. 对新图像进行推理

让我们在接下来的几节中逐一讨论这些问题。

获取和准备数据

我们将在谷歌在 https://storage.googleapis.com/openimg/web/index.html 提供的开放图像数据集中(包含数百万张图像及其注释)处理这些图像。

在这部分代码中,我们将学习如何只获取所需的图像,而不是整个数据集。请注意,这一步是必需的,因为数据集大小会阻止可能没有大量资源的典型用户构建模型:

以下代码在本书的 GitHub 知识库的Chapter10文件夹中以Multi_object_segmentation.ipynb的形式提供-【https://tinyurl.com/mcvp-packt 代码包含下载数据的 URL,长度适中。我们强烈建议您在 GitHub 中执行 notebook 以重现结果,同时理解执行的步骤和对文本中各种代码组件的解释。

  1. 安装所需的软件包:
!pip install -qU openimages torch_snippets
  1. 下载所需的注释文件:
from torch_snippets import *
!wget -O train-annotations-object-segmentation.csv -q https://storage.googleapis.com/openimg/v5/train-annotations-object-segmentation.csv
!wget -O classes.csv -q \
 https://raw.githubusercontent.com/openimg/dataset/master/dict.csv 
  1. 指定我们希望模型预测的类(您可以访问 Open Images 网站查看所有类的列表):
required_classes = 'person,dog,bird,car,elephant,football,\
jug,laptop,Mushroom,Pizza,Rocket,Shirt,Traffic sign,\
Watermelon,Zebra'
required_classes = [c.lower() for c in \
                        required_classes.lower().split(',')]

classes = pd.read_csv('classes.csv', header=None)
classes.columns = ['class','class_name']
classes = classes[classes['class_name'].map(lambda x: x \
                                        in required_classes)]
  1. 获取对应于required_classes的图像 id 和遮罩:
from torch_snippets import *
df = pd.read_csv('train-annotations-object-segmentation.csv')

data = pd.merge(df, classes, left_on='LabelName', 
                right_on='class')

subset_data = data.groupby('class_name').agg( \
                        {'ImageID': lambda x: list(x)[:500]})
subset_data = flatten(subset_data.ImageID.tolist())
subset_data = data[data['ImageID'].map(lambda x: x \
                                       in subset_data)]
subset_masks = subset_data['MaskPath'].tolist()

考虑到大量的数据,我们在subset_data中每个类只获取 500 张图片。这取决于你是否为每个类获取一个更小或更大的文件集和唯一类列表(required_classes)。

到目前为止,我们只有对应于一个图像的ImageIdMaskPath值。在接下来的步骤中,我们将继续从open-images下载实际的图像和蒙版。

  1. 现在我们有了要下载的掩码数据子集,让我们开始下载。Open Images 有 16 个用于训练遮罩的 ZIP 文件。每个 ZIP 文件只有几个来自subset_masks的掩码,所以在将所需的掩码移到一个单独的文件夹后,我们将删除其余的。这个下载 - > 移动 - > 删除动作会保持内存占用相对较小。我们必须为 16 个文件中的每个文件运行一次该步骤:
!mkdir -p masks
for c in Tqdm('0123456789abcdef'):
    !wget -q \
     https://storage.googleapis.com/openimg/v5/train-masks/train-masks-{c}.zip
    !unzip -q train-masks-{c}.zip -d tmp_masks
    !rm train-masks-{c}.zip
    tmp_masks = Glob('tmp_masks', silent=True)
    items = [(m,fname(m)) for m in tmp_masks]
    items = [(i,j) for (i,j) in items if j in subset_masks]
    for i,j in items:
        os.rename(i, f'masks/{j}')
    !rm -rf tmp_masks
  1. 下载ImageId对应的图片:
masks = Glob('masks')
masks = [fname(mask) for mask in masks]

subset_data = subset_data[subset_data['MaskPath'].map(lambda \
                                              x: x in masks)]
subset_imageIds = subset_data['ImageID'].tolist()

from openimages.download import _download_images_by_id
!mkdir images
_download_images_by_id(subset_imageIds, 'train', './img/')
  1. 压缩所有图像、遮罩和基本事实并保存它们——以防您的会话崩溃,保存和检索文件对以后的培训很有帮助。创建 ZIP 文件后,请确保将文件保存在您的驱动器中或下载该文件。文件大小最终约为 2.5 GB:
import zipfile
files = Glob('images') + Glob('masks') + \
['train-annotations-object-segmentation.csv', 'classes.csv']
with zipfile.ZipFile('data.zip','w') as zipme:
    for file in Tqdm(files):
        zipme.write(file, compress_type=zipfile.ZIP_DEFLATED)

最后,将数据移动到一个目录中:

!mkdir -p train/
!mv images train/myData2020
!mv masks train/annotations

鉴于目标检测代码中有如此多的移动组件,作为一种标准化的方式,Detectron 接受一种严格的数据格式进行训练。虽然可以编写数据集定义并将其提供给 Detectron,但以 COCO 格式保存整个训练数据更容易(也更有利)。这样,你可以利用其他的训练算法,比如detectron transformers(DETR),而不需要改变任何数据。首先,我们将从定义类的类别开始。

  1. 在 COCO 格式中定义所需的类别:
!pip install \
 git+git://github.com/waspinator/pycococreator.git@0.2.0
import datetime

INFO = {
    "description": "MyData2020",
    "url": "None",
    "version": "1.0",
    "year": 2020,
    "contributor": "sizhky",
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

LICENSES = [
    {
        "id": 1,
        "name": "MIT"
    }
]

CATEGORIES = [{'id': id+1, 'name': name.replace('/',''), \
               'supercategory': 'none'} \
              for id,(_,(name, clss_name)) in \
              enumerate(classes.iterrows())]

在前面的代码中,在CATEGORIES的定义中,我们创建了一个名为supercategory的新键。为了理解supercategory,让我们看一个例子:ManWoman类属于Person超类别。在我们的例子中,假设我们对超级类别不感兴趣,我们将把它指定为none

  • 导入相关的包并创建一个空字典,其中包含保存 COCO JSON 文件所需的键:
!pip install pycocotools
from pycococreatortools import pycococreatortools
from os import listdir
from os.path import isfile, join
from PIL import Image

coco_output = {
    "info": INFO,
    "licenses": LICENSES,
    "categories": CATEGORIES,
    "images": [],
    "annotations": []
}
  • 在适当的位置设置几个包含图像位置和注释文件位置信息的变量:
ROOT_DIR = "train"
IMAGE_DIR, ANNOTATION_DIR = 'train/myData2020/', \
                            'train/annotations/'
image_files = [f for f in listdir(IMAGE_DIR) if \
               isfile(join(IMAGE_DIR, f))]
annotation_files = [f for f in listdir(ANNOTATION_DIR) if \
                    isfile(join(ANNOTATION_DIR, f))]
  • 遍历每个图像文件名,并在coco_output字典中填充images键:
image_id = 1
# go through each image
for image_filename in Tqdm(image_files):
    image = Image.open(IMAGE_DIR + '/' + image_filename)
    image_info = pycococreatortools\
                    .create_image_info(image_id, \
                os.path.basename(image_filename), image.size)
    coco_output["images"].append(image_info)
    image_id = image_id + 1
  1. 遍历每个分段注释,并在coco_output字典中填充annotations键:
segmentation_id = 1
for annotation_filename in Tqdm(annotation_files):
    image_id = [f for f in coco_output['images'] if \
                stem(f['file_name']) == \
                annotation_filename.split('_')[0]][0]['id']
    class_id = [x['id'] for x in CATEGORIES \
                if x['name'] in annotation_filename][0]
    category_info = {'id': class_id, \
                    'is_crowd': 'crowd' in image_filename}
    binary_mask = np.asarray(Image.open(f'{ANNOTATION_DIR}/\
{annotation_filename}').convert('1')).astype(np.uint8)

    annotation_info = pycococreatortools\
                    .create_annotation_info( \
                    segmentation_id, image_id, category_info, 
                    binary_mask, image.size, tolerance=2)

    if annotation_info is not None:
        coco_output["annotations"].append(annotation_info)
        segmentation_id = segmentation_id + 1
  1. coco_output保存在一个 JSON 文件中:
coco_output['categories'] = [{'id': id+1, 'name':clss_name, \
                              'supercategory': 'none'} for \
                             id,(_,(name, clss_name)) in \
                             enumerate(classes.iterrows())]

import json
with open('images.json', 'w') as output_json_file:
    json.dump(coco_output, output_json_file)

这样,我们就有了 COCO 格式的文件,可以很容易地使用 Detectron2 框架来训练我们的模型。

训练模型进行实例分割

使用 Detectron2 进行培训可以通过几个步骤完成:

  1. 安装所需的 Detectron2 软件包。在安装正确的软件包之前,您应该检查您的 CUDA 和 PyTorch 版本。截至撰写本书时,Colab 包含 PyTorch 1.7 和 CUDA 10.1,因此我们将使用相应的文件:
!pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.7/index.html
!pip install pyyaml==5.1 pycocotools>=2.0.1

在进行下一步之前,重新启动 Colab。

  1. 导入相关的detectron2包:
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.engine import DefaultTrainer
  • 假设我们已经重启了 Colab,让我们重新获取所需的类:
from torch_snippets import *
required_classes= 'person,dog,bird,car,elephant,football,jug,\
laptop,Mushroom,Pizza,Rocket,Shirt,Traffic sign,\
Watermelon,Zebra'
required_classes = [c.lower() for c in \
                    required_classes.lower().split(',')]

classes = pd.read_csv('classes.csv', header=None)
classes.columns = ['class','class_name']
classes = classes[classes['class_name'].map(lambda \
                                x: x in required_classes)]
  1. 使用register_coco_instances注册创建的数据集:
from detectron2.data.datasets import register_coco_instances
register_coco_instances("dataset_train", {}, \
                        "images.json", "train/myData2020")
  1. 定义cfg配置文件中的所有参数。

Configuration ( cfg)是一个特殊的 Detectron 对象,它保存了用于训练模型的所有相关信息:

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-\ InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("dataset_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-\ InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") # pretrained 
# weights
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR
cfg.SOLVER.MAX_ITER = 5000 # instead of epochs, we train on 
# 5000 batches
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(classes) 

正如您在前面的代码中看到的,您可以设置训练模型所需的所有主要超参数。merge_from_file正在从一个预先存在的配置文件中导入所有核心参数,该配置文件用于以FPN为骨干的mask_rccnn的预训练。这还将包含关于预训练实验的附加信息,例如优化器和损失函数。为了我们的目的,在cfg中设置的超参数是不言自明的。

  1. 训练模型:
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

使用前面的代码行,我们可以训练一个模型来预测类、边界框以及属于自定义数据集中已定义类的对象的分段。

  • 将模型保存在文件夹中:
!cp output/model_final.pth output/trained_model.pth

至此,我们已经训练好了模型。在下一节中,我们将对一个新图像进行推理。

对新图像进行推断

为了对新图像执行推断,我们加载路径,设置概率阈值,并通过DefaultPredictor方法传递它,如下所示:

  1. 用训练好的模型加载权重。使用相同的cfg并加载模型权重,如以下代码所示:
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \
                                 "trained_model.pth")
  1. 为对象属于某个类别的概率设置阈值:
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.25
  1. 定义predictor方法:
predictor = DefaultPredictor(cfg)
  1. 对感兴趣的图像执行分割并将其可视化:

在下面的代码中,我们随机绘制了 30 幅训练图像(注意,我们还没有创建验证数据;我们将此作为练习留给您),但是您也可以加载自己的图像路径来代替choose(files):

from detectron2.utils.visualizer import ColorMode
files = Glob('train/myData2020')
for _ in range(30):
    im = cv2.imread(choose(files))
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1], scale=0.5, \
                    metadata=MetadataCatalog.get(\
                              "dataset_train"), \
                    instance_mode=ColorMode.IMAGE_BW 
# remove the colors of unsegmented pixels. 
# This option is only available for segmentation models
    )

    out = v.draw_instance_predictions(\
                         outputs["instances"].to("cpu"))
    show(out.get_image())

Visualizer是 Detectron2 绘制对象实例的方式。鉴于预测(出现在outputs变量中)仅仅是一个张量字典,Visualizer将它们转换成像素信息并绘制在图像上。

让我们看看每个输入的含义:

  • 我们想要可视化的图像。
  • scale:打印时图像的大小。在这里,我们要求它将图像缩小到 50%。
  • metadata:我们需要数据集的类级信息,主要是索引到类的映射,这样当我们发送原始张量作为要绘制的输入时,类会将它们解码成实际的人类可读的类。
  • 我们要求模型只突出显示分割的像素。

最后,一旦创建了类(在我们的例子中,它是v),我们可以要求它绘制来自模型的实例预测并显示图像。

上述代码给出了以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

从前面的输出中,您可以看到我们能够相当准确地识别对应于大象的像素。

现在我们已经了解了如何利用 detector 2 来识别图像中与类别相对应的像素,在下一节中,我们将了解如何利用 detector 2 来执行图像中人物的姿势检测。

人体姿态检测

在上一节中,我们学习了如何检测多个对象并对其进行分割。在本节中,我们将学习如何检测图像中的多人,以及如何使用 Detectron2 检测图像中人物的各个身体部位的关键点。在多种用例中,检测关键点非常方便。例如在体育分析和安全领域。

在本练习中,我们将利用配置文件中提供的预训练关键点模型:

下面的代码在 GitHub 知识库的Chapter10文件夹中以Human_pose_detection.ipynb的形式提供-【https://tinyurl.com/mcvp-packt 代码包含了下载数据的 URL。我们强烈建议您在 GitHub 中执行 notebook 以重现结果,同时理解执行的步骤和对文本中各种代码组件的解释。

  1. 安装上一节所示的所有要求:
!pip install detectron2 -f \
  https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.7/index.html
!pip install torch_snippets
!pip install pyyaml==5.1 pycocotools>=2.0.1

from torch_snippets import *
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
  1. 获取配置文件并加载 Detectron2 中预先训练的关键点检测模型:
cfg = get_cfg() # get a fresh new config
cfg.merge_from_file(model_zoo.get_config_file("COCO-\ Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
  1. 指定配置参数:
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold 
# for this model
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-\ Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
  1. 加载我们想要预测的图像:
from torch_snippets import read, resize
!wget -q https://i.imgur.com/ldzGSHk.jpg -O image.png
im = read('image.png',1)
im = resize(im, 0.5) # resize image to half its dimensions
  1. 预测图像并绘制关键点:
outputs = predictor(im)
v = Visualizer(im[:,:,::-1], \
               MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), \
               scale=1.2)
out = v.draw_instance_predictions(\
                outputs["instances"].to("cpu"))
import matplotlib.pyplot as plt
%matplotlib inline
plt.imshow(out.get_image())

前面的代码给出如下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

从前面的输出中,我们可以看到该模型能够准确地识别与图像中的人相对应的各种关键点。

在本节中,我们学习了如何使用 Detectron2 平台执行关键点检测。在下一节中,我们将从头开始学习实现一个修改的 VGG 架构,以估计图像中存在的人数。

人群计数

想象一个场景,给你一张人群的照片,要求你估计照片中的人数。在这种情况下,人群计数模型就派上了用场。在我们继续构建一个模型来执行人群计数之前,让我们先了解可用的数据和模型架构。

为了训练预测图像中人数的模型,我们必须首先加载图像。图像应该构成图像中出现的所有人的头部中心的位置。输入图像的样本和图像中各个人物的头部中心位置如下(来源:上海科技数据集(github.com/desenzhou/ShanghaiTechDataset)):

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

在前面的示例中,表示地面实况的图像(右边的图像–图像中出现的人的头部中心)非常稀疏。正好有 N 个白色像素,其中 N 是图像中的人数。让我们放大到图像的左上角,再次看到相同的地图:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

在下一步中,我们将地面实况稀疏图像转换为密度图,该密度图表示图像中该区域的人数:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

同一作物的最终输入输出对看起来像这样:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

整个图像看起来也是如此:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

请注意,在前面的图像中,当两个人彼此靠近时,像素强度较高。然而,当一个人远离休息时,对应于该人的像素密度更均匀地散开,导致对应于远离休息的人的像素密度更低。本质上,热图是以这样一种方式生成的,即像素值的总和等于图像中出现的人数。

现在,我们已经能够接受输入图像和图像中人的头部中心的位置(图像被处理以获取地面实况输出热图),我们将利用标题为 CSRNet:用于理解高度拥堵场景的扩展卷积神经网络的论文中详细描述的架构来预测图像中出现的人数。

模型架构(arxiv.org/pdf/1802.10062.pdf)如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

在模型架构的上述结构中,我们在首先将图像通过标准的 VGG-16 主干网之后,再将它通过四个额外的卷积层。该输出通过四种配置之一,并最终通过 1 x 1 x 1 卷积层。我们将使用 A 配置,因为它是最小的。

接下来,我们对输出图像执行均方误差 ( MSE )损失最小化,以达到最佳权重值,同时使用 MAE 跟踪实际人群计数。

该架构的另一个细节是,作者使用了扩展卷积而不是普通卷积。

典型的扩张卷积如下(图像来源:arxiv.org/pdf/1802.10062.pdf):

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

在前面的内容中,左边的图表代表了我们目前为止一直在做的一个典型的内核。第二个和第三个图表示扩展的内核,在各个像素之间有一个间隙。这样,内核就有了更大的感受野。一个大的感受野可以派上用场,因为我们需要了解一个给定的人附近的人数,以便估计对应于这个人的像素密度。我们使用一个膨胀的核(有九个参数)而不是普通的核(有 49 个参数,相当于三个核的膨胀率)来用更少的参数获取更多的信息。

了解了如何构建模型之后,让我们继续编写模型代码来执行下一节中的人群计数。(对于那些希望了解工作细节的人,我们建议你浏览这里的文件:arxiv.org/pdf/1802.10062.pdf。我们将在下一节中培训的模型是受本文的启发。)

编码人群计数

我们将采用以下策略进行人群计数:

  1. 导入相关的包和数据集。

  2. 我们将要处理的数据集 ShanghaiTech 数据集——已经将人脸中心转换为基于高斯过滤器密度的分布,因此我们无需再次执行。使用网络映射输入图像和输出高斯密度图。

  3. 定义一个函数来执行扩张卷积。

  4. 定义网络模型,并对批量数据进行训练,以最小化 MSE。

让我们将我们的策略编码如下:

以下代码在本书的 GitHub 知识库的Chapter 10文件夹中以crowd_counting.ipynb的形式提供-【https://tinyurl.com/mcvp-packt 代码包含下载数据的 URL,长度适中。我们强烈建议您在 GitHub 中执行 notebook 以重现结果,同时理解执行的步骤和对文本中各种代码组件的解释。

  1. 导入包并下载数据集:
%%time
import os
if not os.path.exists('CSRNet-pytorch/'):
    !pip install -U scipy torch_snippets torch_summary
    !git clone https://github.com/sizhky/CSRNet-pytorch.git
    from google.colab import files
    files.upload() # upload kaggle.json
    !mkdir -p ~/.kaggle
    !mv kaggle.json ~/.kaggle/
    !ls ~/.kaggle
    !chmod 600 /root/.kaggle/kaggle.json
    print('downloading data...')
    !kaggle datasets download -d \
        tthien/shanghaitech-with-people-density-map/
    print('unzipping data...')
    !unzip -qq shanghaitech-with-people-density-map.zip

%cd CSRNet-pytorch
!ln -s ../shanghaitech_with_people_density_map
from torch_snippets import *
import h5py
from scipy import io
  • 提供图像(image_folder)、地面实况(gt_folder)和热图文件夹(heatmap_folder)的位置:
part_A = Glob('shanghaitech_with_people_density_map/\
ShanghaiTech/part_A/train_data/');

image_folder = 'shanghaitech_with_people_density_map/\
ShanghaiTech/part_A/train_data/img/'
heatmap_folder = 'shanghaitech_with_people_density_map/\
ShanghaiTech/part_A/train_data/ground-truth-h5/'
gt_folder = 'shanghaitech_with_people_density_map/\
ShanghaiTech/part_A/train_data/ground-truth/'
  1. 定义训练和验证数据集以及数据加载器:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
tfm = T.Compose([
    T.ToTensor()
])

class Crowds(Dataset):
    def __init__(self, stems):
        self.stems = stems

    def __len__(self):
        return len(self.stems)

    def __getitem__(self, ix):
        _stem = self.stems[ix]
        image_path = f'{image_folder}/{_stem}.jpg'
        heatmap_path = f'{heatmap_folder}/{_stem}.h5'
        gt_path = f'{gt_folder}/GT_{_stem}.mat'

        pts = io.loadmat(gt_path)
        pts = len(pts['image_info'][0,0][0,0][0])

        image = read(image_path, 1)
        with h5py.File(heatmap_path, 'r') as hf:
            gt = hf['density'][:]
        gt = resize(gt, 1/8)*64
        return image.copy(), gt.copy(), pts

    def collate_fn(self, batch):
        ims, gts, pts = list(zip(*batch))
        ims = torch.cat([tfm(im)[None] for im in \
                            ims]).to(device)
        gts = torch.cat([tfm(gt)[None] for gt in \
                            gts]).to(device)
        return ims, gts, torch.tensor(pts).to(device)

    def choose(self):
        return self[randint(len(self))]

from sklearn.model_selection import train_test_split
trn_stems, val_stems = train_test_split(\
            stems(Glob(image_folder)), random_state=10)

trn_ds = Crowds(trn_stems)
val_ds = Crowds(val_stems)

trn_dl = DataLoader(trn_ds, batch_size=1, shuffle=True, \
                    collate_fn=trn_ds.collate_fn)
val_dl = DataLoader(val_ds, batch_size=1, shuffle=True, \
                    collate_fn=val_ds.collate_fn)

请注意,到目前为止,我们编写的典型数据集类的唯一附加内容是前面代码中以粗体显示的代码行。我们正在调整地面真实的大小,因为我们网络的输出将被缩小到原始大小的 1/8 ^(th) ,因此我们将地图乘以 64,以便图像像素的总和将被缩放回原始的人群计数。

  1. 定义网络架构:
  • 定义启用扩展卷积的函数(make_layers):
import torch.nn as nn
import torch
from torchvision import models
from utils import save_net,load_net

def make_layers(cfg, in_channels = 3, batch_norm=False, 
                dilation = False):
    if dilation:
        d_rate = 2
    else:
        d_rate = 1
    layers = []
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels,v,kernel_size=3,\
                               padding=d_rate, dilation=d_rate)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), \
                           nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)
  • 定义网络架构-CSRNet:
class CSRNet(nn.Module):
    def __init__(self, load_weights=False):
        super(CSRNet, self).__init__()
        self.seen = 0
        self.frontend_feat = [64, 64, 'M', 128, 128, 'M',256,
                                256, 256, 'M', 512, 512, 512]
        self.backend_feat = [512, 512, 512, 256, 128, 64]
        self.frontend = make_layers(self.frontend_feat)
        self.backend = make_layers(self.backend_feat,
                          in_channels = 512,dilation = True)
        self.output_layer = nn.Conv2d(64, 1, kernel_size=1)
        if not load_weights:
            mod = models.vgg16(pretrained = True)
            self._initialize_weights()
            items = list(self.frontend.state_dict().items())
            _items = list(mod.state_dict().items())
            for i in range(len(self.frontend.state_dict()\
                               .items())):
                items[i][1].data[:] = _items[i][1].data[:]
    def forward(self,x):
        x = self.frontend(x)
        x = self.backend(x)
        x = self.output_layer(x)
        return x
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
  1. 定义对一批数据进行训练和验证的函数:
def train_batch(model, data, optimizer, criterion):
    model.train()
    optimizer.zero_grad()
    ims, gts, pts = data
    _gts = model(ims)
    loss = criterion(_gts, gts)
    loss.backward()
    optimizer.step()
    pts_loss = nn.L1Loss()(_gts.sum(), gts.sum())
    return loss.item(), pts_loss.item()

@torch.no_grad()
def validate_batch(model, data, criterion):
    model.eval()
    ims, gts, pts = data
    _gts = model(ims)
    loss = criterion(_gts, gts)
    pts_loss = nn.L1Loss()(_gts.sum(), gts.sum())
    return loss.item(), pts_loss.item()
  1. 在不断增加的时期内训练模型:
model = CSRNet().to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-6)
n_epochs = 20

log = Report(n_epochs)
for ex in range(n_epochs):
    N = len(trn_dl)
    for bx, data in enumerate(trn_dl):
        loss,pts_loss=train_batch(model, data, optimizer, \
                                        criterion)
        log.record(ex+(bx+1)/N, trn_loss=loss, 
                           trn_pts_loss=pts_loss, end='\r')

    N = len(val_dl)
    for bx, data in enumerate(val_dl):
        loss, pts_loss = validate_batch(model, data, \
                                        criterion)
        log.record(ex+(bx+1)/N, val_loss=loss, 
                    val_pts_loss=pts_loss, end='\r')

    log.report_avgs(ex+1)
    if ex == 10: optimizer = optim.Adam(model.parameters(), \
                                        lr=1e-7)

前面的代码导致了训练和验证损失的变化(这里,损失是群体计数的 MAE),如下所示:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

从前面的图中,我们可以看到我们的预测误差了大约 150 人。我们可以从以下两个方面改进该模型:

  • 通过对原始图像的裁剪使用数据扩充和训练
  • 通过使用更大的网络(我们使用 A 配置,而 B、C 和 D 更大)。
  1. 对新图像进行推断:
  • 获取测试图像并使其正常化:
from matplotlib import cm as c
from torchvision import datasets, transforms
from PIL import Image
transform=transforms.Compose([
                 transforms.ToTensor(),transforms.Normalize(\
                          mean=[0.485, 0.456, 0.406],\
                          std=[0.229, 0.224, 0.225]),\
                  ])

test_folder = 'shanghaitech_with_people_density_map/\
ShanghaiTech/part_A/test_data/'
imgs = Glob(f'{test_folder}/images')
f = choose(imgs)
print(f)
img = transform(Image.open(f).convert('RGB')).to(device)
  • 将图像传递给训练好的模型:
output = model(img[None])
print("Predicted Count : ", int(output.detach().cpu()\
                                      .sum().numpy()))
temp = np.asarray(output.detach().cpu()\
                    .reshape(output.detach().cpu()\
                    .shape[2],output.detach()\
                    .cpu().shape[3]))
plt.imshow(temp,cmap = c.jet)
plt.show()

上述代码会生成输入图像(左图)的热图(右图):

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

从前面的输出中,我们可以看到该模型相当准确地预测了热图,并且预测的人数接近实际值。

在下一节中,我们将利用 U-Net 架构来给图像着色。

图像彩色化

想象一个场景,给你一堆黑白图像,要求你把它们变成彩色图像。你会怎么解决这个问题?解决这个问题的一种方法是使用伪监督管道,我们获取原始图像,将其转换为黑白图像,并将它们视为输入输出对。我们将通过利用 CIFAR-10 数据集对图像执行彩色化来演示这一点。

当我们编码图像彩色化网络时,我们将采用的策略如下:

  1. 获取训练数据集中的原始彩色图像,并将其转换为灰度,以获取输入(灰度)和输出(原始彩色图像)组合。
  2. 标准化输入和输出。
  3. 构建一个 U-Net 架构。
  4. 在不断增加的时期内训练模型。

有了前面的策略,让我们继续编写模型代码,如下所示:

下面的代码可以在本书的 GitHub 库【https://tinyurl.com/mcvp-packtChapter 10文件夹中找到Image colorization.ipynb

  1. 安装所需的软件包并导入它们:
!pip install torch_snippets
from torch_snippets import *
device = 'cuda' if torch.cuda.is_available() else 'cpu'
  1. 下载数据集并定义训练和验证数据集以及数据加载器:
  • 下载数据集:
from torchvision import datasets
import torch
data_folder = '~/cifar10/cifar/' 
datasets.CIFAR10(data_folder, download=True)
  • 定义训练和验证数据集以及数据加载器:
class Colorize(torchvision.datasets.CIFAR10):
    def __init__(self, root, train):
        super().__init__(root, train)

    def __getitem__(self, ix):
        im, _ = super().__getitem__(ix)
        bw = im.convert('L').convert('RGB')
        bw, im = np.array(bw)/255., np.array(im)/255.
        bw, im = [torch.tensor(i).permute(2,0,1)\
                  .to(device).float() for i in [bw,im]]
        return bw, im

trn_ds = Colorize('~/cifar10/cifar/', train=True)
val_ds = Colorize('~/cifar10/cifar/', train=False)

trn_dl = DataLoader(trn_ds, batch_size=256, shuffle=True)
val_dl = DataLoader(val_ds, batch_size=256, shuffle=False)

输入和输出图像的示例如下:

a,b = trn_ds[0]
subplots([a,b], nc=2)

上述代码会产生以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

请注意,CIFAR-10 图像的形状为 32 x 32。

  1. 定义网络架构:
class Identity(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return x

class DownConv(nn.Module):
    def __init__(self, ni, no, maxpool=True):
        super().__init__()
        self.model = nn.Sequential(
            nn.MaxPool2d(2) if maxpool else Identity(),
            nn.Conv2d(ni, no, 3, padding=1),
            nn.BatchNorm2d(no),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(no, no, 3, padding=1),
            nn.BatchNorm2d(no),
            nn.LeakyReLU(0.2, inplace=True),
        )
    def forward(self, x):
        return self.model(x)

class UpConv(nn.Module):
    def __init__(self, ni, no, maxpool=True):
        super().__init__()
        self.convtranspose = nn.ConvTranspose2d(ni, no, \
                                                2, stride=2)
        self.convlayers = nn.Sequential(
            nn.Conv2d(no+no, no, 3, padding=1),
            nn.BatchNorm2d(no),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(no, no, 3, padding=1),
            nn.BatchNorm2d(no),
            nn.LeakyReLU(0.2, inplace=True),
        )

    def forward(self, x, y):
        x = self.convtranspose(x)
        x = torch.cat([x,y], axis=1)
        x = self.convlayers(x)
        return x

class UNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.d1 = DownConv( 3, 64, maxpool=False)
        self.d2 = DownConv( 64, 128)
        self.d3 = DownConv( 128, 256)
        self.d4 = DownConv( 256, 512)
        self.d5 = DownConv( 512, 1024)
        self.u5 = UpConv (1024, 512)
        self.u4 = UpConv ( 512, 256)
        self.u3 = UpConv ( 256, 128)
        self.u2 = UpConv ( 128, 64)
        self.u1 = nn.Conv2d(64, 3, kernel_size=1, stride=1)

    def forward(self, x):
        x0 = self.d1( x) # 32
        x1 = self.d2(x0) # 16
        x2 = self.d3(x1) # 8
        x3 = self.d4(x2) # 4
        x4 = self.d5(x3) # 2
        X4 = self.u5(x4, x3)# 4
        X3 = self.u4(X4, x2)# 8
        X2 = self.u3(X3, x1)# 16
        X1 = self.u2(X2, x0)# 32
        X0 = self.u1(X1) # 3
        return X0
  1. 定义模型、优化器和损失函数:
def get_model():
    model = UNet().to(device)
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.MSELoss()
    return model, optimizer, loss_fn
  1. 定义对一批数据进行训练和验证的函数:
def train_batch(model, data, optimizer, criterion):
    model.train()
    x, y = data
    _y = model(x)
    optimizer.zero_grad()
    loss = criterion(_y, y)
    loss.backward()
    optimizer.step()
    return loss.item()

@torch.no_grad()
def validate_batch(model, data, criterion):
    model.eval()
    x, y = data
    _y = model(x)
    loss = criterion(_y, y)
    return loss.item()
  1. 在不断增加的时期内训练模型:
model, optimizer, criterion = get_model()
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, \
                                    step_size=10, gamma=0.1)

_val_dl = DataLoader(val_ds, batch_size=1, shuffle=True)

n_epochs = 100
log = Report(n_epochs)
for ex in range(n_epochs):
    N = len(trn_dl)
    for bx, data in enumerate(trn_dl):
        loss = train_batch(model, data, optimizer, criterion)
        log.record(ex+(bx+1)/N, trn_loss=loss, end='\r')
        if (bx+1)%50 == 0:
            for _ in range(5):
                a,b = next(iter(_val_dl))
                _b = model(a)
                subplots([a[0], b[0], _b[0]], nc=3, \
                          figsize=(5,5))

    N = len(val_dl)
    for bx, data in enumerate(val_dl):
        loss = validate_batch(model, data, criterion)
        log.record(ex+(bx+1)/N, val_loss=loss, end='\r')

    exp_lr_scheduler.step()
    if (ex+1) % 5 == 0: log.report_avgs(ex+1)

    for _ in range(5):
        a,b = next(iter(_val_dl))
        _b = model(a)
        subplots([a[0], b[0], _b[0]], nc=3, figsize=(5,5))

log.plot_epochs()

上述代码生成如下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

从前面的输出中,我们可以看到该模型能够很好地为灰度图像着色。

到目前为止,我们已经了解了如何利用 Detectron2 进行分割和关键点检测、人群计数中的膨胀卷积以及图像着色中的 U-Net。在下一部分,我们将了解如何利用 YOLO 进行 3D 目标检测。

基于点云的三维目标检测

到目前为止,我们已经学习了如何使用具有锚定框核心基本概念的算法来预测 2D 图像上的边界矩形。我们现在将学习如何将相同的概念扩展到预测物体周围的 3D 边界框。

在自动驾驶汽车中,如果不了解环境,行人/障碍物检测和路线规划等任务就无法完成。预测 3D 对象的位置及其方向成为一项重要的任务。障碍物周围的 2D 边界框不仅重要,而且知道障碍物与对象的距离、高度、宽度和方向对于在 3D 世界中安全导航也至关重要。

在本节中,我们将了解如何使用 YOLO 来预测现实数据集上汽车和行人的 3D 方向和位置。

下载数据、训练和测试集的说明都在这个 GitHub repo 中给出:GitHub . com/sizhky/Complex-yolov 4-py torch/blob/master/readme . MD # training-instructions。鉴于很少有公开可用的 3D 数据集,我们选择了最常用的数据集来做这个练习,您仍然需要注册下载。我们也在前面的链接中提供了注册说明。

理论

收集实时 3D 数据的一个众所周知的传感器是激光雷达 ( 光探测和测距)。它是一个安装在旋转装置上的激光器,每秒钟发射数百次激光束。另一个传感器接收来自周围物体的激光反射,并计算激光在遇到障碍物之前已经行进了多远。在汽车的所有方向上这样做将产生反映环境本身的距离的 3D 点云。在我们将要学习的数据集中,我们已经从称为velodyne的特定硬件获得了 3D 点云。让我们了解如何为 3D 目标检测编码输入和输出。

输入编码

我们的原始输入将是以.bin文件的形式呈现给我们的 3D 点云。每一个都可以使用np.fromfile(<filepath>)作为 NumPy 数组加载,下面是样本文件的数据外观(按照 GitHub repo 指令下载并移动原始文件后,这些文件位于dataset/.../training/velodyne目录中):

files = Glob('training/velodyne')
F = choose(files)
pts = np.fromfile(F, dtype=np.float32).reshape(-1, 4)
pts

上述代码给出了以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

这可以想象如下:

# take the points and remove faraway points
x,y,z = np.clip(pts[:,0], 0, 50), 
        np.clip(pts[:,1], -25, 25), 
        np.clip(pts[:,2],-3, 1.27)

fig = go.Figure(data=[go.Scatter3d(\
        x=x, y=y, z=z, mode='markers',
        marker=dict(
            size=2,
            color=z, # set color to a list of desired values
            colorscale='Viridis', # choose a colorscale
            opacity=0.8
        )
    )])

fig.update_layout(margin=dict(l=0, r=0, b=0, t=0))
fig.show()

上述代码会产生以下输出:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

通过执行以下步骤,我们可以将这些信息转换成鸟瞰图。

  1. 将 3D 点云投影到 XY 平面(地面)上,分割成一个网格,每个网格单元的分辨率为 8 cm ² 。
  2. 对于每个单元格,计算以下内容,并将它们与指定的通道相关联:
  • 红色通道:网格中最高点的高度
  • 绿色通道:网格中最高点的亮度
  • 蓝色通道:网格中的点数除以 64(这是一个归一化因子)

例如,重建的云的俯视图可能如下所示:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

可以清楚地看到图像中的“阴影”,表示有障碍物。

这就是我们如何从激光雷达点云数据创建图像。

我们将 3D 点云作为原始输入,并获得鸟瞰图像作为输出。这是创建图像所必需的预处理步骤,该图像将作为 YOLO 模型的输入。

输出编码

现在我们已经将鸟瞰图像(3D 点云的)作为模型的输入,模型需要预测以下真实世界的特征:

  • 图像中出现的对象()是什么
  • 物体在东西轴( x )上离汽车有多远(以米为单位)
  • 物体在南北轴( y )上离汽车有多远(以米为单位)
  • 物体的方位()是什么**
    ***** 物体有多大(物体的
    长度宽度**,单位为米)****

****可以预测(鸟瞰图像的)像素坐标系中的边界框。但是它没有任何现实意义,因为预测仍然在像素空间中(在鸟瞰图中)。在这种情况下,我们需要将这些像素坐标(鸟瞰图的)边界框预测转换为以米为单位的真实坐标。为了避免后处理过程中的额外步骤,我们直接预测真实值。

此外,在现实场景中,对象可以朝向任何方向。如果我们只计算长度和宽度,将不足以描述紧密包围盒。这种情况的一个例子如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

为了得到物体的紧密包围盒,我们还需要障碍物朝向哪个方向的信息,因此我们还需要额外的偏航参数。从形式上来说,它是具有南北轴的物体所做的定向。

首先,考虑到汽车的 dashcam(以及激光雷达)视图的宽度大于高度,YOLO 模型使用了 32 x 64 单元的锚定网格(宽度大于高度)。该模型对任务使用了两个损失。第一个是正常的 YOLO 损失(它负责预测我们在第八章、高级物体探测中了解到的 xylw 类),另一个损失称为欧拉损失,它专门预测偏航。形式上,从模型输出预测最终边界框的方程组如下:

b[x]=σ(t[x]??)+c[x
b[y]=σ(t[y])+c[y] b[w]= p[w]e^(t[w])]

**这里,b[x]??、b[y]??、b[w]??、 b [l]b [φ] 分别是障碍物的 xz 坐标值、宽度、长度和偏航。
t [x]t [y]t [w]t [l]t [Im]t[Re]是预测的六个回归值
c [x]c [y] 是 32×64 矩阵中网格单元的中心位置,而 p [w]p [l] 是通过取汽车和行人的平均宽度和长度选择的预定义先验。此外,在实现中有五个先验(锚盒)。

同一类的每个对象的高度被假定为一个固定的数字。

参考这里给出的插图,它以图片的形式展示了这一点(图片来源:arxiv.org/pdf/1803.06199.pdf):

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

总损失计算如下:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

你已经从上一章知道了损失 [YOLO] (使用t[x]??、 t [y]t [w]t [l] 作为目标)。另外,请注意以下事项:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

现在我们已经了解了 3D 目标检测的基本原理如何与 2D 目标检测的基本原理保持相同(但是有更多的参数要预测)以及该任务的输入输出对,让我们利用现有的 GitHub repo 来训练我们的模型。

关于 3D 目标检测的更多细节,请参考在 https://arxiv.org/pdf/1803.06199.pdf的论文复合体-YOLO

训练用于 3D 目标检测的 YOLO 模型

由于标准化的代码,编码工作在很大程度上远离了用户。很像 Detectron2,我们可以通过确保数据在正确的位置以正确的格式来训练和测试算法。一旦确保了这一点,我们就可以用最少的代码行来训练和测试代码。

我们需要首先克隆Complex-YOLOv4-Pytorch存储库:

$ git clone https://github.com/sizhky/Complex-YOLOv4-Pytorch

按照README.md文件中的说明下载数据集并将其移动到正确的位置。

The instructions for downloading the data, training, and testing sets are all given in this GitHub repo: github.com/sizhky/Complex-YOLOv4-Pytorch/blob/master/README.md#training-instructions.

Given that there are very few openly available 3D datasets, we have chosen the most-used dataset for this exercise, which you still need to register for download. We also give the instructions for registration in the preceding link.

数据格式

在这个练习中,我们可以使用任何带有地面实况的 3D 点云数据。有关如何下载和移动数据的更多说明,请参考 GitHub repo 上的README文件。数据需要以下列格式存储在根目录中:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

三个对我们来说是新的文件夹是velodynecaliblabel_2:

  • velodyne包含一个.bin文件列表,这些文件对存在于image_2文件夹中的相应图像的 3D 点云信息进行编码。

  • calib包含每个点云对应的校准文件。通过使用calib文件夹中每个文件中的 3 x 4 投影矩阵,可以将激光雷达点云坐标系中的 3D 坐标投影到相机坐标系上,即图像上。实际上,激光雷达传感器捕捉的点与相机捕捉的点略有偏差。这种偏移是由于两个传感器彼此相隔几英寸安装。了解正确的偏移量将有助于我们正确地将边界框和 3D 点投影到来自相机的图像上。

  • label_2包含每幅图像的基本事实(每行一个基本事实),以 15 个值的形式,如下表所述:

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

请注意,我们的目标列是这里看到的类型(class)、 wlxzry (偏航)。我们将忽略该任务的其余值。

数据检查

我们可以通过运行以下命令来验证数据是否正确下载:

$ cd Complex-YOLOv4-Pytorch/src/data_process
$ python kitti_dataloader.py --output-width 600

前面的代码显示了多个图像,一次一个图像。下面就是这样一个例子(图片来源:arxiv.org/pdf/1803.06199.pdf):

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

既然我们已经能够下载并查看一些图像,在下一节中,我们将学习如何训练模型来预测 3D 边界框。

培养

训练代码封装在单个 Python 文件中,可以按如下方式调用:

$ cd Complex-YOLOv4-Pytorch/src
$ python train.py --gpu_idx 0 --batch_size 2 --num_workers 4 \
                  --num_epochs 5

默认的历元数是 300,但是从第五个历元开始,结果是相当合理的。在 GTX 1070 GPU 上,每个历元需要 30 到 45 分钟。如果无法一次性完成训练,您可以使用--resume_path恢复训练。代码每五个时期保存一个新的检查点。

测试

就像在数据检查部分一样,可以用下面的代码测试训练好的模型:

$ cd Complex-YOLOv4-Pytorch/src
$ python test.py --gpu_idx 0 --pretrained_path ../checkpoints/complexer_yolo/Model_complexer_yolo_epoch_5.pth --cfgfile ./config/cfg/complex_yolov4.cfg --show_image 

代码的主要输入是检查点路径和模型配置路径。给出它们并运行代码后,弹出如下输出(图片来源:arxiv.org/pdf/1803.06199.pdf):

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

由于模型的简单性,我们可以使用普通 GPU 在实时场景中使用它,每秒钟可以获得大约 15-20 次预测。

摘要

在这一章中,我们学习了处理对象定位和分割的各种实际方面。具体来说,我们了解了如何利用 Detectron2 平台来执行图像分割和检测以及关键点检测。此外,当我们从开放图像数据集中获取图像时,我们还了解了处理大型数据集所涉及的一些复杂性。接下来,我们致力于利用 VGG 和 U-Net 架构分别进行人群计数和图像着色。最后,我们了解了使用点云图像进行三维目标检测的理论和实现步骤。正如您从所有这些示例中看到的,底层基础与前面章节中描述的相同,只是在网络的输入/输出方面进行了修改,以适应手头的任务。

在下一章中,我们将转换话题,学习图像编码,这有助于识别相似的图像以及生成新的图像。******

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值