mask-rcnn原理与实战

nice-wyh

已于 2024-03-15 14:32:40 修改

阅读量1.5k

点赞数 11

分类专栏： pytorch目标检测文章标签：人工智能

于 2023-11-15 19:44:32 首次发布

本文链接：https://blog.csdn.net/qq_46454669/article/details/134426740

版权

pytorch目标检测专栏收录该内容

5 篇文章 1 订阅

订阅专栏

一、Mask R-CNN是什么，可以做哪些任务？

Mask R-CNN是一个实例分割（Instance segmentation）算法，可以用来做“目标检测”、“目标实例分割”、“目标关键点检测”。

1. 实例分割（Instance segmentation）和语义分割（Semantic segmentation）的区别与联系

联系：语义分割和实例分割都是目标分割中的两个小的领域，都是用来对输入的图片做分割处理；

区别：

1. 通常意义上的目标分割指的是语义分割，语义分割已经有很长的发展历史，已经取得了很好地进展，目前有很多的学者在做这方面的研究；然而实例分割是一个从目标分割领域独立出来的一个小领域，是最近几年才发展起来的，与前者相比，后者更加复杂，当前研究的学者也比较少，是一个有研究空间的热门领域，如图1所示，这是一个正在探索中的领域；

2. 观察图3中的c和d图，c图是对a图进行语义分割的结果，d图是对a图进行实例分割的结果。两者最大的区别就是图中的"cube对象"，在语义分割中给了它们相同的颜色，而在实例分割中却给了不同的颜色。即实例分割需要在语义分割的基础上对同类物体进行更精细的分割。

注：很多博客中都没有完全理解清楚这个问题，很多人将这个算法看做语义分割，其实它是一个实例分割算法。

2. Mask R-CNN可以完成的任务

图4 Mask R-CNN进行目标检测与实例分割

图5 Mask R-CNN进行人体姿态识别

总之，Mask R-CNN是一个非常灵活的框架，可以增加不同的分支完成不同的任务，可以完成目标分类、目标检测、语义分割、实例分割、人体姿势识别等多种任务，真不愧是一个好算法！

3. Mask R-CNN预期达到的目标

高速
高准确率（高的分类准确率、高的检测准确率、高的实例分割准确率等）
简单直观
易于使用

4. 如何实现这些目标

高速和高准确率：为了实现这个目的，作者选用了经典的目标检测算法Faster-rcnn和经典的语义分割算法FCN。Faster-rcnn可以既快又准的完成目标检测的功能；FCN可以精准的完成语义分割的功能，这两个算法都是对应领域中的经典之作。Mask R-CNN比Faster-rcnn复杂，但是最终仍然可以达到5fps的速度，这和原始的Faster-rcnn的速度相当。由于发现了ROI Pooling中所存在的像素偏差问题，提出了对应的ROIAlign策略，加上FCN精准的像素MASK，使得其可以获得高准确率。

简单直观：整个Mask R-CNN算法的思路很简单，就是在原始Faster-rcnn算法的基础上面增加了FCN来产生对应的MASK分支。即Faster-rcnn + FCN，更细致的是 RPN + ROIAlign + Fast-rcnn + FCN。

易于使用：整个Mask R-CNN算法非常的灵活，可以用来完成多种任务，包括目标分类、目标检测、语义分割、实例分割、人体姿态识别等多个任务，这将其易于使用的特点展现的淋漓尽致。我很少见到有哪个算法有这么好的扩展性和易用性，值得我们学习和借鉴。除此之外，我们可以更换不同的backbone architecture和Head Architecture来获得不同性能的结果。

二、Mask R-CNN框架解析

1. Mask R-CNN算法步骤

首先，输入一幅你想处理的图片，然后进行对应的预处理操作，或者预处理后的图片；
然后，将其输入到一个预训练好的神经网络中（ResNeXt等）获得对应的feature map；
接着，对这个feature map中的每一点设定预定个的ROI，从而获得多个候选ROI；
接着，将这些候选的ROI送入RPN网络进行二值分类（前景或背景）和BB回归，过滤掉一部分候选的ROI；
接着，对这些剩下的ROI进行ROIAlign操作（即先将原图和feature map的pixel对应起来，然后将feature map和固定的feature对应起来）；
最后，对这些ROI进行分类（N类别分类）、BB回归和MASK生成（在每一个ROI里面进行FCN操作）。

2. Mask R-CNN架构分解

在这里，我将Mask R-CNN分解为如下的3个模块，Faster-rcnn、ROIAlign和FCN。然后分别对这3个模块进行讲解，这也是该算法的核心。

3. Faster-rcnn（该算法请参考该链接，我进行了详细的分析）

4. FCN

