DETR论文笔记

End-to-End Object Detection with Transformers

论文笔记及与相关代码阅读

---

---

一、论文阅读

摘要：提出了一种将目标检测看作直接集预测问题的新方法。我们的方法简化了检测流程，有效地消除了对许多手工设计的组件的需要，例如非最大抑制nms或achor生成，它们明确地编码了我们关于任务的先验知识。新框架的主要组成部分，称为DETR，是一个基于集合的全局损失，通过二部匹配以及Transformers编码器-解码器架构强制唯一的预测。给定一个固定的小范围学习对象查询集，DETR根据对象与全局图像上下文的关系，直接并行输出最终的预测集。与许多其他检测算法不同，新模型在概念上很简单，不需要专门的库。在具有挑战性的COCO目标检测数据集上，DETR的准确度和运行时性能与成熟且高度优化的Faster RCNN基线相当。此外，DETR可以很容易地推广到产生一个统一的方式全景分割。我们发现，它显著优于竞争基线。代码

1 介绍

  目标检测的目标是为每个感兴趣的目标预测一组边界框和类别标签。现代检测器通过在一大组建议[37,5]、锚[23]或窗口中心[53,46]上定义替代回归和分类问题，间接地解决了这一集合预测任务。它们的性能受到后处理步骤（瓦解接近重复的预测）、锚集的设计以及将目标框分配给锚的启发式算法的显著影响[52]。为了简化这些流程，我们提出了一种直接集预测方法来绕过代理任务。这种端到端的思想在复杂的结构化预测任务（如机器翻译或语音识别）方面取得了重大进展，但在目标检测方面还没有取得进展：以前的尝试[43,16,4,39]要么添加了其他形式的先验知识，要么在具有挑战性的基准上没有被证明与强大的基线具有竞争力。本文旨在填补这一空白。
  我们将目标检测视为一个直接集预测问题来简化训练。我们采用了一种基于transformers的编解码器架构[47]，这是一种流行的序列预测架构。transformers的自注意机制可以显式地对序列中元素之间的所有成对交互进行建模，这使得这些体系结构特别适用于集预测的特定约束，例如消除重复预测。
  我们的检测变换器（DETR，见图1）一次预测所有对象，并使用一个集损函数进行端到端训练，该函数在预测对象和地面真实对象之间执行二部匹配。DETR通过删除编码先验知识的多个手工设计的组件（如anchor或非最大抑制nms），简化了检测流程。与大多数现有的检测方法不同，DETR不需要任何定制的层，因此可以在任何包含标准CNN和transformer类的框架中轻松复制。
  与以往大多数关于直接集预测的工作相比，DETR的主要特点是二部匹配损耗和transformer（非自回归）并行译码的结合[29,12,10,8]。相比之下，以前的工作主要集中在使用RNNs的自回归解码上[43,41,30,36,42]。我们的匹配损失函数唯一地将预测分配给地面真值对象，并且对预测对象的排列是不变的，因此我们可以并行地发射它们。
  我们在一个最流行的目标检测数据集COCO[24]上评估了DETR，并与一个非常有竞争力的Faster R-CNN基线进行了比较[37]。Faster R-CNN经历了多次设计迭代，自最初发布以来，其性能得到了极大的提高。我们的实验表明，我们的新模型达到了相当的性能。更准确地说，DETR在大型对象上表现出明显更好的性能，这一结果很可能是由transformer的非局部计算实现的。但是，它在小对象上的性能较低。我们期望未来的工作能改善这一方面，像FPN[22]的发展Faster R-CNN一样。
  DETR的设计理念很容易扩展到更复杂的任务。在我们的实验中，我们证明了一个简单的分割头训练在一个预先训练好的DETR上比全景分割的竞争基线更具竞争力[19]，这是一个具有挑战性的像素级识别任务，最近得到了广泛的应用。

