fix: 修复 NMSBoxes 索引错误
This commit is contained in:
+17
-9
@@ -189,19 +189,27 @@ def _raw_detect(det, data):
|
||||
scores.append(float(max_score))
|
||||
class_ids.append(class_id)
|
||||
|
||||
indices = cv2.dnn.NMSBoxes(boxes, scores, 0.1, 0.45)
|
||||
# 没有检测到任何类别,直接返回空
|
||||
if not boxes:
|
||||
return []
|
||||
|
||||
# 转成 numpy array,OpenCV 的 NMSBoxes 需要
|
||||
boxes_arr = np.array(boxes, dtype=np.float32)
|
||||
scores_arr = np.array(scores, dtype=np.float32)
|
||||
|
||||
indices = cv2.dnn.NMSBoxes(boxes_arr.tolist(), scores_arr.tolist(), 0.1, 0.45)
|
||||
|
||||
detections = []
|
||||
for i in indices:
|
||||
idx = i if isinstance(i, int) else i[0]
|
||||
detections.append({
|
||||
"class": __labels[class_ids[idx]],
|
||||
"score": scores[idx],
|
||||
"box": [int(boxes[idx][0]), int(boxes[idx][1]), int(boxes[idx][2]), int(boxes[idx][3])],
|
||||
})
|
||||
if indices is not None and len(indices) > 0:
|
||||
for i in np.array(indices).flatten():
|
||||
idx = int(i)
|
||||
detections.append({
|
||||
"class": __labels[class_ids[idx]],
|
||||
"score": scores[idx],
|
||||
"box": [int(boxes[idx][0]), int(boxes[idx][1]), int(boxes[idx][2]), int(boxes[idx][3])],
|
||||
})
|
||||
return detections
|
||||
|
||||
|
||||
# NudeNet 的标签表(和 nudenet.nudenet.__labels 一致)
|
||||
__labels = [
|
||||
"FEMALE_GENITALIA_COVERED", "FACE_FEMALE", "BUTTOCKS_EXPOSED",
|
||||
|
||||
Reference in New Issue
Block a user