FCN算法是一个经典的语义分割算法，可以对图片中的目标进行准确的分割。其总体架构如上图所示，它是一个端到端的网络，主要的模快包括卷积和去卷积，即先对图像进行卷积和池化，使其feature map的大小不断减小；然后进行反卷积操作，即进行插值操作，不断的增大其feature map，最后对每一个像素值进行分类。从而实现对输入图像的准确分割。具体的细节请参考该链接。

5. ROIPooling和ROIAlign的分析与比较

如图所示，ROI Pooling和ROIAlign最大的区别是：前者使用了两次量化操作，而后者并没有采用量化操作，使用了线性插值算法，具体的解释如下所示。

如上图所示，为了得到固定大小（7X7）的feature map，我们需要做两次量化操作：1）图像坐标 — feature map坐标，2）feature map坐标 — ROI feature坐标。我们来说一下具体的细节，如图我们输入的是一张800x800的图像，在图像中有两个目标（猫和狗），狗的BB大小为665x665，经过VGG16网络后，我们可以获得对应的feature map，如果我们对卷积层进行Padding操作，我们的图片经过卷积层后保持原来的大小，但是由于池化层的存在，我们最终获得feature map 会比原图缩小一定的比例，这和Pooling层的个数和大小有关。在该VGG16中，我们使用了5个池化操作，每个池化操作都是2Pooling，因此我们最终获得feature map的大小为800/32 x 800/32 = 25x25（是整数），但是将狗的BB对应到feature map上面，我们得到的结果是665/32 x 665/32 = 20.78 x 20.78，结果是浮点数，含有小数，但是我们的像素值可没有小数，那么作者就对其进行了量化操作（即取整操作），即其结果变为20 x 20，在这里引入了第一次的量化误差；然而我们的feature map中有不同大小的ROI，但是我们后面的网络却要求我们有固定的输入，因此，我们需要将不同大小的ROI转化为固定的ROI feature，在这里使用的是7x7的ROI feature，那么我们需要将20 x 20的ROI映射成7 x 7的ROI feature，其结果是 20 /7 x 20/7 = 2.86 x 2.86，同样是浮点数，含有小数点，我们采取同样的操作对其进行取整吧，在这里引入了第二次量化误差。其实，这里引入的误差会导致图像中的像素和特征中的像素的偏差，即将feature空间的ROI对应到原图上面会出现很大的偏差。原因如下：比如用我们第二次引入的误差来分析，本来是2,86，我们将其量化为2，这期间引入了0.86的误差，看起来是一个很小的误差呀，但是你要记得这是在feature空间，我们的feature空间和图像空间是有比例关系的，在这里是1:32，那么对应到原图上面的差距就是0.86 x 32 = 27.52。这个差距不小吧，这还是仅仅考虑了第二次的量化误差。这会大大影响整个检测算法的性能，因此是一个严重的问题。好的，应该解释清楚了吧，好累！

如上图所示，为了得到为了得到固定大小（7X7）的feature map，ROIAlign技术并没有使用量化操作，即我们不想引入量化误差，比如665 / 32 = 20.78，我们就用20.78，不用什么20来替代它，比如20.78 / 7 = 2.97，我们就用2.97，而不用2来代替它。这就是ROIAlign的初衷。那么我们如何处理这些浮点数呢，我们的解决思路是使用“双线性插值”算法。双线性插值是一种比较好的图像缩放算法，它充分的利用了原图中虚拟点（比如20.56这个浮点数，像素位置都是整数值，没有浮点值）四周的四个真实存在的像素值来共同决定目标图中的一个像素值，即可以将20.56这个虚拟的位置点对应的像素值估计出来。厉害哈。如图11所示，蓝色的虚线框表示卷积后获得的feature map，黑色实线框表示ROI feature，最后需要输出的大小是2x2，那么我们就利用双线性插值来估计这些蓝点（虚拟坐标点，又称双线性插值的网格点）处所对应的像素值，最后得到相应的输出。这些蓝点是2x2Cell中的随机采样的普通点，作者指出，这些采样点的个数和位置不会对性能产生很大的影响，你也可以用其它的方法获得。然后在每一个橘红色的区域里面进行max pooling或者average pooling操作，获得最终2x2的输出结果。我们的整个过程中没有用到量化操作，没有引入误差，即原图中的像素和feature map中的像素是完全对齐的，没有偏差，这不仅会提高检测的精度，同时也会有利于实例分割。这么细心，做科研就应该关注细节，细节决定成败。

6. LOSS计算与分析

由于增加了mask分支，每个ROI的Loss函数如下所示：

