tf.map_fn合并yolov3 feature map

tf.map_fn合并yolov3 feature map

微信公众号：幼儿园的学霸
个人的学习笔记，关于OpenCV,关于机器学习, …。问题或建议，请公众号留言;

在代码yolov3_tensorflow中，网络输出为3个feature map，如果要利用该网络进行预测，那么还需要额外对3个feature map进行合并、提取结果，比较麻烦。有必要对网络输出进行处理，以使网络的输出结果为目标检测框[x_min,y_min,x_max,y_max,score,id]的形式。

目录

函数介绍

函数声明如下:

map_fn(fn, #一个可调用函数，用来对elems中的元素进行处理
    elems, #要处理的tensor，tf会从第0维(最外层)展开，进行map操作，如一批图片[batchsiz,H,W,C]，则是对其中的每一张图片调用函数fn进行处理
    dtype=None, # fn函数的返回类型(对elems处理后的返回类型)，如果fn返回的类型和输入elems中的不通，此时需要显式指定fn的返回类型
    parallel_iterations=None, # 允许并行运行的迭代次数
    back_prop=True,
    swap_memory=False, 
    infer_shape=True, #对一致输出形状的测试，建议选择False
    name=None)#返回的张量的名称前缀

例如，一个tensor A具有这样的shape:[?, 13,13,255],其中?表示batchsize，该维度和训练神经网络时的输入有关，我希望tensorflow能够自动对A batch内的张量进行操作(对[13,13,255]进行数据处理)，即输入不含?(batchsize)这个维度，那么就可以采用map_fn函数，参数elems= A,此时fn接收的输入即为[13,13,255]

简单示例

1. map_fn 的最简单版本是反复地将可调用的fn应用于从第一个到最后一个的元素序列。这些元素由 elems解压缩的张量构成。如：

import tensorflow as tf
tf.enable_eager_execution()
import numpy as np

elems = np.array([1, 2, 3, 4, 5, 6],dtype=np.int32)

def squareElements(arg):
    res = arg * arg
    res = tf.cast(res,dtype=tf.float32)
    return res*1.0

squares = tf.map_fn(squareElements, elems,dtype=tf.float32)

print(squares)
# 输出结果
# tf.Tensor([ 1.  4.  9. 16. 25. 36.], shape=(6,), dtype=float32)

这种用法类似于对tensor中的元素进行一次遍历操作，输入和输入元素的shape是一致的，因此infer_shape参数为默认值True，但是数据类型不一致，所以需要显式指明fn的返回类型

该示例的一种应用场景：在目标检测任务中，神经网络以[x_min,y_min,x_max,y_max,score,cls_id]的形式输出了一副图像中的检测结果，但是这个结果是没有经过处理的，同一类别，存在重合的box，我们需要对这些重合的box进行NMS过滤处理，对每一类别进行NMS处理时，就可以首先获取检测结果中的所有目标类别，然后对每一类别进行过滤。

代码如下:

# ===============nms过滤=======================#
# bboxes 具有这样的shape: [xmin,ymin,xmax,ymax,prob,classid]
# self.per_cls_maxboxes : 每一类别最多输出的检测框box数量，例如:50,或者169(13*13=169)，等等，根据自己情况而定

def nms_map_fn(args):
    '''

    :param args:
    :return:
    '''

    cls = args
    cls = tf.cast(cls, dtype=tf.int32)

    _bboxes = tf.cast(bboxes[:, 5], dtype=tf.int32)  # 类别ID
    cls_mask = tf.equal(_bboxes, cls)
    cls_bboxes = bboxes[cls_mask]  # ID为cls的目标框

    # 拆分得到boxes，scores，以便调用tf.image.non_max_suppression
    # nms之后再来合并
    boxes = cls_bboxes[:, 0:4]
    scores = cls_bboxes[:, 4]
    _maxbox = tf.shape(scores)[0]  # nms操作最多输出多少个目标

    selected_indices = tf.image.non_max_suppression(boxes=boxes, scores=scores,
                                                    iou_threshold=iou_threshold,
                                                    max_output_size=_maxbox)
    selected_boxes = tf.gather(boxes, selected_indices)
    seclected_scores = tf.gather(scores, selected_indices)
    classes = tf.ones_like(seclected_scores, dtype=tf.int32) * cls
    classes = tf.cast(calsses, dtype=tf.float32)

    selected_bboxes = tf.concat([selected_boxes,
                                 seclected_scores[:, tf.newaxis],
                                 classes[:, tf.newaxis]],
                                axis=-1)  # [xmin,ymin,xmax,ymax,prob,classid]

    objnum = tf.shape(selected_boxes)[0]  # nms得到的目标数量

    # 根据概率降序排序
    # indices = tf.argsort(tf.cast(selected_bboxes[:, 4] * 1000, dtype=tf.int32),
    #                      direction='DESCENDING')  # 索引 # tf 1.14.0有该函数
    indices = tf.nn.top_k(tf.cast(selected_bboxes[:, 4] * 1000, dtype=tf.int32), k=objnum).indices
    
    selected_bboxes = tf.gather(selected_bboxes, indices)

    selected_bboxes = selected_bboxes[:self.per_cls_maxboxes]

    def add_boxes():
        temp_bboxes = tf.fill([self.per_cls_maxboxes - objnum, 6], -1)  # 创建一个常量
        temp_bboxes = tf.to_float(temp_bboxes)
        _selected_bboxes = tf.concat([selected_bboxes, temp_bboxes], axis=0)
        return _selected_bboxes

    def ori_boxes():
        return selected_bboxes

    selected_bboxes = tf.cond(objnum < self.per_cls_maxboxes, true_fn=add_boxes, false_fn=ori_boxes)

    return selected_bboxes


