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()