其中Lcls和Lbox和Faster r-cnn中定义的相同。对于每一个ROI，mask分支有Km*m维度的输出，其对K个大小为m*m的mask进行编码，每一个mask有K个类别。我们使用了per-pixel sigmoid，并且将Lmask定义为the average binary cross-entropy loss 。对应一个属于GT中的第k类的ROI，Lmask仅仅在第k个mask上面有定义（其它的k-1个mask输出对整个Loss没有贡献）。我们定义的Lmask允许网络为每一类生成一个mask，而不用和其它类进行竞争；我们依赖于分类分支所预测的类别标签来选择输出的mask。这样将分类和mask生成分解开来。这与利用FCN进行语义分割的有所不同，它通常使用一个per-pixel sigmoid和一个multinomial cross-entropy loss ，在这种情况下mask之间存在竞争关系；而由于我们使用了一个per-pixel sigmoid 和一个binary loss ，不同的mask之间不存在竞争关系。经验表明，这可以提高实例分割的效果。

一个mask对一个目标的输入空间布局进行编码，与类别标签和BB偏置不同，它们通常需要通过FC层而导致其以短向量的形式输出。我们可以通过由卷积提供的像素和像素的对应关系来获得mask的空间结构信息。具体的来说，我们使用FCN从每一个ROI中预测出一个m*m大小的mask，这使得mask分支中的每个层能够明确的保持m×m空间布局，而不将其折叠成缺少空间维度的向量表示。和以前用fc层做mask预测的方法不同的是，我们的实验表明我们的mask表示需要更少的参数，而且更加准确。这些像素到像素的行为需要我们的ROI特征，而我们的ROI特征通常是比较小的feature map，其已经进行了对其操作，为了一致的较好的保持明确的单像素空间对应关系，我们提出了ROIAlign操作。

三、Mask R-CNN细节分析

1. Head Architecture

如上图所示，为了产生对应的Mask，文中提出了两种架构，即左边的Faster R-CNN/ResNet和右边的Faster R-CNN/FPN。对于左边的架构，我们的backbone使用的是预训练好的ResNet，使用了ResNet倒数第4层的网络。输入的ROI首先获得7x7x1024的ROI feature，然后将其升维到2048个通道（这里修改了原始的ResNet网络架构），然后有两个分支，上面的分支负责分类和回归，下面的分支负责生成对应的mask。由于前面进行了多次卷积和池化，减小了对应的分辨率，mask分支开始利用反卷积进行分辨率的提升，同时减少通道的个数，变为14x14x256，最后输出了14x14x80的mask模板。而右边使用到的backbone是FPN网络，这是一个新的网络，通过输入单一尺度的图片，最后可以对应的特征金字塔，如果想要了解它的细节，请参考该链接。得到证实的是，该网络可以在一定程度上面提高检测的精度，当前很多的方法都用到了它。由于FPN网络已经包含了res5，可以更加高效的使用特征，因此这里使用了较少的filters。该架构也分为两个分支，作用于前者相同，但是分类分支和mask分支和前者相比有很大的区别。可能是因为FPN网络可以在不同尺度的特征上面获得许多有用信息，因此分类时使用了更少的滤波器。而mask分支中进行了多次卷积操作，首先将ROI变化为14x14x256的feature，然后进行了5次相同的操作（不清楚这里的原理，期待着你的解释），然后进行反卷积操作，最后输出28x28x80的mask。即输出了更大的mask，与前者相比可以获得更细致的mask。

如上图所示，图像中红色的BB表示检测到的目标，我们可以用肉眼可以观察到检测结果并不是很好，即整个BB稍微偏右，左边的一部分像素并没有包括在BB之内，但是右边显示的最终结果却很完美。

2. Equivariance in Mask R-CNN

Equivariance 指随着输入的变化输出也会发生变化。

图 Equivariance 1

即全卷积特征（Faster R-CNN网络）和图像的变换具有同变形，即随着图像的变换，全卷积的特征也会发生对应的变化；

图 Equivariance2

在ROI上面的全卷积操作（FCN网络）和在ROI中的变换具有同变性；

图 Equivariance3

ROIAlign操作保持了ROI变换前后的同变性；

图 ROI中的全卷积

图 ROIAlign的尺度同变性

图 Mask R-CNN中的同变性总结

3. 算法实现细节

图算法实现细节

观察上图，我们可以得到以下的信息：

Mask R-CNN中的超参数都是用了Faster r-cnn中的值，机智，省时省力，效果还好，别人已经替你调节过啦，哈哈哈；
使用到的预训练网络包括ResNet50、ResNet101、FPN，都是一些性能很好地网络，尤其是FPN，后面会有分析；
对于过大的图片，它会将其裁剪成800x800大小，图像太大的话会大大的增加计算量的；
利用8个GPU同时训练，开始的学习率是0.01，经过18k次将其衰减为0.001，ResNet50-FPN网络训练了32小时，ResNet101-FPN训练了44小时；
在Nvidia Tesla M40 GPU上面的测试时间是195ms/张；
使用了MS COCO数据集，将120k的数据集划分为80k的训练集、35k的验证集和5k的测试集；

四、性能比较

1. 定量结果分析

表 ROI Pool和ROIAlign性能比较

