# -*- 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}")