该模型首先是经历一个CNN提取特征，然后得到的特征进入transformer, 最后将transformer输出的结果转化为class和box.输出固定个数的预测结果。最后将得到的检测结果和GroundTruth进行二部图匹配计算loss

2 相关工作

我们的工作建立在几个领域的先前工作之上：集合预测的二部匹配损失、基于transformer的编码器-解码器结构、并行解码和目标检测方法。

2.1 Set Prediction

  没有一个规范的深度学习模型可以直接预测集合。基本的集合预测任务是多标签分类（参见[40,33]以获取计算机视觉背景下的参考文献），对于该分类，基线方法“一对一”不适用于诸如在元素之间存在底层结构（即，接近相同的框）的检测等问题。这些任务的第一个困难是避免近乎重复的任务。目前大多数检测器采用非最大值抑制等后处理方法来解决这一问题，而直接集预测是无需后处理的。它们需要全局推理方案来模拟所有预测元素之间的交互，以避免冗余。对于恒定大小集预测，密集的完全连接网络[9]是足够的，但成本高昂。一般的方法是使用自回归序列模型，如递归神经网络[48]。在所有情况下，损失函数都应通过预测的排列保持不变。通常的解决方案是基于匈牙利算法[20]设计一个损耗，在地面真值和预测之间找到一个二部匹配。这强制了置换不变性，并保证每个目标元素具有唯一的匹配。我们遵循二部匹配损失方法。然而，与大多数先前的工作相比，我们远离自回归模型，使用具有并行解码的transformer，我们将在下面描述。

2.2 transformer与并行译码

  Vaswani等人[47]引入了Transformers，作为机器翻译的一个新的基于注意力的构建块。注意机制[2]是从整个输入序列聚合信息的神经网络层。Transformers引入了自我注意层，与非局部神经网络[49]类似，它扫描序列的每个元素，并通过聚集整个序列的信息来更新。基于注意的模型的一个主要优点是其全局计算能力和完美的记忆能力，这使得它比RNNs更适合于长序列。在自然语言处理、语音处理和计算机视觉的许多问题中，Transformers正在取代RNN[8,27,45,34,31]。
  Transformers首先用于自回归模型，继早期的序列到序列模型[44]之后，一个接一个地生成输出标记。然而，在音频[29]、机器翻译[12,10]、词表示学习[8]和最近的语音识别[6]等领域，禁止性的推理成本（与输出长度成正比，且难以批量）导致了并行序列生成的发展。我们还结合Transformers和并行解码之间的适当权衡计算成本和能力，以执行集预测所需的全局计算。

2.3 目标检测

  基于集合的损失。一些目标探测器[9,25,35]使用了二部匹配损失。然而，在这些早期的深度学习模型中，不同的预测之间的关系仅用卷积层或完全连接层来建模，手工设计的NMS后处理可以提高它们的性能。最近的检测器[37,23,53]在真值和预测以及NMS之间使用非唯一分配规则。
  可学习的NMS方法[16,4]和关系网络[17]显式地建模了不同预测之间的关系。使用直接设置损耗，它们不需要任何后处理步骤。然而，这些方法采用了额外的手工构建的上下文特征，如建议框坐标，来有效地建立检测之间的关系模型，同时我们寻找减少模型中编码的先验知识的解决方案。
  循环检测器。最接近我们的方法是目标检测的端到端集预测[43]和实例分割[41,30,36,42]。与我们类似，他们使用基于CNN激活的编码器-解码器结构的二部匹配损失来直接生成一组边界框。然而，这些方法只在小数据集上进行了评估，没有对照现代基线进行评估。特别是，它们基于自回归模型（更准确地说是RNN），因此它们不会利用最近的Transformers进行并行解码。

3 DETR模型

  在检测中，直接集合预测有两个要素是必不可少的：（1）集合预测损失，强制预测和地面真值框之间的唯一匹配；（2）预测（在一次传递中）一组对象并对其关系建模的体系结构。我们在图2中详细描述了我们的体系结构。