由前面的分析，我们就可以定性的得到一个结论，ROIAlign会使得目标检测的效果有很大的性能提升。根据上表，我们进行定量的分析，结果表明，ROIAlign使得mask的AP值提升了10.5个百分点，使得box的AP值提升了9.5个百分点。

表 Multinomial和Binary loss比较

根据上表的分析，我们知道Mask R-CNN利用两个分支将分类和mask生成解耦出来，然后利用Binary Loss代替Multinomial Loss，使得不同类别的mask之间消除了竞争。依赖于分类分支所预测的类别标签来选择输出对应的mask。使得mask分支不需要进行重新的分类工作，使得性能得到了提升。

表 MLP与FCN mask性能比较

如上表所示，MLP即利用FC来生成对应的mask，而FCN利用Conv来生成对应的mask，仅仅从参数量上来讲，后者比前者少了很多，这样不仅会节约大量的内存空间，同时会加速整个训练过程（因此需要进行推理、更新的参数更少啦）。除此之外，由于MLP获得的特征比较抽象，使得最终的mask中丢失了一部分有用信息，我们可以直观的从右边看到差别。从定性角度来讲，FCN使得mask AP值提升了2.1个百分点。

表实例分割的结果

表目标检测的结果

观察目标检测的表格，我们可以发现使用了ROIAlign操作的Faster R-CNN算法性能得到了0.9个百分点，Mask R-CNN比最好的Faster R-CNN高出了2.6个百分点。

2. 定性结果分析

图实例分割结果1

图实例分割结果2

图人体姿势识别结果

图失败检测案例1

图失败检测案例2

五、总结

Mask R-CNN论文的主要贡献包括以下几点：

分析了ROI Pool的不足，提升了ROIAlign，提升了检测和实例分割的效果；
将实例分割分解为分类和mask生成两个分支，依赖于分类分支所预测的类别标签来选择输出对应的mask。同时利用Binary Loss代替Multinomial Loss，消除了不同类别的mask之间的竞争，生成了准确的二值mask；
并行进行分类和mask生成任务，对模型进行了加速。

六、代码解析

根据pytorch官方提供的源码链接，进行适当修改和精简后，可以得到一份简易好用的代码，现在对他进行简单的解析。

1.train.py

首先定义了MaskRCNN网络模型，他使用resnet50+fpn作为backbone，并且调用了resnet50的预训练模型进行迁移学习，并在训练时冻结了BN层，同时只训练后三层（即layer4，layer3，layer2）。

def create_model(num_classes, load_pretrain_weights=True):
    # 如果GPU显存很小，batch_size不能设置很大，建议将norm_layer设置成FrozenBatchNorm2d(默认是nn.BatchNorm2d)
    # FrozenBatchNorm2d的功能与BatchNorm2d类似，但参数无法更新
    # trainable_layers包括['layer4', 'layer3', 'layer2', 'layer1', 'conv1']， 5代表全部训练
    # backbone = resnet50_fpn_backbone(norm_layer=FrozenBatchNorm2d,
    #                                  trainable_layers=3)
    # resnet50 imagenet weights url: https://download.pytorch.org/models/resnet50-0676ba61.pth
    backbone = resnet50_fpn_backbone(pretrain_path="resnet50.pth", trainable_layers=3)

    model = MaskRCNN(backbone, num_classes=num_classes)

    if load_pretrain_weights:
        # coco weights url: "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth"
        weights_dict = torch.load("./maskrcnn_resnet50_fpn_coco.pth", map_location="cpu")
        for k in list(weights_dict.keys()):
            if ("box_predictor" in k) or ("mask_fcn_logits" in k):
                del weights_dict[k]

        print(model.load_state_dict(weights_dict, strict=False))

    return model

在main函数下，会先指定设备，然后加载数据集，将图片相似高宽比采样图片组成batch，减小训练时所需GPU显存，通过这样的方式来读取图片。接下来设置学习率、参数更新、和学习率更新方法等；同时如果传入了resume参数接着上次训练的进度继续训练，这里也给出了方法；这些都设置好后，就会开始迭代epoch，开始训练，并把每轮的训练相关参数信息保存到两个txt文件里，以便后续复查，最后保存权重文件，绘制loss、lr、和map的变化图像。

