def split_coco(imgs_path,
annotaions_path,
dst_dir,
num_catetory=20,
num_per_category=18):
"""
:param origin_path:
:param split_ratio:
:return:
"""
dataset = json.load(open(annotaions_path, 'r'))
sub_annotations_path = os.path.join(dst_dir, 'Annotations')
sub_img_path = os.path.join(dst_dir, 'Images')
anns, cats, imgs, img_anns, cate_imgs = create_index(dataset)
img_id_list, category_id_list = get_img_per_categorise(
cate_imgs, num_catetory, num_per_category)
img_name_dict = {}
for i, img_id in enumerate(img_id_list):
img_name_dict[img_id] = '0' * (
12 - len(str(img_id))) + '{0}.jpg'.format(img_id)
#----------------------------write annotaion info-----------------------------------
images_list, annotations_list = get_images_annotaion_info(
img_id_list, imgs, img_anns, category_id_list)
new_dataset = defaultdict(list)
new_dataset['info'] = dataset['info']
new_dataset['licenses'] = dataset['licenses']
new_dataset['images'] = images_list
new_dataset['annotations'] = annotations_list
new_dataset['categories'] = dataset['categories']
makedir(sub_annotations_path)
json_path = os.path.join(sub_annotations_path, 'instances.json')
with open(json_path, 'w') as fw:
json.dump(new_dataset, fw)
print(
'Successful write the number of {0} annotations respect to {1} images to {2}'
.format(len(new_dataset['annotations']), len(new_dataset['images']),
json_path))
#---------------------------------remove image---------------------------------------
makedir(sub_img_path)
num_samples = 0
for img_id, img_name in img_name_dict.items():
shutil.copy(os.path.join(imgs_path, img_name),
os.path.join(sub_img_path, '{0}.jpg'.format(img_id)))
num_samples += 1
view_bar("split coco:", num_samples, len(img_name_dict))
print('Successful copy the number of {0} images to {1}'.format(
len(img_name_dict), sub_img_path))
def show_save_image_grid(images, save_dir=None, batch_size=128, id=None):
fig = plt.figure(figsize=(8, batch_size / 32)) # width height
fig.suptitle("Epoch {}".format(id))
gs = plt.GridSpec(nrows=int(batch_size / 16), ncols=16)
gs.update(wspace=0.05, hspace=0.05)
for i, image in enumerate(images):
ax = plt.subplot(gs[i])
plt.axis('off')
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_aspect('equal')
plt.imshow(image * 127.5 + 127.5, cmap='Greys_r')
makedir(save_dir)
plt.savefig(os.path.join(save_dir, 'epoch_{:04d}.png'.format(id)))
plt.show()
def generate_cls_bbox(all_boxes, img_name_list, detect_bbox_save_path):
"""
:param all_boxes: is a list. each item reprensent the detections of a img.
the detections is a array. shape is [-1, 6]. [category, score, xmin, ymin, xmax, ymax]
Note that: if none detections in this img. that the detetions is : []
:param img_name_list:
:param detect_bbox_save_path:
:return:
"""
# create path
makedir(detect_bbox_save_path)
for cls_name, cls_id in NAME_LABEL_MAP.items():
if cls_id == 0:
continue
print("Writing {} VOC detect bbox file".format(cls_name))
cls_save_path = os.path.join(detect_bbox_save_path, cls_name + '.txt')
with open(cls_save_path, 'wt') as fw:
for index, img_name in enumerate(img_name_list):
detection_per_img = all_boxes[index]
detection_cls_per_img = detection_per_img[
detection_per_img[:, 0].astype(np.int32) == cls_id]
if detection_cls_per_img.shape[
0] == 0: # this cls has none detections in this image
continue
# write box to txt
for detect_box in detection_cls_per_img:
fw.write(
'{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.format(
# [img_name, score, xmin, ymin, xmax, ymax]
img_name,
detect_box[1],
detect_box[2],
detect_box[3],
detect_box[4],
detect_box[5]))
def split_pascal(origin_path, dst_path, split_rate=0.8):
"""
split pascal dataset
:param origin_path:
:return:
"""
image_path = os.path.join(origin_path, 'JPEGImages')
xml_path = os.path.join(origin_path, 'Annotations')
image_train_path = os.path.join(dst_path, 'train', 'JPEGImages')
xml_train_path = os.path.join(dst_path, 'train', 'Annotations')
image_val_path = os.path.join(dst_path, 'val', 'JPEGImages')
xml_val_path = os.path.join(dst_path, 'val', 'Annotations')
makedir(image_train_path)
makedir(xml_train_path)
makedir(image_val_path)
makedir(xml_val_path)
image_list = os.listdir(image_path)
image_name = [image.split('.')[0] for image in image_list]
image_name = np.random.permutation(image_name)
train_image_name = image_name[:int(math.ceil(len(image_name) *
split_rate))]
val_image_name = image_name[int(math.ceil(len(image_name) * split_rate)):]
for n, image in enumerate(train_image_name):
shutil.copy(os.path.join(image_path, image + '.jpg'),
os.path.join(image_train_path, image + '.jpg'))
shutil.copy(os.path.join(xml_path, image + '.xml'),
os.path.join(xml_train_path, image + '.xml'))
view_bar(message="split train dataset:",
num=n,
total=len(train_image_name))
print('Total of {0} data split to {1}'.format(
len(train_image_name), os.path.dirname(image_train_path)))
for n, image in enumerate(val_image_name):
shutil.copy(os.path.join(image_path, image + '.jpg'),
os.path.join(image_val_path, image + '.jpg'))
shutil.copy(os.path.join(xml_path, image + '.xml'),
os.path.join(xml_val_path, image + '.xml'))
view_bar(message="split val dataset:",
num=n,
total=len(val_image_name))
print('Total of {0} data split to {1}'.format(
len(val_image_name), os.path.dirname(image_val_path)))
def train():
faster_rcnn = models.FasterRCNN(base_network_name=cfgs.NET_NAME,
is_training=True)
with tf.name_scope('get_batch'):
img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \
dataset_tfrecord(batch_size=cfgs.BATCH_SIZE,
shortside_len=cfgs.IMG_SHORT_SIDE_LEN,
length_limitation=cfgs.IMG_MAX_LENGTH,
record_file=cfgs.TFRECORD_DIR,
is_training=True)
# construct net work
faster_rcnn.inference()
# ----------------------------------------------------------------------------------------------------build loss
# weight_decay_loss = tf.add_n(slim.losses.get_regularization_losses())
# weight_decay_loss = tf.add_n(tf.losses.get_regularization_losses())
rpn_location_loss = faster_rcnn.loss_dict['rpn_loc_loss']
rpn_cls_loss = faster_rcnn.loss_dict['rpn_cls_loss']
rpn_total_loss = rpn_location_loss + rpn_cls_loss
fastrcnn_cls_loss = faster_rcnn.loss_dict['fastrcnn_cls_loss']
fastrcnn_loc_loss = faster_rcnn.loss_dict['fastrcnn_loc_loss']
fastrcnn_total_loss = fastrcnn_cls_loss + fastrcnn_loc_loss
total_loss = rpn_total_loss + fastrcnn_total_loss
# ____________________________________________________________________________________________________build loss
# gtboxes_in_img = show_box_in_tensor.draw_boxes_with_categories(img_batch=img_batch,
# boxes=gtboxes_and_label[:, :-1],
# labels=gtboxes_and_label[:, -1])
# if cfgs.ADD_BOX_IN_TENSORBOARD:
# detections_in_img = show_box_in_tensor.draw_boxes_with_categories_and_scores(img_batch=img_batch,
# boxes=final_bbox,
# labels=final_category,
# scores=final_scores)
# tf.summary.image('Compare/final_detection', detections_in_img)
# tf.summary.image('Compare/gtboxes', gtboxes_in_img)
# ___________________________________________________________________________________________________add summary
global_step = slim.get_or_create_global_step()
lr = tf.train.piecewise_constant(
global_step,
boundaries=[
np.int64(cfgs.DECAY_STEP[0]),
np.int64(cfgs.DECAY_STEP[1])
],
values=[cfgs.LR, cfgs.LR / 10., cfgs.LR / 100.])
tf.summary.scalar('learning_rate', lr)
optimizer = tf.train.MomentumOptimizer(lr, momentum=cfgs.MOMENTUM)
# ---------------------------------------------------------------------------------------------compute gradients
# -----------------------------------------computer gradient-------------------------------------------------------
gradients = faster_rcnn.get_gradients(optimizer, total_loss)
# enlarge_gradients for bias
if cfgs.MUTILPY_BIAS_GRADIENT:
gradients = faster_rcnn.enlarge_gradients_for_bias(gradients)
if cfgs.GRADIENT_CLIPPING_BY_NORM:
with tf.name_scope('clip_gradients_YJR'):
gradients = slim.learning.clip_gradient_norms(
gradients, cfgs.GRADIENT_CLIPPING_BY_NORM)
# +++++++++++++++++++++++++++++++++++++++++start train+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# train_op
train_op = optimizer.apply_gradients(grads_and_vars=gradients,
global_step=global_step)
summary_op = tf.summary.merge_all()
restorer, restore_ckpt = faster_rcnn.get_restorer()
saver = tf.train.Saver(max_to_keep=30)
# support growth train
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
with tf.Session(config=config) as sess:
sess.run(init_op)
if not restorer is None:
restorer.restore(sess, save_path=restore_ckpt)
print(
'*' * 80 +
'\nSuccessful restore model from {0}\n'.format(restore_ckpt) +
'*' * 80)
# model_variables = slim.get_model_variables()
# for var in model_variables:
# print(var.name, var.shape)
# build summary write
summary_path = os.path.join(cfgs.SUMMARY_PATH, cfgs.VERSION)
summary_writer = tf.summary.FileWriter(summary_path, graph=sess.graph)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess, coord)
try:
if not coord.should_stop():
# ++++++++++++++++++++++++++++++++++++++++start training+++++++++++++++++++++++++++++++++++++++++++++++++++++
for step in range(cfgs.MAX_ITERATION):
training_time = time.strftime('%Y-%m-%d %H:%M:%S',
time.localtime(time.time()))
img_name, image, gtboxes_and_label, num_objects = \
sess.run([img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch])
feed_dict = faster_rcnn.fill_feed_dict(
image_feed=image, gtboxes_feed=gtboxes_and_label)
if step % cfgs.SHOW_TRAIN_INFO_INTE != 0 and step % cfgs.SMRY_ITER != 0:
_, globalStep = sess.run([train_op, global_step],
feed_dict=feed_dict)
else:
if step % cfgs.SHOW_TRAIN_INFO_INTE == 0 and step % cfgs.SMRY_ITER != 0:
start_time = time.time()
_, globalStep, rpnLocLoss, rpnClsLoss, rpnTotalLoss, \
fastrcnnLocLoss, fastrcnnClsLoss, fastrcnnTotalLoss, totalLoss = \
sess.run([train_op, global_step, rpn_location_loss, rpn_cls_loss, rpn_total_loss,
fastrcnn_loc_loss, fastrcnn_cls_loss, fastrcnn_total_loss, total_loss],
feed_dict=feed_dict)
end_time = time.time()
print(""" {}: step {}\t\timage_name:{} |\t
rpn_loc_loss:{} |\t rpn_cla_loss:{} |\t rpn_total_loss:{} |
fast_rcnn_loc_loss:{} |\t fast_rcnn_cla_loss:{} |\t fast_rcnn_total_loss:{} |
total_loss:{} |\t per_cost_time:{}s""" \
.format(training_time, globalStep, str(img_name[0]), rpnLocLoss, rpnClsLoss,
rpnTotalLoss, fastrcnnLocLoss, fastrcnnClsLoss, fastrcnnTotalLoss, totalLoss,
(end_time - start_time)))
else:
if step % cfgs.SMRY_ITER == 0:
_, globalStep, summary_str = sess.run(
[train_op, global_step, summary_op],
feed_dict=feed_dict)
summary_writer.add_summary(
summary_str, globalStep)
summary_writer.flush()
if (step > 0 and step % cfgs.SAVE_WEIGHTS_INTE
== 0) or (step == cfgs.MAX_ITERATION - 1):
save_dir = os.path.join(cfgs.TRAINED_CKPT,
cfgs.VERSION)
makedir(save_dir)
save_ckpt = os.path.join(
save_dir, 'voc_' + str(globalStep) + '_model.ckpt')
saver.save(sess, save_ckpt)
print(' weights had been saved')
except Exception as e:
# Report exceptions to the coordinator.
coord.request_stop(e)
finally:
coord.request_stop()
coord.join(threads)
print('all threads are asked to stop!')
def exucute_detect(self, image_path, save_path):
"""
execute object detect
:param detect_net:
:param image_path:
:return:
"""
input_image = tf.placeholder(dtype=tf.uint8,
shape=(None, None, 3),
name='inputs_images')
resize_img = self.image_process(input_image)
# expend dimension
image_batch = tf.expand_dims(input=resize_img,
axis=0) # (1, None, None, 3)
self.detect_net.images_batch = image_batch
# img_shape = tf.shape(inputs_img)
# load detect network
detection_boxes, detection_scores, detection_category = self.detect_net.inference(
)
# restore pretrain weight
restorer, restore_ckpt = self.detect_net.get_restorer()
# config gpu to growth train
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
with tf.Session(config=config) as sess:
sess.run(init_op)
if restorer is not None:
restorer.restore(sess, save_path=restore_ckpt)
print('Successful restore model from {0}'.format(restore_ckpt))
# construct image path list
format_list = ('.jpg', '.png', '.jpeg', '.tif', '.tiff')
if os.path.isfile(image_path):
image_name_list = [image_path]
else:
image_name_list = [
img_name for img_name in os.listdir(image_path)
if img_name.endswith(format_list)
and os.path.isfile(os.path.join(image_path, img_name))
]
assert len(image_name_list) != 0
print(
"test_dir has no imgs there. Note that, we only support img format of {0}"
.format(format_list))
#+++++++++++++++++++++++++++++++++++++start detect+++++++++++++++++++++++++++++++++++++++++++++++++++++=++
makedir(save_path)
fw = open(os.path.join(save_path, 'detect_bbox.txt'), 'w')
for index, img_name in enumerate(image_name_list):
detect_dict = {}
bgr_img = cv.imread(os.path.join(image_path, img_name))
rgb_img = cv.cvtColor(
bgr_img, cv.COLOR_BGR2RGB
) # convert channel from BGR to RGB (cv is BGR)
start_time = time.perf_counter()
# image resize and white process
# construct feed_dict
feed_dict = {input_image: rgb_img}
resized_img, detected_boxes, detected_scores, detected_categories = \
sess.run([resize_img, detection_boxes, detection_scores, detection_category],
feed_dict=feed_dict)
end_time = time.perf_counter()
# select object according to threshold
object_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
object_scores = detected_scores[object_indices]
object_boxes = detected_boxes[object_indices]
object_categories = detected_categories[object_indices]
final_detections_img = draw_box_in_img.draw_boxes_with_label_and_scores(
resized_img,
boxes=object_boxes,
labels=object_categories,
scores=object_scores)
final_detections_img = cv.cvtColor(final_detections_img,
cv.COLOR_RGB2BGR)
cv.imwrite(os.path.join(save_path, img_name),
final_detections_img)
# resize boxes and image according to raw input image
raw_h, raw_w = rgb_img.shape[0], rgb_img.shape[1]
resized_h, resized_w = resized_img.shape[1], resized_img.shape[
2]
x_min, y_min, x_max, y_max = object_boxes[:, 0], object_boxes[:, 1], object_boxes[:, 2], \
object_boxes[:, 3]
x_min = x_min * raw_w / resized_w
y_min = y_min * raw_h / resized_h
x_max = x_max * raw_w / resized_w
y_max = y_max * raw_h / resized_h
object_boxes = np.stack([x_min, y_min, x_max, y_max], axis=1)
# final_detections= cv.resize(final_detections[:, :, ::-1], (raw_w, raw_h))
# recover to raw size
detect_dict['score'] = object_scores
detect_dict['boxes'] = object_boxes
detect_dict['categories'] = object_categories
# convert from RGB to BG
fw.write(f'\n{img_name}')
for score, boxes, categories in zip(object_scores,
object_boxes,
object_categories):
fw.write('\n\tscore:' + str(score))
fw.write('\tbboxes:' + str(boxes))
fw.write('\tcategories:' + str(categories))
view_bar(
'{} image cost {} second'.format(img_name,
(end_time - start_time)),
index + 1, len(image_name_list))
fw.close()
tf.app.flags.DEFINE_string('xml_dir', 'Annotations', 'xml dir')
tf.app.flags.DEFINE_string('image_dir', 'JPEGImages', 'image dir')
tf.app.flags.DEFINE_string('save_name', 'train', 'save name')
tf.app.flags.DEFINE_string('year', '2007,2012', 'Desired challenge year.')
tf.app.flags.DEFINE_string('save_dir', tfrecord_dir, 'save name')
tf.app.flags.DEFINE_string('img_format', 'jpg', 'format of image')
tf.app.flags.DEFINE_string('dataset', 'car', 'dataset')
FLAGS = tf.app.flags.FLAGS
try:
if os.path.exists(original_dataset_dir) is False:
raise IOError('dataset is not exist please check the path')
except FileNotFoundError as e:
print(e)
finally:
makedir(tfrecord_dir)
def _int64_feature(value):
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def _bytes_feature(value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
def read_xml_gtbox_and_label(xml_path):
"""
read gtbox(ground truth) and label from xml
:param xml_path: the path of voc xml
:return: a list contains gtboxes and labels, shape is [num_of_gtboxes, 5],
def train():
# -----------------step 1 Get the SSD network and its anchors.----------------
SSDNet.default_params._replace(num_classes=cfgs.NUM_CLASS+1)
# construct ssd net
ssd_net = SSDNet()
ssd_shape = ssd_net.params.img_shape
global_step = ssd_net.global_step
ssd_anchors = ssd_net.make_anchors(ssd_shape)
# ----------------- step 2 Create a dataset provider and batches.---------------
with tf.name_scope(cfgs.DATASET_NAME + '_data_provider'):
anchor_encoder_fn = lambda gt_labels, gb_bboxes: ssd_net.bboxes_encode(gt_labels, gb_bboxes, ssd_anchors,
scope="anchor_encode")
batch_shape = [1] * 3 + [len(ssd_anchors)] * 3
image_batch, filename_batch, shape_batch, labels_batch, bboxes_batch, scores_batch = dataset_tfrecord(
dataset_dir=cfgs.TFRECORD_DIR, split_name='train', batch_size=cfgs.BATCH_SIZE,
anchor_encoder_fn=anchor_encoder_fn, batch_shape=batch_shape, num_threads=cfgs.NUM_THREADS,
is_training=True)
# batch_queue = slim.prefetch_queue.prefetch_queue(
# tf_utils.reshape_list([image_batch, filename_batch, shape_batch, labels_batch, bboxes_batch, scores_batch]),
# capacity=2)
#
# image_batch, filename_batch, shape_batch, labels_batch, bboxes_batch, scores_batch = \
# tf_utils.reshape_list(batch_queue.dequeue(), batch_shape)
# -------------------step 3 construct foward network-----------------------------
# Construct SSD network.
arg_scope = ssd_net.arg_scope(weight_decay=cfgs.WEIGHT_DECAY, data_format=cfgs.DATA_FORMAT)
with slim.arg_scope(arg_scope):
predictions, localisations, logits, end_points = ssd_net.net(image_batch, is_training=True)
# Add loss function.
ssd_net.losses(logits, localisations,
labels_batch, bboxes_batch, scores_batch,
match_threshold=cfgs.MATCH_THRESHOLD,
negative_ratio=cfgs.NEGATIVE_RATIO,
alpha=cfgs.LOSS_ALPHA,
label_smoothing=cfgs.LABELS_SMOOTH)
# Gather initial summaries.
# summaries = set(tf.get_collection(tf.GraphKeys.SUMMARIES))
# =================================================================== #
# Add summaries from first clone.
# =================================================================== #
# Gather update_ops from the first clone. These contain, for example,
# the updates for the batch_norm variables created by network_fn.
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
# Add summaries for end_points.
# Add summaries for losses and extra losses.
# for loss in tf.get_collection(tf.GraphKeys.LOSSES):
# tf.summary.scalar(loss.op.name, loss)
# Add summaries for variables.
for variable in slim.get_model_variables():
tf.summary.histogram(variable.op.name, variable)
# =================================================================== #
# Configure the moving averages.
# =================================================================== #
if cfgs.MOVING_AVERATE_DECAY:
moving_average_variables = slim.get_model_variables()
variable_averages = tf.train.ExponentialMovingAverage(
cfgs.MOVING_AVERATE_DECAY, ssd_net.global_step)
else:
moving_average_variables, variable_averages = None, None
# =================================================================== #
# Configure the optimization procedure.
# =================================================================== #
learning_rate = ssd_net.learning_rate(boundaries=[cfgs.WARM_UP_STEP, cfgs.DECAY_STEP[0], cfgs.DECAY_STEP[1]],
rates=[cfgs.WARM_UP_LEARING_RATE, cfgs.LEARING_RATE_BASE,
cfgs.LEARING_RATE_BASE / 10., cfgs.LEARING_RATE_BASE / 100.],
global_step=ssd_net.global_step,
warmup=True)
truncated_learning_rate = tf.maximum(learning_rate, tf.constant(cfgs.END_LEARNING_RATE, dtype=learning_rate.dtype), name='learning_rate')
tf.summary.scalar('learning_rate', truncated_learning_rate)
#-
optimizer = tf.train.MomentumOptimizer(truncated_learning_rate, momentum=cfgs.MOMENTUM)
if cfgs.MOVING_AVERATE_DECAY:
with tf.name_scope("weight_decay"):
# Update ops executed locally by trainer.
update_ops.append(variable_averages.apply(moving_average_variables))
# and returns a train_tensor and summary_op
# with tf.name_scope("first_stage_train"):
# first_stage_trainable_var_list = ssd_net.get_train_variable(scopes=cfgs.BACKBONE_SCOPE)
# print("*" * 40, "first stage trainable variable:", "*" * 40)
# for var in first_stage_trainable_var_list:
# print(var.op.name, var.shape)
# total_loss, gradients = ssd_net.optimize_gradient(optimizer=optimizer, var_list=first_stage_trainable_var_list)
# # Create gradient updates.
# grad_updates = optimizer.apply_gradients(gradients,
# global_step=global_step)
# update_ops.append(grad_updates)
# update_ops.append(global_step)
# train_op_with_frozen = control_flow_ops.with_dependencies(update_ops, total_loss, name='train_op')
#
# with tf.name_scope("second_stage_train"):
# second_stage_trainable_var_list = ssd_net.get_train_variable(scopes=None)
# print("*" * 40, "second stage trainable variable:", "*" * 40)
# for var in second_stage_trainable_var_list:
# print(var.op.name, var.shape)
# total_loss, gradients = ssd_net.optimize_gradient(optimizer=optimizer, var_list=second_stage_trainable_var_list)
# # Create gradient updates.
# grad_updates = optimizer.apply_gradients(gradients, global_step=global_step)
# update_ops.append(grad_updates)
# update_ops.append(global_step)
# train_op_with_all_variable = control_flow_ops.with_dependencies(update_ops, total_loss, name='train_op')
total_loss, gradients = ssd_net.optimize_gradient(optimizer=optimizer)
# Create gradient updates.
grad_updates = optimizer.apply_gradients(gradients, global_step=global_step)
update_ops.append(grad_updates)
update_ops.append(global_step)
train_op = control_flow_ops.with_dependencies(update_ops, total_loss, name='train_op')
# Add total_loss to summary.
tf.summary.scalar('total_loss', total_loss)
# Add the summaries from the first clone. These contain the summaries
summary_op = tf.summary.merge_all()
# =================================================================== #
# Kicks off the training.
# =================================================================== #
saver = tf.train.Saver(max_to_keep=30,
keep_checkpoint_every_n_hours=2.0,
write_version=2)
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=cfgs.GPU_MEMORY_FRACTION,
allow_growth=True)
config = tf.ConfigProto(log_device_placement=False,
gpu_options=gpu_options)
init_op = tf.group(
tf.global_variables_initializer(),
tf.local_variables_initializer()
)
with tf.Session(config=config) as sess:
sess.run(init_op)
# if not restorer is None:
# restorer.restore(sess, save_path=restore_ckpt)
restore_ckpt = ssd_net.restore_ckpt(sess)
print('*' * 80 + '\nSuccessful restore model from {0}\n'.format(restore_ckpt) + '*' * 80)
# model_variables = slim.get_model_variables()
# for var in model_variables:
# print(var.name, var.shape)
# build summary write
summary_writer = tf.summary.FileWriter(cfgs.SUMMARY_PATH, graph=sess.graph)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess, coord)
# ++++++++++++++++++++++++++++++++++++++++start training+++++++++++++++++++++++++++++++++++++++++++++++++++++
try:
if not coord.should_stop():
for step in range(cfgs.MAX_ITERATION):
# if step < cfgs.FIRST_STAGE_STEP:
# train_op = train_op_with_frozen
# train_stage = 1
# else:
# train_op = train_op_with_all_variable
# train_stage = 2
train_stage = 0
# image, labels, bboxes, scores = \
# sess.run([image_batch, labels_batch, bboxes_batch, scores_batch])
# feed_dict = ssd_net.fill_feed_dict(image, labels, bboxes, scores)
training_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
if step % cfgs.SHOW_TRAIN_INFO_INTE != 0 and step % cfgs.SMRY_ITER != 0:
_, globalStep = sess.run([train_op, global_step])
else:
if step % cfgs.SHOW_TRAIN_INFO_INTE == 0 and step % cfgs.SMRY_ITER != 0:
start_time = time.time()
_, globalStep, totalLoss = sess.run([train_op, global_step, total_loss])
end_time = time.time()
print(""" {} train stage-{}: step {}\t | total_loss:{} |\t per_cost_time:{}s""" \
.format(training_time, train_stage, globalStep, totalLoss,(end_time - start_time)))
else:
if step % cfgs.SMRY_ITER == 0:
_, globalStep, summary_str = sess.run([train_op, global_step, summary_op])
summary_writer.add_summary(summary_str, globalStep)
summary_writer.flush()
if (step > 0 and step % cfgs.SAVE_WEIGHTS_INTE == 0) or (step == cfgs.MAX_ITERATION - 1):
save_dir = os.path.join(cfgs.TRAINED_CKPT, cfgs.VERSION)
makedir(save_dir)
save_ckpt = os.path.join(save_dir, 'voc_' + str(globalStep) + '_model.ckpt')
saver.save(sess, save_ckpt)
print(' weights had been saved')
except Exception as e:
# Report exceptions to the coordinator.
coord.request_stop(e)
finally:
coord.request_stop()
coord.join(threads)
print('all threads are asked to stop!')
def exucute_detect(self, image_path, save_path):
"""
execute object detect
:param detect_net:
:param image_path:
:return:
"""
# load detect network
pred_sbbox_batch, pred_mbbox_batch, pred_lbbox_batch = self.detector.pred_sbbox, self.detector.pred_mbbox, self.detector.pred_lbbox
# TensorFlow session: grow memory when needed. TF, DO NOT USE ALL MY GPU MEMORY!!!
gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options)
init_op = tf.group(
tf.global_variables_initializer(),
tf.local_variables_initializer()
)
with tf.Session(config=config) as sess:
sess.run(init_op)
# restore pretrain weight
if self.ckpt_path is not None:
restorer = tf.train.Saver()
restorer.restore(sess, self.ckpt_path)
else:
restorer, ckpt_path = self.detector.get_restorer(is_training=False)
restorer.restore(sess, ckpt_path)
print('*'*80 +'\nSuccessful restore model from {0}\n'.format(self.ckpt_path) + '*'*80)
# construct image path list
format_list = ('.jpg', '.png', '.jpeg', '.tif', '.tiff')
if os.path.isfile(image_path):
image_name_list = [image_path]
else:
image_name_list = [img_name for img_name in os.listdir(image_path)
if img_name.endswith(format_list) and os.path.isfile(os.path.join(image_path, img_name))]
assert len(image_name_list) != 0
print("test_dir has no imgs there. Note that, we only support img format of {0}".format(format_list))
#+++++++++++++++++++++++++++++++++++++start detect+++++++++++++++++++++++++++++++++++++++++++++++++++++=++
makedir(save_path)
fw = open(os.path.join(save_path, 'detect_bbox.txt'), 'w')
for index, img_name in enumerate(image_name_list):
detect_dict = {}
original_image, image_batch, original_size = self.image_process(img_path=os.path.join(image_path, img_name))
start_time = time.perf_counter()
# image resize and white process
# construct feed_dict
# Run SSD network.]
feed_dict = {self.input_data: image_batch,
self.trainable: False}
pred_sbbox, pred_mbbox, pred_lbbox = sess.run([pred_sbbox_batch, pred_mbbox_batch, pred_lbbox_batch],
feed_dict=feed_dict)
pred_bbox = np.concatenate([np.reshape(pred_sbbox, (-1, 5 + self.num_classes)),
np.reshape(pred_mbbox, (-1, 5 + self.num_classes)),
np.reshape(pred_lbbox, (-1, 5 + self.num_classes))], axis=0)
bboxes = box_utils.postprocess_boxes(pred_bbox, original_size, self.input_size[0], self.score_threshold)
bboxes = box_utils.nms(bboxes, self.num_threshold, method='nms')
end_time = time.perf_counter()
image = draw_box_in_image.draw_bbox(original_image, bboxes, classes=self.class_name)
image = Image.fromarray(image)
image.save(os.path.join(save_path, img_name))
# resize boxes and image according to raw input image
# final_detections= cv.resize(final_detections[:, :, ::-1], (raw_w, raw_h))
# recover to raw size
bboxes = np.array(bboxes)
rbboxes = bboxes[:, :4]
rscores = bboxes[:, 4]
rclasses = bboxes[:, 5]
# convert from RGB to BG
fw.write(f'\n{img_name}')
for score, boxes, categories in zip(rscores, rbboxes, rclasses):
fw.write('\n\tscore:' + str(score))
fw.write('\tbboxes:' + str(boxes))
fw.write('\tcategories:' + str(int(categories)))
view_bar('{} image cost {} second'.format(img_name, (end_time - start_time)), index + 1,
len(image_name_list))
fw.close()
def convert_coco_to_tfrecord(src_path,
save_path,
record_capacity=2000,
raw_coco=True):
"""
:param src_path:
:param save_path:
:param record_capacity:
:param raw_coco:
:return:
"""
imgs_path = os.path.join(src_path, FLAGS.image_dir)
anns_path = os.path.join(src_path, FLAGS.anns_dir)
# img_name_list = glob.glob(os.path.join(img_path,'*'+FLAGS.img_format))
annotation_list = glob.glob(os.path.join(anns_path, '*.json'))
anns, cats, imgs, img_anns, cate_imgs = create_index(annotation_list[0])
image_id_list = [img_id for img_id in img_anns.keys()]
remainder_num = len(image_id_list) % record_capacity
if remainder_num == 0:
num_record = int(len(image_id_list) / record_capacity)
else:
num_record = int(len(image_id_list) / record_capacity) + 1
for index in range(num_record):
makedir(save_path)
record_filename = os.path.join(save_path, f'{index}.record')
write = tf.io.TFRecordWriter(record_filename)
if index < num_record - 1:
sub_img_id_list = image_id_list[index *
record_capacity:(index + 1) *
record_capacity]
else:
sub_img_id_list = image_id_list[(index * record_capacity):(
index * record_capacity + remainder_num)]
num_samples = 0
for index, img_id in enumerate(sub_img_id_list):
try:
# get gtbox_label
gtbox_label = read_json_gtbox_label(img_anns[img_id])
# get image name
if raw_coco:
img_name = '0' * (
12 - len(str(img_id))) + f'{img_id}.{FLAGS.img_format}'
else:
img_name = '{0}.jpg'.format(img_id)
img_path = os.path.join(imgs_path, img_name)
# load image
bgr_image = cv.imread(img_path)
# BGR TO RGB
rgb_image = cv.cvtColor(bgr_image, cv.COLOR_BGR2RGB)
img_height = rgb_image.shape[0]
img_width = rgb_image.shape[1]
image_record = serialize_example(image=rgb_image,
img_height=img_height,
img_width=img_width,
img_depth=3,
filename=img_name,
gtbox_label=gtbox_label)
write.write(record=image_record)
num_samples += 1
view_bar(message='\nConversion progress',
num=num_samples,
total=len(img_anns))
except Exception as e:
print(e)
continue
write.close()
print('There are {0} samples convert to {1}'.format(
num_samples, save_path))
str(class_ind)])
print(annotation)
f.write(annotation + "\n")
index += 1
return len(image_inds)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--data_path", default=cfgs.DATASET_DIR)
parser.add_argument("--train_annotation", default=cfgs.TRAIN_ANNOT_PATH)
parser.add_argument("--test_annotation", default=cfgs.TEST_ANNOT_PATH)
flags = parser.parse_args()
makedir(os.path.dirname(flags.train_annotation))
if os.path.exists(flags.train_annotation):
os.remove(flags.train_annotation)
if os.path.exists(flags.test_annotation): os.remove(flags.test_annotation)
train_path_1 = os.path.join(flags.data_path, 'train', 'VOC2007')
train_path_2 = os.path.join(flags.data_path, 'train', 'VOC2012')
test_path_1 = os.path.join(flags.data_path, 'test', 'VOC2007')
num1 = convert_voc_annotation(train_path_1, 'trainval',
flags.train_annotation, False)
num2 = convert_voc_annotation(train_path_2, 'trainval',
flags.train_annotation, False)
num3 = convert_voc_annotation(test_path_1, 'test', flags.test_annotation,
False)
def train(self):
self.sess.run(tf.global_variables_initializer())
try:
print('=> Restoring weights from: {0} ... '.format(
self.checkpoint_path))
self.loader.restore(self.sess, self.checkpoint_path)
except:
print('=> {0}does not exist !!!'.format(self.checkpoint_path))
print('=> Now it starts to train YOLOV3 from scratch ...')
self.first_stage_epochs = 0
for epoch in range(
1, 1 + self.first_stage_epochs + self.second_stage_epochs):
if epoch <= self.first_stage_epochs:
train_op = self.train_op_with_frozen_variables
else:
train_op = self.train_op_with_all_variables
pbar = tqdm(self.train_dataset)
train_epoch_loss, test_epoch_loss = [], []
for train_data in pbar:
_, summary, train_step_loss, global_step_val = self.sess.run(
[train_op, self.write_op, self.loss, self.global_step],
feed_dict={
self.input_data: train_data[0],
self.label_sbbox: train_data[1],
self.label_mbbox: train_data[2],
self.label_lbbox: train_data[3],
self.true_sbboxes: train_data[4],
self.true_mbboxes: train_data[5],
self.true_lbboxes: train_data[6],
self.trainable: True,
})
train_epoch_loss.append(train_step_loss)
self.summary_writer.add_summary(summary, global_step_val)
pbar.set_description(
"train loss: {:.2f}".format(train_step_loss))
for test_data in self.test_dataset:
test_step_loss = self.sess.run(
self.loss,
feed_dict={
self.input_data: test_data[0],
self.label_sbbox: test_data[1],
self.label_mbbox: test_data[2],
self.label_lbbox: test_data[3],
self.true_sbboxes: test_data[4],
self.true_mbboxes: test_data[5],
self.true_lbboxes: test_data[6],
self.trainable: False,
})
test_epoch_loss.append(test_step_loss)
train_epoch_loss, test_epoch_loss = np.mean(
train_epoch_loss), np.mean(test_epoch_loss)
save_dir = os.path.join(cfgs.TRAINED_CKPT, cfgs.VERSION)
tools.makedir(save_dir)
ckpt_file = os.path.join(
save_dir, "yolov3_loss={:.4f}.ckpt".format(test_epoch_loss))
log_time = time.strftime('%Y-%m-%d %H:%M:%S',
time.localtime(time.time()))
print(
"=> Epoch: {0} Time: {1} Train loss: {2:.4f} Test loss: {3:.4f} Saving {4} ..."
.format(epoch, log_time, train_epoch_loss, test_epoch_loss,
ckpt_file))
self.saver.save(self.sess, ckpt_file, global_step=epoch)
def imsave(images, size, path):
makedir(os.path.dirname(path))
return plt.imsave(path, merge(images, size))
def execute(self):
predicted_dir_path = os.path.join(cfgs.TEST_SAVE_MAP_PATH, 'predicted')
ground_truth_dir_path = os.path.join(cfgs.TEST_SAVE_MAP_PATH, 'truth')
if os.path.exists(predicted_dir_path):
shutil.rmtree(predicted_dir_path)
if os.path.exists(ground_truth_dir_path):
shutil.rmtree(ground_truth_dir_path)
if os.path.exists(self.write_image_path):
shutil.rmtree(self.write_image_path)
tools.makedir(predicted_dir_path)
tools.makedir(ground_truth_dir_path)
tools.makedir(self.write_image_path)
with open(self.annotation_path, 'r') as annotation_file:
for num, line in enumerate(annotation_file):
annotation = line.strip().split()
image_path = annotation[1]
image_name = image_path.split('/')[-1]
image = cv2.imread(image_path)
bbox_data_gt = np.array(
[list(map(int, box.split(','))) for box in annotation[4:]])
if len(bbox_data_gt) == 0:
bboxes_gt = []
classes_gt = []
else:
bboxes_gt, classes_gt = bbox_data_gt[:, :
4], bbox_data_gt[:, 4]
ground_truth_path = os.path.join(ground_truth_dir_path,
str(num) + '.txt')
print('=> ground truth of {0}:'.format(image_name))
num_bbox_gt = len(bboxes_gt)
with open(ground_truth_path, 'w') as f:
for i in range(num_bbox_gt):
class_name = self.classes[classes_gt[i]]
xmin, ymin, xmax, ymax = list(map(str, bboxes_gt[i]))
bbox_mess = ' '.join(
[class_name, xmin, ymin, xmax, ymax]) + '\n'
f.write(bbox_mess)
print('\t' + str(bbox_mess).strip())
print('=> predict result of {0}:'.format(image_name))
predict_result_path = os.path.join(predicted_dir_path,
str(num) + '.txt')
bboxes_pr = self.predict(image)
if self.write_image:
image = draw_box_in_image.draw_bbox(
image,
bboxes_pr,
classes=self.classes,
show_label=self.show_label)
cv2.imwrite(
os.path.join(self.write_image_path, image_name), image)
with open(predict_result_path, 'w') as f:
for bbox in bboxes_pr:
coor = np.array(bbox[:4], dtype=np.int32)
score = bbox[4]
class_ind = int(bbox[5])
class_name = self.classes[class_ind]
score = '%.4f' % score
xmin, ymin, xmax, ymax = list(map(str, coor))
bbox_mess = ' '.join(
[class_name, score, xmin, ymin, xmax, ymax]) + '\n'
f.write(bbox_mess)
print('\t' + str(bbox_mess).strip())
def exucute_detect(self, image_path, save_path):
"""
execute object detect
:param detect_net:
:param image_path:
:return:
"""
input_image = tf.placeholder(dtype=tf.uint8,
shape=(None, None, 3),
name='inputs_images')
image_pre, labels_pre, bboxes_pre = self.image_process(
input_image, img_shape=self.net_shape, img_format=self.data_format)
# expend dimension
image_batch = tf.expand_dims(input=image_pre,
axis=0) # (1, None, None, 3)
# img_shape = tf.shape(inputs_img)
# load detect network
reuse = True if 'ssd_net' in locals() else None
with slim.arg_scope(
self.ssd_net.arg_scope(data_format=self.data_format)):
detection_category, detection_bbox, _, _ = self.ssd_net.net(
image_batch, is_training=False, reuse=reuse)
# restore pretrain weight
restorer = tf.train.Saver()
# TensorFlow session: grow memory when needed. TF, DO NOT USE ALL MY GPU MEMORY!!!
gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(log_device_placement=False,
gpu_options=gpu_options)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
with tf.Session(config=config) as sess:
sess.run(init_op)
if self.ckpt_path is not None:
restorer.restore(sess, self.ckpt_path)
else:
self.ckpt_path = self.ssd_net.restore_ckpt(sess)
print('*' * 80 + '\nSuccessful restore model from {0}\n'.format(
self.ckpt_path) + '*' * 80)
# construct image path list
format_list = ('.jpg', '.png', '.jpeg', '.tif', '.tiff')
if os.path.isfile(image_path):
image_name_list = [image_path]
else:
image_name_list = [
img_name for img_name in os.listdir(image_path)
if img_name.endswith(format_list)
and os.path.isfile(os.path.join(image_path, img_name))
]
assert len(image_name_list) != 0
print(
"test_dir has no imgs there. Note that, we only support img format of {0}"
.format(format_list))
#+++++++++++++++++++++++++++++++++++++start detect+++++++++++++++++++++++++++++++++++++++++++++++++++++=++
makedir(save_path)
fw = open(os.path.join(save_path, 'detect_bbox.txt'), 'w')
for index, img_name in enumerate(image_name_list):
detect_dict = {}
bgr_img = cv.imread(os.path.join(image_path, img_name))
rgb_img = cv.cvtColor(
bgr_img, cv.COLOR_BGR2RGB
) # convert channel from BGR to RGB (cv is BGR)
start_time = time.perf_counter()
# image resize and white process
# construct feed_dict
# Run SSD network.]
feed_dict = {input_image: rgb_img}
image, category, bbox = sess.run(
[image_batch, detection_category, detection_bbox],
feed_dict=feed_dict)
# Get classes and bboxes from the net outputs.
rclasses, rscores, rbboxes = np_methods.ssd_bboxes_select(
category,
bbox,
self.ssd_anchors,
select_threshold=self.select_threshold,
img_shape=self.net_shape,
num_classes=self.num_classes,
decode=True)
rbboxes = np_methods.bboxes_clip(self.bbox_image, rbboxes)
rclasses, rscores, rbboxes = np_methods.bboxes_sort(rclasses,
rscores,
rbboxes,
top_k=400)
rclasses, rscores, rbboxes = np_methods.bboxes_nms(
rclasses,
rscores,
rbboxes,
nms_threshold=self.nms_threshold)
# Resize bboxes to original image shape. Note: useless for Resize.WARP!
rbboxes = np_methods.bboxes_resize(self.bbox_image, rbboxes)
end_time = time.perf_counter()
rbboxes = np_methods.bboxes_recover(rbboxes, rgb_img)
final_detections_img = draw_box_in_image.draw_boxes_with_label_and_scores(
rgb_img, rbboxes, rclasses, rscores)
final_detections_img = cv.cvtColor(final_detections_img,
cv.COLOR_RGB2BGR)
cv.imwrite(os.path.join(save_path, img_name),
final_detections_img)
# resize boxes and image according to raw input image
# final_detections= cv.resize(final_detections[:, :, ::-1], (raw_w, raw_h))
# recover to raw size
detect_dict['score'] = rscores
detect_dict['boxes'] = rbboxes
detect_dict['categories'] = rclasses
# convert from RGB to BG
fw.write(f'\n{img_name}')
for score, boxes, categories in zip(rscores, rbboxes,
rclasses):
fw.write('\n\tscore:' + str(score))
fw.write('\tbboxes:' + str(boxes))
fw.write('\tcategories:' + str(int(categories)))
view_bar(
'{} image cost {} second'.format(img_name,
(end_time - start_time)),
index + 1, len(image_name_list))
fw.close()
def __init__(self):
self.anchor_per_scale = cfgs.ANCHOR_PER_SCALE
self.classes = tools.read_class_names(cfgs.CLASSES)
self.num_classes = len(self.classes)
self.learn_rate_init = cfgs.LEARNING_RATE_INIT
self.learn_rate_end = cfgs.LEARNING_RATE_END
self.first_stage_epochs = cfgs.FIRST_STAGE_EPOCHS
self.second_stage_epochs = cfgs.SECOND_STAGE_EPOCHS
self.warmup_periods = cfgs.WARMUP_EPOCHS
self.time = time.strftime('%Y-%m-%d-%H-%M-%S',
time.localtime(time.time()))
self.moving_ave_decay = cfgs.MOVING_AVE_DECAY
self.max_bbox_per_scale = 150
self.log_dir = os.path.join(cfgs.SUMMARY_PATH, cfgs.VERSION)
self.train_dataset = Dataset(is_training=True)
self.test_dataset = Dataset(is_training=False)
self.steps_per_period = self.train_dataset.num_steps_per_epoches
self.test_steps_per_period = self.train_dataset.num_steps_per_epoches
self.train_data_batch = self.train_dataset.dataset_tfrecord(
batch_size=cfgs.TRAIN_BATCH_SIZE, is_training=True)
self.test_data_batch = self.test_dataset.dataset_tfrecord(
batch_size=cfgs.TEST_BATCH_SIZE, is_training=False)
self.sess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True,
gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.8,
allow_growth=True)))
with tf.name_scope('define_input'):
self.input_data = tf.placeholder(dtype=tf.float32,
name='input_data')
self.label_sbbox = tf.placeholder(dtype=tf.float32,
name='label_sbbox')
self.label_mbbox = tf.placeholder(dtype=tf.float32,
name='label_mbbox')
self.label_lbbox = tf.placeholder(dtype=tf.float32,
name='label_lbbox')
self.true_sbboxes = tf.placeholder(dtype=tf.float32,
name='sbboxes')
self.true_mbboxes = tf.placeholder(dtype=tf.float32,
name='mbboxes')
self.true_lbboxes = tf.placeholder(dtype=tf.float32,
name='lbboxes')
self.trainable = tf.placeholder(dtype=tf.bool, name='training')
with tf.name_scope("define_loss"):
self.model = YOLOV3(self.input_data, self.trainable)
# self.net_var = tf.global_variables()
# get loader and saver
self.loader, self.checkpoint_path = self.model.get_restorer(
is_training=True)
# self.global_step = tf.train.get_or_create_global_step()
self.giou_loss, self.conf_loss, self.prob_loss = self.model.compute_loss(
self.label_sbbox, self.label_mbbox, self.label_lbbox,
self.true_sbboxes, self.true_mbboxes, self.true_lbboxes)
self.loss = self.giou_loss + self.conf_loss + self.prob_loss
with tf.name_scope('learning_rate'):
self.global_step = tf.Variable(1.0,
dtype=tf.float64,
trainable=False,
name='global_step')
# # self.global_step = self.model.global_step
# # self.global_step = tf.train.get_or_create_global_step()
# warmup_steps = tf.constant(self.warmup_periods * self.steps_per_period,
# dtype=tf.float64, name='warmup_steps')
# train_steps = tf.constant( (self.first_stage_epochs + self.second_stage_epochs)* self.steps_per_period,
# dtype=tf.float64, name='train_steps')
# self.learn_rate = tf.cond(
# pred=self.global_step < warmup_steps,
# true_fn=lambda: self.global_step / warmup_steps * self.learn_rate_init,
# false_fn=lambda: self.learn_rate_end + 0.5 * (self.learn_rate_init - self.learn_rate_end) *
# (1 + tf.cos(
# (self.global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))
# )
warmup_steps = int(self.warmup_periods * self.steps_per_period)
total_train_step = int(
(self.first_stage_epochs + self.second_stage_epochs) *
self.steps_per_period)
self.learning_rate = self.model.cosine_decay_with_warmup(
learning_rate_base=self.learn_rate_init,
learning_rate_end=self.learn_rate_end,
total_decay_steps=total_train_step,
warmup_steps=warmup_steps,
global_step=self.global_step)
global_step_update = tf.assign_add(self.global_step, 1.0)
with tf.name_scope("define_weight_decay"):
moving_ave = tf.train.ExponentialMovingAverage(
self.moving_ave_decay).apply(tf.trainable_variables())
with tf.name_scope("define_first_stage_train"):
self.first_stage_trainable_var_list = []
for var in tf.trainable_variables():
var_name = var.op.name
var_name_mess = str(var_name).split('/')
if var_name_mess[1] in [
'conv_sbbox', 'conv_mbbox', 'conv_lbbox'
]:
self.first_stage_trainable_var_list.append(var)
first_stage_optimizer = tf.compat.v1.train.AdamOptimizer(
self.learning_rate).minimize(
self.loss, var_list=self.first_stage_trainable_var_list)
with tf.control_dependencies(
tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
with tf.control_dependencies(
[first_stage_optimizer, global_step_update]):
with tf.control_dependencies([moving_ave]):
self.train_op_with_frozen_variables = tf.no_op()
with tf.name_scope("define_second_stage_train"):
second_stage_trainable_var_list = tf.trainable_variables()
second_stage_optimizer = tf.compat.v1.train.AdamOptimizer(
self.learning_rate).minimize(
self.loss, var_list=second_stage_trainable_var_list)
with tf.control_dependencies(
tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
with tf.control_dependencies(
[second_stage_optimizer, global_step_update]):
with tf.control_dependencies([moving_ave]):
self.train_op_with_all_variables = tf.no_op()
with tf.name_scope('loader_and_saver'):
# self.loader = tf.train.Saver(self.net_var)
self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=10)
with tf.name_scope('summary'):
tf.summary.scalar("learning_rate", self.learning_rate)
tf.summary.scalar("giou_loss", self.giou_loss)
tf.summary.scalar("conf_loss", self.conf_loss)
tf.summary.scalar("prob_loss", self.prob_loss)
tf.summary.scalar("total_loss", self.loss)
if os.path.exists(self.log_dir): shutil.rmtree(self.log_dir)
tools.makedir(self.log_dir)
self.write_op = tf.summary.merge_all()
self.summary_writer = tf.summary.FileWriter(self.log_dir,
graph=self.sess.graph)
def train(self):
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
self.sess.run(init_op)
try:
print('=> Restoring weights from: {0} ... '.format(
self.checkpoint_path))
self.loader.restore(self.sess, self.checkpoint_path)
except:
print('=> {0}does not exist !!!'.format(self.checkpoint_path))
print('=> Now it starts to train YOLOV3 from scratch ...')
self.first_stage_epochs = 0
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=self.sess, coord=coord)
# ++++++++++++++++++++++++++++++++++++++++start training+++++++++++++++++++++++++++++++++++++++++++++++++++++
try:
if not coord.should_stop():
for epoch in range(
1, 1 + self.first_stage_epochs +
self.second_stage_epochs):
if epoch <= self.first_stage_epochs:
train_op = self.train_op_with_frozen_variables
else:
train_op = self.train_op_with_all_variables
train_epoch_loss, test_epoch_loss = [], []
train_bar = tqdm(range(self.steps_per_period))
for _ in train_bar:
train_data = self.sess.run(self.train_data_batch)
_, summary, train_step_loss, global_step_val = self.sess.run(
[
train_op, self.write_op, self.loss,
self.global_step
],
feed_dict={
self.input_data: train_data[0],
self.label_sbbox: train_data[1],
self.label_mbbox: train_data[2],
self.label_lbbox: train_data[3],
self.true_sbboxes: train_data[4],
self.true_mbboxes: train_data[5],
self.true_lbboxes: train_data[6],
self.trainable: True,
})
train_epoch_loss.append(train_step_loss)
self.summary_writer.add_summary(
summary, global_step_val)
train_bar.set_description(
"train loss: {:.2f}".format(train_step_loss))
for _ in range(self.test_steps_per_period):
test_data = self.sess.run(self.test_data_batch)
test_step_loss = self.sess.run(self.loss,
feed_dict={
self.input_data:
test_data[0],
self.label_sbbox:
test_data[1],
self.label_mbbox:
test_data[2],
self.label_lbbox:
test_data[3],
self.true_sbboxes:
test_data[4],
self.true_mbboxes:
test_data[5],
self.true_lbboxes:
test_data[6],
self.trainable:
False,
})
test_epoch_loss.append(test_step_loss)
train_epoch_loss, test_epoch_loss = np.mean(
train_epoch_loss), np.mean(test_epoch_loss)
save_dir = os.path.join(cfgs.TRAINED_CKPT, cfgs.VERSION)
tools.makedir(save_dir)
ckpt_file = os.path.join(
save_dir,
"yolov3_loss={:.4f}.ckpt".format(test_epoch_loss))
log_time = time.strftime('%Y-%m-%d %H:%M:%S',
time.localtime(time.time()))
print(
"=> Epoch: {0} Time: {1} Train loss: {2:.4f} Test loss: {3:.4f} Saving {4} ..."
.format(epoch, log_time, train_epoch_loss,
test_epoch_loss, ckpt_file))
self.saver.save(self.sess, ckpt_file, global_step=epoch)
except Exception as e:
# Report exceptions to the coordinator.
coord.request_stop(e)
finally:
coord.request_stop()
coord.join(threads)
print('all threads are asked to stop!')