(一)test_single_image.py
默认输入图片尺寸为[416,416]。
# coding: utf-8
from __future__ import division, print_function
import tensorflow as tf
import numpy as np
import argparse
import cv2
from utils.misc_utils import parse_anchors, read_class_names
from utils.nms_utils import gpu_nms
from utils.plot_utils import get_color_table, plot_one_box
from model import yolov3
# 设置命令行参数,具体可参见每一个命令行参数的含义
parser = argparse.ArgumentParser(description="YOLO-V3 test single image test procedure.")
parser.add_argument("input_image", type=str,
help="The path of the input image.")
parser.add_argument("--anchor_path", type=str, default="./data/yolo_anchors.txt",
help="The path of the anchor txt file.")
parser.add_argument("--new_size", nargs='*', type=int, default=[416, 416],
help="Resize the input image with `new_size`, size format: [width, height]")
parser.add_argument("--class_name_path", type=str, default="./data/coco.names",
help="The path of the class names.")
parser.add_argument("--restore_path", type=str, default="./data/darknet_weights/yolov3.ckpt",
help="The path of the weights to restore.")
args = parser.parse_args()
# 处理anchors,这些anchors是通过数据聚类获得,一共9个,shape为:[9, 2]。
# 需要注意的是,最后一个维度的顺序是[width, height]
args.anchors = parse_anchors(args.anchor_path)
# 处理classes, 这里是将所有的class的名称提取了出来,组成了一个列表
args.classes = read_class_names(args.class_name_path)
# 类别的数目
args.num_class = len(args.classes)
# 根据类别的数目为每一个类别分配不同的颜色,以便展示
color_table = get_color_table(args.num_class)
# 读取图片
img_ori = cv2.imread(args.input_image)
# 获取图片的尺寸
height_ori, width_ori = img_ori.shape[:2]
# resize,根据之前设定的尺寸值进行resize,默认是[416, 416],还是[width, height]的顺序
img = cv2.resize(img_ori, tuple(args.new_size))
# 对图片像素进行一定的数据处理
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
# TF会话
with tf.Session() as sess:
# 输入的placeholder,用于输入图片
input_data = tf.placeholder(tf.float32, [1, args.new_size[1], args.new_size[0], 3], name='input_data')
# 定义一个YOLOv3的类,在后面可以用来做模型建立以及loss计算等操作,参数分别是类别的数目和anchors
yolo_model = yolov3(args.num_class, args.anchors)
with tf.variable_scope('yolov3'):
# 对图片进行正向传播,返回多张特征图
pred_feature_maps = yolo_model.forward(input_data, False)
# 对这些特征图进行处理,获得计算出的bounding box以及属于前景的概率已经每一个类别的概率分布
pred_boxes, pred_confs, pred_probs = yolo_model.predict(pred_feature_maps)
# 将两个概率值分别相乘就可以获得最终的概率值
pred_scores = pred_confs * pred_probs
# 对这些bounding boxes和概率值进行非最大抑制(NMS)就可以获得最后的bounding boxes和与其对应的概率值以及标签
boxes, scores, labels = gpu_nms(pred_boxes, pred_scores, args.num_class, max_boxes=30, score_thresh=0.4, nms_thresh=0.5)
# Saver类,用以保存和恢复模型
saver = tf.train.Saver()
# 恢复模型参数
saver.restore(sess, args.restore_path)
# 运行graph,获得对应tensors的具体数值,这里是[boxes, scores, labels],对应于NMS之后获得的结果
boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={
input_data: img})
# rescale the coordinates to the original image
# 将坐标重新映射到原始图片上,因为前面的计算都是在resize之后的图片上进行的,所以需要进行映射
boxes_[:, 0] *= (width_ori/float(args.new_size[0]))
boxes_[:, 2] *= (width_ori/float(args.new_size[0]))
boxes_[:, 1] *= (height_ori/float(args.new_size[1]))
boxes_[:, 3] *= (height_ori/float(args.new_size[1]))
# 输出
print("box coords:")
print(boxes_)
print('*' * 30)
print("scores:")
print(scores_)
print('*' * 30)
print("labels:")
print(labels_)
# 绘制并展示,保存最后的结果
for i in range(len(boxes_)):
x0, y0, x1, y1 = boxes_[i]
plot_one_box(img_ori, [x0, y0, x1, y1], label=args.classes[labels_[i]], color=color_table[labels_[i]])
cv2.imshow('Detection result', img_ori)
cv2.imwrite('detection_result.jpg', img_ori)
cv2.waitKey(0)
(二)get_kmeans.py
这里函数的主要作用是使用kmeans聚类产生若干个anchors中心,在训练的时候使用这些作为一种先验条件。这里的聚类主要是对目标检测框的尺寸进行聚类。
# coding: utf-8
# This script is modified from https://github.com/lars76/kmeans-anchor-boxes
from __future__ import division, print_function
import numpy as np
# 计算IOU,box一个长度为2的数组,表示box的尺寸,clusters表示的是若干集群的中心,同样也是尺寸。
def iou(box, clusters):
"""
Calculates the Intersection over Union (IoU) between a box and k clusters.
param:
box: tuple or array, shifted to the origin (i. e. width and height)
clusters: numpy array of shape (k, 2) where k is the number of clusters
return:
numpy array of shape (k, 0) where k is the number of clusters
"""
x = np.minimum(clusters[:, 0], box[0])
y = np.minimum(clusters[:, 1], box[1])
if np.count_nonzero(x == 0) > 0 or np.count_nonzero(y == 0) > 0:
raise ValueError("Box has no area")
intersection = x * y
box_area = box[0] * box[1]
cluster_area = clusters[:, 0] * clusters[:, 1]
iou_ = intersection / (box_area + cluster_area - intersection + 1e-10)
return iou_
def avg_iou(boxes, clusters):
"""
Calculates the average Intersection over Union (IoU) between a numpy array of boxes and k clusters.
param:
boxes: numpy array of shape (r, 2), where r is the number of rows
clusters: numpy array of shape (k, 2) where k is the number of clusters
return:
average IoU as a single float
"""
# 计算平均IOU
return np.mean([np.max(iou(boxes[i], clusters)) for i in range(boxes.shape[0])])
# 这个函数并未在任何地方被使用
def translate_boxes(boxes):
"""
Translates all the boxes to the origin.
param:
boxes: numpy array of shape (r, 4)
return:
numpy array of shape (r, 2)
"""
new_boxes = boxes.copy()
for row in range(new_boxes.shape[0]):
new_boxes[row][2] = np.abs(new_boxes[row][2] - new_boxes[row][0])
new_boxes[row][3] = np.abs(new_boxes[row][3] - new_boxes[row][1])
return np.delete(new_boxes, [0, 1], axis=1)
def kmeans(boxes, k, dist=np.median):
"""
Calculates k-means clustering with the Intersection over Union (IoU) metric.
param:
boxes: numpy array of shape (r, 2), where r is the number of rows
k: number of clusters
dist: distance function
return:
numpy array of shape (k, 2)
"""
# rows表示的是数据集中一共有多少个标注框
rows = boxes.shape[0]
# 初始化统计距离的矩阵和每一个标注框的所属集群编号,
# 这里使用last cluster记录下一轮循环开始时标注框的集群编号,如果在这某一轮的迭代中不发生改变则算法已经收敛。
distances = np.empty((rows, k))
last_clusters = np.zeros((rows,))
np.random.seed()
# the Forgy method will fail if the whole array contains the same rows
# 随机选择几个数据作为初始的集群中心
clusters = boxes[np.random.choice(rows, k, replace=False)]
# 循环
while True:
# 对每一个标注框,计算其与每个集群中心的距离,这里的距离采用的是(1 - 标注框与集群中心的IOU)来表示,
# IOU数值越大, 则(1- IOU)越小, 则表示距离越接近.
for row in range(rows):
distances[row] = 1 - iou(boxes[row], clusters)
# 对每个标注框选择与其距离最接近的集群中心的标号作为所属类别的编号。
nearest_clusters = np.argmin(distances, axis=1)
# 如果在这轮循环中所有的标注框的所属类别不再变化,则说明算法已经收敛,可以跳出循环。
if (last_clusters == nearest_clusters).all():
break
# 对每一类集群,取出所有属于该集群的数据,并按照给定的方法计算集群的中心,
# 这里默认采用中位数的方法来计算集群中心
for cluster in range(k):
clusters[cluster] = dist(boxes[nearest_clusters == cluster], axis=0)
# 更新每一个标注框所属的集群类别。
last_clusters = nearest_clusters
# 返回所有的集群中心
return clusters
def parse_anno(annotation_path):
# 打开数据标记的文件
anno = open(annotation_path, 'r')
# 用以储存最后的提取出的所有的高度和宽度的结果,
result = []
# 对每一个标记图片
for line in anno:
# 根据空格将数据行进行分割
s = line.strip().split(' ')
# 按照数据的标记规则,每一行的第一个数据是编号,第二个数据是图片地址,从第三个开始才是标记框的信息。
s = s[2:]
# 当前图片的标记框的数目,每个标记框包含五个信息,四个坐标信息和一个类别信息
box_cnt = len(s) // 5
# 分别处理每一个标记框的信息,并提取标记框的高度和宽度,存入result 列表。
for i in range(box_cnt):
x_min, y_min, x_max, y_max = float(s[i*5+1]), float(s[i*5+2]), float(s[i*5+3]), float(s[i*5+4])
width = x_max - x_min
height = y_max - y_min
assert width > 0
assert height > 0
result.append([width, height])
# 将list变为numpy的数组
result = np.asarray(result)
# 返回
return result
def get_kmeans(anno, cluster_num=9):
# 使用kmeans算法计算需要的anchors
anchors = kmeans(anno, cluster_num)
# 计算平均IOU
ave_iou = avg_iou(anno, anchors)
# 格式化为int类型
anchors = anchors.astype('int').tolist()
# 按照面积大小排序,
anchors = sorted(anchors, key=lambda x: x[0] * x[1])
# 返回
return anchors, ave_iou
if __name__ == '__main__':
annotation_path = "./data/my_data/train.txt"
anno_result = parse_anno(annotation_path)
anchors, ave_iou = get_kmeans(anno_result, 9)
# 格式化输出anchors数据
anchor_string = ''
for anchor in anchors:
anchor_string += '{},{}, '.format(anchor[0], anchor[1])
anchor_string = anchor_string[:-2]
print('anchors are:')
print(anchor_string)
print('the average iou is:')
print(ave_iou)
(三)model.py
这里函数和类的主要作用是对YOLO模型进行封装,类中的函数主要包括:
- 模型的简历
- 特征图信息和anchors的联合使用
- loss的计算
# coding=utf-8
# for better understanding about yolov3 architecture, refer to this website (in Chinese):
# https://blog.csdn.net/leviopku/article/details/82660381
from __future__ import division, print_function
import tensorflow as tf
slim = tf.contrib.slim
from utils.layer_utils import conv2d, darknet53_body, yolo_block, upsample_layer
class yolov3(object):
def __init__(self,
class_num,
anchors,
use_label_smooth=False,
use_focal_loss=False,
batch_norm_decay=0.999,
weight_decay=5e-4):
"""
yolov3 class
:param class_num: 类别数目
:param anchors: anchors,一般来说是9个anchors
:param use_label_smooth: 是否使用label smooth,默认为False
:param use_focal_loss: 是否使用focal loss,默认为False
:param batch_norm_decay: BN的衰减系数
:param weight_decay: 权重衰减系数
"""
# self.anchors = [[10, 13], [16, 30], [33, 23],
# [30, 61], [62, 45], [59, 119],
# [116, 90], [156, 198], [373,326]]
self.class_num = class_num
self.anchors = anchors
self.batch_norm_decay = batch_norm_decay
self.use_label_smooth = use_label_smooth
self.use_focal_loss = use_focal_loss
self.weight_decay = weight_decay
def forward(self, inputs, is_training=False, reuse=False):
"""
进行正向传播,返回的是若干特征图
:param inputs: shape: [N, height, width, channel]
:param is_training:
:param reuse:
:return:
"""
# 获取输入图片的高度height和宽度width
# the input img_size, form: [height, width]
# it will be used later
self.img_size = tf.shape(inputs)[1:3]
# batch normalization的相关参数
# set batch norm params
batch_norm_params = {
'decay': self.batch_norm_decay,
'epsilon': 1e-05,
'scale': True,
'is_training': is_training,
'fused': None, # Use fused batch norm if possible.
}
# slim的arg scope,可以简化代码的编写,共用一套参数设置
with slim.arg_scope([slim.conv2d, slim.batch_norm], 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=0.1),
weights_regularizer=slim.l2_regularizer(self.weight_decay)):
# DarkNet 的主体部分,主要作用是提取图片中的各种特征信息。
# 这里可以获取三张特征图,分别取自DarkNet的三个不同的阶段。
# 每一个阶段对应于不同的特征粒度,结合更多的特征可以增强模型的表达能力。
# 理论上来说特征提取网络也可以采用其他的网络结构,但是效果可能会有所差异。
# 如果输入图片的尺寸为[416, 416],则三张特征图的尺寸分别为
# route_1 : [1, 52, 52, 256]
# route_2 : [1, 26, 26, 512]
# route_3 : [1, 13, 13, 1024]
with tf.variable_scope('darknet53_body'):
route_1, route_2, route_3 = darknet53_body(inputs)
# 根据前面的特征图,进行特征融合操作,这样可以提供更多的信息。
with tf.variable_scope('yolov3_head'):
# 使用YOLO_block函数来处理得到的特征图,并返回两张特征图。
# 本质上,YOLO_block函数仅仅包含若干层卷积层。
# 其中,inter1的作用是用来后续进行特征融合,net的主要作用是用以计算后续的坐标和概率等信息。
inter1, net = yolo_block(route_3