def main(args):
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    print("Using {} device training.".format(device.type))

    # 用来保存coco_info的文件
    now = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    det_results_file = f"det_results{now}.txt"
    seg_results_file = f"seg_results{now}.txt"

    data_transform = {
        "train": transforms.Compose([transforms.ToTensor(),
                                     transforms.RandomHorizontalFlip(0.5)]),
        "val": transforms.Compose([transforms.ToTensor()])
    }

    data_root = args.data_path

    # load train data set
    # coco2017 -> annotations -> instances_train2017.json
    train_dataset = CocoDetection(data_root, "train", data_transform["train"])
    # VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
    # train_dataset = VOCInstances(data_root, year="2012", txt_name="train.txt", transforms=data_transform["train"])
    train_sampler = None

    # 是否按图片相似高宽比采样图片组成batch
    # 使用的话能够减小训练时所需GPU显存，默认使用
    if args.aspect_ratio_group_factor >= 0:
        train_sampler = torch.utils.data.RandomSampler(train_dataset)
        # 统计所有图像高宽比例在bins区间中的位置索引
        group_ids = create_aspect_ratio_groups(train_dataset, k=args.aspect_ratio_group_factor)
        # 每个batch图片从同一高宽比例区间中取
        train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size)

    # 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
    batch_size = args.batch_size
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
    print('Using %g dataloader workers' % nw)

    if train_sampler:
        # 如果按照图片高宽比采样图片，dataloader中需要使用batch_sampler
        train_data_loader = torch.utils.data.DataLoader(train_dataset,
                                                        batch_sampler=train_batch_sampler,
                                                        pin_memory=True,
                                                        num_workers=nw,
                                                        collate_fn=train_dataset.collate_fn)
    else:
        train_data_loader = torch.utils.data.DataLoader(train_dataset,
                                                        batch_size=batch_size,
                                                        shuffle=True,
                                                        pin_memory=True,
                                                        num_workers=nw,
                                                        collate_fn=train_dataset.collate_fn)

    # load validation data set
    # coco2017 -> annotations -> instances_val2017.json
    val_dataset = CocoDetection(data_root, "val", data_transform["val"])
    # VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
    # val_dataset = VOCInstances(data_root, year="2012", txt_name="val.txt", transforms=data_transform["val"])
    val_data_loader = torch.utils.data.DataLoader(val_dataset,
                                                  batch_size=1,
                                                  shuffle=False,
                                                  pin_memory=True,
                                                  num_workers=nw,
                                                  collate_fn=train_dataset.collate_fn)

    # create model num_classes equal background + classes
    model = create_model(num_classes=args.num_classes + 1, load_pretrain_weights=args.pretrain)
    model.to(device)

    train_loss = []
    learning_rate = []
    val_map = []

    # define optimizer
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    scaler = torch.cuda.amp.GradScaler() if args.amp else None

    # learning rate scheduler
    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
                                                        milestones=args.lr_steps,
                                                        gamma=args.lr_gamma)
    # 如果传入resume参数，即上次训练的权重地址，则接着上次的参数训练
    if args.resume:
        # If map_location is missing, torch.load will first load the module to CPU
        # and then copy each parameter to where it was saved,
        # which would result in all processes on the same machine using the same set of devices.
        checkpoint = torch.load(args.resume, map_location='cpu')  # 读取之前保存的权重文件(包括优化器以及学习率策略)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
        args.start_epoch = checkpoint['epoch'] + 1
        if args.amp and "scaler" in checkpoint:
            scaler.load_state_dict(checkpoint["scaler"])

    for epoch in range(args.start_epoch, args.epochs):
        # train for one epoch, printing every 50 iterations
        mean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader,
                                              device, epoch, print_freq=50,
                                              warmup=True, scaler=scaler)
        train_loss.append(mean_loss.item())
        learning_rate.append(lr)

        # update the learning rate
        lr_scheduler.step()

        # evaluate on the test dataset
        det_info, seg_info = utils.evaluate(model, val_data_loader, device=device)

        # write detection into txt
        with open(det_results_file, "a") as f:
            # 写入的数据包括coco指标还有loss和learning rate
            result_info = [f"{i:.4f}" for i in det_info + [mean_loss.item()]] + [f"{lr:.6f}"]
            txt = "epoch:{} {}".format(epoch, '  '.join(result_info))
            f.write(txt + "\n")

        # write seg into txt
        with open(seg_results_file, "a") as f:
            # 写入的数据包括coco指标还有loss和learning rate
            result_info = [f"{i:.4f}" for i in seg_info + [mean_loss.item()]] + [f"{lr:.6f}"]
            txt = "epoch:{} {}".format(epoch, '  '.join(result_info))
            f.write(txt + "\n")

        val_map.append(det_info[1])  # pascal mAP

        # save weights
        save_files = {
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'lr_scheduler': lr_scheduler.state_dict(),
            'epoch': epoch}
        if args.amp:
            save_files["scaler"] = scaler.state_dict()
        torch.save(save_files, "./save_weights/model_{}.pth".format(epoch))

    # plot loss and lr curve
    if len(train_loss) != 0 and len(learning_rate) != 0:
        from plot_curve import plot_loss_and_lr
        plot_loss_and_lr(train_loss, learning_rate)

    # plot mAP curve
    if len(val_map) != 0:
        from plot_curve import plot_map
        plot_map(val_map)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description=__doc__)

    # 训练设备类型
    parser.add_argument('--device', default='cuda:0', help='device')
    # 训练数据集的根目录
    parser.add_argument('--data-path', default='/data/coco2017', help='dataset')
    # 检测目标类别数(不包含背景)
    parser.add_argument('--num-classes', default=90, type=int, help='num_classes')
    # 文件保存地址
    parser.add_argument('--output-dir', default='./save_weights', help='path where to save')
    # 若需要接着上次训练，则指定上次训练保存权重文件地址
    parser.add_argument('--resume', default='', type=str, help='resume from checkpoint')
    # 指定接着从哪个epoch数开始训练
    parser.add_argument('--start_epoch', default=0, type=int, help='start epoch')
    # 训练的总epoch数
    parser.add_argument('--epochs', default=26, type=int, metavar='N',
                        help='number of total epochs to run')
    # 学习率
    parser.add_argument('--lr', default=0.004, type=float,
                        help='initial learning rate, 0.02 is the default value for training '
                             'on 8 gpus and 2 images_per_gpu')
    # SGD的momentum参数
    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                        help='momentum')
    # SGD的weight_decay参数
    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                        metavar='W', help='weight decay (default: 1e-4)',
                        dest='weight_decay')
    # 针对torch.optim.lr_scheduler.MultiStepLR的参数
    parser.add_argument('--lr-steps', default=[16, 22], nargs='+', type=int,
                        help='decrease lr every step-size epochs')
    # 针对torch.optim.lr_scheduler.MultiStepLR的参数
    parser.add_argument('--lr-gamma', default=0.1, type=float, help='decrease lr by a factor of lr-gamma')
    # 训练的batch size(如果内存/GPU显存充裕，建议设置更大)
    parser.add_argument('--batch_size', default=2, type=int, metavar='N',
                        help='batch size when training.')
    parser.add_argument('--aspect-ratio-group-factor', default=3, type=int)
    parser.add_argument("--pretrain", type=bool, default=True, help="load COCO pretrain weights.")
    # 是否使用混合精度训练(需要GPU支持混合精度)
    parser.add_argument("--amp", default=False, help="Use torch.cuda.amp for mixed precision training")

    args = parser.parse_args()
    print(args)

    # 检查保存权重文件夹是否存在，不存在则创建
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    main(args)

