def test(ssd, data, default_boxes):
""" Execute test on one image
then perform NMS algorithm
then rescale boxes back to normal size
Args:
ssd: trained SSD model
img: Torch array of shape (1, 3, H, W)
default_boxes: Torch array of shape (num_default, 4)
Returns:
img: numpy array of (3, H, W)
all_boxes: final boxes, numpy array of (num_boxes, 4)
all_scores: final scores, numpy array of (num_boxes,)
all_names: final class names, numpy array of (num_boxes,)
"""
img, _, _ = data
img = img.to(device)
default_boxes = default_boxes.to(device)
with torch.no_grad():
out_confs, out_locs = ssd(img)
out_confs = out_confs.squeeze(0)
out_locs = out_locs.squeeze(0)
out_boxes = decode(default_boxes, out_locs)
out_labels = F.softmax(out_confs, dim=1)
all_boxes = []
all_scores = []
all_names = []
for c in range(1, NUM_CLASSES):
cls_scores = out_labels[:, c]
score_idx = cls_scores > args.score_thresh
cls_boxes = out_boxes[score_idx]
cls_scores = cls_scores[score_idx]
box_idx = compute_nms(cls_boxes, cls_scores, args.nms_thresh,
args.max_num_boxes_per_class)
cls_boxes = cls_boxes[box_idx]
cls_scores = cls_scores[box_idx]
cls_names = [c] * cls_boxes.size(0)
all_boxes.append(cls_boxes)
all_scores.append(cls_scores)
all_names.extend(cls_names)
all_boxes = torch.cat(all_boxes, dim=0)
all_scores = torch.cat(all_scores, dim=0)
img = img.squeeze(0).cpu().numpy()
all_boxes *= img.shape[-1]
all_boxes = all_boxes.cpu().numpy()
all_scores = all_scores.cpu().numpy()
all_names = np.array(all_names)
return img, all_boxes, all_scores, all_names
def post_process(confs, locs, scores, default_boxes, mode=1):
# confs = tf.squeeze(confs, 0)
# locs = tf.squeeze(locs, 0)
# i have o return the locs back
# print(scores)
# print(confs)
# print(locs)
print(tf.math.reduce_max(locs), tf.math.reduce_min(locs))
newres = decode(default_boxes, locs).numpy()
if (mode == 2):
confs = tf.math.softmax(confs, axis=-1)
classes = tf.math.argmax(confs, axis=-1)
scores = tf.math.reduce_max(confs, axis=-1)
out_boxes = []
out_labels = []
out_scores = []
# print(confs.shape,classes.shape, scores.shape, boxes.shape)
for c in range(1, NUM_CLASSES):
if (mode == 1):
cls_scores = np.zeros(np.shape(confs))
cls_scores[confs == c] = confs[confs == c]
else:
cls_scores = confs[:, c]
score_idx = cls_scores > 0.5
# cls_boxes = tf.boolean_mask(boxes, score_idx)
# cls_scores = tf.boolean_mask(cls_scores, score_idx)
cls_boxes = newres[score_idx]
cls_scores = cls_scores[score_idx]
nms_idx = compute_nms(cls_boxes, cls_scores, 0.35, 200)
cls_boxes = tf.gather(cls_boxes, nms_idx)
cls_scores = tf.gather(cls_scores, nms_idx)
cls_labels = [c] * cls_boxes.shape[0]
out_boxes.append(cls_boxes)
out_labels.extend(cls_labels)
out_scores.append(cls_scores)
out_boxes = tf.concat(out_boxes, axis=0)
out_scores = tf.concat(out_scores, axis=0)
boxes = tf.clip_by_value(out_boxes, 0.0, 1.0).numpy()
classes = np.array(out_labels)
scores = out_scores.numpy()
return boxes, classes, scores
def predict(imgs, default_boxes):
confs, locs = ssd(imgs)
confs = tf.squeeze(confs, 0)
locs = tf.squeeze(locs, 0)
confs = tf.math.softmax(confs, axis=-1)
classes = tf.math.argmax(confs, axis=-1)
scores = tf.math.reduce_max(confs, axis=-1)
boxes = decode(default_boxes, locs)
out_boxes = []
out_labels = []
out_scores = []
for c in range(1, NUM_CLASSES):
cls_scores = confs[:, c]
score_idx = cls_scores > 0.6
# cls_boxes = tf.boolean_mask(boxes, score_idx)
# cls_scores = tf.boolean_mask(cls_scores, score_idx)
cls_boxes = boxes[score_idx]
cls_scores = cls_scores[score_idx]
nms_idx = compute_nms(cls_boxes, cls_scores, 0.45, 200)
cls_boxes = tf.gather(cls_boxes, nms_idx)
cls_scores = tf.gather(cls_scores, nms_idx)
cls_labels = [c] * cls_boxes.shape[0]
out_boxes.append(cls_boxes)
out_labels.extend(cls_labels)
out_scores.append(cls_scores)
out_boxes = tf.concat(out_boxes, axis=0)
out_scores = tf.concat(out_scores, axis=0)
boxes = tf.clip_by_value(out_boxes, 0.0, 1.0).numpy()
classes = np.array(out_labels)
scores = out_scores.numpy()
return boxes, classes, scores
def predict(confs, locs, thresh, default_boxes):
confs = tf.squeeze(confs, 0)
locs = tf.squeeze(locs, 0)
confs = tf.math.softmax(confs, axis=-1)
boxes = decode(default_boxes, locs)
out_boxes = []
out_labels = []
out_scores = []
print('confs shape =', np.shape(confs))
print('confs max =', np.amax(confs))
print('confs min =', np.amin(confs))
print('confs max index =',
np.where(np.array(confs) == np.array(confs).max()))
for c in range(1, num_classes):
cls_scores = confs[:, c]
score_idx = cls_scores > thresh
cls_boxes = boxes[score_idx]
cls_scores = cls_scores[score_idx]
nms_idx = compute_nms(cls_boxes, cls_scores, 0.45, 200)
cls_boxes = tf.gather(cls_boxes, nms_idx)
cls_scores = tf.gather(cls_scores, nms_idx)
cls_labels = [c] * cls_boxes.shape[0]
out_boxes.append(cls_boxes)
out_labels.extend(cls_labels)
out_scores.append(cls_scores)
out_boxes = tf.concat(out_boxes, axis=0)
out_scores = tf.concat(out_scores, axis=0)
boxes = tf.clip_by_value(out_boxes, 0.0, 1.0).numpy()
classes = np.array(out_labels)
scores = out_scores.numpy()
return boxes, classes, scores
