faster_rcnn.py 文件源码

def _rcn_head(self, inputs, image_shape, nms_threshold, rpn_thresholds,
                  rcn_batch, batch_size, name='rcn_head', **kwargs):
        anchors_labels = self.anchors_placeholders['labels']
        feature_maps, rpn_reg, rpn_cls = inputs
        n_anchors = self.n_anchors

        with tf.variable_scope(name):
            rcn_input_indices = non_max_suppression(rpn_reg, rpn_cls, batch_size, n_anchors,
                                                    iou_threshold=nms_threshold,
                                                    score_threshold=rpn_thresholds[1],
                                                    nonempty=True)

            rcn_input_indices = tf.cond(self.is_training,
                                        lambda: self.create_bbox_batch(rcn_input_indices, rcn_batch),
                                        lambda: rcn_input_indices)

            rcn_input_rois, rcn_input_labels = self._get_rois_and_labels(rpn_reg, anchors_labels, rcn_input_indices)

            for tensor in rcn_input_rois:
                tf.add_to_collection('roi', tensor)
            for tensor in rcn_input_labels:
                tf.add_to_collection('targets', tensor)
            roi_factor = np.array(self.map_shape/image_shape)

            rcn_input_rois = self.stop_gradient_tuple(rcn_input_rois)
            rcn_input_labels = self.stop_gradient_tuple(rcn_input_labels)

            roi_cropped = roi_pooling_layer(feature_maps, rcn_input_rois,
                                            factor=roi_factor, shape=(7, 7), data_format=kwargs['data_format'])
            indices, roi_cropped, rcn_input_labels = self._stack_tuple(roi_cropped, rcn_input_labels) # pylint: disable=unbalanced-tuple-unpacking
            rcn_clsf = conv_block(roi_cropped, 'f', units=10, name='output_conv', **kwargs)

            loss = self.rcn_loss(rcn_clsf, rcn_input_labels)

            rcn_clsf = tf.argmax(rcn_clsf, axis=-1)
            rcn_clsf = self._unstack_tuple(rcn_clsf, indices)
            rcn_clsf = tf.tuple(rcn_clsf, name='clsf')
            for tensor in rcn_clsf:
                tf.add_to_collection('rcn_output', tensor)
            loss = tf.identity(loss, 'loss')

        return rcn_clsf, loss