2.predict.py

对一张图片进行实例分割

import os
import time
import json

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import torch
from torchvision import transforms

from network_files import MaskRCNN
from backbone import resnet50_fpn_backbone
from draw_box_utils import draw_objs


def create_model(num_classes, box_thresh=0.5):
    backbone = resnet50_fpn_backbone()
    model = MaskRCNN(backbone,
                     num_classes=num_classes,
                     rpn_score_thresh=box_thresh,
                     box_score_thresh=box_thresh)

    return model


def time_synchronized():
    torch.cuda.synchronize() if torch.cuda.is_available() else None
    return time.time()


def main():
    num_classes = 90  # 不包含背景
    box_thresh = 0.5
    weights_path = "./save_weights/mask_rcnn_weights.pth"
    img_path = "./1.jpg"
    label_json_path = './coco91_indices.json'

    # get devices
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    # create model
    model = create_model(num_classes=num_classes + 1, box_thresh=box_thresh)

    # load train weights
    assert os.path.exists(weights_path), "{} file dose not exist.".format(weights_path)
    weights_dict = torch.load(weights_path, map_location='cpu')
    weights_dict = weights_dict["model"] if "model" in weights_dict else weights_dict
    model.load_state_dict(weights_dict)
    model.to(device)

    # read class_indict
    assert os.path.exists(label_json_path), "json file {} dose not exist.".format(label_json_path)
    with open(label_json_path, 'r') as json_file:
        category_index = json.load(json_file)

    # load image
    assert os.path.exists(img_path), f"{img_path} does not exits."
    original_img = Image.open(img_path).convert('RGB')

    # from pil image to tensor, do not normalize image
    data_transform = transforms.Compose([transforms.ToTensor()])
    img = data_transform(original_img)
    # expand batch dimension
    img = torch.unsqueeze(img, dim=0)

    model.eval()  # 进入验证模式
    with torch.no_grad():
        # init
        img_height, img_width = img.shape[-2:]
        init_img = torch.zeros((1, 3, img_height, img_width), device=device)
        model(init_img)

        t_start = time_synchronized()
        predictions = model(img.to(device))[0]
        t_end = time_synchronized()
        print("inference+NMS time: {}".format(t_end - t_start))

        predict_boxes = predictions["boxes"].to("cpu").numpy()
        predict_classes = predictions["labels"].to("cpu").numpy()
        predict_scores = predictions["scores"].to("cpu").numpy()
        predict_mask = predictions["masks"].to("cpu").numpy()
        predict_mask = np.squeeze(predict_mask, axis=1)  # [batch, 1, h, w] -> [batch, h, w]

        if len(predict_boxes) == 0:
            print("没有检测到任何目标!")
            return

        plot_img = draw_objs(original_img,
                             boxes=predict_boxes,
                             classes=predict_classes,
                             scores=predict_scores,
                             masks=predict_mask,
                             category_index=category_index,
                             line_thickness=3,
                             font='arial.ttf',
                             font_size=20)
        plt.imshow(plot_img)
        plt.show()
        # 保存预测的图片结果
        plot_img.save("test_result.jpg")