3.1 目标检测集预测损失

  DETR在通过解码器的一次过程中推断N个预测的固定大小集，其中N被设置为显著大于图像中对象的典型数目。训练的主要困难之一是根据地面真实情况对预测对象（类别、位置、大小）进行评分。我们的损失产生一个最佳的二部匹配之间的预测和地面真相的对象，然后优化特定对象（边界框）的损失。
  首先将gt box也变成长度为N的序列以便和网络输出进行匹配，不够长度的用 ∅ 补充，然后对这个序列排列组合，找到和预测的N个序列损失最小的序列来优化。这样就可以得到一对一的关系，也不用后处理操作NMS。

3.2 DETR架构

  它包含三个主要组件，我们将在下面描述：一个用于提取紧凑特征表示的CNN主干、一个编码器-解码器转换器和一个用于进行最终检测预测的简单前馈网络（FFN）。

backbone就是传统的CNN结构，用于提取图像的2D信息，网络一开始是使用Backbone（比如ResNet）提取一些feature，然后降维到d×HW，backbone：[3, H, W] 变为[2048, H/32, W/32]。prediction heads用于对decoder的输出进行分类。
  Transformer encoder。首先，1x1卷积将高级激活图f的通道维数从C降低到更小的维数d。创建一个新的特征映射z0∈Rd×H×W。编码器期望一个序列作为输入，因此我们将z0的空间维数压缩为一维，得到一个d×HW特征映射。每个编码器层都有一个标准的体系结构，由一个多头自注意模块和一个前馈网络（FFN）组成。由于transformer结构是置换不变的，因此我们用固定的位置编码[31,3]来补充它，这些编码被添加到每个注意层的输入中。我们根据补充材料对该体系结构的详细定义，该定义遵循[47]中所述的定义。
  Transformer decoder. 解码器遵循转换器的标准架构，使用多头自和编码器-解码器注意机制转换大小为d的N个嵌入。与原始转换器的不同之处在于，我们的模型在每个解码器层并行解码N个对象，而Vaswani等人[47]使用自回归模型一次预测一个元素的输出序列。我们建议不熟悉这些概念的读者参阅补充材料。由于解码器也是置换不变的，N个输入嵌入必须不同才能产生不同的结果。这些输入嵌入是学习的位置编码，我们称之为对象查询，类似于编码器，我们将它们添加到每个注意层的输入中。解码器将N个对象查询转化为一个嵌入输出。然后，它们被前馈网络（在下一小节中描述）独立地解码成方框坐标和类标签，从而得到N个最终预测。该模型利用自身和编解码器对这些嵌入的关注，利用对象之间的成对关系，对所有对象进行全局推理，同时能够使用整个图像作为上下文。

1. 和NLP中的transformer不太一样，DETR中在decoder部分将Object queries一起作为输入，而非序列化地一个一个输入。
2. 只用image features进行编码是不够的，需要融入位置信息，也就是上图中的Spatial positional encoding部分。
     Prediction feed-forward networks (FFNs) 最终预测是由具有ReLU激活功能且具有隐藏层的3层感知器和线性层计算的。 FFN预测框的标准化中心坐标，高度和宽度， 输入图像，然后线性层使用softmax函数预测类标签。 由于我们预测了一组固定大小的N个边界框，其中N通常比图像中感兴趣的对象的实际数量大得多，因此使用了一个额外的特殊类标签∅来表示在未检测到任何对象。 此类在标准对象检测方法中与“背景”类具有相似的作用。
     

二、源码阅读

pytorch版本的detr推理过程只有50行代码。给定固定的学习对象查询集，则DETR会考虑对象与全局图像上下文之间的关系，以直接并行并行输出最终的预测集。 由于这种并行性质，DETR非常快速和高效。
python中的transformer

1.a simplified version of DETR

a demo of DETR (Detection Transformer), with slight differences with the baseline model in the paper.We show how to define the model, load pretrained weights and visualize bounding box and class predictions.基于预训练权重实现detr的简化版，并测试一张图像。

