def build_detection_graph():
input_data = tf.placeholder(dtype=tf.uint8, shape=[FLAGS.size, FLAGS.size, 3], name='input_data')
input_data = tf.expand_dims(input_data, 0)
input_data = tf.cast(input_data, tf.float32)
input_data = input_data / 255.
if FLAGS.tiny:
model = yolo_v3_tiny.yolo_v3_tiny
elif FLAGS.spp:
model = yolo_v3.yolo_v3_spp
else:
model = yolo_v3.yolo_v3
classes = load_coco_names(FLAGS.class_names)
# yolo_model = model(cfgs.class_num, cfgs.anchors)
with tf.variable_scope('detector'):
detections = model(input_data, len(classes), data_format=FLAGS.data_format)
print(detections.get_shape().as_list())
boxes, pred_confs, pred_probs = tf.split(detections, [4, 1, len(classes)], axis=-1)
center_x, center_y, width, height = tf.split(boxes, [1, 1, 1, 1], axis=-1)
x_min = center_x - width / 2
y_min = center_y - height / 2
x_max = center_x + width / 2
y_max = center_y + height / 2
pred_boxes = tf.concat([x_min, y_min, x_max, y_max], axis=-1)
pred_scores = pred_confs * pred_probs
boxes, scores, labels = gpu_nms(pred_boxes, pred_scores, len(classes), max_boxes=20, score_thresh=0.3,
nms_thresh=0.4)
boxes = tf.identity(boxes, name='boxes')
scores = tf.identity(scores, name='scores')
labels = tf.identity(labels, name='labels')
return boxes, scores, labels
def compute_mAP_nms_all_class(self):
gt_boxes = []
pred_boxes = []
with tqdm(total=len(self.testDataset), ncols=80) as t:
for img_id in range(len(self.testDataset)):
img_org, img_norm, labels_org, boxes_org, ratio, dh, dw = self.testDataset[
img_id]
# [3, 416, 416] -> [1, 3, 416, 416]
img_norm = img_norm.unsqueeze(0)
img_in = img_norm.to(self.opt_test.device)
boxes, confs, probs = self.trainer.predict(img_in)
boxes = boxes.squeeze(0) # [13*13*5, 4] / Tensor
confs = confs.squeeze(0) # [13*13*5, 1] / Tensor
probs = probs.squeeze(0) # [13*13*5, 20] / Tensor
cls_scores = (confs * probs).squeeze(
0) # [13*13*5, 20] / Tensor
cls_scores, labels = cls_scores.max(
dim=-1, keepdim=True) # [13*13*5, 1] / Tensor
labels = labels.squeeze()
# box_out: [xmin, ymin, xmax, ymax]
score_thresh = self.opt_test.nms_score_thresh
iou_thresh = self.opt_test.nms_iou_thresh
max_box_num = self.opt_test.nms_max_box_num
keep_index = gpu_nms(boxes, cls_scores, score_thresh,
iou_thresh, max_box_num)
cls_scores = cls_scores.squeeze()
keep_boxes = boxes[keep_index].detach().cpu().numpy()
keep_scores = cls_scores[keep_index].detach().cpu().numpy()
keep_labels = labels[keep_index].detach().cpu().numpy()
keep_boxes = inverse_letter_resize(keep_boxes, ratio, dh, dw)
keep_index = self.filter_small_box(keep_boxes, 16)
keep_boxes = keep_boxes[keep_index]
keep_labels = keep_labels[keep_index]
keep_scores = keep_scores[keep_index]
if len(keep_boxes) > 0:
pred_map_boxe = np.hstack(
[keep_boxes,
np.expand_dims(keep_scores, axis=-1)])
else:
pred_map_boxe = np.zeros(shape=[1, 5])
gt_boxes.append(boxes_org)
pred_boxes.append(pred_map_boxe)
t.update(1)
pickle.dump(
{
'gt_boxes': gt_boxes,
'pred_boxes': pred_boxes
},
open('/home/dk/ML/V2/VOC/result/yolov3_loss_nms_all_class_mAP.pkl',
'wb'))
map = mAP(predict=pred_boxes, ground_truth=gt_boxes, iou_threshold=0.5)
print(map.elevenPointAP * 100)
def nms_all_class(boxes, scores, score_threshold, iou_threshold,
max_box_num):
"""
:param boxes: [13*13*5, 4]
:param scores: [13*13*5, 20]
:param score_threshold: 0.3
:param iou_threshold: 0.45
:param max_box_num:
:return:
boxes_output shape: [X, 4]
scores_output shape: [X,]
labels_output shape: [X,]
"""
assert boxes.dim() == 2 and scores.dim() == 2
boxes = boxes.clamp(0., opt.img_size)
boxes_output = []
scores_output = []
labels_output = []
# [13*13*5, 20] -> [13*13*5, 1]
scores_mask = scores.max(dim=-1)
labels = scores_mask[1]
max_scores = scores_mask[0]
# [13*13*5, 1]
valid_mask = max_scores.ge(score_threshold)
# do nms for all class
if valid_mask.sum() != 0:
valid_boxes = boxes[valid_mask] # [M, 4]
valid_scores = max_scores[valid_mask] # [M, 1]
valid_labels = labels[valid_mask]
keep_index = gpu_nms(valid_boxes, valid_scores, iou_threshold)
for keep_box in valid_boxes[keep_index]:
boxes_output.append(keep_box)
scores_output.extend(valid_scores[keep_index])
labels_output.extend(valid_labels[keep_index])
assert len(boxes_output) == len(scores_output) == len(labels_output)
num_out = len(labels_output)
if num_out == 0:
return torch.tensor([], device=opt.device), torch.tensor(
[], device=opt.device), torch.tensor([], device=opt.device)
else:
boxes_output = torch.stack(boxes_output, dim=0)
scores_output = torch.tensor(scores_output)
labels_output = torch.tensor(labels_output)
assert boxes_output.dim() == 2
assert labels_output.dim() == scores_output.dim() == 1
assert boxes_output.size(
0) == scores_output.numel() == labels_output.numel()
if num_out > max_box_num:
descend_order_index = torch.argsort(scores_output)[::-1]
output_index = descend_order_index[:max_box_num]
else:
output_index = torch.arange(num_out)
return boxes_output[output_index], scores_output[
output_index], labels_output[output_index]
def nms_each_class(boxes, scores, score_threshold, iou_threshold,
max_box_num):
"""
:param boxes: [13*13*5, 4]
:param scores: [13*13*5, 20]
:param score_threshold:
:param iou_threshold:
:param max_box_num:
:return:
boxes_output shape: [X, 4]
scores_output shape: [X,]
labels_output shape: [X,]
"""
assert boxes.dim() == 2 and scores.dim() == 2
boxes = boxes.clamp(0., opt.img_size)
boxes_output = []
scores_output = []
labels_output = []
# [13*13*5, 20]
score_mask = scores.ge(score_threshold)
# do nms for each class
for k in range(opt.coco_class_num):
valid_mask = score_mask[:, k] # [M, 20]
if valid_mask.sum() == 0:
continue
else:
valid_boxes = boxes[valid_mask] # [M, 4]
valid_scores = scores[:, k][valid_mask] # [M, 1]
keep_index = gpu_nms(valid_boxes, valid_scores, iou_threshold)
for keep_box in valid_boxes[keep_index]:
boxes_output.append(keep_box)
scores_output.extend(valid_scores[keep_index])
labels_output.extend([k for _ in range(len(keep_index))])
assert len(boxes_output) == len(scores_output) == len(labels_output)
num_out = len(labels_output)
if num_out == 0:
return torch.tensor([], device=opt.device), torch.tensor(
[], device=opt.device), torch.tensor([], device=opt.device)
else:
boxes_output = torch.stack(boxes_output, dim=0)
scores_output = torch.tensor(scores_output)
labels_output = torch.tensor(labels_output)
assert boxes_output.dim() == 2
assert labels_output.dim() == scores_output.dim() == 1
assert boxes_output.size(
0) == scores_output.numel() == labels_output.numel()
if num_out > max_box_num:
descend_order_index = torch.argsort(scores_output)[::-1]
output_index = descend_order_index[:max_box_num]
else:
output_index = torch.arange(num_out)
return boxes_output[output_index], scores_output[
output_index], labels_output[output_index]
def nms_all_class(boxes, scores, labels, score_threshold, iou_threshold, max_box_num):
"""
:param boxes: [300, 4]
:param scores: [300,]
:param score_threshold: 0.3
:param iou_threshold: 0.45
:param max_box_num:
:return:
boxes_output shape: [X, 4]
scores_output shape: [X,]
labels_output shape: [X,]
"""
# assert boxes.dim() == 2 and scores.dim() == 2
boxes_output = []
scores_output = []
labels_output = []
max_scores = scores
# [13*13*5, 1]
valid_mask = max_scores.ge(score_threshold)
# do nms for all class
if valid_mask.sum() != 0:
valid_boxes = boxes[valid_mask] # [M, 4]
valid_scores = max_scores[valid_mask] # [M, 1]
valid_labels = labels[valid_mask]
keep_index = gpu_nms(valid_boxes, valid_scores, iou_threshold)
boxes_output = valid_boxes[keep_index]
scores_output.extend(valid_scores[keep_index])
labels_output.extend(valid_labels[keep_index])
num_out = len(boxes_output)
if num_out == 0:
return torch.tensor([], device='cuda'), torch.tensor([], device='cuda'), torch.tensor([], device='cuda')
else:
boxes_output = torch.tensor(boxes_output)
scores_output = torch.tensor(scores_output)
labels_output = torch.tensor(labels_output)
if num_out > max_box_num:
descend_order_index = torch.argsort(scores_output)[::-1]
output_index = descend_order_index[:max_box_num]
else:
output_index = torch.arange(num_out)
return boxes_output[output_index], scores_output[output_index], labels_output[output_index]
def eval(self, epoch, step):
test_imgs = np.random.randint(low=0,
high=len(self.testDataset),
size=7)
for img_id in test_imgs:
_, _, img_org, labels, boxes = self.testDataset.get_example_for_testing(
img_id)
img_resized, _, ratio, dh, dw = letterbox_resize(
img_org, boxes, [416, 416])
img = self.normailze(img_resized)
# [3, 416, 416] -> [1, 3, 416, 416]
img = torch.unsqueeze(img, dim=0)
img = img.to(self.opt_train.device)
boxes, confs, probs = self.trainer.predict(img)
boxes = boxes.squeeze(0) # [13*13*5, 4] / Tensor
confs = confs.squeeze(0) # [13*13*5, 1] / Tensor
probs = probs.squeeze(0) # [13*13*5, 20] / Tensor
cls_scores = (confs * probs).squeeze(0) # [13*13*5, 20] / Tensor
cls_scores, labels = cls_scores.max(
dim=-1, keepdim=True) # [13*13*5, 1] / Tensor
labels = labels.squeeze()
# box_out: [xmin, ymin, xmax, ymax]
score_thresh = self.opt_train.nms_score_thresh
iou_thresh = self.opt_train.nms_iou_thresh
max_box_num = self.opt_train.nms_max_box_num
keep_index = gpu_nms(boxes, cls_scores, score_thresh, iou_thresh,
max_box_num)
cls_scores = cls_scores.squeeze()
keep_boxes = boxes[keep_index].detach().cpu().numpy()
keep_scores = cls_scores[keep_index].detach().cpu().numpy()
keep_labels = labels[keep_index].detach().cpu().numpy()
keep_boxes = inverse_letter_resize(keep_boxes, ratio, dh, dw)
keep_index = self.fliter_small_box(keep_boxes, 16)
keep_boxes = keep_boxes[keep_index]
keep_labels = keep_labels[keep_index]
keep_scores = keep_scores[keep_index]
imgname = self.opt_train.result_img_dir / f'E{epoch}_S{step}_img{img_id}.png'
cv2_savefig(img_org, keep_boxes, keep_labels, keep_scores, imgname)
def eval(self):
with tqdm(total=len(self.testDataset), ncols=80) as t:
for img_id in range(len(self.testDataset)):
img_org, img_norm, labels_org, boxes_org, ratio, dh, dw = self.testDataset[
img_id]
# [3, 416, 416] -> [1, 3, 416, 416]
img_norm = img_norm.unsqueeze(0)
img_in = img_norm.to(self.opt_test.device)
boxes, confs, probs = self.trainer.predict(img_in)
boxes = boxes.squeeze(0) # [13*13*5, 4] / Tensor
confs = confs.squeeze(0) # [13*13*5, 1] / Tensor
probs = probs.squeeze(0) # [13*13*5, 20] / Tensor
cls_scores = (confs * probs).squeeze(
0) # [13*13*5, 20] / Tensor
cls_scores, labels = cls_scores.max(
dim=-1, keepdim=True) # [13*13*5, 1] / Tensor
labels = labels.squeeze()
# box_out: [xmin, ymin, xmax, ymax]
score_thresh = self.opt_test.nms_score_thresh
iou_thresh = self.opt_test.nms_iou_thresh
max_box_num = self.opt_test.nms_max_box_num
keep_index = gpu_nms(boxes, cls_scores, score_thresh,
iou_thresh, max_box_num)
cls_scores = cls_scores.squeeze()
keep_boxes = boxes[keep_index].detach().cpu().numpy()
keep_scores = cls_scores[keep_index].detach().cpu().numpy()
keep_labels = labels[keep_index].detach().cpu().numpy()
keep_boxes = inverse_letter_resize(keep_boxes, ratio, dh, dw)
keep_index = self.filter_small_box(keep_boxes, 16)
keep_boxes = keep_boxes[keep_index]
keep_labels = keep_labels[keep_index]
keep_scores = keep_scores[keep_index]
imgname = Path('/home/dk/Desktop/Yolov2/v3loss/nms_all_class'
) / f'{img_id}.png'
cv2_savefig(img_org, keep_boxes, keep_labels, keep_scores,
imgname)
t.update(1)
sess = tf.Session(graph=graph)
inputs = tf.placeholder(
tf.float32,
[1, image_h, image_w, 3]) # placeholder for detector inputs
print("=>", inputs)
with tf.variable_scope('yolov3-tiny'):
feature_map = model.forward(inputs, is_training=False)
boxes, confs, probs = model.predict(feature_map)
scores = confs * probs
print("=>", boxes.name[:-2], scores.name[:-2])
cpu_out_node_names = [boxes.name[:-2], scores.name[:-2]]
boxes, scores, labels = utils.gpu_nms(boxes,
scores,
num_classes,
score_thresh=0.5,
iou_thresh=0.5)
print("=>", boxes.name[:-2], scores.name[:-2], labels.name[:-2])
gpu_out_node_names = [
boxes.name[:-2], scores.name[:-2], labels.name[:-2]
]
feature_map_1, feature_map_2, feature_map_3 = feature_map
saver = tf.train.Saver(var_list=tf.global_variables(
scope='yolov3-tiny'))
saver.restore(sess, ckpt_file)
print('=> checkpoint file restored from ', ckpt_file)
utils.freeze_graph(sess, './pb/yolov3_cpu_nms_tiny_mix_v2_low_q_5k.pb',
cpu_out_node_names)
utils.freeze_graph(sess, './pb/yolov3_gpu_nms_tiny_mix_v2_low_q_5k.pb',