if __name__ == '__main__':
    main()

代码首先加载已训练权重文件、预测图片、和存有类别信息的、json文件，对图像进行简单处理后开始预测，并最终把结果绘制在图片上

3.predict2.py

为了实现对多张图片实现预测，我对代码进行了改进，实现了对文件夹下多张图片预测的功能

import os
import time
import json

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import torch
from torchvision import transforms

from network_files import MaskRCNN
from backbone import resnet50_fpn_backbone
from draw_box_utils import draw_objs


def create_model(num_classes, box_thresh=0.5):
    backbone = resnet50_fpn_backbone()
    model = MaskRCNN(backbone,
                     num_classes=num_classes,
                     rpn_score_thresh=box_thresh,
                     box_score_thresh=box_thresh)

    return model


def time_synchronized():
    torch.cuda.synchronize() if torch.cuda.is_available() else None
    return time.time()


# 预测多张图片
def main():
    num_classes = 90  # 不包含背景
    box_thresh = 0.5
    weights_path = "./save_weights/mask_rcnn_weights.pth"

    # get devices
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))
    
    files = os.listdir('./images/')
    for file in files:
        img_name = file.split('.')
        img_path = './images/' + img_name[0] + ".jpg"
        print(img_path)
        label_json_path = './coco91_indices.json'


        # create model
        model = create_model(num_classes=num_classes + 1, box_thresh=box_thresh)

        # load train weights
        assert os.path.exists(weights_path), "{} file dose not exist.".format(weights_path)
        weights_dict = torch.load(weights_path, map_location='cpu')
        weights_dict = weights_dict["model"] if "model" in weights_dict else weights_dict
        model.load_state_dict(weights_dict)
        model.to(device)

        # read class_indict
        assert os.path.exists(label_json_path), "json file {} dose not exist.".format(label_json_path)
        with open(label_json_path, 'r') as json_file:
            category_index = json.load(json_file)

        # load image
        assert os.path.exists(img_path), f"{img_path} does not exits."
        original_img = Image.open(img_path).convert('RGB')

        # from pil image to tensor, do not normalize image
        data_transform = transforms.Compose([transforms.ToTensor()])
        img = data_transform(original_img)
        # expand batch dimension
        img = torch.unsqueeze(img, dim=0)

        model.eval()  # 进入验证模式
        with torch.no_grad():
            # init
            img_height, img_width = img.shape[-2:]
            init_img = torch.zeros((1, 3, img_height, img_width), device=device)
            model(init_img)

            t_start = time_synchronized()
            predictions = model(img.to(device))[0]
            t_end = time_synchronized()
            print("inference+NMS time: {}".format(t_end - t_start))

            predict_boxes = predictions["boxes"].to("cpu").numpy()
            predict_classes = predictions["labels"].to("cpu").numpy()
            predict_scores = predictions["scores"].to("cpu").numpy()
            predict_mask = predictions["masks"].to("cpu").numpy()
            predict_mask = np.squeeze(predict_mask, axis=1)  # [batch, 1, h, w] -> [batch, h, w]

            if len(predict_boxes) == 0:
                print("没有检测到任何目标!")
                return

            plot_img = draw_objs(original_img,
                                boxes=predict_boxes,
                                classes=predict_classes,
                                scores=predict_scores,
                                masks=predict_mask,
                                category_index=category_index,
                                line_thickness=3,
                                font='arial.ttf',
                                font_size=20)
            # plt.imshow(plot_img)
            # plt.show()
            # 保存预测的图片结果
            plot_img.save("./result/" + img_name[0] + ".jpg")


if __name__ == '__main__':
    main()

代码会读取images文件夹下的所有图片，并对他们逐一预测，最终将输出的结果保存在result文件夹下

4.predict3.py

只能处理图片还不够，所以我又修改了代码，使其可以调用摄像头，对摄像头的每帧图片实时预测，代码如下

import os
import time
import json
import cv2

import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

import torch
from torchvision import transforms

from network_files import MaskRCNN
from backbone import resnet50_fpn_backbone
from draw_box_utils import draw_objs


def create_model(num_classes, box_thresh=0.5):
    backbone = resnet50_fpn_backbone()
    model = MaskRCNN(backbone,
                     num_classes=num_classes,
                     rpn_score_thresh=box_thresh,
                     box_score_thresh=box_thresh)

    return model


def time_synchronized():
    torch.cuda.synchronize() if torch.cuda.is_available() else None
    return time.time()