得益于transformer的代表性功能，DETR架构非常简单。有两个主要组成部分：
backbone:采用Resnet50
Transformer：使用默认的PyTorch nn.Transformer

模型代码如下：

'''
detr简单应用demo推理一张图像的目标检测
'''

from PIL import Image
import requests
import matplotlib.pyplot as plt
import torch
from torch import nn
from torchvision.models import resnet50
import torchvision.transforms as T
torch.set_grad_enabled(False)  # 前向推理，不需要梯度反传

# 简化的DETR实现
class DETRdemo(nn.Module):
    """
    Demo DETR implementation.

    Demo implementation of DETR in minimal number of lines, with the
    following differences wrt DETR in the paper:
    * learned positional encoding (instead of sine)
    * positional encoding is passed at input (instead of attention)
    * fc bbox predictor (instead of MLP)
    The model achieves ~40 AP on COCO val5k and runs at ~28 FPS on Tesla V100.
    Only batch size 1 supported.
    """

    def __init__(self, num_classes, hidden_dim=256, nheads=8, num_encoder_layers=6, num_decoder_layers=6):
        # 6个编码器/解码器,nheads为多头注意力的头数
        super().__init__()

        # create ResNet-50 backbone
        self.backbone = resnet50()
        del self.backbone.fc

        # create conversion layer （输入transformer之前降低特征通道数）
        self.conv = nn.Conv2d(2048, hidden_dim, 1)  # 把通道数从2048变成hidden_dim，卷积核大小为1

        # create a default PyTorch transformer
        self.transformer = nn.Transformer(
            hidden_dim, nheads, num_encoder_layers, num_decoder_layers)
        # PyTorch将Transformer相关的模型分为nn.TransformerEncoderLayer、nn.TransformerDecoderLayer、nn.LayerNorm等几个部分
        '''
        pytorch中有5类：
        1.torch.nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=2048, 
        dropout=0.1, activation='relu', custom_encoder=None, custom_decoder=None)
        transformer模型，基于论文 Attention Is All You Need
        d_model –编码器/解码器输入大小（默认 512）
        nhead –多头注意力模型的头数（默认为8）
        num_encoder_layers –编码器中子编码器层的数量（默认为6）
        num_decoder_layers –解码器中子解码器层的数量（默认为6）
        dim_feedforward –前馈网络模型的中间层维度（默认= 2048）
        activation–编码器/解码器中间层的激活函数，relu或gelu（默认值= relu）
        custom_encoder –自定义编码器（默认=None）
        custom_decoder –自定义解码器（默认=None）

        2.torch.nn.TransformerEncoder(encoder_layer, num_layers, norm=None)
        TransformerEncoder是N个编码器层的堆叠
        encoder_layer – TransformerEncoderLayer（）的实例（必需）
        num_layers –编码器中的子编码器层数（必填）。

        3.torch.nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu')
        TransformerEncoderLayer 由self-attn和feedforward组成
        d_model – 编码器/解码器输入中预期词向量的大小.
        nhead – 多头注意力模型中的头数.
        dim_feedforward – 前馈网络模型的尺寸(默认值= 2048).
        activation – 编码器/解码器中间层，激活函数relu或gelu(默认=relu).
        '''

        # prediction heads, one extra class for predicting non-empty slots
        # note that in baseline DETR linear_bbox layer is 3-layer MLP
        # bbox含框的中心的坐标和宽度高度，即（x,y,w,h）四维
        self.linear_class = nn.Linear(hidden_dim, num_classes + 1)
        self.linear_bbox = nn.Linear(hidden_dim, 4)

        # output positional encodings (object queries)
        # 一张图片最多检测100个物体
        self.query_pos = nn.Parameter(torch.rand(100, hidden_dim))

        # spatial positional encodings
        # note that in baseline DETR we use sine positional encodings
        # 编码器的位置编码，对特征图行和列分别进行位置编码，后面会将两者拼接起来故维度为一半
        # 特征图尺寸不超过50*50  50个行，50个列最多
        self.row_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))  # 行
        self.col_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))  # 列
        # nn.Parameter参数可学习的，这个tensor是Parameter类

    def forward(self, inputs):
        # propagate inputs through ResNet-50 up to avg-pool layer
        # inputs为[3,H,W]
        x = self.backbone.conv1(inputs)  # resnet50的conv1图像变成[64,H/2,W/2]  网络输入图像深度为3(正好是彩色图片的3个通道)，输出深度为64，滤波器为7*7，步长为2，填充3层，特征图的h,w缩小1/2
        x = self.backbone.bn1(x)
        x = self.backbone.relu(x)
        x = self.backbone.maxpool(x)   # 此时，特征图的尺寸已成为输入的1/4, [64,H/4,W/4]

        x = self.backbone.layer1(x)  # [256,H/4,W/4]
        x = self.backbone.layer2(x)  # [512,H/8,W/8]
        x = self.backbone.layer3(x)  # [1024,H/16,W/16]
        x = self.backbone.layer4(x)  # [2048,H/32,W/32]

        # convert from 2048 to 256 feature planes for the transformer
        h = self.conv(x)    # [hidden_dim,H/32,W/32]  论文中的[d,H0/32,W0/32]

        # construct positional encodings
        H, W = h.shape[-2:]  # img.shape[:2] 取彩色图片的长、宽；img.shape[:3] 则取彩色图片的长、宽、通道；img.shape[0]：图像的垂直尺寸（高度）;pytorch tensor matrix: (N, C, H, W)
        pos = torch.cat([
            self.col_embed[:W].unsqueeze(0).repeat(H, 1, 1),  # repeat(H, 1, 1)沿着特定的维度重复这个张量；unsqueeze(0)在0维上增加一维，[H,W,hidden_dim//2]
            self.row_embed[:H].unsqueeze(1).repeat(1, W, 1),  # [H,W,hidden_dim//2]
        ], dim=-1).flatten(0, 1).unsqueeze(1)  # （H×W，1，hidden_dim） 对应论文中d×HW的特征图
        # torch.cat是将两个张量（tensor）拼接在一起，torch.cat((A,B),0)就表示按维数0（行）拼接A和B，也就是竖着拼接，A上B下；torch.cat((A,B),1)就表示按维数1（列）拼接A和B，也就是横着拼接，A左B右
        # dim=-1 表示倒数第一维即hiddden_dim//2上
        # torch.flatten(input, start_dim, end_dim=) 把input从第start_dim起到end_dim止展平
        # 变成（H×W，1，hidden_dim）,1为batchsize

        # propagate through the transformer
        h = self.transformer(pos + 0.1 * h.flatten(2).permute(2, 0, 1),
                             self.query_pos.unsqueeze(1)).transpose(0, 1)
        # 现在query_pos大小为（100, 1，hidden_dim），得到的h大小为（batch, 100，hidden_dim）

        # finally project transformer outputs to class labels and bounding boxes
        return {
   'pred_logits': self.linear_class(h),
                'pred_boxes': self.linear_bbox(h).sigmoid()}

