Недавно я играл с детектором объектов YOLO v3 в Tensorflow. Я не нашел на GitHub подходящей реализации, поэтому решил преобразовать этот код, написанный на PyTorch, в Tensorflow. Оригинальную конфигурацию YOLO v3, опубликованную вместе с бумагой, можно найти в репозитории Darknet на GitHub здесь.
Я хотел бы поделиться своим кодом вместе с решениями некоторых проблем, с которыми я боролся при его реализации.
Я не буду особо останавливаться на аспектах, не связанных с реализацией. Я предполагаю, что вы знакомы с CNN, обнаружением объектов, архитектурой YOLO v3 и т. Д., А также с фреймворком Tensorflow и TF-Slim. Если нет, возможно, лучше начать с соответствующих документов / руководств. Я не буду объяснять, что делает каждая отдельная строка, а представлю рабочий код с объяснениями некоторых проблем, на которые я наткнулся.
Весь код, необходимый для запуска этого детектора, и некоторые демонстрации доступны в моем репозитории на GitHub. Я тестировал Ubuntu 16.04, Tensorflow 1.8.0 и CUDA 9.0.
Этот пост организован следующим образом:
- Настраивать
- Реализация слоев Darknet-53
- Реализация слоев детектирования YOLO v3
- Преобразование предварительно обученных весов COCO
- Реализация алгоритмов постобработки
- Резюме
1. Настройка
Я хочу организовать код аналогично тому, как он организован в репозитории Tensorflow models. Я использую TF-Slim, потому что он позволяет нам определять общие аргументы, такие как функция активации, параметры пакетной нормализации и т. Д., Как глобальные, что значительно ускоряет определение нейронных сетей.
Мы начнем с yolo_v3.py файла, в который мы поместим функции, которые инициализируют сеть, а также функции для загрузки предварительно обученных весов.
Добавьте необходимые константы (настроенные авторами YOLO) где-нибудь поверх файла.
_BATCH_NORM_DECAY = 0.9 _BATCH_NORM_EPSILON = 1e-05 _LEAKY_RELU = 0.1
YOLO v3 нормализует вход в диапазоне 0..1. Большинство слоев детектора выполняют пакетную нормализацию сразу после свертки, не имеют смещений и используют активацию Leaky ReLU. Для использования в таких случаях удобно определить slim arg scope. В слоях, которые не используют BN и LReLU, нам нужно будет неявно определить это.
# transpose the inputs to NCHW
if data_format == 'NCHW':
inputs = tf.transpose(inputs, [0, 3, 1, 2])
# normalize values to range [0..1]
inputs = inputs / 255
# set batch norm params
batch_norm_params = {
'decay': _BATCH_NORM_DECAY,
'epsilon': _BATCH_NORM_EPSILON,
'scale': True,
'is_training': is_training,
'fused': None, # Use fused batch norm if possible.
}
# Set activation_fn and parameters for conv2d, batch_norm.
with slim.arg_scope([slim.conv2d, slim.batch_norm, _fixed_padding], data_format=data_format, reuse=reuse):
with slim.arg_scope([slim.conv2d], normalizer_fn=slim.batch_norm, normalizer_params=batch_norm_params,
biases_initializer=None, activation_fn=lambda x: tf.nn.leaky_relu(x, alpha=_LEAKY_RELU)):
with tf.variable_scope('darknet-53'):
inputs = darknet53(inputs)
Теперь мы готовы определить слои Darknet-53.
2. Реализация слоев Darknet-53.
В статье YOLO v3 авторы представляют новую, более глубокую архитектуру экстрактора функций под названием Darknet-53. Как следует из названия, он содержит 53 сверточных слоя, за каждым из которых следует слой пакетной нормализации и активация Leaky ReLU. Понижение разрешения выполняется с помощью сверточных слоев с stride=2.
Прежде чем мы определим сверточные слои, мы должны понять, что реализация авторов использует фиксированное заполнение независимо от размера ввода. Чтобы добиться такого же поведения, мы можем использовать функцию ниже (я немного изменил код, найденный здесь).
@tf.contrib.framework.add_arg_scope
def _fixed_padding(inputs, kernel_size, *args, mode='CONSTANT', **kwargs):
"""
Pads the input along the spatial dimensions independently of input size.
Args:
inputs: A tensor of size [batch, channels, height_in, width_in] or
[batch, height_in, width_in, channels] depending on data_format.
kernel_size: The kernel to be used in the conv2d or max_pool2d operation.
Should be a positive integer.
data_format: The input format ('NHWC' or 'NCHW').
mode: The mode for tf.pad.
Returns:
A tensor with the same format as the input with the data either intact
(if kernel_size == 1) or padded (if kernel_size > 1).
"""
pad_total = kernel_size - 1
pad_beg = pad_total // 2
pad_end = pad_total - pad_beg
if kwargs['data_format'] == 'NCHW':
padded_inputs = tf.pad(inputs, [[0, 0], [0, 0],
[pad_beg, pad_end], [pad_beg, pad_end]], mode=mode)
else:
padded_inputs = tf.pad(inputs, [[0, 0], [pad_beg, pad_end],
[pad_beg, pad_end], [0, 0]], mode=mode)
return padded_inputs
_fixed_padding дополняет ввод по высоте и ширине соответствующим количеством нулей (когда mode='CONSTANT'). Позже мы также будем использовать mode='SYMMETRIC'.
Теперь мы можем определить функцию _conv2d_fixed_padding:
def _conv2d_fixed_padding(inputs, filters, kernel_size, strides=1):
if strides > 1:
inputs = _fixed_padding(inputs, kernel_size)
inputs = slim.conv2d(inputs, filters, kernel_size, stride=strides, padding=('SAME' if strides == 1 else 'VALID'))
return inputs
Модель Darknet-53 построена из некоторого количества блоков с двумя сверточными слоями и ярлыком соединения, за которым следует слой понижающей дискретизации. Чтобы избежать шаблонного кода, мы определяем функцию _darknet_block:
def _darknet53_block(inputs, filters):
shortcut = inputs
inputs = _conv2d_fixed_padding(inputs, filters, 1)
inputs = _conv2d_fixed_padding(inputs, filters * 2, 3)
inputs = inputs + shortcut
return inputs
Наконец, у нас есть все необходимые строительные блоки для модели Darknet-53:
def darknet53(inputs):
"""
Builds Darknet-53 model.
"""
inputs = _conv2d_fixed_padding(inputs, 32, 3)
inputs = _conv2d_fixed_padding(inputs, 64, 3, strides=2)
inputs = _darknet53_block(inputs, 32)
inputs = _conv2d_fixed_padding(inputs, 128, 3, strides=2)
for i in range(2):
inputs = _darknet53_block(inputs, 64)
inputs = _conv2d_fixed_padding(inputs, 256, 3, strides=2)
for i in range(8):
inputs = _darknet53_block(inputs, 128)
inputs = _conv2d_fixed_padding(inputs, 512, 3, strides=2)
for i in range(8):
inputs = _darknet53_block(inputs, 256)
inputs = _conv2d_fixed_padding(inputs, 1024, 3, strides=2)
for i in range(4):
inputs = _darknet53_block(inputs, 512)
return inputs
Изначально после последнего блока есть слой global avg pool и softmax, но они не используются YOLO v3 (так что фактически у нас 52 слоя вместо 53;))
3. Реализация слоев детектирования YOLO v3.
Характеристики, извлеченные Darknet-53, направлены на уровни обнаружения. Модуль обнаружения состоит из некоторого количества сверточных слоев, сгруппированных в блоки, слоев с повышающей дискретизацией и трех сверточных слоев с функцией линейной активации, что позволяет выполнять обнаружение в 3 различных масштабах. Начнем с написания вспомогательной функции _yolo_block:
def _yolo_block(inputs, filters):
inputs = _conv2d_fixed_padding(inputs, filters, 1)
inputs = _conv2d_fixed_padding(inputs, filters * 2, 3)
inputs = _conv2d_fixed_padding(inputs, filters, 1)
inputs = _conv2d_fixed_padding(inputs, filters * 2, 3)
inputs = _conv2d_fixed_padding(inputs, filters, 1)
route = inputs
inputs = _conv2d_fixed_padding(inputs, filters * 2, 3)
return route, inputs
Затем активации из 5-го слоя в блоке направляются на другой сверточный слой и подвергаются повышающей дискретизации, в то время как активации из 6-го уровня переходят в _detection_layer, который мы собираемся определить сейчас:
def _detection_layer(inputs, num_classes, anchors, img_size, data_format):
num_anchors = len(anchors)
predictions = slim.conv2d(inputs, num_anchors * (5 + num_classes), 1, stride=1, normalizer_fn=None,
activation_fn=None, biases_initializer=tf.zeros_initializer())
shape = predictions.get_shape().as_list()
grid_size = _get_size(shape, data_format)
dim = grid_size[0] * grid_size[1]
bbox_attrs = 5 + num_classes
if data_format == 'NCHW':
predictions = tf.reshape(predictions, [-1, num_anchors * bbox_attrs, dim])
predictions = tf.transpose(predictions, [0, 2, 1])
predictions = tf.reshape(predictions, [-1, num_anchors * dim, bbox_attrs])
stride = (img_size[0] // grid_size[0], img_size[1] // grid_size[1])
anchors = [(a[0] / stride[0], a[1] / stride[1]) for a in anchors]
box_centers, box_sizes, confidence, classes = tf.split(predictions, [2, 2, 1, num_classes], axis=-1)
box_centers = tf.nn.sigmoid(box_centers)
confidence = tf.nn.sigmoid(confidence)
grid_x = tf.range(grid_size[0], dtype=tf.float32)
grid_y = tf.range(grid_size[1], dtype=tf.float32)
a, b = tf.meshgrid(grid_x, grid_y)
x_offset = tf.reshape(a, (-1, 1))
y_offset = tf.reshape(b, (-1, 1))
x_y_offset = tf.concat([x_offset, y_offset], axis=-1)
x_y_offset = tf.reshape(tf.tile(x_y_offset, [1, num_anchors]), [1, -1, 2])
box_centers = box_centers + x_y_offset
box_centers = box_centers * stride
anchors = tf.tile(anchors, [dim, 1])
box_sizes = tf.exp(box_sizes) * anchors
box_sizes = box_sizes * stride
detections = tf.concat([box_centers, box_sizes, confidence], axis=-1)
classes = tf.nn.sigmoid(classes)
predictions = tf.concat([detections, classes], axis=-1)
return predictions
Этот слой преобразует необработанные прогнозы в соответствии со следующими уравнениями. Поскольку YOLO v3 на каждом масштабе обнаруживает объекты разных размеров и соотношений сторон, передается аргумент anchors, который представляет собой список из 3 кортежей (высота, ширина) для каждого масштаба. Якоря должны быть адаптированы для набора данных (в этом руководстве мы будем использовать якоря для набора данных COCO). Просто добавьте эту константу где-нибудь поверх yolo_v3.py файла.
_ANCHORS = [(10, 13), (16, 30), (33, 23), (30, 61), (62, 45), (59, 119), (116, 90), (156, 198), (373, 326)]
Нам нужна одна небольшая вспомогательная функция _get_size, которая возвращает высоту и ширину ввода:
def _get_size(shape, data_format):
if len(shape) == 4:
shape = shape[1:]
return shape[1:3] if data_format == 'NCHW' else shape[0:2]
Как упоминалось ранее, последний строительный блок, который нам нужен для реализации YOLO v3, - это уровень передискретизации. Детектор YOLO использует метод билинейной передискретизации. Почему мы не можем просто использовать стандартный tf.image.resize_bilinear метод из Tensorflow API? Причина в том, что на сегодняшний день (версия TF 1.8.0) все методы повышающей дискретизации используют режим constant pad. Стандартный метод заполнения в репозитории авторов YOLO и в PyTorch - edge (хорошее сравнение режимов заполнения можно найти здесь). Это незначительное различие существенно влияет на обнаружение (и стоило мне пары часов отладки).
Чтобы обойти это, мы вручную добавим к входам 1 пиксель и mode='SYMMETRIC', что эквивалентно режиму edge.
# we just need to pad with one pixel, so we set kernel_size = 3 inputs = _fixed_padding(inputs, 3, 'NHWC', mode='SYMMETRIC')
Весь _upsample код функции выглядит следующим образом:
def _upsample(inputs, out_shape, data_format='NCHW'):
# we need to pad with one pixel, so we set kernel_size = 3
inputs = _fixed_padding(inputs, 3, mode='SYMMETRIC')
# tf.image.resize_bilinear accepts input in format NHWC
if data_format == 'NCHW':
inputs = tf.transpose(inputs, [0, 2, 3, 1])
if data_format == 'NCHW':
height = out_shape[3]
width = out_shape[2]
else:
height = out_shape[2]
width = out_shape[1]
# we padded with 1 pixel from each side and upsample by factor of 2, so new dimensions will be
# greater by 4 pixels after interpolation
new_height = height + 4
new_width = width + 4
inputs = tf.image.resize_bilinear(inputs, (new_height, new_width))
# trim back to desired size
inputs = inputs[:, 2:-2, 2:-2, :]
# back to NCHW if needed
if data_format == 'NCHW':
inputs = tf.transpose(inputs, [0, 3, 1, 2])
inputs = tf.identity(inputs, name='upsampled')
return inputs
ОБНОВЛЕНИЕ: Благодаря Srikanth Vidapanakal я проверил исходный код даркнета и обнаружил, что метод апсемплинга является ближайшим соседом, а не билинейным. Нам больше не нужно дополнять изображение. Обновленный код уже доступен в моем репо.
Код функции Fixed _upsample выглядит следующим образом:
def _upsample(inputs, out_shape, data_format='NCHW'):
# tf.image.resize_nearest_neighbor accepts input in format NHWC
if data_format == 'NCHW':
inputs = tf.transpose(inputs, [0, 2, 3, 1])
if data_format == 'NCHW':
new_height = out_shape[3]
new_width = out_shape[2]
else:
new_height = out_shape[2]
new_width = out_shape[1]
inputs = tf.image.resize_nearest_neighbor(inputs, (new_height, new_width))
# back to NCHW if needed
if data_format == 'NCHW':
inputs = tf.transpose(inputs, [0, 3, 1, 2])
inputs = tf.identity(inputs, name='upsampled')
return inputs
Активации с повышенной дискретизацией объединяются по оси каналов с активациями из слоев Darknet-53. Вот почему нам нужно вернуться к функции darknet53 и вернуть активации из сверточных слоев перед 4-м и 5-м уровнями понижающей дискретизации.
def darknet53(inputs):
"""
Builds Darknet-53 model.
"""
inputs = _conv2d_fixed_padding(inputs, 32, 3)
inputs = _conv2d_fixed_padding(inputs, 64, 3, strides=2)
inputs = _darknet53_block(inputs, 32)
inputs = _conv2d_fixed_padding(inputs, 128, 3, strides=2)
for i in range(2):
inputs = _darknet53_block(inputs, 64)
inputs = _conv2d_fixed_padding(inputs, 256, 3, strides=2)
for i in range(8):
inputs = _darknet53_block(inputs, 128)
route1 = inputs
inputs = _conv2d_fixed_padding(inputs, 512, 3, strides=2)
for i in range(8):
inputs = _darknet53_block(inputs, 256)
route2 = inputs
inputs = _conv2d_fixed_padding(inputs, 1024, 3, strides=2)
for i in range(4):
inputs = _darknet53_block(inputs, 512)
return route1, route2, inputs
Теперь мы готовы определить модуль детектора. Вернемся к функции yolo_v3 и добавим следующие строки под узкой областью аргументов:
with tf.variable_scope('darknet-53'):
route_1, route_2, inputs = darknet53(inputs)
with tf.variable_scope('yolo-v3'):
route, inputs = _yolo_block(inputs, 512)
detect_1 = _detection_layer(inputs, num_classes, _ANCHORS[6:9], img_size, data_format)
detect_1 = tf.identity(detect_1, name='detect_1')
inputs = _conv2d_fixed_padding(route, 256, 1)
upsample_size = route_2.get_shape().as_list()
inputs = _upsample(inputs, upsample_size, data_format)
inputs = tf.concat([inputs, route_2], axis=1 if data_format == 'NCHW' else 3)
route, inputs = _yolo_block(inputs, 256)
detect_2 = _detection_layer(inputs, num_classes, _ANCHORS[3:6], img_size, data_format)
detect_2 = tf.identity(detect_2, name='detect_2')
inputs = _conv2d_fixed_padding(route, 128, 1)
upsample_size = route_1.get_shape().as_list()
inputs = _upsample(inputs, upsample_size, data_format)
inputs = tf.concat([inputs, route_1], axis=1 if data_format == 'NCHW' else 3)
_, inputs = _yolo_block(inputs, 128)
detect_3 = _detection_layer(inputs, num_classes, _ANCHORS[0:3], img_size, data_format)
detect_3 = tf.identity(detect_3, name='detect_3')
detections = tf.concat([detect_1, detect_2, detect_3], axis=1)
return detections
4. Преобразование предварительно обученных весов COCO
Мы определили архитектуру детектора. Чтобы использовать его, мы должны либо обучить его на собственном наборе данных, либо использовать предварительно обученные веса. Веса, предварительно обученные на наборе данных COCO, доступны для публичного использования. Мы можем скачать его с помощью этой команды:
wget https://pjreddie.com/media/files/yolov3.weights
Структура этого двоичного файла следующая:
Первые 3 значения int32 - это информация заголовка: основной номер версии, дополнительный номер версии, номер подрывной версии, за которым следует значение int64: количество изображений, просмотренных сетью во время обучения. После них идут 62 001 757 значений float32, которые представляют собой веса каждого слоя с нормами свертки и пакета. Важно помнить, что они сохраняются в формате основных строк, который противоположен формату, используемому Tensorflow (основные столбцы).
Итак, как нам читать веса из этого файла?
Начнем с первого свёрточного слоя. За большинством сверточных слоев сразу следует слой пакетной нормализации. В этом случае нам нужно сначала прочитать 4 * num_filters весов слоя пакетных норм: гамма, бета, скользящее среднее и скользящее отклонение, затемkernel_size[0] * kernel_size[1] * num_filters * input_channels веса сверточного слоя.
В противоположном случае, когда за сверточным слоем не следует слой пакетной нормы, вместо чтения параметров пакетной нормы нам нужно читатьnum_filters веса смещения.
Приступим к написанию кода функции load_weights. Он принимает 2 аргумента: список переменных в нашем графике и имя двоичного файла.
Мы начинаем с открытия файла, пропуская первые 5 значений int32 и читая все остальное в виде списка:
def load_weights(var_list, weights_file):
with open(weights_file, "rb") as fp:
_ = np.fromfile(fp, dtype=np.int32, count=5)
weights = np.fromfile(fp, dtype=np.float32)
Затем мы будем использовать два указателя, первый для перебора списка переменных var_list, а второй для перебора списка с загруженными переменными weights. Нам нужно проверить тип слоя, следующего за текущим обрабатываемым, и прочитать соответствующее количество значений. В коде i будет повторять var_list, а ptr будет повторять weights. Мы вернем список tf.assign операций. Я проверяю тип слоя, просто сравнивая его имя. (Я согласен с тем, что это немного некрасиво, но я не знаю лучшего способа сделать это. Этот подход, кажется, работает для меня.)
ptr = 0
i = 0
assign_ops = []
while i < len(var_list) - 1:
var1 = var_list[i]
var2 = var_list[i + 1]
# do something only if we process conv layer
if 'Conv' in var1.name.split('/')[-2]:
# check type of next layer
if 'BatchNorm' in var2.name.split('/')[-2]:
# load batch norm params
gamma, beta, mean, var = var_list[i + 1:i + 5]
batch_norm_vars = [beta, gamma, mean, var]
for var in batch_norm_vars:
shape = var.shape.as_list()
num_params = np.prod(shape)
var_weights = weights[ptr:ptr + num_params].reshape(shape)
ptr += num_params
assign_ops.append(tf.assign(var, var_weights, validate_shape=True))
# we move the pointer by 4, because we loaded 4 variables
i += 4
elif 'Conv' in var2.name.split('/')[-2]:
# load biases
bias = var2
bias_shape = bias.shape.as_list()
bias_params = np.prod(bias_shape)
bias_weights = weights[ptr:ptr + bias_params].reshape(bias_shape)
ptr += bias_params
assign_ops.append(tf.assign(bias, bias_weights, validate_shape=True))
# we loaded 2 variables
i += 1
# we can load weights of conv layer
shape = var1.shape.as_list()
num_params = np.prod(shape)
var_weights = weights[ptr:ptr + num_params].reshape((shape[3], shape[2], shape[0], shape[1]))
# remember to transpose to column-major
var_weights = np.transpose(var_weights, (2, 3, 1, 0))
ptr += num_params
assign_ops.append(tf.assign(var1, var_weights, validate_shape=True))
i += 1
return assign_ops
Вот и все! Теперь мы можем восстановить веса модели, выполнив строки кода, подобные этим:
with tf.variable_scope('model'):
model = yolo_v3(inputs, 80)
model_vars = tf.global_variables(scope='model')
assign_ops = load_variables(model_vars, 'yolov3.weights')
sess = tf.Session()
sess.run(assign_ops)
Для будущего использования, вероятно, будет намного проще экспортировать веса с помощью tf.train.Saver и загрузить с контрольной точки.
5. Реализация алгоритмов постобработки.
Наша модель возвращает тензор формы:
batch_size x 10647 x (num_classes + 5 bounding box attrs)
Число 10647 равно сумме 507 +2028 + 8112, которые представляют собой количество возможных объектов, обнаруженных на каждой шкале. 5 значений, описывающих атрибуты ограничивающей рамки, обозначают center_x, center_y, width, height. В большинстве случаев проще работать с координатами двух точек: верхнего левого и нижнего правого. Преобразуем вывод детектора в этот формат.
Функция, которая это делает, довольно проста:
def detections_boxes(detections):
center_x, center_y, width, height, attrs = tf.split(detections, [1, 1, 1, 1, -1], axis=-1)
w2 = width / 2
h2 = height / 2
x0 = center_x - w2
y0 = center_y - h2
x1 = center_x + w2
y1 = center_y + h2
boxes = tf.concat([x0, y0, x1, y1], axis=-1)
detections = tf.concat([boxes, attrs], axis=-1)
return detections
Обычно наш детектор обнаруживает один и тот же объект несколько раз (с немного разными центрами и размерами). В большинстве случаев мы не хотим сохранять все эти обнаружения, которые отличаются только небольшим количеством пикселей. Стандартное решение этой проблемы - подавление без максимума. Хорошее описание этого метода доступно здесь.
Почему мы не используем функцию tf.image.non_max_suppression из Tensorflow API? Есть две основные причины. Во-первых, на мой взгляд, намного лучше выполнять NMS для каждого класса, потому что у нас может быть ситуация, когда объекты из двух разных классов сильно перекрываются, и глобальная NMS подавит один из блоков. Во-вторых, некоторые люди жалуются, что эта функция работает медленно, потому что она еще не оптимизирована.
Давайте реализуем алгоритм NMS. Сначала нам нужна функция для вычисления IoU (пересечение по объединению) двух ограничивающих прямоугольников:
def _iou(box1, box2):
b1_x0, b1_y0, b1_x1, b1_y1 = box1
b2_x0, b2_y0, b2_x1, b2_y1 = box2
int_x0 = max(b1_x0, b2_x0)
int_y0 = max(b1_y0, b2_y0)
int_x1 = min(b1_x1, b2_x1)
int_y1 = min(b1_y1, b2_y1)
int_area = (int_x1 - int_x0) * (int_y1 - int_y0)
b1_area = (b1_x1 - b1_x0) * (b1_y1 - b1_y0)
b2_area = (b2_x1 - b2_x0) * (b2_y1 - b2_y0)
iou = int_area / (b1_area + b2_area - int_area + 1e-05)
return iou
Теперь мы можем написать код функции non_max_suppression. Я использую библиотеку NumPy для быстрых векторных операций.
def non_max_suppression(predictions_with_boxes, confidence_threshold, iou_threshold=0.4):
"""
Applies Non-max suppression to prediction boxes.
:param predictions_with_boxes: 3D numpy array, first 4 values in 3rd dimension are bbox attrs, 5th is confidence
:param confidence_threshold: the threshold for deciding if prediction is valid
:param iou_threshold: the threshold for deciding if two boxes overlap
:return: dict: class -> [(box, score)]
"""
Он принимает 3 аргумента: выходные данные детектора YOLO v3, порог достоверности и порог IoU. Тело этой функции выглядит следующим образом:
conf_mask = np.expand_dims((predictions_with_boxes[:, :, 4] > confidence_threshold), -1)
predictions = predictions_with_boxes * conf_mask
result = {}
for i, image_pred in enumerate(predictions):
shape = image_pred.shape
non_zero_idxs = np.nonzero(image_pred)
image_pred = image_pred[non_zero_idxs]
image_pred = image_pred.reshape(-1, shape[-1])
bbox_attrs = image_pred[:, :5]
classes = image_pred[:, 5:]
classes = np.argmax(classes, axis=-1)
unique_classes = list(set(classes.reshape(-1)))
for cls in unique_classes:
cls_mask = classes == cls
cls_boxes = bbox_attrs[np.nonzero(cls_mask)]
cls_boxes = cls_boxes[cls_boxes[:, -1].argsort()[::-1]]
cls_scores = cls_boxes[:, -1]
cls_boxes = cls_boxes[:, :-1]
while len(cls_boxes) > 0:
box = cls_boxes[0]
score = cls_scores[0]
if not cls in result:
result[cls] = []
result[cls].append((box, score))
cls_boxes = cls_boxes[1:]
ious = np.array([_iou(box, x) for x in cls_boxes])
iou_mask = ious < iou_threshold
cls_boxes = cls_boxes[np.nonzero(iou_mask)]
cls_scores = cls_scores[np.nonzero(iou_mask)]
return result
Вот и все. Мы реализовали все функции, необходимые для работы YOLO v3.
6. Резюме
В учебном репо вы можете найти код и демонстрационный скрипт для запуска детекторов. Детектор работает с форматами данных NHWC и NCHW, поэтому вы можете легко выбрать, какой формат быстрее работает на вашем компьютере.
Если у вас есть вопросы, не стесняйтесь обращаться ко мне.
Я планирую написать следующую часть этого руководства, в которой я покажу, как обучать (настраивать) YOLO v3 на пользовательских наборах данных.
Спасибо за прочтение. Пожалуйста, дайте мне знать, понравился ли он вам, хлопнув в ладоши и / или поделившись им! :)