# 预测相机流
def main():
    num_classes = 90  # 不包含背景
    box_thresh = 0.5

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"using {device} device.")

    weights_path = "./save_weights/mask_rcnn_weights.pth"
    label_json_path = './coco91_indices.json'

    # create model
    model = create_model(num_classes=num_classes + 1, box_thresh=box_thresh)

    # load train weights
    assert os.path.exists(weights_path), "{} file dose not exist.".format(weights_path)
    weights_dict = torch.load(weights_path, map_location='cpu')
    weights_dict = weights_dict["model"] if "model" in weights_dict else weights_dict
    model.load_state_dict(weights_dict)
    model.to(device)

    # read class_indict
    assert os.path.exists(label_json_path), "json file {} dose not exist.".format(label_json_path)
    with open(label_json_path, 'r') as json_file:
        category_index = json.load(json_file)

    
    cap = cv2.VideoCapture(0)
    while 1:
        ret, frame = cap.read()

        # cv2.imshow("frame", frame)    

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        original_img = Image.fromarray(frame)

        # from pil image to tensor, do not normalize image
        data_transform = transforms.Compose([transforms.ToTensor()])
        img = data_transform(original_img)
        # expand batch dimension
        img = torch.unsqueeze(img, dim=0)

        model.eval()  # 进入验证模式
        with torch.no_grad():
            # init
            img_height, img_width = img.shape[-2:]
            init_img = torch.zeros((1, 3, img_height, img_width), device=device)
            model(init_img)

            t_start = time_synchronized()
            predictions = model(img.to(device))[0]
            t_end = time_synchronized()
            print("inference+NMS time: {}".format(t_end - t_start))

            predict_boxes = predictions["boxes"].to("cpu").numpy()
            predict_classes = predictions["labels"].to("cpu").numpy()
            predict_scores = predictions["scores"].to("cpu").numpy()
            predict_mask = predictions["masks"].to("cpu").numpy()
            predict_mask = np.squeeze(predict_mask, axis=1)  # [batch, 1, h, w] -> [batch, h, w]
            
            if len(predict_boxes) == 0:
                print("没有检测到任何目标!")
                return

            plot_img = draw_objs(original_img,
                                boxes=predict_boxes,
                                classes=predict_classes,
                                scores=predict_scores,
                                masks=predict_mask,
                                category_index=category_index,
                                line_thickness=3,
                                font='arial.ttf',
                                font_size=20)
            # print(type(plot_img))
            plot_img = np.array(plot_img)
            # print(type(plot_img))
            converted_img = cv2.cvtColor(plot_img, cv2.COLOR_RGB2BGR)
            cv2.imshow("frame", converted_img)  # 修改这里
            if cv2.waitKey(1) & 0xff == ord("q"):  # 修改这里
                break     
 

if __name__ == '__main__':
    main()

七.模型pth文件解读

1.读取基本信息

新建一个pth.py文件

查看pth文件的基本信息

import torch

weight_path = "./save_weights/mask_rcnn_weights.pth"
net = torch.load(weight_path)

print(len(net))
for i in net.keys():
    print(i)

print(net["model"])

print(type(net))输出的为dict，表明模型是一个键值对，而len(net)输出为1，他的key为model，因此我们打印了model的内容。

import torch

weight_path = "./save_weights/mask_rcnn_weights.pth"
net = torch.load(weight_path)

print(len(net))
# for i in net.keys():
#     print(i)

# print(net["model"])

for key, value in net["model"].items():
    print(key, value.size(),sep="  ")

输出为

建model所对应的值是一个OrderedDict，而这个OrderedDict字典里面存储着每一层的参数名称以及对应的参数

2.加载模型

import torch
from network_files import MaskRCNN
from backbone import resnet50_fpn_backbone

weight_path = "./save_weights/mask_rcnn_weights.pth"
weight_dict = torch.load(weight_path)

# 创建相应的模型结构
model = MaskRCNN(resnet50_fpn_backbone(), 91, 0.5, 0.5)

# 读取网络中的model参数
weight = weight_dict["model"]

# 加载模型权重
model.load_state_dict(weight)

整个代码链接：https://pan.baidu.com/s/1XIsn-LJSAFuGFGhd9FRAtQ提取码: e3ma 复制这段内容后打开百度网盘手机App，操作更方便哦
参考链接：https://blog.csdn.net/WZZ18191171661/article/details/79453780

nice-wyh

关注

11
点赞
踩
19

收藏

觉得还不错? 一键收藏
0
评论
mask-rcnn原理与实战

Mask R-CNN论文的主要贡献包括以下几点：分析了ROI Pool的不足，提升了ROIAlign，提升了检测和实例分割的效果；将实例分割分解为分类和mask生成两个分支，依赖于分类分支所预测的类别标签来选择输出对应的mask。同时利用Binary Loss代替Multinomial Loss，消除了不同类别的mask之间的竞争，生成了准确的二值mask；并行进行分类和mask生成任务，对模型进行了加速。
复制链接

扫一扫

专栏目录