def tagging(self, text, sen_words, id2tag, len_sen):
if text:
max_len = self.cfg.getint('get_pkl', 'max_len')
y_pred = []
tags = []
text_len = len(text)
if text_len == 1:
X_batch = [text[0]]
else:
X_batch = np.squeeze(text, axis=(1, ))
fetches = [self.y_pred]
feed_dict = {
self.X_inputs: X_batch,
self._model.lr: 0.0,
self._model.batch_size: len(text),
self._model.keep_prob: 1.0
}
_y_pred = self.session.run(fetches, feed_dict)
_y_pred = np.squeeze(_y_pred, axis=(0, ))
for i in range(text_len):
tags.append(
decode(self.cfg,
_y_pred[i * max_len:i * max_len + len_sen[i]],
id2tag))
return zip(sen_words, tags)
def __init__(self):
self.input_size = cfg.input_image_h
self.classes = utils.read_class_names(cfg.classes_file)
self.num_classes = len(self.classes)
self.score_threshold = cfg.score_threshold
self.iou_threshold = cfg.nms_thresh
self.moving_ave_decay = cfg.moving_ave_decay
self.annotation_path = cfg.test_data_file
self.weight_file = cfg.weight_file
self.write_image = cfg.write_image
self.write_image_path = cfg.write_image_path
self.show_label = cfg.show_label
self.input_data = tf.placeholder(
shape=[1, cfg.input_image_h, cfg.input_image_w, 3],
dtype=tf.float32)
self.sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
self.sess.run(tf.global_variables_initializer())
model = CenterNet(self.input_data, False)
saver = tf.train.Saver()
saver.restore(
self.sess,
'./checkpoint/2021_02_24-centernet_test_loss=0.5797.ckpt-80')
self.hm = model.pred_hm
self.wh = model.pred_wh
self.reg = model.pred_reg
self.det = decode(self.hm, self.wh, self.reg, K=cfg.max_objs)
def _decode(self, images):
outputs = self.forward(images).detach()
offset = outputs[:, :2].tanh().detach()
wh = outputs[:, 2:4].exp().detach()
kpt = outputs[:, 4:].sigmoid().detach()
bboxes, categs, scores = decode(kpt, wh, offset, 100)
bboxes *= self.net.stride
return bboxes, categs, scores
from utils.image import get_affine_transform, affine_transform
from utils.utils import image_preporcess, py_nms, post_process, bboxes_draw_on_img, read_class_names
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
ckpt_path = './checkpoint/'
sess = tf.Session()
inputs = tf.placeholder(shape=[None, None, None, 3], dtype=tf.float32)
model = CenterNet(inputs, False, net_name=cfg.net)
saver = tf.train.Saver()
saver.restore(sess, tf.train.latest_checkpoint(ckpt_path))
hm = model.pred_hm
wh = model.pred_wh
reg = model.pred_reg
det = decode(hm, wh, reg, K=cfg.max_objs)
class_names = read_class_names(cfg.classes_file)
img_names = os.listdir(
'/home/pcl/tf_work/TF_CenterNet/VOC/test/VOCdevkit/VOC2007/JPEGImages')
for img_name in img_names:
img_path = '/home/pcl/tf_work/TF_CenterNet/VOC/test/VOCdevkit/VOC2007/JPEGImages/' + img_name
print(img_path)
original_image = cv2.imread(img_path)
original_image_size = original_image.shape[:2]
image_data = image_preporcess(np.copy(original_image),
[cfg.input_image_h, cfg.input_image_w])
image_data = image_data[np.newaxis, ...]
t0 = time.time()
detections = sess.run(det, feed_dict={inputs: image_data})
def train():
# define dataset
num_train_imgs = len(open(cfg.train_data_file, 'r').readlines())
num_train_batch = int(math.ceil(float(num_train_imgs) / cfg.batch_size))
num_test_imgs = len(open(cfg.test_data_file, 'r').readlines())
num_test_batch = int(math.ceil(float(num_test_imgs) / 1))
train_dataset = tf.data.TextLineDataset(cfg.train_data_file)
train_dataset = train_dataset.shuffle(num_train_imgs)
train_dataset = train_dataset.batch(cfg.batch_size)
train_dataset = train_dataset.map(lambda x: tf.py_func(get_data, inp=[x, True],
Tout=[tf.float32, tf.float32, tf.float32, tf.float32,
tf.float32, tf.float32, tf.int32, tf.int32]),
num_parallel_calls=6)
train_dataset = train_dataset.prefetch(3)
test_dataset = tf.data.TextLineDataset(cfg.test_data_file)
test_dataset = test_dataset.batch(1)
test_dataset = test_dataset.map(lambda x: tf.py_func(get_data, inp=[x, False],
Tout=[tf.float32, tf.float32, tf.float32, tf.float32,
tf.float32, tf.float32, tf.int32, tf.int32]),
num_parallel_calls=1)
test_dataset = test_dataset.prefetch(1)
iterator = tf.data.Iterator.from_structure(train_dataset.output_types, train_dataset.output_shapes)
trainset_init_op = iterator.make_initializer(train_dataset)
testset_init_op = iterator.make_initializer(test_dataset)
input_data, batch_hm, batch_wh, batch_reg, batch_reg_mask, batch_ind, batch_img_size, batch_id = iterator.get_next()
input_data.set_shape([None, None, None, 3])
batch_hm.set_shape([None, None, None, None])
batch_wh.set_shape([None, None, None])
batch_reg.set_shape([None, None, None])
batch_reg_mask.set_shape([None, None])
batch_ind.set_shape([None, None])
batch_img_size.set_shape([None, None])
batch_id.set_shape([None])
# training flag
is_training = tf.placeholder(dtype=tf.bool, name='is_training')
# difine model and loss
model = CenterNet(input_data, is_training, "dla_34")
hm = model.pred_hm
wh = model.pred_wh
reg = model.pred_reg
from utils.decode import decode
det = decode(hm, wh, reg, K=cfg.max_objs)
with tf.variable_scope('loss'):
# hm_loss, wh_loss, reg_loss = model.compute_loss(batch_hm, batch_wh, batch_reg, batch_reg_mask, batch_ind)
hm_loss, wh_loss, reg_loss = model.compute_loss_pcl(batch_hm, batch_wh, batch_reg, batch_reg_mask, batch_ind)
total_loss = hm_loss + wh_loss + reg_loss
# define train op
if cfg.lr_type == "CosineAnnealing":
learning_rate = 0.0001
global_step = tf.Variable(1.0, dtype=tf.float64, trainable=False, name='global_step')
# warmup_steps = tf.constant(cfg.warm_up_epochs * num_train_batch, dtype=tf.float64, name='warmup_steps')
# train_steps = tf.constant(cfg.epochs * num_train_batch, dtype=tf.float64, name='train_steps')
# learning_rate = tf.cond(
# pred=global_step < warmup_steps,
# true_fn=lambda: global_step / warmup_steps * cfg.init_lr,
# false_fn=lambda: cfg.end_lr + 0.5 * (cfg.init_lr - cfg.end_lr) *
# (1 + tf.cos(
# (global_step - warmup_steps) / (train_steps - warmup_steps) * np.pi))
# )
global_step_update = tf.assign_add(global_step, 1.0)
optimizer = tf.train.AdamOptimizer(learning_rate).minimize(total_loss)
with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
with tf.control_dependencies([optimizer, global_step_update]):
train_op = tf.no_op()
else:
global_step = tf.Variable(0, trainable=False)
if cfg.lr_type == "exponential":
learning_rate = 0.0001
# learning_rate = tf.train.exponential_decay(cfg.lr,
# global_step,
# cfg.lr_decay_steps,
# cfg.lr_decay_rate,
# staircase=True)
elif cfg.lr_type == "piecewise":
learning_rate = 0.0001
# learning_rate = tf.train.piecewise_constant(global_step, cfg.lr_boundaries, cfg.lr_piecewise)
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)
saver = tf.train.Saver(tf.global_variables(), max_to_keep=10)
with tf.Session() as sess:
with tf.name_scope('summary'):
tf.summary.scalar("learning_rate", learning_rate)
tf.summary.scalar("hm_loss", hm_loss)
tf.summary.scalar("wh_loss", wh_loss)
tf.summary.scalar("reg_loss", reg_loss)
tf.summary.scalar("total_loss", total_loss)
logdir = "./log_dla/"
if os.path.exists(logdir): shutil.rmtree(logdir)
os.mkdir(logdir)
write_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(logdir, graph=sess.graph)
# train
sess.run(tf.global_variables_initializer())
# if cfg.dla_pretrain:
# saver.restore(sess, './checkpoint/centernet_test_loss=3.1386.ckpt-79')
for epoch in range(1, 1 + cfg.epochs):
pbar = tqdm(range(num_train_batch))
train_epoch_loss, test_epoch_loss = [], []
sess.run(trainset_init_op)
for i in pbar:
_, summary, train_step_loss, global_step_val, _hm_loss, _wh_loss, _reg_loss = sess.run(
[train_op, write_op, total_loss, global_step, hm_loss, wh_loss, reg_loss],
feed_dict={is_training: True})
train_epoch_loss.append(train_step_loss)
summary_writer.add_summary(summary, global_step_val)
pbar.set_description("train loss: %.2f" % train_step_loss)
if i % 20 == 0:
print("train loss: %.2f hm_loss: %.2f wh_loss:%2f reg_loss:%2f learning_rate:%f" % (train_step_loss, _hm_loss, _wh_loss, _reg_loss, learning_rate) )
print("begining test")
sess.run(testset_init_op)
val_preds = []
# for j in range(num_test_batch ):
# test_step_loss = sess.run(total_loss, feed_dict={is_training: False})
# test_epoch_loss.append(test_step_loss)
train_epoch_loss = np.mean(train_epoch_loss)
# train_epoch_loss, test_epoch_loss = np.mean(train_epoch_loss), np.mean(test_epoch_loss)
ckpt_file = "./checkpoint/centernet_train_loss=%.4f.ckpt" % train_epoch_loss
log_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
print("=> Epoch: %2d Time: %s Train loss: %.2f Saving %s ..."
% (epoch, log_time, train_epoch_loss, ckpt_file))
saver.save(sess, ckpt_file, global_step=epoch)
if epoch % cfg.eval_epoch == 0 and epoch > 0:
sess.run(testset_init_op)
for j in range(num_test_batch):
detections, _batch_img_size, _batch_id = sess.run([det, batch_img_size, batch_id], feed_dict={is_training: False})
ori_h = _batch_img_size[0][1]
ori_w = _batch_img_size[0][0]
detect_post = post_process(detections, (ori_h, ori_w), [cfg.input_image_h, cfg.input_image_w],
cfg.down_ratio, cfg.score_threshold)
id = _batch_id[0]
detect_per_img = get_preds_gpu(detect_post, id)
val_preds.extend(detect_per_img)
rec_total, prec_total, ap_total = AverageMeter(), AverageMeter(), AverageMeter()
info = ""
gt_dict = parse_gt_rec(cfg.test_data_file, [cfg.input_image_h, cfg.input_image_w], cfg.letterbox_resize)
for ii in range(cfg.num_classes):
from utils.utils import voc_eval
npos, nd, rec, prec, ap = voc_eval(gt_dict, val_preds, ii, iou_thres=cfg.score_threshold,
use_07_metric=cfg.use_voc_07_metric)
info += 'EVAL: Class {}: Recall: {:.4f}, Precision: {:.4f}, AP: {:.4f}\n'.format(ii, rec, prec, ap)
rec_total.update(rec, npos)
prec_total.update(prec, nd)
ap_total.update(ap, 1)
mAP = ap_total.average
info += 'EVAL: Recall: {:.4f}, Precison: {:.4f}, mAP: {:.4f}\n'.format(rec_total.average,
prec_total.average, mAP)
print(info)
def train():
# define dataset
num_train_imgs = len(open(cfg.train_data_file, 'r').readlines())
num_train_batch = int(math.ceil(float(num_train_imgs) / cfg.batch_size))
num_test_imgs = len(open(cfg.test_data_file, 'r').readlines())
num_test_batch = int(math.ceil(float(num_test_imgs) / 1))
train_dataset = tf.data.TextLineDataset(cfg.train_data_file)
train_dataset = train_dataset.shuffle(num_train_imgs)
train_dataset = train_dataset.batch(cfg.batch_size)
train_dataset = train_dataset.map(lambda x: tf.py_func(
get_data,
inp=[x, True],
Tout=[
tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.
float32, tf.int32, tf.int32
]),
num_parallel_calls=6)
train_dataset = train_dataset.prefetch(3)
test_dataset = tf.data.TextLineDataset(cfg.test_data_file)
test_dataset = test_dataset.batch(1)
test_dataset = test_dataset.map(lambda x: tf.py_func(
get_data,
inp=[x, False],
Tout=[
tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.
float32, tf.int32, tf.int32
]),
num_parallel_calls=1)
test_dataset = test_dataset.prefetch(1)
iterator = tf.data.Iterator.from_structure(train_dataset.output_types,
train_dataset.output_shapes)
trainset_init_op = iterator.make_initializer(train_dataset)
testset_init_op = iterator.make_initializer(test_dataset)
input_data, hm, wh, reg, reg_mask, ind, img_size, id = iterator.get_next()
batch_input_data = []
batch_hm = []
batch_wh = []
batch_reg = []
batch_reg_mask = []
batch_ind = []
batch_img_size = []
batch_id = []
for i in range(cfg.NUM_GPU):
start = i * (cfg.batch_size // cfg.NUM_GPU)
end = (i + 1) * (cfg.batch_size // cfg.NUM_GPU)
single_input_data = input_data[start:end, :, :, :]
single_hm = hm[start:end, :, :, :]
single_wh = wh[start:end, :, :]
single_reg = reg[start:end, :, :]
single_reg_mask = reg_mask[start:end, :]
single_ind = ind[start:end, :]
single_img_size = img_size[start:end:, :]
single_id = id[start:end, :]
batch_input_data.append(single_input_data)
batch_hm.append(single_hm)
batch_wh.append(single_wh)
batch_reg.append(single_reg)
batch_reg_mask.append(single_reg_mask)
batch_ind.append(single_ind)
batch_img_size.append(single_img_size)
batch_id.append(single_id)
batch_input_data[i].set_shape([None, None, None, 3])
batch_hm[i].set_shape([None, None, None, None])
batch_wh[i].set_shape([None, None, None])
batch_reg[i].set_shape([None, None, None])
batch_reg_mask[i].set_shape([None, None])
batch_ind[i].set_shape([None, None])
batch_img_size[i].set_shape([None, None])
batch_id[i].set_shape([None])
# difine model and loss
with tf.device('/cpu:0'):
tower_grads = []
hm_loss = []
wh_loss = []
reg_loss = []
total_loss = []
hm_pred_list = []
wh_pred_list = []
reg_pred_list = []
pred_det = []
# training flag
is_training = tf.placeholder(dtype=tf.bool, name='is_training')
# with tf.variable_scope(tf.get_variable_scope()):
with tf.variable_scope(tf.get_variable_scope(),
reuse=tf.AUTO_REUSE) as scope:
for i in range(cfg.NUM_GPU):
print("current gpu is", i)
with tf.device('/gpu:%d' % i):
model = CenterNet(batch_input_data[i], is_training,
"dla_34")
hm_pred = model.pred_hm
wh_pred = model.pred_wh
reg_pred = model.pred_reg
hm_pred_list.append(hm_pred)
wh_pred_list.append(wh_pred)
reg_pred_list.append(reg_pred)
det = decode(hm_pred_list[i],
wh_pred_list[i],
reg_pred_list[i],
K=cfg.max_objs)
pred_det.append(det)
with tf.variable_scope('loss'):
l2_loss = tf.losses.get_regularization_loss()
# hm_loss[i], wh_loss[i], reg_loss[i] = model.compute_loss(batch_hm[i], batch_wh[i], batch_reg[i], batch_reg_mask[i], batch_ind[i])
hm_loss_single, wh_loss_single, reg_loss_single = model.compute_loss(
batch_hm[i], batch_wh[i], batch_reg[i],
batch_reg_mask[i], batch_ind[i])
hm_loss.append(hm_loss_single)
wh_loss.append(wh_loss_single)
reg_loss.append(reg_loss_single)
total_loss_single = hm_loss[i] + wh_loss[i] + reg_loss[
i] + l2_loss
total_loss.append(total_loss_single)
# define train op
global_step = tf.Variable(0, trainable=False)
if cfg.lr_type == "exponential":
learning_rate = tf.train.exponential_decay(
cfg.lr,
global_step,
cfg.lr_decay_steps,
cfg.lr_decay_rate,
staircase=True)
elif cfg.lr_type == "piecewise":
learning_rate = tf.train.piecewise_constant(
global_step, cfg.lr_boundaries, cfg.lr_piecewise)
optimizer = tf.train.AdamOptimizer(learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
grads = optimizer.compute_gradients(total_loss[i])
clip_grad_var = [
gv if gv[0] is None else
[tf.clip_by_norm(gv[0], 100.), gv[1]]
for gv in grads
]
tower_grads.append(clip_grad_var)
last_loss = tf.reduce_mean(total_loss)
if len(tower_grads) > 1:
clip_grad_var = sum_gradients(tower_grads)
else:
clip_grad_var = tower_grads[0]
train_op = optimizer.apply_gradients(clip_grad_var,
global_step=global_step)
saver = tf.train.Saver(tf.global_variables(), max_to_keep=10)
config = tf.ConfigProto(allow_soft_placement=True)
with tf.Session(config=config) as sess:
with tf.name_scope('summary'):
tf.summary.scalar("learning_rate", learning_rate)
tf.summary.scalar("hm_loss", tf.reduce_mean(hm_loss))
tf.summary.scalar("wh_loss", tf.reduce_mean(wh_loss))
tf.summary.scalar("reg_loss", tf.reduce_mean(reg_loss))
tf.summary.scalar("total_loss", tf.reduce_mean(total_loss))
logdir = "./log_dla/"
if os.path.exists(logdir): shutil.rmtree(logdir)
os.mkdir(logdir)
write_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(logdir,
graph=sess.graph)
sess.run(tf.global_variables_initializer())
if cfg.pre_train:
saver.restore(
sess,
'./checkpoint/centernet_train_epoch_loss=313.7357.ckpt-6'
)
for epoch in range(1, 1 + cfg.epochs):
pbar = tqdm(range(num_train_batch))
train_epoch_loss, test_epoch_loss = [], []
sess.run(trainset_init_op)
for i in pbar:
_, summary, train_step_loss, global_step_val = sess.run(
[train_op, write_op, last_loss, global_step],
feed_dict={is_training: True})
train_epoch_loss.append(train_step_loss)
summary_writer.add_summary(summary, global_step_val)
pbar.set_description("train loss: %.2f" %
train_step_loss)
# sess.run(testset_init_op)
# for j in range(num_test_batch ):
# test_step_loss = sess.run(last_loss, feed_dict={is_training: False})
# test_epoch_loss.append(test_step_loss)
# train_epoch_loss, test_epoch_loss = np.mean(train_epoch_loss), np.mean(test_epoch_loss)
train_epoch_loss = np.mean(train_epoch_loss)
ckpt_file = "./checkpoint/centernet_train_epoch_loss=%.4f.ckpt" % train_epoch_loss
log_time = time.strftime('%Y-%m-%d %H:%M:%S',
time.localtime(time.time()))
print(
"=> Epoch: %2d Time: %s Train loss: %.2f Saving %s ..."
% (epoch, log_time, train_epoch_loss, ckpt_file))
saver.save(sess, ckpt_file, global_step=epoch)
if epoch % cfg.eval_epoch == 0 and epoch > 0:
print("begining test")
sess.run(testset_init_op)
val_preds = []
for j in tqdm(range(num_test_batch)):
detections, _batch_img_size, _batch_id = sess.run(
[pred_det[0], img_size, id],
feed_dict={is_training: False})
# print("detecttiion is", detections)
# print("_batch_img_size is", _batch_img_size)
# print("id is", _batch_id)
ori_h = _batch_img_size[0][1]
ori_w = _batch_img_size[0][0]
detect_post = post_process(
detections, (ori_h, ori_w),
[cfg.input_image_h, cfg.input_image_w],
cfg.down_ratio, cfg.score_threshold)
id_test = _batch_id[0]
detect_per_img = get_preds_gpu(
detect_post, id_test)
val_preds.extend(detect_per_img)
rec_total, prec_total, ap_total = AverageMeter(
), AverageMeter(), AverageMeter()
info = ""
gt_dict = parse_gt_rec(
cfg.test_data_file,
[cfg.input_image_h, cfg.input_image_w],
cfg.letterbox_resize)
for ii in range(cfg.num_classes):
from utils.utils import voc_eval
npos, nd, rec, prec, ap = voc_eval(
gt_dict,
val_preds,
ii,
iou_thres=cfg.score_threshold,
use_07_metric=cfg.use_voc_07_metric)
info += 'EVAL: Class {}: Recall: {:.4f}, Precision: {:.4f}, AP: {:.4f}\n'.format(
ii, rec, prec, ap)
rec_total.update(rec, npos)
prec_total.update(prec, nd)
ap_total.update(ap, 1)
mAP = ap_total.average
info += 'EVAL: Recall: {:.4f}, Precison: {:.4f}, mAP: {:.4f}\n'.format(
rec_total.average, prec_total.average, mAP)
print(info)
def evaluate_net(args):
logging.basicConfig()
logger = logging.getLogger()
logger.setLevel(logging.INFO)
ctx = [mx.gpu(int(i)) for i in args.gpus.split(',') if i.strip()]
ctx = [mx.cpu()] if not ctx else ctx
ctx = ctx[0]
cfg['num_ctx'] = cfg['batch_size']
sym = CornerNet(is_train=False, cfgs=cfg)
valset = args.valset
imdb = eval(args.dataset)(valset, args.root_dir, args.image_dir,
args.result_path)
val_roidb = imdb.gt_roidb()
if args.DEBUG:
val_roidb = val_roidb[0:100]
eval_iter = TestDetRecordIter(val_roidb, cfg, 1, False)
if args.DEBUG and 0:
mod = mx.mod.Module(symbol=sym, context=ctx, data_names=['data'])
mod.bind(data_shapes=[('data', (2, 3, 511, 767))], label_shapes=None)
_, params, auxs = mx.model.load_checkpoint(
'model/pretrained_cornernet', 0)
mod.set_params(params, auxs, allow_missing=False, force_init=True)
it = np.load('image.npy')
mod.forward(mx.io.DataBatch(data=[mx.nd.array(it)]), is_train=False)
out = mod.get_outputs()[0]
mod = mx.mod.Module(symbol=sym,
context=ctx,
data_names=eval_iter.data_names,
label_names=None)
mod.bind(data_shapes=eval_iter.provide_data,
label_shapes=None,
for_training=False)
_, params, auxs = mx.model.load_checkpoint(args.prefix, args.epoch)
mod.set_params(params, auxs, allow_missing=False, force_init=True)
num_classes = cfg['num_classes']
all_boxes = [[[] for _ in range(len(val_roidb))]
for _ in range(num_classes + 1)]
eval_iter.reset()
for idx, (info, it) in enumerate(eval_iter):
print('predicting image %d' % idx)
mod.forward(it, is_train=False)
out = mod.get_outputs()
bboxes = decode(*out, info, cfg['test_scales'])
for i in range(num_classes):
all_boxes[i + 1][idx] = bboxes[i + 1]
if args.DEBUG:
for i in range(1):
im = it.data[0].asnumpy()
im = im[i].transpose((1, 2, 0))
im *= 70
im += 110
im_map = im.copy()
im = im.clip(0, 255).astype(np.uint8)
h, w, c = im.shape
tl_map = mx.nd.sigmoid(out[0]).asnumpy()
tl_map = tl_map[i].transpose(1, 2, 0).sum(axis=-1) * 255
tl_map = cv2.resize(tl_map, (w, h))
br_map = mx.nd.sigmoid(out[1]).asnumpy()
br_map = br_map[i].transpose(1, 2, 0).sum(axis=-1) * 255
br_map = cv2.resize(br_map, (w, h))
im_map[:, :, 0] += br_map
im_map[:, :, -1] += tl_map
cv2.imwrite('images/heat_%04d.jpg' % idx, im_map)
for i in range(1, 81):
cat_name = coco_dict[i - 1]
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).astype(np.int32).tolist()
for bbox in bboxes[i]:
b = bbox[:4]
b[0:4:2] += info[0][0, 2]
b[1:4:2] += info[0][0, 0]
b = b.astype(np.int)
if bbox[-1] < 0.5:
continue
if b[1] - cat_size[1] - 2 < 0:
cv2.rectangle(
im, (b[0], b[1] + 2),
(b[0] + cat_size[0], b[1] + cat_size[1] + 2),
color, -1)
cv2.putText(im,
cat_name,
(b[0], b[1] + cat_size[1] + 2),
cv2.FONT_HERSHEY_SIMPLEX,
0.5, (0, 0, 0),
thickness=1)
else:
cv2.rectangle(im, (b[0], b[1] - cat_size[1] - 2),
(b[0] + cat_size[0], b[1] - 2),
color, -1)
cv2.putText(im,
cat_name, (b[0], b[1] - 2),
cv2.FONT_HERSHEY_SIMPLEX,
0.5, (0, 0, 0),
thickness=1)
cv2.rectangle(im, (b[0], b[1]), (b[2], b[3]), color, 2)
cv2.imwrite('images/image_%04d.jpg' % idx, im)
if not args.DEBUG:
info_str = imdb.evaluate_detections(all_boxes)
pprint.pprint(info_str)