Spaces:
Runtime error
Runtime error
| from typing import Optional | |
| import numpy as np | |
| def w_np_non_max_suppression( | |
| prediction, | |
| conf_thresh: float = 0.25, | |
| iou_thresh: float = 0.45, | |
| class_agnostic: bool = False, | |
| max_detections: int = 300, | |
| max_candidate_detections: int = 3000, | |
| timeout_seconds: Optional[int] = None, | |
| num_masks: int = 0, | |
| box_format: str = "xywh", | |
| ): | |
| """Applies non-maximum suppression to predictions. | |
| Args: | |
| prediction (np.ndarray): Array of predictions. Format for single prediction is | |
| [bbox x 4, max_class_confidence, (confidence) x num_of_classes, additional_element x num_masks] | |
| conf_thresh (float, optional): Confidence threshold. Defaults to 0.25. | |
| iou_thresh (float, optional): IOU threshold. Defaults to 0.45. | |
| class_agnostic (bool, optional): Whether to ignore class labels. Defaults to False. | |
| max_detections (int, optional): Maximum number of detections. Defaults to 300. | |
| max_candidate_detections (int, optional): Maximum number of candidate detections. Defaults to 3000. | |
| timeout_seconds (Optional[int], optional): Timeout in seconds. Defaults to None. | |
| num_masks (int, optional): Number of masks. Defaults to 0. | |
| box_format (str, optional): Format of bounding boxes. Either 'xywh' or 'xyxy'. Defaults to 'xywh'. | |
| Returns: | |
| list: List of filtered predictions after non-maximum suppression. Format of a single result is: | |
| [bbox x 4, max_class_confidence, max_class_confidence, id_of_class_with_max_confidence, | |
| additional_element x num_masks] | |
| """ | |
| num_classes = prediction.shape[2] - 5 - num_masks | |
| np_box_corner = np.zeros(prediction.shape) | |
| if box_format == "xywh": | |
| np_box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2 | |
| np_box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2 | |
| np_box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2 | |
| np_box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2 | |
| prediction[:, :, :4] = np_box_corner[:, :, :4] | |
| elif box_format == "xyxy": | |
| pass | |
| else: | |
| raise ValueError( | |
| "box_format must be either 'xywh' or 'xyxy', got {}".format(box_format) | |
| ) | |
| batch_predictions = [] | |
| for np_image_i, np_image_pred in enumerate(prediction): | |
| filtered_predictions = [] | |
| np_conf_mask = (np_image_pred[:, 4] >= conf_thresh).squeeze() | |
| np_image_pred = np_image_pred[np_conf_mask] | |
| if np_image_pred.shape[0] == 0: | |
| batch_predictions.append(filtered_predictions) | |
| continue | |
| np_class_conf = np.max(np_image_pred[:, 5 : num_classes + 5], 1) | |
| np_class_pred = np.argmax(np_image_pred[:, 5 : num_classes + 5], 1) | |
| np_class_conf = np.expand_dims(np_class_conf, axis=1) | |
| np_class_pred = np.expand_dims(np_class_pred, axis=1) | |
| np_mask_pred = np_image_pred[:, 5 + num_classes :] | |
| np_detections = np.append( | |
| np.append( | |
| np.append(np_image_pred[:, :5], np_class_conf, axis=1), | |
| np_class_pred, | |
| axis=1, | |
| ), | |
| np_mask_pred, | |
| axis=1, | |
| ) | |
| np_unique_labels = np.unique(np_detections[:, 6]) | |
| if class_agnostic: | |
| np_detections_class = sorted( | |
| np_detections, key=lambda row: row[4], reverse=True | |
| ) | |
| filtered_predictions.extend( | |
| non_max_suppression_fast(np.array(np_detections_class), iou_thresh) | |
| ) | |
| else: | |
| for c in np_unique_labels: | |
| np_detections_class = np_detections[np_detections[:, 6] == c] | |
| np_detections_class = sorted( | |
| np_detections_class, key=lambda row: row[4], reverse=True | |
| ) | |
| filtered_predictions.extend( | |
| non_max_suppression_fast(np.array(np_detections_class), iou_thresh) | |
| ) | |
| filtered_predictions = sorted( | |
| filtered_predictions, key=lambda row: row[4], reverse=True | |
| ) | |
| batch_predictions.append(filtered_predictions[:max_detections]) | |
| return batch_predictions | |
| # Malisiewicz et al. | |
| def non_max_suppression_fast(boxes, overlapThresh): | |
| """Applies non-maximum suppression to bounding boxes. | |
| Args: | |
| boxes (np.ndarray): Array of bounding boxes with confidence scores. | |
| overlapThresh (float): Overlap threshold for suppression. | |
| Returns: | |
| list: List of bounding boxes after non-maximum suppression. | |
| """ | |
| # if there are no boxes, return an empty list | |
| if len(boxes) == 0: | |
| return [] | |
| # if the bounding boxes integers, convert them to floats -- | |
| # this is important since we'll be doing a bunch of divisions | |
| if boxes.dtype.kind == "i": | |
| boxes = boxes.astype("float") | |
| # initialize the list of picked indexes | |
| pick = [] | |
| # grab the coordinates of the bounding boxes | |
| x1 = boxes[:, 0] | |
| y1 = boxes[:, 1] | |
| x2 = boxes[:, 2] | |
| y2 = boxes[:, 3] | |
| conf = boxes[:, 4] | |
| # compute the area of the bounding boxes and sort the bounding | |
| # boxes by the bottom-right y-coordinate of the bounding box | |
| area = (x2 - x1 + 1) * (y2 - y1 + 1) | |
| idxs = np.argsort(conf) | |
| # keep looping while some indexes still remain in the indexes | |
| # list | |
| while len(idxs) > 0: | |
| # grab the last index in the indexes list and add the | |
| # index value to the list of picked indexes | |
| last = len(idxs) - 1 | |
| i = idxs[last] | |
| pick.append(i) | |
| # find the largest (x, y) coordinates for the start of | |
| # the bounding box and the smallest (x, y) coordinates | |
| # for the end of the bounding box | |
| xx1 = np.maximum(x1[i], x1[idxs[:last]]) | |
| yy1 = np.maximum(y1[i], y1[idxs[:last]]) | |
| xx2 = np.minimum(x2[i], x2[idxs[:last]]) | |
| yy2 = np.minimum(y2[i], y2[idxs[:last]]) | |
| # compute the width and height of the bounding box | |
| w = np.maximum(0, xx2 - xx1 + 1) | |
| h = np.maximum(0, yy2 - yy1 + 1) | |
| # compute the ratio of overlap | |
| overlap = (w * h) / area[idxs[:last]] | |
| # delete all indexes from the index list that have | |
| idxs = np.delete( | |
| idxs, np.concatenate(([last], np.where(overlap > overlapThresh)[0])) | |
| ) | |
| # return only the bounding boxes that were picked using the | |
| # integer data type | |
| return boxes[pick].astype("float") | |