Select boxes that contain human faces Args: width: original image width height: original image height confidences (N, 2): confidence array boxes (N, 4): boxes array in corner-form iou_threshold: intersection over union threshold. top_k: keep t
(width, height, confidences, boxes, prob_threshold, iou_threshold=0.5, top_k=-1)
| 64 | return box_scores[picked, :] |
| 65 | |
| 66 | def predict(width, height, confidences, boxes, prob_threshold, iou_threshold=0.5, top_k=-1): |
| 67 | """ |
| 68 | Select boxes that contain human faces |
| 69 | Args: |
| 70 | width: original image width |
| 71 | height: original image height |
| 72 | confidences (N, 2): confidence array |
| 73 | boxes (N, 4): boxes array in corner-form |
| 74 | iou_threshold: intersection over union threshold. |
| 75 | top_k: keep top_k results. If k <= 0, keep all the results. |
| 76 | Returns: |
| 77 | boxes (k, 4): an array of boxes kept |
| 78 | labels (k): an array of labels for each boxes kept |
| 79 | probs (k): an array of probabilities for each boxes being in corresponding labels |
| 80 | """ |
| 81 | boxes = boxes[0] |
| 82 | confidences = confidences[0] |
| 83 | #print(boxes) |
| 84 | #print(confidences) |
| 85 | |
| 86 | picked_box_probs = [] |
| 87 | picked_labels = [] |
| 88 | for class_index in range(1, confidences.shape[1]): |
| 89 | #print(confidences.shape[1]) |
| 90 | probs = confidences[:, class_index] |
| 91 | #print(probs) |
| 92 | mask = probs > prob_threshold |
| 93 | probs = probs[mask] |
| 94 | |
| 95 | if probs.shape[0] == 0: |
| 96 | continue |
| 97 | subset_boxes = boxes[mask, :] |
| 98 | #print(subset_boxes) |
| 99 | box_probs = np.concatenate([subset_boxes, probs.reshape(-1, 1)], axis=1) |
| 100 | box_probs = hard_nms(box_probs, |
| 101 | iou_threshold=iou_threshold, |
| 102 | top_k=top_k, |
| 103 | ) |
| 104 | picked_box_probs.append(box_probs) |
| 105 | picked_labels.extend([class_index] * box_probs.shape[0]) |
| 106 | if not picked_box_probs: |
| 107 | return np.array([]), np.array([]), np.array([]) |
| 108 | picked_box_probs = np.concatenate(picked_box_probs) |
| 109 | picked_box_probs[:, 0] *= width |
| 110 | picked_box_probs[:, 1] *= height |
| 111 | picked_box_probs[:, 2] *= width |
| 112 | picked_box_probs[:, 3] *= height |
| 113 | return picked_box_probs[:, :4].astype(np.int32), np.array(picked_labels), picked_box_probs[:, 4] |
no test coverage detected