def convert(is_tiny=False):
if is_tiny:
anchors = np.array([[1, 1]] * 6)
weight_path = join('model_data', 'yolov3-tiny.weights')
save_path = join('logs', 'cnn_tiny', 'cnn_tiny_model')
else:
anchors = np.array([[1, 1]] * 9)
weight_path = join('model_data', 'yolov3.weights')
save_path = join('logs', 'cnn_full', 'cnn_full_model')
if not exists(split(save_path)[0]):
makedirs(split(save_path)[0])
input_data = tf.placeholder(dtype=tf.float32, shape=(1, 416, 416, 3))
model(input_data, 80, anchors, 'cnn', True, False)
model_vars_ = tf.global_variables()
assert weight_path.endswith(
'.weights'), '{} is not a .weights files'.format(weight_path)
assign_ops_ = load_weight(model_vars_, weight_path)
t0 = time.time()
print("start loading weights")
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(assign_ops_)
saver.save(sess, save_path, write_meta_graph=False, write_state=False)
t1 = time.time()
print("convert weights is over, cost {0:.4f}s".format(t1 - t0))
def __init__(self, config):
self.config = config
net_type, tiny = split(self.config.weight_path)[-1].split('_')[:2]
if tiny == 'tiny':
self.anchor_path = join('model_data', 'yolo_anchors_tiny.txt')
else:
self.anchor_path = join('model_data', 'yolo_anchors.txt')
self.classes = self._get_classes()
self.anchors = self._get_anchors()
self.hw = [416, 416]
self.batch_size = 64
self.ious_thres = [0.5, 0.75]
self.test_path = "model_data/test.txt"
with open(self.test_path) as f:
self.test_data = f.readlines()
if tiny == 'tiny':
assert 6 == len(
self.anchors
), 'the model type does not match with anchors, check anchors or type param'
else:
assert 9 == len(
self.anchors
), 'the model type does not match with anchors, check anchors or type param'
self.input = tf.placeholder(tf.float32,
[self.batch_size] + self.hw + [3])
self.is_training = tf.placeholder(tf.bool, shape=[])
self.pred = model(self.input, len(self.classes), self.anchors,
net_type, self.is_training, False)
print('start load net_type: {}_{}_model'.format(net_type, tiny))
# load weights
conf = tf.ConfigProto()
conf.gpu_options.allow_growth = True
# change fraction according to your GPU
# conf.gpu_options.per_process_gpu_memory_fraction = 0.05
self.sess = tf.Session(config=conf)
saver = tf.train.Saver()
saver.restore(self.sess, self.config.weight_path)
self.color_table = get_color_table(len(self.classes))
def train(self):
# pred, losses, op = self.create_model()
pred = model(self.input, len(self.classes), self.anchors,
config.net_type, True)
grid_shape = [g.get_shape().as_list() for g in pred[2]]
s = sum([g[2] * g[1] for g in grid_shape])
self.label = tf.placeholder(
tf.float32, [self.batch_size, s, 3, 9 + len(self.classes)])
# for data in self.generate_data(grid_shape):
# print()
losses = loss(pred, self.label, self.hw, self.lambda_coord,
self.lambda_noobj, self.lambda_cls, self.iou_threshold,
config.debug)
opt = tf.train.AdamOptimizer(self.learn_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
var_list = tf.global_variables()
with tf.control_dependencies(update_ops):
op = opt.minimize(losses)
# summary
writer = tf.summary.FileWriter(self.log_path)
img_tensor = tf.placeholder(tf.float32,
[2 * self.batch_size] + self.hw + [3])
train_loss_tensor = tf.placeholder(tf.float32)
val_loss_tensor = tf.placeholder(tf.float32)
tf.summary.scalar('train_loss', train_loss_tensor)
tf.summary.scalar('val_loss', val_loss_tensor)
tf.summary.image('img', img_tensor, 2 * self.batch_size)
summary = tf.summary.merge_all()
conf = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
sess = tf.Session(config=conf)
# sess = tf_debug.LocalCLIDebugWrapperSession(sess)
# sess = tf_debug.TensorBoardDebugWrapperSession(sess, "PC-DAIXILI:6001")
saver = tf.train.Saver(var_list=var_list, max_to_keep=1)
# init
init = tf.global_variables_initializer()
sess.run(init)
if len(self.pretrain_path):
flag = 0
try:
print('try to restore the whole graph')
saver.restore(sess, self.pretrain_path)
except:
print('failed to restore the whole graph')
flag = 1
if flag:
try:
print('try to restore the graph body')
restore_weights = [
v for v in var_list if 'yolo_head' not in v.name
]
sv = tf.train.Saver(var_list=restore_weights)
sv.restore(sess, self.pretrain_path)
except Exception:
raise Exception(
'restore body failed, please check the pretained weight'
)
total_step = int(np.ceil(
len(self.train_data) / self.batch_size)) * self.epoch
print(
'train on {} samples, val on {} samples, batch size {}, total {} epoch'
.format(len(self.train_data), len(self.val_data), self.batch_size,
self.epoch))
step = 0
epoch = 0
t0 = time.time()
for data in self.generate_data(grid_shape):
step += 1
img, label, idx = data
pred_, losses_, _ = sess.run([pred, losses, op], {
self.input: img,
self.label: label
})
t1 = time.time()
print('step:{:<d}/{} epoch:{} loss:{:< .3f} ETA:{}'.format(
step, total_step, epoch, losses_,
sec2time((t1 - t0) * (total_step / step - 1))))
if idx == 0:
# for visual
raw_, boxes, grid = pred_
vis_img = []
for b in range(self.batch_size):
picked_boxes = pick_box(boxes[b], 0.3, self.hw,
self.classes)
per_img = np.array(img[b] * 255, dtype=np.uint8)
# draw pred
per_img_ = per_img.copy()
per_img_ = plot_rectangle(per_img_, picked_boxes, 1, 1,
self.color_table, self.classes)
vis_img.append(per_img_)
# draw gts
per_img_ = per_img.copy()
per_label = label[b]
picked_boxes = pick_box(per_label[..., 4:], 0.3, self.hw,
self.classes)
per_img_ = plot_rectangle(per_img_, picked_boxes, 1, 1,
self.color_table, self.classes,
True)
vis_img.append(per_img_)
# cal vaild_loss
val_loss_ = 0
val_step = 0
for val_data in self.generate_data(grid_shape, is_val=True):
img, label = val_data
_, losses__ = sess.run([pred, losses], {
self.input: img,
self.label: label
})
val_loss_ += losses__
val_step += self.batch_size
if val_step >= len(self.val_data):
break
val_loss_ /= (val_step / self.batch_size)
ss = sess.run(summary,
feed_dict={
img_tensor: np.array(vis_img),
train_loss_tensor: losses_,
val_loss_tensor: val_loss_
})
writer.add_summary(ss, epoch)
saver.save(sess,
join(self.log_path,
split(self.log_path)[-1] + '_model'),
write_meta_graph=False,
write_state=False)
print('epoch:{} train_loss:{:< .3f} val_loss:{:< .3f}'.format(
epoch, losses_, val_loss_))
epoch += 1
if epoch >= self.epoch:
break
def train(self):
# pred, losses, op = self.create_model()
pred = model(self.input, len(self.classes), self.anchors,
self.config.net_type, self.is_training, True)
grid_shape = [g.get_shape().as_list() for g in pred[2]]
s = sum([g[2] * g[1] for g in grid_shape])
self.label = tf.placeholder(
tf.float32, [self.batch_size, s, 3, 9 + len(self.classes)])
# for data in self.generate_data(grid_shape):
# print()
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
var_list = tf.global_variables()
losses = loss(pred, self.label, self.hw, self.lambda_coord,
self.lambda_noobj, self.lambda_cls, self.iou_threshold,
self.config.debug)
opt = tf.train.AdamOptimizer(self.learn_rate)
with tf.control_dependencies(update_ops):
op = opt.minimize(losses)
# summary
writer = tf.summary.FileWriter(self.log_path, max_queue=-1)
img_tensor = tf.placeholder(tf.float32,
[2 * self.batch_size] + self.hw + [3])
with tf.name_scope('loss'):
train_loss_tensor = tf.placeholder(tf.float32)
val_loss_tensor = tf.placeholder(tf.float32)
tf.summary.scalar('train_loss', train_loss_tensor)
tf.summary.scalar('val_loss', val_loss_tensor)
with tf.name_scope('mAP'):
for iou in self.ious_thres:
with tf.name_scope('iou{}'.format(iou)):
exec('map_with_iou{} = tf.placeholder(tf.float32)'.format(
int(iou * 100)))
exec('tf.summary.scalar("mAP", map_with_iou{})'.format(
int(iou * 100)))
with tf.name_scope('per_class_AP'):
for iou in self.ious_thres:
with tf.name_scope('iou{}'.format(iou)):
for per_cls in self.classes:
per_cls = per_cls.replace(' ', '_')
exec('ap_{}_with_iou{} = tf.placeholder(tf.float32)'.
format(per_cls, int(iou * 100)))
exec(
'tf.summary.scalar("{}", ap_{}_with_iou{})'.format(
per_cls, per_cls, int(iou * 100)))
tf.summary.image('img', img_tensor, 2 * self.batch_size)
summary = tf.summary.merge_all()
conf = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))
sess = tf.Session(config=conf)
# sess = tf_debug.LocalCLIDebugWrapperSession(sess)
# sess = tf_debug.TensorBoardDebugWrapperSession(sess, "PC-DAIXILI:6001")
saver = tf.train.Saver(var_list=var_list, max_to_keep=5)
# saver = tf.train.Saver()
# init
init = tf.global_variables_initializer()
sess.run(init)
if len(self.pretrain_path):
flag = 0
try:
print('try to restore the whole graph')
saver.restore(sess, self.pretrain_path)
print('successfully restore the whole graph')
except:
print('failed to restore the whole graph')
flag = 1
if flag:
try:
print('try to restore the graph body')
restore_weights = [
v for v in var_list if 'yolo_head' not in v.name
]
sv = tf.train.Saver(var_list=restore_weights)
sv.restore(sess, self.pretrain_path)
print('successfully restore the graph body')
except Exception:
raise Exception(
'restore body failed, please check the pretained weight'
)
total_step = int(np.ceil(
len(self.train_data) / self.batch_size)) * self.epoch
print(
'train on {} samples, val on {} samples, batch size {}, total {} epoch'
.format(len(self.train_data), len(self.val_data), self.batch_size,
self.epoch))
step = 0
epoch = 0
t0 = time.time()
DETECTION = defaultdict(lambda: defaultdict(list))
FP_TP = defaultdict(lambda: defaultdict(list))
GT_NUMS = defaultdict(int)
for data in self.generate_data(grid_shape):
step += 1
img, label, idx = data
pred_, losses_, _ = sess.run([pred, losses, op], {
self.input: img,
self.label: label,
self.is_training: True
})
t1 = time.time()
print('step:{:<d}/{} epoch:{} loss:{:< .3f} ETA:{}'.format(
step, total_step, epoch, losses_,
sec2time((t1 - t0) * (total_step / step - 1))))
if idx == 0:
# for training visual
raw_, boxes, grid = pred_
vis_img = []
for b in range(self.batch_size):
picked_boxes = pick_box(boxes[b], 0.3, 0.3, self.hw,
self.classes)
per_img = np.array(img[b] * 255, dtype=np.uint8)
# draw pred
per_img_ = per_img.copy()
per_img_ = plot_img(per_img_, picked_boxes,
self.color_table, self.classes)
vis_img.append(per_img_)
# draw gts
per_img_ = per_img.copy()
per_label = label[b]
picked_boxes = pick_box(per_label[..., 4:], 0.3, 0.3,
self.hw, self.classes)
per_img_ = plot_img(per_img_, picked_boxes,
self.color_table, self.classes, True)
vis_img.append(per_img_)
# cal valid_loss
val_loss_ = 0
val_step = 0
cnt = 0
GT_NUMS.clear()
DETECTION.clear()
FP_TP.clear()
for val_data in self.generate_data(grid_shape, is_val=True):
cnt += self.batch_size
print("valid data: {}/{}".format(cnt, len(self.val_data)),
end='\n')
img, label, GTS = val_data
pred_, losses__ = sess.run(
[pred, losses], {
self.input: img,
self.label: label,
self.is_training: False
})
_, boxes_, _ = pred_
for b in range(self.batch_size):
DETECTION[b] = defaultdict(list)
picked_boxes = pick_box(boxes_[b], 0.01, 0.5, self.hw,
self.classes) # NMS
for picked_box in picked_boxes:
DETECTION[b][self.classes[int(
picked_box[5])]].append(
picked_box[:5].tolist())
# cal FP TP
# import pdb
# pdb.set_trace()
cal_fp_fn_tp_tn(DETECTION, GTS, FP_TP, GT_NUMS,
self.classes, self.ious_thres)
val_loss_ += losses__
val_step += self.batch_size
DETECTION.clear()
if val_step >= len(self.val_data):
break
APs, mAPs = cal_mAP(FP_TP, GT_NUMS, self.classes,
self.ious_thres)
print(APs)
print(mAPs)
# import pdb
# pdb.set_trace()
val_loss_ /= (val_step / self.batch_size)
feed_dict = {
img_tensor: np.array(vis_img),
train_loss_tensor: losses_,
val_loss_tensor: val_loss_
}
for iou in self.ious_thres:
exec('feed_dict[map_with_iou{0}] = mAPs[{1}] '.format(
int(iou * 100), iou))
for per_cls in self.classes:
per_clses = per_cls.replace(' ', '_')
exec(
'feed_dict[ap_{0}_with_iou{1}] = APs[{2}]["{3}"] '.
format(per_clses, int(iou * 100), iou, per_cls))
ss = sess.run(summary, feed_dict=feed_dict)
writer.add_summary(ss, epoch)
saver.save(sess,
join(
self.log_path,
split(self.log_path)[-1] +
'_model_epoch_{}'.format(epoch)),
write_meta_graph=False,
write_state=False)
print('epoch:{} train_loss:{:< .3f} val_loss:{:< .3f}'.format(
epoch, losses_, val_loss_))
epoch += 1
if epoch >= self.epoch:
break