def predict_image():
image_path = "/home/chenwei/HDD/Project/datasets/object_detection/FDDB2016/convert/images/2002_07_19_big_img_130.jpg"
image = cv2.imread(image_path)
image_size = image.shape[:2]
input_shape = [model_params['input_height'], model_params['input_width']]
image_data = pre_process(image, input_shape)
image_data = image_data[np.newaxis, ...]
input = tf.placeholder(shape=[1, None, None, 3], dtype=tf.float32)
network = Network(is_train=False)
logits = network.build_network(input)
output = network.reorg_layer(logits, model_params['anchors'])
checkpoints = "./checkpoints/model.ckpt-128"
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, checkpoints)
bboxes, obj_probs, class_probs = sess.run(
output, feed_dict={input: image_data})
bboxes, scores, class_id = postprocess(bboxes,
obj_probs,
class_probs,
image_shape=image_size,
input_shape=input_shape)
img_detection = visualization(image, bboxes, scores, class_id,
model_params["classes"])
cv2.imshow("result", img_detection)
cv2.waitKey(0)
def predict_video():
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
capture = cv2.VideoCapture(0)
input = tf.placeholder(shape=[1, None, None, 3], dtype=tf.float32)
network = Network(is_train=False)
logits = network.build_network(input)
output = network.reorg_layer(logits, model_params['anchors'])
checkpoints = "./checkpoints/model.ckpt-128"
saver = tf.train.Saver()
with tf.Session(config=config) as sess:
saver.restore(sess, checkpoints)
while (True):
ref, image = capture.read()
image_size = image.shape[:2]
input_shape = [
model_params['input_height'], model_params['input_width']
]
image_data = pre_process(image, input_shape)
image_data = image_data[np.newaxis, ...]
bboxes, obj_probs, class_probs = sess.run(
output, feed_dict={input: image_data})
bboxes, scores, class_id = postprocess(bboxes,
obj_probs,
class_probs,
image_shape=image_size,
input_shape=input_shape)
img_detection = visualization(image, bboxes, scores, class_id,
model_params["classes"])
cv2.imshow("result", img_detection)
cv2.waitKey(1)
cv2.destroyAllWindows()
def predict_image():
image_path = "/home/chenwei/HDD/Project/datasets/object_detection/VOCdevkit/VOC2007/JPEGImages/000066.jpg"
image = cv2.imread(image_path)
image_size = image.shape[:2]
input_shape = [model_params['input_height'], model_params['input_width']]
image_data = preporcess(image, input_shape)
image_data = image_data[np.newaxis, ...]
input = tf.placeholder(shape=[1, input_shape[0], input_shape[1], 3],
dtype=tf.float32)
model = Network(len(model_params['classes']),
model_params['anchors'],
is_train=False)
with tf.variable_scope('yolov3'):
logits = model.build_network(input)
output = model.inference(logits)
checkpoints = "/home/chenwei/HDD/Project/YOLOv3/weights/yolov3.ckpt"
saver = tf.train.Saver()
with tf.Session() as sess:
saver.restore(sess, checkpoints)
bboxes, obj_probs, class_probs = sess.run(
output, feed_dict={input: image_data})
bboxes, scores, class_max_index = postprocess(bboxes,
obj_probs,
class_probs,
image_shape=image_size,
input_shape=input_shape)
resize_ratio = min(input_shape[1] / image_size[1],
input_shape[0] / image_size[0])
dw = (input_shape[1] - resize_ratio * image_size[1]) / 2
dh = (input_shape[0] - resize_ratio * image_size[0]) / 2
bboxes[:, [0, 2]] = (bboxes[:, [0, 2]] - dw) / resize_ratio
bboxes[:, [1, 3]] = (bboxes[:, [1, 3]] - dh) / resize_ratio
img_detection = visualization(image, bboxes, scores, class_max_index,
model_params["classes"])
cv2.imshow("result", img_detection)
cv2.waitKey(0)
def remove_optimizers_params():
ckpt_path = ''
class_num = 2
save_dir = 'shrinked_ckpt'
if not os.path.exists(save_dir):
os.makedirs(save_dir)
anchors = [[676, 197], [763, 250], [684, 283], [868, 231], [745, 273],
[544, 391], [829, 258], [678, 316, 713, 355]]
image = tf.placeholder(tf.float32, [1, 416, 416, 3])
model = Network(class_num, anchors, False)
with tf.variable_scope('yolov3'):
feature_maps = model.build_network(image)
saver_to_restore = tf.train.Saver()
saver_to_save = tf.train.Saver()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
saver_to_restore.restore(sess, ckpt_path)
saver_to_save.save(sess, save_dir + '/shrinked')
def weights_to_ckpt():
num_class = 80
image_size = 416
anchors = [[676, 197], [763, 250], [684, 283], [868, 231], [745, 273],
[544, 391], [829, 258], [678, 316, 713, 355]]
weight_path = '../weights/yolov3.weights'
save_path = '../weights/yolov3.ckpt'
model = Network(num_class, anchors, False)
with tf.Session() as sess:
inputs = tf.placeholder(tf.float32, [1, image_size, image_size, 3])
with tf.variable_scope('yolov3'):
feature_maps = model.build_network(inputs)
saver = tf.train.Saver(var_list=tf.global_variables(scope='yolov3'))
load_ops = load_weights(tf.global_variables(scope='yolov3'),
weight_path)
sess.run(load_ops)
saver.save(sess, save_path=save_path)
print('TensorFlow model checkpoint has been saved to {}'.format(
save_path))
def train():
start_step = 0
log_step = solver_params['log_step']
restore = solver_params['restore']
checkpoint_dir = path_params['checkpoints_dir']
checkpoints_name = path_params['checkpoints_name']
tfrecord_dir = path_params['tfrecord_dir']
tfrecord_name = path_params['train_tfrecord_name']
log_dir = path_params['logs_dir']
batch_size = solver_params['batch_size']
# 配置GPU
gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(gpu_options=gpu_options)
# 解析得到训练样本以及标注
data = tfrecord.TFRecord()
train_tfrecord = os.path.join(tfrecord_dir, tfrecord_name)
data_num = total_sample(train_tfrecord)
batch_num = int(math.ceil(float(data_num) / batch_size))
dataset = data.create_dataset(train_tfrecord,
batch_num,
batch_size=batch_size,
is_shuffle=True)
iterator = dataset.make_one_shot_iterator()
images, y_true = iterator.get_next()
images.set_shape([None, 416, 416, 3])
y_true.set_shape([None, 13, 13, 5, 6])
# 构建网络
network = Network(is_train=True)
logits = network.build_network(images)
# 计算损失函数
total_loss, diou_loss, confs_loss, class_loss = network.calc_loss(
logits, y_true)
global_step = tf.Variable(0, trainable=False)
learning_rate = tf.train.exponential_decay(solver_params['lr'],
global_step,
solver_params['decay_steps'],
solver_params['decay_rate'],
staircase=True)
optimizer = tf.train.AdamOptimizer(learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer.minimize(total_loss, global_step=global_step)
# 配置tensorboard
tf.summary.scalar("learning_rate", learning_rate)
tf.summary.scalar('total_loss', total_loss)
tf.summary.scalar("diou_loss", diou_loss)
tf.summary.scalar("confs_loss", confs_loss)
tf.summary.scalar("class_loss", class_loss)
# 配置tensorboard
summary_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(log_dir,
graph=tf.get_default_graph(),
flush_secs=60)
# 模型保存
save_variable = tf.global_variables()
saver = tf.train.Saver(save_variable, max_to_keep=50)
with tf.Session(config=config) as sess:
sess.run([
tf.global_variables_initializer(),
tf.local_variables_initializer()
])
if restore == True:
ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
if ckpt and ckpt.model_checkpoint_path:
stem = os.path.basename(ckpt.model_checkpoint_path)
restore_step = int(stem.split('.')[0].split('-')[-1])
start_step = restore_step
sess.run(global_step.assign(restore_step))
saver.restore(sess, ckpt.model_checkpoint_path)
print('Restoreing from {}'.format(ckpt.model_checkpoint_path))
else:
print("Failed to find a checkpoint")
if solver_params['pre_train']:
pretrained = np.load(path_params['pretrain_weights'],
allow_pickle=True).item()
for variable in tf.trainable_variables():
for key in pretrained.keys():
key2 = variable.name.rstrip(':0')
if (key == key2):
sess.run(tf.assign(variable, pretrained[key]))
summary_writer.add_graph(sess.graph)
print('\n----------- start to train -----------\n')
for epoch in range(start_step + 1, solver_params['epoches']):
train_epoch_loss, train_epoch_diou_loss, train_epoch_confs_loss, train_epoch_class_loss = [], [], [], []
for index in tqdm(range(batch_num)):
_, summary_, loss_, diou_loss_, confs_loss_, class_loss_, global_step_, lr = sess.run(
[
train_op, summary_op, total_loss, diou_loss,
confs_loss, class_loss, global_step, learning_rate
])
train_epoch_loss.append(loss_)
train_epoch_diou_loss.append(diou_loss_)
train_epoch_confs_loss.append(confs_loss_)
train_epoch_class_loss.append(class_loss_)
summary_writer.add_summary(summary_, global_step_)
train_epoch_loss, train_epoch_diou_loss, train_epoch_confs_loss, train_epoch_class_loss = np.mean(
train_epoch_loss), np.mean(train_epoch_diou_loss), np.mean(
train_epoch_confs_loss), np.mean(train_epoch_class_loss)
print(
"Epoch: {}, global_step: {}, lr: {:.8f}, total_loss: {:.3f}, diou_loss: {:.3f}, confs_loss: {:.3f}, class_loss: {:.3f}"
.format(epoch, global_step_, lr, train_epoch_loss,
train_epoch_diou_loss, train_epoch_confs_loss,
train_epoch_class_loss))
saver.save(sess,
os.path.join(checkpoint_dir, checkpoints_name),
global_step=epoch)
sess.close()
