71 lines
2.1 KiB
Python
71 lines
2.1 KiB
Python
# -*- coding: utf-8 -*-
|
|
import cv2
|
|
import numpy as np
|
|
from nudenet import NudeDetector
|
|
|
|
# 测试图片路径
|
|
img_path = r"E:\已测试分拣men\数据源\xxxx\图片_横屏_2K\dark_25\e0d12fd3a8c5436f801f766fad4a22fb.jpg"
|
|
|
|
print(f"测试图片: {img_path}")
|
|
|
|
# 加载模型
|
|
det = NudeDetector()
|
|
|
|
# 读取图片
|
|
with open(img_path, 'rb') as f:
|
|
img_bytes = f.read()
|
|
|
|
print(f"文件大小: {len(img_bytes)} bytes")
|
|
|
|
# 解码图片
|
|
mat = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), -1)
|
|
if mat is None:
|
|
print("图片解码失败")
|
|
exit(1)
|
|
|
|
print(f"图片尺寸: {mat.shape}")
|
|
|
|
# 预处理
|
|
mat_c3 = cv2.cvtColor(mat, cv2.COLOR_RGBA2BGR)
|
|
ow, oh = mat_c3.shape[1], mat_c3.shape[0]
|
|
max_size = max(ow, oh)
|
|
x_pad, y_pad = max_size - ow, max_size - oh
|
|
mat_pad = cv2.copyMakeBorder(mat_c3, 0, y_pad, 0, x_pad, cv2.BORDER_CONSTANT)
|
|
blob = cv2.dnn.blobFromImage(mat_pad, 1/255.0, (320, 320), (0, 0, 0), swapRB=True, crop=False)
|
|
|
|
# 推理
|
|
outputs = det.onnx_session.run(None, {det.input_name: blob})
|
|
outs = np.transpose(np.squeeze(outputs[0]))
|
|
|
|
# 标签列表
|
|
labels = [
|
|
'FEMALE_GENITALIA_COVERED', 'FACE_FEMALE', 'BUTTOCKS_EXPOSED',
|
|
'FEMALE_BREAST_EXPOSED', 'FEMALE_GENITALIA_EXPOSED', 'MALE_BREAST_EXPOSED',
|
|
'ANUS_EXPOSED', 'FEET_EXPOSED', 'BELLY_COVERED', 'FEET_COVERED',
|
|
'ARMPITS_COVERED', 'ARMPITS_EXPOSED', 'FACE_MALE', 'BELLY_EXPOSED',
|
|
'MALE_GENITALIA_EXPOSED', 'ANUS_COVERED', 'FEMALE_BREAST_COVERED',
|
|
'BUTTOCKS_COVERED'
|
|
]
|
|
|
|
# 收集所有得分 > 0.01 的检测结果
|
|
all_scores = []
|
|
for i in range(outs.shape[0]):
|
|
cls_scores = outs[i][4:]
|
|
max_idx = int(np.argmax(cls_scores))
|
|
max_score = float(cls_scores[max_idx])
|
|
if max_score > 0.01:
|
|
all_scores.append((labels[max_idx], max_score))
|
|
|
|
# 按得分排序
|
|
all_scores.sort(key=lambda x: -x[1])
|
|
|
|
print(f"\n检测到 {len(all_scores)} 个得分 > 0.01 的区域")
|
|
print("\nTop 20 最高分:")
|
|
for cls, sc in all_scores[:20]:
|
|
print(f" {cls}: {sc:.4f}")
|
|
|
|
# 使用标准 detect 方法对比
|
|
print("\n--- 使用标准 detect 方法 ---")
|
|
result = det.detect(img_bytes)
|
|
print(f"标准检测结果: {result}")
|