# 获取检测结果中的所有目标类别
classes_in_img, idx = tf.unique(tf.cast(bboxes[:, 5], tf.int32))
# 对每一类别进行NMS过滤处理
best_bboxes = tf.map_fn(nms_map_fn, classes_in_img, infer_shape=False, dtype=tf.float32)

代码中首先利用tf.unique()函数获取检测结果中的所有类别，然后遍历该类别列表，对每一类别进行处理。

由于类别列表类型为tf.int32，而fn返回的是每一类别的box，类型为tf.float32,输入和返回类型不一致，因此需要显式指定返回类型； 同时，输入的shape可能为[85,](85个类别)， 而各类别返回的结果的shape可能是[50,5](50个box，每个box具有5个属性) ，显然，返回结果的shape和输入的[85,]是不一致的，因此需要禁用对一致输出形状的测试，即infer_shape=False；最后，对每一类别而言，NMS处理后返回的box数量很大概率上都是不一样的，如类别0返回的shape是[30,5],类别2返回的shape是[3,5]，而map_fn要求fn函数返回值具有相同的shape，这导致map_fn不能将各类的结果组合在一起，形成最终的结果，因此代码中对NMS过滤后的结果进行了补齐(填充box数量)，以及裁剪(只选取固定数量的box)

yolov3 feature map 合并输出

在代码yolov3_tensorflow(文后附地址及说明)中，yolo网络输出了3个feature map，各feature map存储有各自的检测结果，但是这检测结果是不完整的，对这3个feature map的检测结果进行合并，然后处理之后得到的结果才能完整的代表这幅图像上的目标检测结果。

以input_data为[None,416,416,3]输入为例，检测目标类别为85类，
3个feature map输出结果为:
self.pred_sbbox:检测小目标，[None, 52, 52, 255]
self.pred_mbbox:检测中等目标,[None, 26, 26, 255]
self.pred_lbbox:检测大目标,[None, 13, 13, 255]

我的目标就是

• 1)要对这3个feature map检测结果进行合并;
• 2)还要对合并后的结果进行处理

由于输入数据具有batch的维度，因此可以按照这样的流程进行处理:

1. 取batch中的1张图片在3个featurmap上的检测结果pred_sbbox,pred_mbbox,pred_lbbox
2. 将这3层的结果合并tf.concat,得到该图片的所有检测结果pred_bbox;
3. 过滤掉pred_bbox中不在图像范围内、置信度比较低的检测结果，得到初步的检测结果bboxes
4. 得到bboxes中的所有类别，对bboxes中的每一类进行上面的NMS处理，通过调用tf.map_fn得到该图片的最终检测结果_best_bboxes
5. 由于步骤1~4是对1个图片进行处理得到检测结果，因此利用tf.map_fn就可以对batch中的所有图片进行处理，得到该batch的检测结果best_bboxes
6. 转换过程中，注意各步骤结果的shape变化，以及我们需要的shape是什么样的形式，及时进行处理

完整代码如下:

def _get_pred_bboxes(self, input_data, score_threshold, iou_threshold):
    '''
    根据置信度和nms阈值，获取该批次数据的预测结果框
    :param input_data NHWC
    :param score_threshold:
    :param iou_threshold:
    :return:
    '''

    # 取出batch中的1个image的检测结果进行处理
    def batch_map_fn(args):
        pred_sbbox, pred_mbbox, pred_lbbox = args

        pred_bbox = tf.concat([tf.reshape(pred_sbbox, (-1, 5 + self.num_class)),
                               tf.reshape(pred_mbbox, (-1, 5 + self.num_class)),
                               tf.reshape(pred_lbbox, (-1, 5 + self.num_class))],
                              axis=0)  # pred_bbox.shape:(?,85)

        pred_xywh = pred_bbox[:, 0:4]  # 4列数据内容为：Center_x,Center_y,width,height(中心点坐标+宽高)
        pred_conf = pred_bbox[:, 4]  # 含有物体的概率
        pred_prob = pred_bbox[:, 5:]  # 各目标的概率

        # # (1) (x, y, w, h) --> (xmin, ymin, xmax, ymax)
        pred_coor = tf.concat([pred_xywh[:, :2] - pred_xywh[:, 2:] * 0.5,
                               pred_xywh[:, :2] + pred_xywh[:, 2:] * 0.5], axis=-1)

        # # (3) clip some boxes those are out of range
        input_image_h = tf.shape(input_data[0])[0]
        input_image_w = tf.shape(input_data[0])[1]

        pred_coor = tf.concat([tf.maximum(pred_coor[:, :2], [0, 0]),
                               tf.minimum(pred_coor[:, 2:], [input_image_w - 1, input_image_h - 1])], axis=-1)
        invalid_mask = tf.logical_or((pred_coor[:, 0] > pred_coor[:, 2]), (pred_coor[:, 1] > pred_coor[:, 3]))
        # pred_coor[invalid_mask] = 0
        # pred_coor1 = tf.where(invalid_mask,[[0,0,0,0]],pred_coor) # 对于mask位置处的坐标值，将值置0，其他位置保留原来的坐标值
        valid_mask = tf.logical_not(invalid_mask)

        # # (4) discard some invalid boxes
        valid_scale = [0, np.inf]
        bboxes_scale = tf.sqrt(
            tf.reduce_prod(pred_coor[:, 2:4] - pred_coor[:, 0:2], -1))  # √((xmax-xmin)*(ymax-ymin))
        scale_mask = tf.logical_and((valid_scale[0] < bboxes_scale), (bboxes_scale < valid_scale[1]))
        scale_mask = tf.logical_and(valid_mask, scale_mask)

        # # (5) discard some boxes with low scores
        classes = tf.argmax(pred_prob, axis=-1)  # 找出概率最大的class索引
        classes = tf.to_float(classes)
        max_value = tf.reduce_max(pred_prob, reduction_indices=[1])  # 找出行上最大值，即找出概率最大的class
        scores = pred_conf * max_value
        score_mask = scores > score_threshold
        mask = tf.logical_and(scale_mask, score_mask)
        # coors, scores, classes = pred_coor[mask], scores[mask], classes[mask]

        # 合并结果
        coors, scores, classes = pred_coor, scores, classes
        bboxes = tf.concat([coors, scores[:, tf.newaxis], classes[:, tf.newaxis]],
                           axis=-1)  # [xmin,ymin,xmax,ymax,prob,classid]
        indices = tf.where(mask)
        indices = tf.squeeze(indices, axis=-1)
        bboxes = tf.gather(bboxes, indices)

        # ===============nms过滤=======================#
        def nms_map_fn(args):
            '''
            :param args:
            :return:
            '''

            cls = args
            cls = tf.cast(cls, dtype=tf.int32)

            _bboxes = tf.cast(bboxes[:, 5], dtype=tf.int32)  # 类别ID
            cls_mask = tf.equal(_bboxes, cls)
            indices = tf.where(cls_mask)
            indices = tf.squeeze(indices, axis=-1)
            cls_bboxes = tf.gather(bboxes, indices)  # ID为cls的目标框

            # 拆分得到boxes，scores，以便调用tf.image.non_max_suppression
            # nms之后再来合并
            boxes = cls_bboxes[:, 0:4]
            scores = cls_bboxes[:, 4]
            _maxbox = tf.shape(scores)[0]  # nms操作最多输出多少个目标

            selected_indices = tf.image.non_max_suppression(boxes=boxes, scores=scores,
                                                            iou_threshold=iou_threshold,
                                                            max_output_size=_maxbox)
            selected_boxes = tf.gather(boxes, selected_indices)
            seclected_scores = tf.gather(scores, selected_indices)
            classes = tf.ones_like(seclected_scores, dtype=tf.int32) * cls
            classes = tf.to_float(classes)

            selected_bboxes = tf.concat([selected_boxes,
                                         seclected_scores[:, tf.newaxis],
                                         classes[:, tf.newaxis]],
                                        axis=-1)  # [xmin,ymin,xmax,ymax,prob,classid]

            objnum = tf.shape(selected_boxes)[0]  # nms得到的目标数量

            # 根据概率降序排序
            # indices = tf.argsort(tf.cast(selected_bboxes[:, 4] * 1000, dtype=tf.int32),
            #                      direction='DESCENDING')  # 索引 # tf 1.14.0有该函数
            indices = tf.nn.top_k(tf.cast(selected_bboxes[:, 4] * 1000, dtype=tf.int32), k=objnum).indices
            selected_bboxes = tf.gather(selected_bboxes, indices)

            selected_bboxes = selected_bboxes[:self.per_cls_maxboxes]

            def add_boxes():
                temp_bboxes = tf.fill([self.per_cls_maxboxes - objnum, 6], -1)  # 创建一个常量
                temp_bboxes = tf.to_float(temp_bboxes)
                _selected_bboxes = tf.concat([selected_bboxes, temp_bboxes], axis=0)
                return _selected_bboxes

            def ori_boxes():
                return selected_bboxes

            selected_bboxes = tf.cond(objnum < self.per_cls_maxboxes, true_fn=add_boxes, false_fn=ori_boxes)

            return selected_bboxes

        classes_in_img, idx = tf.unique(tf.cast(bboxes[:, 5], tf.int32))
        best_bboxes = tf.cond(tf.equal(tf.size(classes_in_img), 0),  # 防止类别为空
                              false_fn=lambda: tf.map_fn(nms_map_fn, classes_in_img,
                                                         infer_shape=False, dtype=tf.float32),
                              true_fn=lambda: tf.to_float(tf.fill([self.num_class, self.per_cls_maxboxes, 6], -1))
                              )

        # 填充行数与类别数一致
        clsnum = tf.shape(best_bboxes)[0]
        best_bboxes = best_bboxes[:self.num_class]

        def add_classes():
            temp_classes = tf.fill([self.num_class - clsnum, self.per_cls_maxboxes, 6], -1)  # 创建一个常量
            temp_classes = tf.to_float(temp_classes)
            _best_bboxes = tf.concat([best_bboxes, temp_classes], axis=0)
            return _best_bboxes

        def ori_classes():
            return best_bboxes

        best_bboxes = tf.cond(clsnum < self.num_class, true_fn=add_classes, false_fn=ori_classes)

        # 给变量一名称
        # best_bboxes = tf.add_n([best_bboxes], name='pred_bboxes')
        return best_bboxes

    best_bboxes = tf.map_fn(batch_map_fn, (self.pred_sbbox, self.pred_mbbox, self.pred_lbbox), dtype=tf.float32,
                            infer_shape=False)
    N = tf.shape(best_bboxes)[0]
    cls = tf.shape(best_bboxes)[1]
    maxbox = tf.shape(best_bboxes)[2]

    best_bboxes = tf.reshape(best_bboxes, [N, cls * maxbox, 6], name='pred_bboxes')

    return best_bboxes

