diff --git a/README.md b/README.md index 8a2c004..e928ac0 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,10 @@ boxes, scores, labels = weighted_boxes_fusion([boxes_list], [scores_list], [labe More examples can be found in [example.py](./example.py) +#### 3D version + +There is support for 3D boxes in WBF method with `weighted_boxes_fusion_3d` function. Check example of usage in [example_3d.py](./example_3d.py) + ## Accuracy and speed comparison Comparison was made for ensemble of 5 different object detection models predictions trained on [Open Images Dataset](https://storage.googleapis.com/openimages/web/index.html) (500 classes). diff --git a/ensemble_boxes/__init__.py b/ensemble_boxes/__init__.py index 383c53b..62ce75e 100644 --- a/ensemble_boxes/__init__.py +++ b/ensemble_boxes/__init__.py @@ -6,3 +6,4 @@ from .ensemble_boxes_nms import nms_method from .ensemble_boxes_nms import nms from .ensemble_boxes_nms import soft_nms +from .ensemble_boxes_wbf_3d import weighted_boxes_fusion_3d \ No newline at end of file diff --git a/ensemble_boxes/ensemble_boxes_wbf_3d.py b/ensemble_boxes/ensemble_boxes_wbf_3d.py new file mode 100644 index 0000000..099a3eb --- /dev/null +++ b/ensemble_boxes/ensemble_boxes_wbf_3d.py @@ -0,0 +1,222 @@ +# coding: utf-8 +__author__ = 'ZFTurbo: https://kaggle.com/zfturbo' + + +import warnings +import numpy as np +from numba import jit + + +@jit(nopython=True) +def bb_intersection_over_union_3d(A, B) -> float: + xA = max(A[0], B[0]) + yA = max(A[1], B[1]) + zA = max(A[2], B[2]) + xB = min(A[3], B[3]) + yB = min(A[4], B[4]) + zB = min(A[5], B[5]) + + interVol = max(0, xB - xA) * max(0, yB - yA) * max(0, zB - zA) + if interVol == 0: + return 0.0 + + # compute the volume of both the prediction and ground-truth rectangular boxes + boxAVol = (A[3] - A[0]) * (A[4] - A[1]) * (A[5] - A[2]) + boxBVol = (B[3] - B[0]) * (B[4] - B[1]) * (B[5] - B[2]) + + iou = interVol / float(boxAVol + boxBVol - interVol) + return iou + + +def prefilter_boxes(boxes, scores, labels, weights, thr): + # Create dict with boxes stored by its label + new_boxes = dict() + + for t in range(len(boxes)): + + if len(boxes[t]) != len(scores[t]): + print('Error. Length of boxes arrays not equal to length of scores array: {} != {}'.format(len(boxes[t]), len(scores[t]))) + exit() + + if len(boxes[t]) != len(labels[t]): + print('Error. Length of boxes arrays not equal to length of labels array: {} != {}'.format(len(boxes[t]), len(labels[t]))) + exit() + + for j in range(len(boxes[t])): + score = scores[t][j] + if score < thr: + continue + label = int(labels[t][j]) + box_part = boxes[t][j] + x1 = float(box_part[0]) + y1 = float(box_part[1]) + z1 = float(box_part[2]) + x2 = float(box_part[3]) + y2 = float(box_part[4]) + z2 = float(box_part[5]) + + # Box data checks + if x2 < x1: + warnings.warn('X2 < X1 value in box. Swap them.') + x1, x2 = x2, x1 + if y2 < y1: + warnings.warn('Y2 < Y1 value in box. Swap them.') + y1, y2 = y2, y1 + if z2 < z1: + warnings.warn('Z2 < Z1 value in box. Swap them.') + z1, z2 = z2, z1 + if x1 < 0: + warnings.warn('X1 < 0 in box. Set it to 0.') + x1 = 0 + if x1 > 1: + warnings.warn('X1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.') + x1 = 1 + if x2 < 0: + warnings.warn('X2 < 0 in box. Set it to 0.') + x2 = 0 + if x2 > 1: + warnings.warn('X2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.') + x2 = 1 + if y1 < 0: + warnings.warn('Y1 < 0 in box. Set it to 0.') + y1 = 0 + if y1 > 1: + warnings.warn('Y1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.') + y1 = 1 + if y2 < 0: + warnings.warn('Y2 < 0 in box. Set it to 0.') + y2 = 0 + if y2 > 1: + warnings.warn('Y2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.') + y2 = 1 + if z1 < 0: + warnings.warn('Z1 < 0 in box. Set it to 0.') + z1 = 0 + if z1 > 1: + warnings.warn('Z1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.') + z1 = 1 + if z2 < 0: + warnings.warn('Z2 < 0 in box. Set it to 0.') + z2 = 0 + if z2 > 1: + warnings.warn('Z2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.') + z2 = 1 + if (x2 - x1) * (y2 - y1) * (z2 - z1) == 0.0: + warnings.warn("Zero volume box skipped: {}.".format(box_part)) + continue + + b = [int(label), float(score) * weights[t], x1, y1, z1, x2, y2, z2] + if label not in new_boxes: + new_boxes[label] = [] + new_boxes[label].append(b) + + # Sort each list in dict by score and transform it to numpy array + for k in new_boxes: + current_boxes = np.array(new_boxes[k]) + new_boxes[k] = current_boxes[current_boxes[:, 1].argsort()[::-1]] + + return new_boxes + + +def get_weighted_box(boxes, conf_type='avg'): + """ + Create weighted box for set of boxes + :param boxes: set of boxes to fuse + :param conf_type: type of confidence one of 'avg' or 'max' + :return: weighted box + """ + + box = np.zeros(8, dtype=np.float32) + conf = 0 + conf_list = [] + for b in boxes: + box[2:] += (b[1] * b[2:]) + conf += b[1] + conf_list.append(b[1]) + box[0] = boxes[0][0] + if conf_type == 'avg': + box[1] = conf / len(boxes) + elif conf_type == 'max': + box[1] = np.array(conf_list).max() + box[2:] /= conf + return box + + +def find_matching_box(boxes_list, new_box, match_iou): + best_iou = match_iou + best_index = -1 + for i in range(len(boxes_list)): + box = boxes_list[i] + if box[0] != new_box[0]: + continue + iou = bb_intersection_over_union_3d(box[2:], new_box[2:]) + if iou > best_iou: + best_index = i + best_iou = iou + + return best_index, best_iou + + +def weighted_boxes_fusion_3d(boxes_list, scores_list, labels_list, weights=None, iou_thr=0.55, skip_box_thr=0.0, conf_type='avg', allows_overflow=False): + ''' + :param boxes_list: list of boxes predictions from each model, each box is 6 numbers. + It has 3 dimensions (models_number, model_preds, 6) + Order of boxes: x1, y1, z1, x2, y2 z2. We expect float normalized coordinates [0; 1] + :param scores_list: list of scores for each model + :param labels_list: list of labels for each model + :param weights: list of weights for each model. Default: None, which means weight == 1 for each model + :param iou_thr: IoU value for boxes to be a match + :param skip_box_thr: exclude boxes with score lower than this variable + :param conf_type: how to calculate confidence in weighted boxes. 'avg': average value, 'max': maximum value + :param allows_overflow: false if we want confidence score not exceed 1.0 + + :return: boxes: boxes coordinates (Order of boxes: x1, y1, z1, x2, y2, z2). + :return: scores: confidence scores + :return: labels: boxes labels + ''' + + if weights is None: + weights = np.ones(len(boxes_list)) + if len(weights) != len(boxes_list): + print('Warning: incorrect number of weights {}. Must be: {}. Set weights equal to 1.'.format(len(weights), len(boxes_list))) + weights = np.ones(len(boxes_list)) + weights = np.array(weights) + + if conf_type not in ['avg', 'max']: + print('Error. Unknown conf_type: {}. Must be "avg" or "max". Use "avg"'.format(conf_type)) + conf_type = 'avg' + + filtered_boxes = prefilter_boxes(boxes_list, scores_list, labels_list, weights, skip_box_thr) + if len(filtered_boxes) == 0: + return np.zeros((0, 6)), np.zeros((0,)), np.zeros((0,)) + + overall_boxes = [] + for label in filtered_boxes: + boxes = filtered_boxes[label] + new_boxes = [] + weighted_boxes = [] + + # Clusterize boxes + for j in range(0, len(boxes)): + index, best_iou = find_matching_box(weighted_boxes, boxes[j], iou_thr) + if index != -1: + new_boxes[index].append(boxes[j]) + weighted_boxes[index] = get_weighted_box(new_boxes[index], conf_type) + else: + new_boxes.append([boxes[j].copy()]) + weighted_boxes.append(boxes[j].copy()) + + # Rescale confidence based on number of models and boxes + for i in range(len(new_boxes)): + if not allows_overflow: + weighted_boxes[i][1] = weighted_boxes[i][1] * min(weights.sum(), len(new_boxes[i])) / weights.sum() + else: + weighted_boxes[i][1] = weighted_boxes[i][1] * len(new_boxes[i]) / weights.sum() + overall_boxes.append(np.array(weighted_boxes)) + + overall_boxes = np.concatenate(overall_boxes, axis=0) + overall_boxes = overall_boxes[overall_boxes[:, 1].argsort()[::-1]] + boxes = overall_boxes[:, 2:] + scores = overall_boxes[:, 1] + labels = overall_boxes[:, 0] + return boxes, scores, labels diff --git a/example_3d.py b/example_3d.py new file mode 100644 index 0000000..fe8d00e --- /dev/null +++ b/example_3d.py @@ -0,0 +1,142 @@ +# coding: utf-8 +__author__ = 'ZFTurbo: https://kaggle.com/zfturbo' + + +import numpy as np +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d.art3d import Poly3DCollection +from ensemble_boxes import * + + +def plot_cube(ax, cube_definition, lbl, thickness): + cube_definition_array = [ + np.array(list(item)) + for item in cube_definition + ] + + points = [] + points += cube_definition_array + vectors = [ + cube_definition_array[1] - cube_definition_array[0], + cube_definition_array[2] - cube_definition_array[0], + cube_definition_array[3] - cube_definition_array[0] + ] + + points += [cube_definition_array[0] + vectors[0] + vectors[1]] + points += [cube_definition_array[0] + vectors[0] + vectors[2]] + points += [cube_definition_array[0] + vectors[1] + vectors[2]] + points += [cube_definition_array[0] + vectors[0] + vectors[1] + vectors[2]] + + points = np.array(points) + + edges = [ + [points[0], points[3], points[5], points[1]], + [points[1], points[5], points[7], points[4]], + [points[4], points[2], points[6], points[7]], + [points[2], points[6], points[3], points[0]], + [points[0], points[2], points[4], points[1]], + [points[3], points[6], points[7], points[5]] + ] + + faces = Poly3DCollection(edges, linewidths=thickness + 1) + if lbl == 0: + faces.set_edgecolor((1, 0, 0)) + else: + faces.set_edgecolor((0, 0, 1)) + faces.set_facecolor((0, 0, 1, 0.1)) + + ax.add_collection3d(faces) + ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0) + + +def show_boxes(boxes_list, scores_list, labels_list, image_size=800): + image = np.zeros((image_size, image_size, 3), dtype=np.uint8) + image[...] = 255 + fig = plt.figure() + ax = fig.add_subplot(111, projection='3d') + + for i in range(len(boxes_list)): + for j in range(len(boxes_list[i])): + x1 = int(image_size * boxes_list[i][j][0]) + y1 = int(image_size * boxes_list[i][j][1]) + z1 = int(image_size * boxes_list[i][j][2]) + x2 = int(image_size * boxes_list[i][j][3]) + y2 = int(image_size * boxes_list[i][j][4]) + z2 = int(image_size * boxes_list[i][j][5]) + lbl = labels_list[i][j] + cube_definition = [ + (x1, y1, z1), (x1, y2, z1), (x2, y1, z1), (x1, y1, z2) + ] + plot_cube(ax, cube_definition, lbl, int(4 * scores_list[i][j])) + + plt.show() + + +def example_wbf_3d_2_models(iou_thr=0.55, draw_image=True): + """ + This example shows how to ensemble boxes from 2 models using WBF_3D method + :return: + """ + + boxes_list = [ + [ + [0.00, 0.51, 0.41, 0.81, 0.91, 0.78], + [0.10, 0.31, 0.45, 0.71, 0.61, 0.85], + [0.01, 0.32, 0.55, 0.83, 0.93, 0.95], + [0.02, 0.53, 0.11, 0.11, 0.94, 0.55], + [0.03, 0.24, 0.34, 0.12, 0.35, 0.67], + ], + [ + [0.04, 0.56, 0.36, 0.84, 0.92, 0.82], + [0.12, 0.33, 0.46, 0.72, 0.64, 0.75], + [0.38, 0.66, 0.55, 0.79, 0.95, 0.90], + [0.08, 0.49, 0.15, 0.21, 0.89, 0.67], + ], + ] + scores_list = [ + [ + 0.9, + 0.8, + 0.2, + 0.4, + 0.7, + ], + [ + 0.5, + 0.8, + 0.7, + 0.3, + ] + ] + labels_list = [ + [ + 0, + 1, + 0, + 1, + 1, + ], + [ + 1, + 1, + 1, + 0, + ] + ] + weights = [2, 1] + if draw_image: + show_boxes(boxes_list, scores_list, labels_list) + + boxes, scores, labels = weighted_boxes_fusion_3d(boxes_list, scores_list, labels_list, weights=weights, iou_thr=iou_thr, skip_box_thr=0.0) + + if draw_image: + show_boxes([boxes], [scores], [labels.astype(np.int32)]) + + print(len(boxes)) + print(boxes) + + +if __name__ == '__main__': + draw_image = True + example_wbf_3d_2_models(iou_thr=0.2, draw_image=draw_image) + diff --git a/setup.py b/setup.py index a497469..5d3cb34 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='ensemble_boxes', - version='1.0.3', + version='1.0.4', author='Roman Solovyev (ZFTurbo)', packages=['ensemble_boxes', ], url='https://github.com/ZFTurbo/Weighted-Boxes-Fusion',