def __init__(self):
self.coco = MSCOCO('minival')
self.net = NetWork()
self.top_k = cfg.top_k
self.ae_threshold = cfg.ae_threshold
self.test_scales = cfg.test_scales
self.weight_exp = cfg.weight_exp
self.merge_bbox = cfg.merge_bbox
self.categories = cfg.categories
self.nms_threshold = cfg.nms_threshold
self.max_per_image = cfg.max_per_image
self.result_dir = cfg.result_dir
def __init__(self):
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
self.net = NetWork()
self.data_train = Image_data('trainval')
self.data_train.inupt_producer()
# self.data_test=Image_data('minival')
# self.data_test.inupt_producer()
self.gpus = [0, 1]
self.batch_i = cfg.batch_size
self.batch_size = self.batch_i * len(self.gpus)
self.save_pre_every = 10
self.num_steps = int(self.save_pre_every * cfg.epoch_num + 1)
self.lr = cfg.learning_rate
self.snapshot_dir = cfg.snapshot_dir
self.snapshot_file = cfg.snapshot_file
self.decay_rate = cfg.decay_rate
self.decay_step = cfg.decay_step
class Test():
def __init__(self):
self.coco = MSCOCO('minival')
self.net = NetWork()
self.top_k = cfg.top_k
self.ae_threshold = cfg.ae_threshold
self.test_scales = cfg.test_scales
self.weight_exp = cfg.weight_exp
self.merge_bbox = cfg.merge_bbox
self.categories = cfg.categories
self.nms_threshold = cfg.nms_threshold
self.max_per_image = cfg.max_per_image
self.result_dir = cfg.result_dir
def test(self, sess):
debug_dir = os.path.join(result_dir, "debug")
if not os.path.exists(debug_dir):
os.makedirs(debug_dir)
img_names = self.coco.get_all_img()
num = len(img_names)
for img_name in tqdm(img_names):
img = self.coco.read_img(img_name)
height, width = img.shape[0:2]
detections = []
for scale in test_scales:
new_height = int(height * scale)
new_width = int(width * scale)
new_center = np.array([new_height // 2, new_width // 2])
inp_height = new_height | 127
inp_width = new_width | 127
images = np.zeros((1, inp_height, inp_width, 3),
dtype=np.float32)
ratios = np.zeros((1, 2), dtype=np.float32)
borders = np.zeros((1, 4), dtype=np.float32)
sizes = np.zeros((1, 2), dtype=np.float32)
out_height, out_width = (inp_height + 1) // 4, (inp_width +
1) // 4
height_ratio = out_height / inp_height
width_ratio = out_width / inp_width
resized_image = cv2.resize(image, (new_width, new_height))
resized_image, border, offset = crop_image(
resized_image, new_center, [inp_height, inp_width])
resized_image = resized_image / 255.
#normalize_(resized_image, db.mean, db.std)
images[0] = resized_image
borders[0] = border
sizes[0] = [int(height * scale), int(width * scale)]
ratios[0] = [height_ratio, width_ratio]
images = np.concatenate((images, images[:, :, ::-1, :]),
axis=0)
images = tf.convert_to_tensor(images)
is_training = tf.convert_to_tensor(False)
outs = self.net.corner_net(images, is_training=is_training)
dets_tensor = self.net.decode(*outs[-6:])
dets = sess.run(dets_tensor)
dets = dets.reshape(2, -1, 8)
dets[1, :, [0, 2]] = out_width - dets[1, :, [2, 0]]
dets = dets.reshape(1, -1, 8)
dets = rescale_dets(dets, ratios, borders, sizes)
dets[:, :, 0:4] /= scale
detections.append(dets)
detections = np.concatenate(detections, axis=1)
classes = detections[..., -1]
classes = classes[0]
detections = detections[0]
# reject detections with negative scores
keep_inds = (detections[:, 4] > -1)
detections = detections[keep_inds]
classes = classes[keep_inds]
top_bboxes[image_id] = {}
for j in range(categories):
keep_inds = (classes == j)
top_bboxes[image_id][j +
1] = detections[keep_inds][:, 0:7].astype(
np.float32)
if merge_bbox:
top_bboxes[image_id][j + 1] = soft_nms_merge(
top_bboxes[image_id][j + 1],
Nt=nms_threshold,
method=2,
weight_exp=weight_exp)
else:
top_bboxes[image_id][j + 1] = soft_nms(
top_bboxes[image_id][j + 1],
Nt=nms_threshold,
method=nms_algorithm)
top_bboxes[image_id][j + 1] = top_bboxes[image_id][j + 1][:,
0:5]
scores = np.hstack([
top_bboxes[image_id][j][:, -1]
for j in range(1, categories + 1)
])
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, categories + 1):
keep_inds = (top_bboxes[image_id][j][:, -1] >= thresh)
top_bboxes[image_id][j] = top_bboxes[image_id][j][
keep_inds]
if debug:
image = self.coco.read_img(img_name)
bboxes = {}
for j in range(1, categories + 1):
keep_inds = (top_bboxes[image_id][j][:, -1] > 0.5)
cat_name = self.coco.class_name(j)
cat_size = cv2.getTextSize(cat_name,
cv2.FONT_HERSHEY_SIMPLEX, 0.5,
2)[0]
color = np.random.random((3, )) * 0.6 + 0.4
color = color * 255
color = color.astype(np.int32).tolist()
for bbox in top_bboxes[image_id][j][keep_inds]:
bbox = bbox[0:4].astype(np.int32)
if bbox[1] - cat_size[1] - 2 < 0:
cv2.rectangle(image, (bbox[0], bbox[1] + 2),
(bbox[0] + cat_size[0],
bbox[1] + cat_size[1] + 2), color,
-1)
cv2.putText(image,
cat_name,
(bbox[0], bbox[1] + cat_size[1] + 2),
cv2.FONT_HERSHEY_SIMPLEX,
0.5, (0, 0, 0),
thickness=1)
else:
cv2.rectangle(image,
(bbox[0], bbox[1] - cat_size[1] - 2),
(bbox[0] + cat_size[0], bbox[1] - 2),
color, -1)
cv2.putText(image,
cat_name, (bbox[0], bbox[1] - 2),
cv2.FONT_HERSHEY_SIMPLEX,
0.5, (0, 0, 0),
thickness=1)
cv2.rectangle(image, (bbox[0], bbox[1]),
(bbox[2], bbox[3]), color, 2)
debug_file = os.path.join(debug_dir, {}.format(img_name))
# result_json = os.path.join(result_dir, "results.json")
# detections = db.convert_to_coco(top_bboxes)
# with open(result_json, "w") as f:
# json.dump(detections, f)
# cls_ids = list(range(1, categories + 1))
# image_ids = [db.image_ids(ind) for ind in db_inds]
# db.evaluate(result_json, cls_ids, image_ids)
return 0
class Test():
def __init__(self):
self.coco=MSCOCO('minival')
self.net=NetWork()
self.top_k=cfg.top_k
self.ae_threshold=cfg.ae_threshold
self.test_scales=cfg.test_scales
self.weight_exp=cfg.weight_exp
self.merge_bbox=cfg.merge_bbox
self.categories=cfg.categories
self.nms_threshold=cfg.nms_threshold
self.max_per_image=cfg.max_per_image
self.result_dir=cfg.result_dir
self.data_train=Image_data('trainval')
self.data_train.inupt_producer()
def test(self,sess):
self.batch_i=cfg.batch_size
images, tags_tl, tags_br,heatmaps_tl, heatmaps_br, tags_mask, offsets_tl, offsets_br,boxes,ratio=self.data_train.get_batch_data(self.batch_i)
outs,test_outs=self.net.corner_net(images,tags_tl,tags_br,is_training=False)
dets_tensor =self.net.decode(*test_outs)
dets=sess.run(dets_tensor)
dets = dets.reshape(2, -1, 8)
dets[1, :, [0, 2]] = out_width - dets[1, :, [2, 0]]
dets = dets.reshape(1, -1, 8)
print(dets)
return
result_dir = os.path.join(self.result_dir, "result")
if not os.path.exists(result_dir):
os.makedirs(result_dir)
img_names=self.coco.get_all_img()
num=len(img_names)
for img_name in tqdm(img_names):
img=self.coco.read_img(img_name)
height, width = img.shape[0:2]
detections=[]
for scale in self.test_scales:
new_height = int(height * scale)
new_width = int(width * scale)
new_center = np.array([new_height // 2, new_width // 2])
inp_height = new_height | 127
inp_width = new_width | 127
images = np.zeros((1, inp_height, inp_width, 3), dtype=np.float32)
ratios = np.zeros((1, 2), dtype=np.float32)
borders = np.zeros((1, 4), dtype=np.float32)
sizes = np.zeros((1, 2), dtype=np.float32)
out_height, out_width = (inp_height + 1) // 4, (inp_width + 1) // 4
height_ratio = out_height / inp_height
width_ratio = out_width / inp_width
resized_image = cv2.resize(images[0], (new_width, new_height))
resized_image, border, offset = crop_image(resized_image, new_center, [inp_height, inp_width])
resized_image = resized_image / 255.
#normalize_(resized_image, db.mean, db.std)
images[0] = resized_image
borders[0] = border
sizes[0] = [int(height * scale), int(width * scale)]
ratios[0] = [height_ratio, width_ratio]
print(images.shape)
images = np.concatenate((images, images[:, :, ::-1, :]), axis=0)
images = tf.convert_to_tensor(images)
is_training=tf.convert_to_tensor(False)
outs=self.net.corner_net(images,gt_tag_tl=None,is_training=is_training)
dets_tensor=self.net.decode(*outs[-6:])
dets=sess.run(dets_tensor)
dets = dets.reshape(2, -1, 8)
dets[1, :, [0, 2]] = out_width - dets[1, :, [2, 0]]
dets = dets.reshape(1, -1, 8)
dets=rescale_dets(dets, ratios, borders, sizes)
dets[:, :, 0:4] /= scale
detections.append(dets)
detections = np.concatenate(detections, axis=1)
classes = detections[..., -1]
classes = classes[0]
detections = detections[0]
# reject detections with negative scores
keep_inds = (detections[:, 4] > -1)
detections = detections[keep_inds]
classes = classes[keep_inds]
top_bboxes[image_id] = {}
for j in range(categories):
keep_inds = (classes == j)
top_bboxes[image_id][j + 1] = detections[keep_inds][:, 0:7].astype(np.float32)
if merge_bbox:
top_bboxes[image_id][j + 1]=soft_nms_merge(top_bboxes[image_id][j + 1], Nt=nms_threshold, method=2, weight_exp=weight_exp)
else:
top_bboxes[image_id][j + 1]=soft_nms(top_bboxes[image_id][j + 1], Nt=nms_threshold, method=nms_algorithm)
top_bboxes[image_id][j + 1] = top_bboxes[image_id][j + 1][:, 0:5]
scores = np.hstack([
top_bboxes[image_id][j][:, -1]
for j in range(1, categories + 1)
])
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, categories + 1):
keep_inds = (top_bboxes[image_id][j][:, -1] >= thresh)
top_bboxes[image_id][j] = top_bboxes[image_id][j][keep_inds]
result_json = os.path.join(result_dir, "results.json")
detections = db.convert_to_coco(top_bboxes)
with open(result_json, "w") as f:
json.dump(detections, f)
cls_ids = list(range(1, categories + 1))
image_ids = [db.image_ids(ind) for ind in db_inds]
db.evaluate(result_json, cls_ids, image_ids)
return 0
class Train():
def __init__(self):
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
self.net = NetWork()
self.data_train = Image_data('trainval')
self.data_train.inupt_producer()
# self.data_test=Image_data('minival')
# self.data_test.inupt_producer()
self.gpus = [0, 1]
self.batch_i = cfg.batch_size
self.batch_size = self.batch_i * len(self.gpus)
self.save_pre_every = int(
self.data_train.num_image / self.batch_size) + 1
self.num_steps = int(self.save_pre_every * cfg.epoch_num + 1)
self.lr = cfg.learning_rate
self.snapshot_dir = cfg.snapshot_dir
self.snapshot_file = cfg.snapshot_file
self.decay_rate = cfg.decay_rate
self.decay_step = cfg.decay_step
def train_mult(self):
coord = tf.train.Coordinator()
images, tags_tl, tags_br, heatmaps_tl, heatmaps_br, tags_mask, offsets_tl, offsets_br, boxes, ratio = self.data_train.get_batch_data(
self.batch_size)
tower_grads = []
steps = tf.Variable(0, name='global_step', trainable=False)
lr = tf.train.exponential_decay(self.lr,
steps,
self.decay_step,
self.decay_rate,
staircase=True,
name='learning_rate')
optim = tf.train.AdamOptimizer(learning_rate=lr)
#optim= tf.train.MomentumOptimizer(0.000025,0.9)
reuse1 = False
for i in range(len(self.gpus)):
with tf.device('/gpu:%d' % i):
with tf.name_scope('Tower_%d' % (i)) as scope:
if i == 0:
reuse1 = False
else:
reuse1 = True
next_imgs = images[i * self.batch_i:(i + 1) * self.batch_i]
next_tags_tl = tags_tl[i * self.batch_i:(i + 1) *
self.batch_i]
next_tags_br = tags_br[i * self.batch_i:(i + 1) *
self.batch_i]
next_heatmaps_tl = heatmaps_tl[i * self.batch_i:(i + 1) *
self.batch_i]
next_heatmaps_br = heatmaps_br[i * self.batch_i:(i + 1) *
self.batch_i]
next_tags_mask = tags_mask[i * self.batch_i:(i + 1) *
self.batch_i]
next_offsets_tl = offsets_tl[i * self.batch_i:(i + 1) *
self.batch_i]
next_offsets_br = offsets_br[i * self.batch_i:(i + 1) *
self.batch_i]
with tf.variable_scope('', reuse=reuse1):
outs, test_outs = self.net.corner_net(next_imgs,
next_tags_tl,
next_tags_br,
is_training=True)
dets_tensor, debug_boxes = self.net.decode(*test_outs)
loss, focal_loss, pull_loss, push_loss, offset_loss = self.net.loss(
outs, [
next_heatmaps_tl, next_heatmaps_br, next_tags_mask,
next_offsets_tl, next_offsets_br
])
trainable_variable = tf.trainable_variables()
grads = optim.compute_gradients(
loss, var_list=trainable_variable)
tower_grads.append(grads)
grads_ave = self.average_gradients(tower_grads)
update = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update):
train_op = optim.apply_gradients(grads_ave, steps)
saver = tf.train.Saver(max_to_keep=100)
#loader = tf.train.Saver()
config = tf.ConfigProto(allow_soft_placement=True,
log_device_placement=False)
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
init = tf.global_variables_initializer()
sess.run(init)
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
print(self.num_steps)
debug = Debug()
epoch = 0
if self.load(saver, sess, self.snapshot_dir):
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
for step in range(self.num_steps):
start = time.time()
#sess.run(update)
_, loss_, focal_loss_, pull_loss_, push_loss_, offset_loss_, lr_ = sess.run(
[
train_op, loss, focal_loss, pull_loss, push_loss,
offset_loss, lr
])
duration = time.time() - start
print(
'step %d, loss %g, focal_loss %g, pull_loss %g, push_loss %g, offset_loss %g, time %g, lr %g'
% (step, loss_, focal_loss_, pull_loss_, push_loss_,
offset_loss_, duration, lr_))
if step % 100 == 0:
dets_, images_, debug_boxes_, boxes_, ratio_ = sess.run(
[dets_tensor, images, debug_boxes, boxes, ratio])
debug.test_debug(images_[0], dets_[0], debug_boxes_[0],
boxes_[0], ratio_[0], self.data_train.coco,
step)
if step % self.save_pre_every == 0 and step > 0:
saver.save(sess, self.snapshot_file, epoch)
epoch += 1
coord.request_stop()
coord.join(threads)
def average_gradients(self, tower_grads):
average_grads = []
for grad_and_vars in zip(*tower_grads):
# Note that each grad_and_vars looks like the following:
# ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
grads = []
for g, _ in grad_and_vars:
# Add 0 dimension to the gradients to represent the tower.
expanded_g = tf.expand_dims(g, 0)
# Append on a 'tower' dimension which we will average over below.
grads.append(expanded_g)
# Average over the 'tower' dimension.
grad = tf.concat(axis=0, values=grads)
grad = tf.reduce_mean(grad, 0)
# Keep in mind that the Variables are redundant because they are shared
# across towers. So .. we will just return the first tower's pointer to
# the Variable.
v = grad_and_vars[0][1]
grad_and_var = (grad, v)
average_grads.append(grad_and_var)
return average_grads
def load(self, saver, sess, ckpt_path):
'''Load trained weights.
Args:
saver: TensorFlow saver object.
sess: TensorFlow session.
ckpt_path: path to checkpoint file with parameters.
'''
ckpt = tf.train.get_checkpoint_state(ckpt_path)
if ckpt and ckpt.model_checkpoint_path:
ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
saver.restore(sess, os.path.join(ckpt_path, ckpt_name))
print("Restored model parameters from {}".format(ckpt_name))
return True
else:
return False
def train_single(self):
#with tf.variable_scope('inputs'):
images, tags_tl, tags_br, heatmaps_tl, heatmaps_br, tags_mask, offsets_tl, offsets_br, boxes, ratio = self.data_train.get_batch_data(
self.batch_i)
#test_images, test_tags_tl, test_tags_br,test_heatmaps_tl, test_heatmaps_br, test_tags_mask, test_offsets_tl, test_offsets_br,test_boxes=self.data_test.get_batch_data(self.batch_size)
#with tf.variable_scope('net'):
#is_training=tf.constant(True)
outs, test_outs = self.net.corner_net(images,
tags_tl,
tags_br,
is_training=True)
dets_tensor, debug_boxes = self.net.decode(*test_outs)
#outs_test=self.net.corner_net(test_images,test_tags_tl,test_tags_br,is_training=False)
loss, focal_loss, pull_loss, push_loss, offset_loss = self.net.loss(
outs,
[heatmaps_tl, heatmaps_br, tags_mask, offsets_tl, offsets_br])
#with tf.variable_scope('train_op'):
steps = tf.Variable(0, name='global_step', trainable=False)
lr = tf.train.exponential_decay(self.lr,
steps,
self.decay_step,
self.decay_rate,
staircase=True,
name='learning_rate')
update = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
#with tf.control_dependencies(update):
train_op = tf.train.AdamOptimizer(learning_rate=lr).minimize(
loss, steps)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.InteractiveSession(config=config)
init = tf.global_variables_initializer()
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
saver = tf.train.Saver(max_to_keep=100)
sess.run(init)
print(self.num_steps)
debug = Debug()
epoch = 5
if self.load(saver, sess, self.snapshot_dir):
print(" [*] Load SUCCESS")
else:
print(" [!] Load failed...")
for step in range(self.num_steps):
start = time.time()
sess.run(update)
_, loss_, focal_loss_, pull_loss_, push_loss_, offset_loss_, lr_ = sess.run(
[
train_op, loss, focal_loss, pull_loss, push_loss,
offset_loss, lr
])
duration = time.time() - start
print(
'step %d, loss %g, focal_loss %g, pull_loss %g, push_loss %g, offset_loss %g, time %g, lr %g'
% (step, loss_, focal_loss_, pull_loss_, push_loss_,
offset_loss_, duration, lr_))
if step % 100 == 0:
dets_, images_, debug_boxes_, boxes_, ratio_ = sess.run(
[dets_tensor, images, debug_boxes, boxes, ratio])
debug.test_debug(images_[0], dets_[0], debug_boxes_[0],
boxes_[0], ratio_[0], self.data_train.coco,
step)
if step % self.save_pre_every == 0 and step > 0:
saver.save(sess, self.snapshot_file, epoch)
epoch += 1
coord.request_stop()
coord.join(threads)