代码偏长，但结构上按照上面分析的流程进行。

至此，就可以让网络输出的结果为唯一的可以直接使用的box，无需进行额外的操作。方便了很多操作：比如，将训练过程中的预测结果绘制在tf.summary中，通过tensorboard进行显示，如下所示：

注意事项

1.openvino不支持tf.map_fn,tf.where,tf.top_k等函数(目前，我的openvino版本为R3.1),这意味着上面合并后的结果不能直接在openvino进行导出操作。解决办法:定义输入数据batchsize=1,同时仅导出合并后的结果，不对结果进行处理即可。具体可见代码中def _get_pred_openvino_bboxes(self, input_data):函数

代码地址及说明

代码地址:https://github.com/leonardohaig/yolov3_tensorflow/tree/dev
1.其中master分支的backbone为darknet53，dev分支被修改为了mobileNetV2结构，这是我目前使用的分支。两个分支的区别主要在于此。所以两分支的代码是通用的（可以说两者代码是完全一样的！）；
2.建议参考dev分支。dev分支是我使用的分支，经过验证，没有问题；而master分支我仅尝试过racoon一类目标的训练，没有发现问题，不能保证多类别时不存在问题，并且dev分支仍在更新中，而master分支已无修复。

---

下面的是我的公众号二维码图片，欢迎关注。