def _filter_detections(scores, labels):
# threshold based on score
indices = backend.where(keras.backend.greater(scores, score_threshold))
if nms:
filtered_boxes = backend.gather_nd(boxes, indices)
filtered_scores = keras.backend.gather(scores, indices)[:, 0]
# perform NMS
nms_indices = backend.non_max_suppression(
filtered_boxes,
filtered_scores,
max_output_size=max_detections,
iou_threshold=nms_threshold)
# filter indices based on NMS
indices = keras.backend.gather(indices, nms_indices)
# add indices to list of all indices
labels = backend.gather_nd(labels, indices)
indices = keras.backend.stack([indices[:, 0], labels], axis=1)
return indices
def filter_detections(boxes,
classification,
other=[],
nms=True,
score_threshold=0.05,
max_detections=300,
nms_threshold=0.5):
""" Filter detections using the boxes and classification values.
Args
boxes : Tensor of shape (num_boxes, 4) containing the boxes in (x1, y1, x2, y2) format.
classification : Tensor of shape (num_boxes, num_classes) containing the classification scores.
other : List of tensors of shape (num_boxes, ...) to filter along with the boxes and classification scores.
nms : Flag to enable/disable non maximum suppression.
score_threshold : Threshold used to prefilter the boxes with.
max_detections : Maximum number of detections to keep.
nms_threshold : Threshold for the IoU value to determine when a box should be suppressed.
Returns
A list of [boxes, scores, labels, other[0], other[1], ...].
boxes is shaped (max_detections, 4) and contains the (x1, y1, x2, y2) of the non-suppressed boxes.
scores is shaped (max_detections,) and contains the scores of the predicted class.
labels is shaped (max_detections,) and contains the predicted label.
other[i] is shaped (max_detections, ...) and contains the filtered other[i] data.
In case there are less than max_detections detections, the tensors are padded with -1's.
"""
all_boxes = []
all_scores = []
all_labels = []
all_other = []
# perform per class filtering
for c in range(int(classification.shape[1])):
scores = classification[:, c]
# threshold based on score
score_indices = backend.where(
keras.backend.greater(scores, score_threshold))
score_indices = keras.backend.cast(score_indices, 'int32')
filtered_boxes = backend.gather_nd(boxes, score_indices)
filtered_scores = keras.backend.gather(scores, score_indices)[:, 0]
filtered_other = [backend.gather_nd(o, score_indices) for o in other]
if nms:
# perform NMS
nms_indices = backend.non_max_suppression(
filtered_boxes,
filtered_scores,
max_output_size=max_detections,
iou_threshold=nms_threshold)
# filter NMS detections
filtered_boxes = keras.backend.gather(filtered_boxes, nms_indices)
filtered_scores = keras.backend.gather(filtered_scores,
nms_indices)
filtered_other = [
keras.backend.gather(o, nms_indices) for o in filtered_other
]
# labels is a vector of the current class label
filtered_labels = c * keras.backend.ones(
(keras.backend.shape(filtered_scores)[0], ), dtype='int32')
# append to lists
all_boxes.append(filtered_boxes)
all_scores.append(filtered_scores)
all_labels.append(filtered_labels)
all_other.append(filtered_other)
# concatenate outputs to single tensors
boxes = keras.backend.concatenate(all_boxes, axis=0)
scores = keras.backend.concatenate(all_scores, axis=0)
labels = keras.backend.concatenate(all_labels, axis=0)
other_ = [
keras.backend.concatenate([o[i] for o in all_other], axis=0)
for i in range(len(other))
]
# select top k
scores, top_indices = backend.top_k(scores,
k=keras.backend.minimum(
max_detections,
keras.backend.shape(scores)[0]))
boxes = keras.backend.gather(boxes, top_indices)
labels = keras.backend.gather(labels, top_indices)
other_ = [keras.backend.gather(o, top_indices) for o in other_]
# zero pad the outputs
pad_size = keras.backend.maximum(
0, max_detections - keras.backend.shape(scores)[0])
boxes = backend.pad(boxes, [[0, pad_size], [0, 0]], constant_values=-1)
scores = backend.pad(scores, [[0, pad_size]], constant_values=-1)
labels = backend.pad(labels, [[0, pad_size]], constant_values=-1)
labels = keras.backend.cast(labels, 'int32')
other_ = [
backend.pad(o, [[0, pad_size]] + [[0, 0]
for _ in range(1, len(o.shape))],
constant_values=-1) for o in other_
]
# set shapes, since we know what they are
boxes.set_shape([max_detections, 4])
scores.set_shape([max_detections])
labels.set_shape([max_detections])
for o, s in zip(other_, [list(keras.backend.int_shape(o)) for o in other]):
o.set_shape([max_detections] + s[1:])
return [boxes, scores, labels] + other_