让我们用80个COCO输出类别+ 1个“无对象”类别构建模型，并加载预训练的权重。 权重以半精度保存，以节省带宽而不损害模型精度。
加载初始化模型实例化一个模型：

detr = DETRdemo(num_classes=91)
state_dict = torch.hub.load_state_dict_from_url(
    url='https://dl.fbaipublicfiles.com/detr/detr_demo-da2a99e9.pth',
    map_location='cpu', check_hash=True)      # 加载预训练模型
detr.load_state_dict(state_dict)  # 模型加载到网络中
detr.eval()  # 测试模型

我们刚刚加载的预训练DETR模型已经在80个COCO类上进行了训练，类索引的范围是1到90（这就是为什么我们在模型构造中考虑了91个类）。 在以下单元格中，我们定义从类索引到名称的映射。

# COCO classes
CLASSES = [
    'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A',
    'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
    'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack',
    'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
    'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
    'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
    'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A',
    'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
    'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A',
    'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
    'toothbrush'
]

# colors for visualization
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]

DETR使用标准ImageNet归一化，并以[xcenter，ycenter，w，h]格式输出相对图像坐标中的框，其中[xcenter，ycenter]是边界框的预测中心，而w，h是边框的宽度和高度。 由于坐标是相对于图像尺寸的并且位于[0,1]之间，因此出于可视化目的，我们将预测转换为绝对图像坐标和[x0，y0，x1，y1]格式。

