Spaces:
Build error
Build error
| # coding: utf-8 | |
| __author__ = 'ZFTurbo: https://kaggle.com/zfturbo' | |
| import warnings | |
| import numpy as np | |
| from numba import jit | |
| 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 | |