def write_voc_results_file(self, all_boxes, test_imgid_list, det_save_dir):
'''
:param all_boxes: is a list. each item reprensent the detections of a img.
the detections is a array. shape is [-1, 7]. [category, score, x, y, w, h, theta]
Note that: if none detections in this img. that the detetions is : []
:param test_imgid_list:
:param det_save_path:
:return:
'''
for cls, cls_id in self.name_label_map.items():
if cls == 'back_ground':
continue
print("Writing {} VOC resutls file".format(cls))
tools.makedirs(det_save_dir)
det_save_path = os.path.join(det_save_dir, "det_" + cls + ".txt")
with open(det_save_path, 'wt') as f:
for index, img_name in enumerate(test_imgid_list):
this_img_detections = all_boxes[index]
this_cls_detections = this_img_detections[
this_img_detections[:, 0] == cls_id]
if this_cls_detections.shape[0] == 0:
continue # this cls has none detections in this img
for a_det in this_cls_detections:
f.write(
'{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.
format(img_name, a_det[1], a_det[2], a_det[3],
a_det[4], a_det[5], a_det[6])
) # that is [img_name, score, x, y, w, h, theta]
def clip_image(file_idx, image, boxes_all, width, height, stride_w, stride_h):
min_pixel = 5
print(file_idx)
boxes_all_5 = backward_convert(boxes_all[:, :8], False)
print(boxes_all[np.logical_or(boxes_all_5[:, 2] <= min_pixel, boxes_all_5[:, 3] <= min_pixel), :])
boxes_all = boxes_all[np.logical_and(boxes_all_5[:, 2] > min_pixel, boxes_all_5[:, 3] > min_pixel), :]
if boxes_all.shape[0] > 0:
shape = image.shape
for start_h in range(0, shape[0], stride_h):
for start_w in range(0, shape[1], stride_w):
boxes = copy.deepcopy(boxes_all)
box = np.zeros_like(boxes_all)
start_h_new = start_h
start_w_new = start_w
if start_h + height > shape[0]:
start_h_new = shape[0] - height
if start_w + width > shape[1]:
start_w_new = shape[1] - width
top_left_row = max(start_h_new, 0)
top_left_col = max(start_w_new, 0)
bottom_right_row = min(start_h + height, shape[0])
bottom_right_col = min(start_w + width, shape[1])
subImage = image[top_left_row:bottom_right_row, top_left_col: bottom_right_col]
box[:, 0] = boxes[:, 0] - top_left_col
box[:, 2] = boxes[:, 2] - top_left_col
box[:, 4] = boxes[:, 4] - top_left_col
box[:, 6] = boxes[:, 6] - top_left_col
box[:, 1] = boxes[:, 1] - top_left_row
box[:, 3] = boxes[:, 3] - top_left_row
box[:, 5] = boxes[:, 5] - top_left_row
box[:, 7] = boxes[:, 7] - top_left_row
box[:, 8] = boxes[:, 8]
center_y = 0.25 * (box[:, 1] + box[:, 3] + box[:, 5] + box[:, 7])
center_x = 0.25 * (box[:, 0] + box[:, 2] + box[:, 4] + box[:, 6])
cond1 = np.intersect1d(np.where(center_y[:] >= 0)[0], np.where(center_x[:] >= 0)[0])
cond2 = np.intersect1d(np.where(center_y[:] <= (bottom_right_row - top_left_row))[0],
np.where(center_x[:] <= (bottom_right_col - top_left_col))[0])
idx = np.intersect1d(cond1, cond2)
if len(idx) > 0 and (subImage.shape[0] > 5 and subImage.shape[1] > 5):
makedirs(os.path.join(save_dir, 'images'))
img = os.path.join(save_dir, 'images',
"%s_%04d_%04d.png" % (file_idx, top_left_row, top_left_col))
cv2.imwrite(img, subImage)
makedirs(os.path.join(save_dir, 'labeltxt'))
xml = os.path.join(save_dir, 'labeltxt',
"%s_%04d_%04d.xml" % (file_idx, top_left_row, top_left_col))
save_to_xml(xml, subImage.shape[0], subImage.shape[1], box[idx, :], class_list)
def test_pb(self, frozen_graph_path, test_dir):
graph = self.load_graph(frozen_graph_path)
print("we are testing ====>>>>", frozen_graph_path)
img = graph.get_tensor_by_name("input_img:0")
dets = graph.get_tensor_by_name("DetResults:0")
with tf.Session(graph=graph) as sess:
for img_path in os.listdir(test_dir):
print(img_path)
a_img = cv2.imread(os.path.join(test_dir, img_path))[:, :, ::-1]
raw_h, raw_w = a_img.shape[0], a_img.shape[1]
short_size, max_len = self.cfgs.IMG_SHORT_SIDE_LEN, cfgs.IMG_MAX_LENGTH
if raw_h < raw_w:
new_h, new_w = short_size, min(int(short_size * float(raw_w) / raw_h), max_len)
else:
new_h, new_w = min(int(short_size * float(raw_h) / raw_w), max_len), short_size
img_resize = cv2.resize(a_img, (new_w, new_h))
dets_val = sess.run(dets, feed_dict={img: img_resize[:, :, ::-1]})
bbox_pred, cls_prob, proposal = dets_val[:, :5], dets_val[:, 5:(5+self.cfgs.CLASS_NUM)], \
dets_val[:, (5+self.cfgs.CLASS_NUM):]
detected_boxes, detected_scores, detected_categories = self.postprocess_detctions(bbox_pred, cls_prob, proposal)
if True:
# detected_indices = det_scores_r_ >= self.cfgs.VIS_SCORE
# detected_scores = det_scores_r_[detected_indices]
# detected_boxes = det_boxes_r_[detected_indices]
# detected_categories = det_category_r_[detected_indices]
drawer = DrawBox(self.cfgs)
det_detections_r = drawer.draw_boxes_with_label_and_scores(img_resize[:, :, ::-1],
boxes=detected_boxes,
labels=detected_categories,
scores=detected_scores,
method=1,
in_graph=True)
save_dir = os.path.join('test_pb', self.cfgs.VERSION, 'pb_img_vis')
tools.makedirs(save_dir)
cv2.imwrite(save_dir + '/{}'.format(img_path),
det_detections_r[:, :, ::-1])
def convert_pascal_to_tfrecord():
xml_path = os.path.join(FLAGS.VOC_dir, FLAGS.xml_dir)
image_path = os.path.join(FLAGS.VOC_dir, FLAGS.image_dir)
save_path = os.path.join(
FLAGS.save_dir, FLAGS.dataset + '_' + FLAGS.save_name + '.tfrecord')
makedirs(FLAGS.save_dir)
# writer_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)
# writer = tf.python_io.TFRecordWriter(path=save_path, options=writer_options)
writer = tf.python_io.TFRecordWriter(path=save_path)
for count, xml in enumerate(glob.glob(xml_path + '/*.xml')):
img_name = xml.split('/')[-1].split('.')[0] + FLAGS.img_format
img_path = image_path + '/' + img_name
if not os.path.exists(img_path):
print('{} is not exist!'.format(img_path))
continue
img_height, img_width, gtbox_label = read_xml_gtbox_and_label(xml)
# if img_height != 600 or img_width != 600:
# continue
img = cv2.imread(img_path)[:, :, ::-1]
feature = tf.train.Features(
feature={
# do not need encode() in linux
'img_name': _bytes_feature(img_name.encode()),
# 'img_name': _bytes_feature(img_name),
'img_height': _int64_feature(img_height),
'img_width': _int64_feature(img_width),
'img': _bytes_feature(img.tostring()),
'gtboxes_and_label': _bytes_feature(gtbox_label.tostring()),
'num_objects': _int64_feature(gtbox_label.shape[0])
})
example = tf.train.Example(features=feature)
writer.write(example.SerializeToString())
view_bar('Conversion progress', count + 1,
len(glob.glob(xml_path + '/*.xml')))
print('\nConversion is complete!')
writer.close()
def test_dota(self, det_net, real_test_img_list, txt_name):
save_path = os.path.join('./test_dota', self.cfgs.VERSION)
nr_records = len(real_test_img_list)
pbar = tqdm(total=nr_records)
gpu_num = len(self.args.gpus.strip().split(','))
nr_image = math.ceil(nr_records / gpu_num)
result_queue = Queue(500)
procs = []
for i, gpu_id in enumerate(self.args.gpus.strip().split(',')):
start = i * nr_image
end = min(start + nr_image, nr_records)
split_records = real_test_img_list[start:end]
proc = Process(target=self.worker,
args=(int(gpu_id), split_records, det_net,
result_queue))
print('process:%d, start:%d, end:%d' % (i, start, end))
proc.start()
procs.append(proc)
for i in range(nr_records):
res = result_queue.get()
if self.args.show_box:
nake_name = res['image_id'].split('/')[-1]
tools.makedirs(os.path.join(save_path, 'dota_img_vis'))
draw_path = os.path.join(save_path, 'dota_img_vis', nake_name)
draw_img = np.array(cv2.imread(res['image_id']), np.float32)
detected_boxes = backward_convert(res['boxes'],
with_label=False)
detected_indices = res['scores'] >= self.cfgs.VIS_SCORE
detected_scores = res['scores'][detected_indices]
detected_boxes = detected_boxes[detected_indices]
detected_categories = res['labels'][detected_indices]
drawer = DrawBox(self.cfgs)
final_detections = drawer.draw_boxes_with_label_and_scores(
draw_img,
boxes=detected_boxes,
labels=detected_categories,
scores=detected_scores,
method=1,
is_csl=True,
in_graph=False)
cv2.imwrite(draw_path, final_detections)
else:
CLASS_DOTA = self.name_label_map.keys()
write_handle = {}
tools.makedirs(os.path.join(save_path, 'dota_res'))
for sub_class in CLASS_DOTA:
if sub_class == 'back_ground':
continue
write_handle[sub_class] = open(
os.path.join(save_path, 'dota_res',
'Task1_%s.txt' % sub_class), 'a+')
for i, rbox in enumerate(res['boxes']):
command = '%s %.3f %.1f %.1f %.1f %.1f %.1f %.1f %.1f %.1f\n' % (
res['image_id'].split('/')[-1].split('.')[0],
res['scores'][i],
rbox[0],
rbox[1],
rbox[2],
rbox[3],
rbox[4],
rbox[5],
rbox[6],
rbox[7],
)
write_handle[self.label_name_map[res['labels'][i]]].write(
command)
for sub_class in CLASS_DOTA:
if sub_class == 'back_ground':
continue
write_handle[sub_class].close()
fw = open(txt_name, 'a+')
fw.write('{}\n'.format(res['image_id'].split('/')[-1]))
fw.close()
pbar.set_description("Test image %s" %
res['image_id'].split('/')[-1])
pbar.update(1)
for p in procs:
p.join()
def convert_pascal_to_tfrecord():
json_file = os.path.join(FLAGS.root_dir, FLAGS.json_file)
image_path = os.path.join(FLAGS.root_dir, FLAGS.image_dir)
save_path = os.path.join(
FLAGS.save_dir, FLAGS.dataset + '_' + FLAGS.save_name + '.tfrecord')
makedirs(FLAGS.save_dir)
# writer_options = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)
# writer = tf.python_io.TFRecordWriter(path=save_path, options=writer_options)
writer = tf.python_io.TFRecordWriter(path=save_path)
with open(json_file, 'r') as fr:
all_gts = json.load(fr)
images = all_gts['images']
annotations = all_gts['annotations']
all_gt_label = {}
for annotation in annotations:
image_id = annotation['image_id']
# print(image_id-1) # 57533
if image_id > len(images):
continue
if images[image_id - 1]['file_name'] in all_gt_label.keys():
# all_gt_label[images[image_id - 1]['file_name']]['gtboxes'].append(annotation['segmentation'])
all_gt_label[images[image_id - 1]['file_name']]['gtboxes'].append(
coordinate_convert_r(annotation['rbbox']))
all_gt_label[images[image_id - 1]['file_name']]['labels'].append(
annotation['category_id'])
else:
all_gt_label[images[image_id - 1]['file_name']] = {
'height': images[image_id - 1]['height'],
'width': images[image_id - 1]['width'],
# 'gtboxes': [annotation['segmentation']],
'gtboxes': [coordinate_convert_r(annotation['rbbox'])],
'labels': [annotation['category_id']]
}
count = 0
for img_name in all_gt_label.keys():
img = cv2.imread(os.path.join(image_path, img_name))
img_height = all_gt_label[img_name]['height']
img_width = all_gt_label[img_name]['width']
gtboxes = np.array(all_gt_label[img_name]['gtboxes']).reshape([-1, 8])
labels = np.array(all_gt_label[img_name]['labels']).reshape([-1, 1])
gtboxes_and_label = np.array(
np.concatenate([gtboxes, labels], axis=-1), np.int32)
feature = tf.train.Features(
feature={
# do not need encode() in linux
'img_name': _bytes_feature(img_name.encode()),
# 'img_name': _bytes_feature(img_name),
'img_height': _int64_feature(img_height),
'img_width': _int64_feature(img_width),
'img': _bytes_feature(img.tostring()),
'gtboxes_and_label': _bytes_feature(
gtboxes_and_label.tostring()),
'num_objects': _int64_feature(gtboxes_and_label.shape[0])
})
example = tf.train.Example(features=feature)
writer.write(example.SerializeToString())
view_bar('Conversion progress', count + 1, len(all_gt_label.keys()))
count += 1
print('\nConversion is complete!')
writer.close()
def eval_with_plac(self, img_dir, det_net, image_ext):
os.environ["CUDA_VISIBLE_DEVICES"] = self.args.gpu
# 1. preprocess img
img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None,
3]) # is RGB. not BGR
img_batch = tf.cast(img_plac, tf.float32)
pretrain_zoo = PretrainModelZoo()
if self.cfgs.NET_NAME in pretrain_zoo.pth_zoo or self.cfgs.NET_NAME in pretrain_zoo.mxnet_zoo:
img_batch = (img_batch / 255 - tf.constant(
self.cfgs.PIXEL_MEAN_)) / tf.constant(self.cfgs.PIXEL_STD)
else:
img_batch = img_batch - tf.constant(self.cfgs.PIXEL_MEAN)
img_batch = tf.expand_dims(img_batch, axis=0)
detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
input_img_batch=img_batch)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
restorer, restore_ckpt = det_net.get_restorer()
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
sess.run(init_op)
if not restorer is None:
restorer.restore(sess, restore_ckpt)
print('restore model')
all_boxes_r = []
imgs = os.listdir(img_dir)
pbar = tqdm(imgs)
for a_img_name in pbar:
a_img_name = a_img_name.split(image_ext)[0]
raw_img = cv2.imread(
os.path.join(img_dir, a_img_name + image_ext))
raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]
det_boxes_r_all, det_scores_r_all, det_category_r_all = [], [], []
img_short_side_len_list = self.cfgs.IMG_SHORT_SIDE_LEN if isinstance(
self.cfgs.IMG_SHORT_SIDE_LEN,
list) else [self.cfgs.IMG_SHORT_SIDE_LEN]
img_short_side_len_list = [
img_short_side_len_list[0]
] if not self.args.multi_scale else img_short_side_len_list
for short_size in img_short_side_len_list:
max_len = self.cfgs.IMG_MAX_LENGTH
if raw_h < raw_w:
new_h, new_w = short_size, min(
int(short_size * float(raw_w) / raw_h), max_len)
else:
new_h, new_w = min(
int(short_size * float(raw_h) / raw_w),
max_len), short_size
img_resize = cv2.resize(raw_img, (new_w, new_h))
resized_img, detected_boxes, detected_scores, detected_categories = \
sess.run(
[img_batch, detection_boxes, detection_scores, detection_category],
feed_dict={img_plac: img_resize[:, :, ::-1]}
)
if detected_boxes.shape[0] == 0:
continue
resized_h, resized_w = resized_img.shape[
1], resized_img.shape[2]
detected_boxes = forward_convert(detected_boxes, False)
detected_boxes[:, 0::2] *= (raw_w / resized_w)
detected_boxes[:, 1::2] *= (raw_h / resized_h)
det_boxes_r_all.extend(detected_boxes)
det_scores_r_all.extend(detected_scores)
det_category_r_all.extend(detected_categories)
det_boxes_r_all = np.array(det_boxes_r_all)
det_scores_r_all = np.array(det_scores_r_all)
det_category_r_all = np.array(det_category_r_all)
box_res_rotate_ = []
label_res_rotate_ = []
score_res_rotate_ = []
if det_scores_r_all.shape[0] != 0:
for sub_class in range(1, self.cfgs.CLASS_NUM + 1):
index = np.where(det_category_r_all == sub_class)[0]
if len(index) == 0:
continue
tmp_boxes_r = det_boxes_r_all[index]
tmp_label_r = det_category_r_all[index]
tmp_score_r = det_scores_r_all[index]
if self.args.multi_scale:
tmp_boxes_r_ = backward_convert(tmp_boxes_r, False)
# try:
# inx = nms_rotate.nms_rotate_cpu(boxes=np.array(tmp_boxes_r_),
# scores=np.array(tmp_score_r),
# iou_threshold=self.cfgs.NMS_IOU_THRESHOLD,
# max_output_size=5000)
# except:
tmp_boxes_r_ = np.array(tmp_boxes_r_)
tmp = np.zeros([
tmp_boxes_r_.shape[0],
tmp_boxes_r_.shape[1] + 1
])
tmp[:, 0:-1] = tmp_boxes_r_
tmp[:, -1] = np.array(tmp_score_r)
# Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms
jitter = np.zeros([
tmp_boxes_r_.shape[0],
tmp_boxes_r_.shape[1] + 1
])
jitter[:, 0] += np.random.rand(
tmp_boxes_r_.shape[0], ) / 1000
inx = rotate_gpu_nms(
np.array(tmp, np.float32) +
np.array(jitter, np.float32),
float(self.cfgs.NMS_IOU_THRESHOLD), 0)
else:
inx = np.arange(0, tmp_score_r.shape[0])
box_res_rotate_.extend(np.array(tmp_boxes_r)[inx])
score_res_rotate_.extend(np.array(tmp_score_r)[inx])
label_res_rotate_.extend(np.array(tmp_label_r)[inx])
if len(box_res_rotate_) == 0:
all_boxes_r.append(np.array([]))
continue
det_boxes_r_ = np.array(box_res_rotate_)
det_scores_r_ = np.array(score_res_rotate_)
det_category_r_ = np.array(label_res_rotate_)
if self.args.draw_imgs:
detected_indices = det_scores_r_ >= self.cfgs.VIS_SCORE
detected_scores = det_scores_r_[detected_indices]
detected_boxes = det_boxes_r_[detected_indices]
detected_categories = det_category_r_[detected_indices]
detected_boxes = backward_convert(detected_boxes, False)
drawer = DrawBox(self.cfgs)
det_detections_r = drawer.draw_boxes_with_label_and_scores(
raw_img[:, :, ::-1],
boxes=detected_boxes,
labels=detected_categories,
scores=detected_scores,
method=1,
in_graph=True)
save_dir = os.path.join('test_hrsc', self.cfgs.VERSION,
'hrsc2016_img_vis')
tools.makedirs(save_dir)
cv2.imwrite(save_dir + '/{}.jpg'.format(a_img_name),
det_detections_r[:, :, ::-1])
det_boxes_r_ = backward_convert(det_boxes_r_, False)
x_c, y_c, w, h, theta = det_boxes_r_[:, 0], det_boxes_r_[:, 1], det_boxes_r_[:, 2], \
det_boxes_r_[:, 3], det_boxes_r_[:, 4]
boxes_r = np.transpose(np.stack([x_c, y_c, w, h, theta]))
dets_r = np.hstack((det_category_r_.reshape(-1, 1),
det_scores_r_.reshape(-1, 1), boxes_r))
all_boxes_r.append(dets_r)
pbar.set_description("Eval image %s" % a_img_name)
# fw1 = open(cfgs.VERSION + '_detections_r.pkl', 'wb')
# pickle.dump(all_boxes_r, fw1)
return all_boxes_r
def test_icdar2015(self, det_net, real_test_img_list, txt_name):
save_path = os.path.join('./test_icdar2015', self.cfgs.VERSION)
tools.makedirs(save_path)
nr_records = len(real_test_img_list)
pbar = tqdm(total=nr_records)
gpu_num = len(self.args.gpus.strip().split(','))
nr_image = math.ceil(nr_records / gpu_num)
result_queue = Queue(500)
procs = []
for i, gpu_id in enumerate(self.args.gpus.strip().split(',')):
start = i * nr_image
end = min(start + nr_image, nr_records)
split_records = real_test_img_list[start:end]
proc = Process(target=self.worker, args=(int(gpu_id), split_records, det_net, result_queue))
print('process:%d, start:%d, end:%d' % (i, start, end))
proc.start()
procs.append(proc)
for i in range(nr_records):
res = result_queue.get()
tools.makedirs(os.path.join(save_path, 'icdar2015_res'))
if res['boxes'].shape[0] == 0:
fw_txt_dt = open(os.path.join(save_path, 'icdar2015_res', 'res_{}.txt'.format(res['image_id'].split('/')[-1].split('.')[0])),
'w')
fw_txt_dt.close()
pbar.update(1)
fw = open(txt_name, 'a+')
fw.write('{}\n'.format(res['image_id'].split('/')[-1]))
fw.close()
continue
x1, y1, x2, y2, x3, y3, x4, y4 = res['boxes'][:, 0], res['boxes'][:, 1], res['boxes'][:, 2], res['boxes'][:, 3],\
res['boxes'][:, 4], res['boxes'][:, 5], res['boxes'][:, 6], res['boxes'][:, 7]
x1, y1 = x1 * res['scales'][0], y1 * res['scales'][1]
x2, y2 = x2 * res['scales'][0], y2 * res['scales'][1]
x3, y3 = x3 * res['scales'][0], y3 * res['scales'][1]
x4, y4 = x4 * res['scales'][0], y4 * res['scales'][1]
boxes = np.transpose(np.stack([x1, y1, x2, y2, x3, y3, x4, y4]))
if self.args.show_box:
boxes = backward_convert(boxes, False)
nake_name = res['image_id'].split('/')[-1]
tools.makedirs(os.path.join(save_path, 'icdar2015_img_vis'))
draw_path = os.path.join(save_path, 'icdar2015_img_vis', nake_name)
draw_img = np.array(cv2.imread(res['image_id']), np.float32)
drawer = DrawBox(self.cfgs)
final_detections = drawer.draw_boxes_with_label_and_scores(draw_img,
boxes=boxes,
labels=res['labels'],
scores=res['scores'],
method=1,
in_graph=False)
cv2.imwrite(draw_path, final_detections)
else:
fw_txt_dt = open(os.path.join(save_path, 'icdar2015_res', 'res_{}.txt'.format(res['image_id'].split('/')[-1].split('.')[0])), 'w')
for box in boxes:
line = '%d,%d,%d,%d,%d,%d,%d,%d\n' % (box[0], box[1], box[2], box[3],
box[4], box[5], box[6], box[7])
fw_txt_dt.write(line)
fw_txt_dt.close()
fw = open(txt_name, 'a+')
fw.write('{}\n'.format(res['image_id'].split('/')[-1]))
fw.close()
pbar.set_description("Test image %s" % res['image_id'].split('/')[-1])
pbar.update(1)
for p in procs:
p.join()
def log_printer(self, deter, optimizer, global_step, tower_grads, total_loss_dict, num_gpu, graph):
for k in total_loss_dict.keys():
tf.summary.scalar('{}/{}'.format(k.split('_')[0], k), total_loss_dict[k])
if len(tower_grads) > 1:
grads = self.sum_gradients(tower_grads)
else:
grads = tower_grads[0]
if self.cfgs.MUTILPY_BIAS_GRADIENT is not None:
final_gvs = []
with tf.variable_scope('Gradient_Mult'):
for grad, var in grads:
scale = 1.
if '/biases:' in var.name:
scale *= self.cfgs.MUTILPY_BIAS_GRADIENT
if 'conv_new' in var.name:
scale *= 3.
if not np.allclose(scale, 1.0):
grad = tf.multiply(grad, scale)
final_gvs.append((grad, var))
apply_gradient_op = optimizer.apply_gradients(final_gvs, global_step=global_step)
else:
apply_gradient_op = optimizer.apply_gradients(grads, global_step=global_step)
variable_averages = tf.train.ExponentialMovingAverage(0.9999, global_step)
variables_averages_op = variable_averages.apply(tf.trainable_variables())
train_op = tf.group(apply_gradient_op, variables_averages_op)
# train_op = optimizer.apply_gradients(final_gvs, global_step=global_step)
summary_op = tf.summary.merge_all()
restorer, restore_ckpt = deter.get_restorer()
saver = tf.train.Saver(max_to_keep=20)
init_op = tf.group(
tf.global_variables_initializer(),
tf.local_variables_initializer()
)
tfconfig = tf.ConfigProto(
allow_soft_placement=True, log_device_placement=False)
tfconfig.gpu_options.allow_growth = True
with tf.Session(config=tfconfig) as sess:
sess.run(init_op)
# sess.run(tf.initialize_all_variables())
tf.local_variables_initializer().run()
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
summary_path = os.path.join(self.cfgs.SUMMARY_PATH, self.cfgs.VERSION)
tools.makedirs(summary_path)
summary_writer = tf.summary.FileWriter(summary_path, graph=sess.graph)
if not restorer is None:
restorer.restore(sess, restore_ckpt)
print('restore model')
self.stats_graph(graph)
for step in range(self.cfgs.MAX_ITERATION // num_gpu):
training_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
if step % self.cfgs.SHOW_TRAIN_INFO_INTE != 0 and step % self.cfgs.SMRY_ITER != 0:
_, global_stepnp = sess.run([train_op, global_step])
else:
if step % self.cfgs.SHOW_TRAIN_INFO_INTE == 0 and step % self.cfgs.SMRY_ITER != 0:
start = time.time()
_, global_stepnp, total_loss_dict_ = \
sess.run([train_op, global_step, total_loss_dict])
end = time.time()
print('***'*24)
print("%s: global_step:%d current_step:%d"
% (training_time, (global_stepnp-1)*num_gpu, step*num_gpu))
seconds = (self.cfgs.MAX_ITERATION - (global_stepnp-1) * num_gpu) * (end - start) / num_gpu
m, s = divmod(seconds, 60)
h, m = divmod(m, 60)
d, h = divmod(h, 24)
print("speed: %.3fs, remaining training time: %02d:%02d:%02d:%02d" % ((end - start) / num_gpu, d, h, m, s))
loss_str = ''
for k in total_loss_dict_.keys():
loss_str += '%s:%.3f\n' % (k, total_loss_dict_[k])
print(loss_str)
if np.isnan(total_loss_dict_['total_losses']):
sys.exit(0)
else:
if step % self.cfgs.SMRY_ITER == 0:
_, global_stepnp, summary_str = sess.run([train_op, global_step, summary_op])
summary_writer.add_summary(summary_str, (global_stepnp-1)*num_gpu)
summary_writer.flush()
if (step > 0 and step % (self.cfgs.SAVE_WEIGHTS_INTE // num_gpu) == 0) or (step >= self.cfgs.MAX_ITERATION // num_gpu - 1):
save_dir = os.path.join(self.cfgs.TRAINED_CKPT, self.cfgs.VERSION)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
save_ckpt = os.path.join(save_dir, '{}_'.format(self.cfgs.DATASET_NAME) +
str((global_stepnp-1)*num_gpu) + 'model.ckpt')
saver.save(sess, save_ckpt)
print('Weights had been saved')
if (global_stepnp-1)*num_gpu > self.cfgs.MAX_ITERATION:
break
print('***' * 24)
print('End of training.')
coord.request_stop()
coord.join(threads)
def test_pb(self, frozen_graph_path, test_dir):
graph = self.load_graph(frozen_graph_path)
print("we are testing ====>>>>", frozen_graph_path)
img = graph.get_tensor_by_name("input_img:0")
dets = graph.get_tensor_by_name("DetResults:0")
with tf.Session(graph=graph) as sess:
for img_path in os.listdir(test_dir):
print(img_path)
a_img = cv2.imread(os.path.join(test_dir,
img_path))[:, :, ::-1]
raw_h, raw_w = a_img.shape[0], a_img.shape[1]
short_size, max_len = self.cfgs.IMG_SHORT_SIDE_LEN, cfgs.IMG_MAX_LENGTH
if raw_h < raw_w:
new_h, new_w = short_size, min(
int(short_size * float(raw_w) / raw_h), max_len)
else:
new_h, new_w = min(int(short_size * float(raw_h) / raw_w),
max_len), short_size
img_resize = cv2.resize(a_img, (new_w, new_h))
dets_val = sess.run(dets,
feed_dict={img: img_resize[:, :, ::-1]})
bbox_pred, cls_prob = dets_val[:, :5], dets_val[:, 5:(
5 + self.cfgs.CLASS_NUM)]
anchor = GenerateAnchors(self.cfgs, 'H')
h1, w1 = math.ceil(new_h / 2), math.ceil(new_w / 2)
h2, w2 = math.ceil(h1 / 2), math.ceil(w1 / 2)
h3, w3 = math.ceil(h2 / 2), math.ceil(w2 / 2)
h4, w4 = math.ceil(h3 / 2), math.ceil(w3 / 2)
h5, w5 = math.ceil(h4 / 2), math.ceil(w4 / 2)
h6, w6 = math.ceil(h5 / 2), math.ceil(w5 / 2)
h7, w7 = math.ceil(h6 / 2), math.ceil(w6 / 2)
h_dict = {'P3': h3, 'P4': h4, 'P5': h5, 'P6': h6, 'P7': h7}
w_dict = {'P3': w3, 'P4': w4, 'P5': w5, 'P6': w6, 'P7': w7}
anchors = anchor.generate_all_anchor_pb(h_dict, w_dict)
anchors = np.concatenate(anchors, axis=0)
x_c = (anchors[:, 2] + anchors[:, 0]) / 2
y_c = (anchors[:, 3] + anchors[:, 1]) / 2
h = anchors[:, 2] - anchors[:, 0] + 1
w = anchors[:, 3] - anchors[:, 1] + 1
theta = -90 * np.ones_like(x_c)
anchors = np.transpose(np.stack([x_c, y_c, w, h, theta]))
detected_boxes, detected_scores, detected_categories = self.postprocess_detctions(
bbox_pred, cls_prob, anchors)
if True:
# detected_indices = det_scores_r_ >= self.cfgs.VIS_SCORE
# detected_scores = det_scores_r_[detected_indices]
# detected_boxes = det_boxes_r_[detected_indices]
# detected_categories = det_category_r_[detected_indices]
drawer = DrawBox(self.cfgs)
det_detections_r = drawer.draw_boxes_with_label_and_scores(
img_resize[:, :, ::-1],
boxes=detected_boxes,
labels=detected_categories,
scores=detected_scores,
method=1,
in_graph=True)
save_dir = os.path.join('test_pb', self.cfgs.VERSION,
'pb_img_vis')
tools.makedirs(save_dir)
cv2.imwrite(save_dir + '/{}'.format(img_path),
det_detections_r[:, :, ::-1])
def test(self, frozen_graph_path, test_dir):
graph = self.load_graph(frozen_graph_path)
print("we are testing ====>>>>", frozen_graph_path)
img = graph.get_tensor_by_name("input_img:0")
dets = graph.get_tensor_by_name("DetResults:0")
with tf.Session(graph=graph) as sess:
for img_path in os.listdir(test_dir):
print(img_path)
a_img = cv2.imread(os.path.join(test_dir,
img_path))[:, :, ::-1]
raw_h, raw_w = a_img.shape[0], a_img.shape[1]
short_size, max_len = self.cfgs.IMG_SHORT_SIDE_LEN, cfgs.IMG_MAX_LENGTH
if raw_h < raw_w:
new_h, new_w = short_size, min(
int(short_size * float(raw_w) / raw_h), max_len)
else:
new_h, new_w = min(int(short_size * float(raw_h) / raw_w),
max_len), short_size
img_resize = cv2.resize(a_img, (new_w, new_h))
dets_val = sess.run(dets,
feed_dict={img: img_resize[:, :, ::-1]})
box_res_rotate_ = []
label_res_rotate_ = []
score_res_rotate_ = []
if dets_val.shape[0] != 0:
for sub_class in range(1, self.cfgs.CLASS_NUM + 1):
index = np.where(dets_val[:, 0] == sub_class)[0]
if len(index) == 0:
continue
tmp_boxes_r = dets_val[:, 2:][index]
tmp_label_r = dets_val[:, 0][index]
tmp_score_r = dets_val[:, 1][index]
# try:
inx = nms_rotate_cpu(
boxes=np.array(tmp_boxes_r),
scores=np.array(tmp_score_r),
iou_threshold=self.cfgs.NMS_IOU_THRESHOLD,
max_output_size=20)
# except:
# tmp_boxes_r_ = np.array(tmp_boxes_r)
# tmp = np.zeros([tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1])
# tmp[:, 0:-1] = tmp_boxes_r_
# tmp[:, -1] = np.array(tmp_score_r)
# # Note: the IoU of two same rectangles is 0, which is calculated by rotate_gpu_nms
# jitter = np.zeros([tmp_boxes_r_.shape[0], tmp_boxes_r_.shape[1] + 1])
# jitter[:, 0] += np.random.rand(tmp_boxes_r_.shape[0], ) / 1000
# inx = rotate_gpu_nms(np.array(tmp, np.float32) + np.array(jitter, np.float32),
# float(self.cfgs.NMS_IOU_THRESHOLD), 0)
box_res_rotate_.extend(np.array(tmp_boxes_r)[inx])
score_res_rotate_.extend(np.array(tmp_score_r)[inx])
label_res_rotate_.extend(np.array(tmp_label_r)[inx])
det_boxes_r_ = np.array(box_res_rotate_)
det_scores_r_ = np.array(score_res_rotate_)
det_category_r_ = np.array(label_res_rotate_)
if True:
detected_indices = det_scores_r_ >= self.cfgs.VIS_SCORE
detected_scores = det_scores_r_[detected_indices]
detected_boxes = det_boxes_r_[detected_indices]
detected_categories = det_category_r_[detected_indices]
drawer = DrawBox(self.cfgs)
det_detections_r = drawer.draw_boxes_with_label_and_scores(
img_resize[:, :, ::-1],
boxes=detected_boxes,
labels=detected_categories,
scores=detected_scores,
method=1,
in_graph=True)
save_dir = os.path.join('test_pb', self.cfgs.VERSION,
'pb_img_vis')
tools.makedirs(save_dir)
cv2.imwrite(save_dir + '/{}'.format(img_path),
det_detections_r[:, :, ::-1])
def clip_image(file_idx, image, boxes_all, width, height, w_overlap,
h_overlap):
print(file_idx)
# fill useless boxes
min_pixel = 5
boxes_all_5 = backward_convert(boxes_all[:, :8], False)
small_boxes = boxes_all[np.logical_or(boxes_all_5[:, 2] <= min_pixel,
boxes_all_5[:, 3] <= min_pixel), :]
cv2.fillConvexPoly(image,
np.reshape(small_boxes, [-1, 2]),
color=(0, 0, 0))
different_boxes = boxes_all[boxes_all[:, 9] == 1]
cv2.fillConvexPoly(image,
np.reshape(different_boxes, [-1, 2]),
color=(0, 0, 0))
boxes_all = boxes_all[np.logical_and(boxes_all_5[:, 2] > min_pixel,
boxes_all_5[:, 3] > min_pixel), :]
boxes_all = boxes_all[boxes_all[:, 9] == 0]
if boxes_all.shape[0] > 0:
imgH = image.shape[0]
imgW = image.shape[1]
if imgH < height:
temp = np.zeros([height, imgW, 3], np.float32)
temp[0:imgH, :, :] = image
image = temp
imgH = height
if imgW < width:
temp = np.zeros([imgH, width, 3], np.float32)
temp[:, 0:imgW, :] = image
image = temp
imgW = width
for hh in range(0, imgH, height - h_overlap):
if imgH - hh - 1 < height:
hh_ = imgH - height
else:
hh_ = hh
for ww in range(0, imgW, width - w_overlap):
if imgW - ww - 1 < width:
ww_ = imgW - width
else:
ww_ = ww
subimg = image[hh_:(hh_ + height), ww_:(ww_ + width), :]
boxes = copy.deepcopy(boxes_all)
box = np.zeros_like(boxes_all)
top_left_row = max(hh_, 0)
top_left_col = max(ww_, 0)
bottom_right_row = min(hh_ + height, imgH)
bottom_right_col = min(ww_ + width, imgW)
box[:, :8:2] = boxes[:, :8:2] - top_left_col
box[:, 1:8:2] = boxes[:, 1:8:2] - top_left_row
box[:, 8:] = boxes[:, 8:]
center_y = 0.25 * (box[:, 1] + box[:, 3] + box[:, 5] +
box[:, 7])
center_x = 0.25 * (box[:, 0] + box[:, 2] + box[:, 4] +
box[:, 6])
cond1 = np.intersect1d(
np.where(center_y[:] >= 0)[0],
np.where(center_x[:] >= 0)[0])
cond2 = np.intersect1d(
np.where(
center_y[:] <= (bottom_right_row - top_left_row))[0],
np.where(
center_x[:] <= (bottom_right_col - top_left_col))[0])
idx = np.intersect1d(cond1, cond2)
if len(idx) > 0:
makedirs(os.path.join(save_dir, 'images'))
img = os.path.join(
save_dir, 'images', "%s_%04d_%04d.png" %
(file_idx, top_left_row, top_left_col))
cv2.imwrite(img, subimg)
makedirs(os.path.join(save_dir, 'labeltxt'))
xml = os.path.join(
save_dir, 'labeltxt', "%s_%04d_%04d.xml" %
(file_idx, top_left_row, top_left_col))
save_to_xml(xml, subimg.shape[0], subimg.shape[1],
box[idx, :], class_list)