# DETR使用标准ImageNet归一化，并以[xcenter，ycenter，w，h]格式输出相对图像坐标中的框，其中[xcenter，ycenter]是边界框的预测中心，
# 而w，h是边框的宽度和高度。 由于坐标是相对于图像尺寸的并且位于[0,1]之间，因此出于可视化目的，我们将预测转换为绝对图像坐标和[x0，y0，x1，y1]格式
# standard PyTorch mean-std input image normalization 标准PyTorch均值输入图像归一化
transform = T.Compose([
    T.Resize(800),
    T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 前面的[0.485, 0.456, 0.406]是RGB三个通道上的均值mean， 后面[0.229, 0.224, 0.225]是三个通道的标准差std
])  # Compose的主要作用是将多个变换组合在一起。


# for output bounding box post-processing
# 将bbox的中心x、中心y坐标，高度宽度转换为左上角右下角坐标
def box_cxcywh_to_xyxy(x):
    x_c, y_c, w, h = x.unbind(1)  # torch.unbind(input, dim=0) → seq，此方法就是将我们的input从dim进行切片，并返回切片的结果，返回的结果里面没有dim这个维度。
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
         (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=1)  # shape为[batch,4]


def rescale_bboxes(out_bbox, size):
    img_w, img_h = size
    b = box_cxcywh_to_xyxy(out_bbox)
    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)  # bbox归一化坐标转为基于图像尺寸的绝对坐标
    return b

封装推断过程，实现输入到输出预测结果：

# Let's put everything together in a detect function:封装整个推断过程获得预测结果
def detect(im, model, transform):
    # mean-std normalize the input image (batch-size: 1)
    img = transform(im).unsqueeze(0)

    # demo model only support by default images with aspect ratio between 0.5 and 2
    # if you want to use images with an aspect ratio outside this range
    # rescale your image so that the maximum size is at most 1333 for best results
    assert img.shape[-2] <= 1600 and img.shape[-1] <= 1600, 'demo model only supports images up to 1600 pixels on each side'

    # propagate through the model
    outputs = model(img)

    # keep only predictions with 0.7+ confidence
    probas = outputs['pred_logits'].softmax(-1)[0, :, :-1]
    keep = probas.max(-1).values > 0.7  # 只取置信度大于0.7的

    # convert boxes from [0; 1] to image scales
    bboxes_scaled = rescale_bboxes(outputs['pred_boxes'][0, keep], im.size)
    return probas[keep], bboxes_scaled

对一张图进行测试

# To try DETRdemo model on your own image just change the URL below.
# 对一张图进行测试
url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
im = Image.open(requests.get(url, stream=True).raw)
scores, boxes = detect(im, detr, transform)

结果可视化

# Let's now visualize the model predictions
def plot_results(pil_img, prob, boxes):
    plt.figure(figsize=(16, 10))
    plt.imshow(pil_img)
    ax = plt.gca()
    for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), COLORS * 100):
        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
                                   fill=False, color=c, linewidth=3))
        cl = p.argmax()
        text = f'{CLASSES[cl]}: {p[cl]:0.2f}'
        ax.text(xmin, ymin, text, fontsize=15,
                bbox=dict<