fix: 修复二次元识别分数计算(添加 sigmoid 激活 + 修正标签格式)
This commit is contained in:
@@ -58,11 +58,11 @@ INPUT_SIZE = 448
|
||||
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
# 我们关注的分级标签
|
||||
# 我们关注的分级标签(wd-tagger 格式:直接是 safe/questionable/explicit)
|
||||
RATING_TAGS = {
|
||||
"rating:safe": "safe",
|
||||
"rating:questionable": "questionable",
|
||||
"rating:explicit": "explicit",
|
||||
"safe": "safe",
|
||||
"questionable": "questionable",
|
||||
"explicit": "explicit",
|
||||
}
|
||||
|
||||
# 模型目录由 Go 侧作为第二个参数传入
|
||||
@@ -225,8 +225,9 @@ def _infer_bytes(session, img_bytes):
|
||||
input_name = session.get_inputs()[0].name
|
||||
outputs = session.run(None, {input_name: inp})
|
||||
|
||||
# 输出是 [1, num_tags] 的分数数组
|
||||
scores = outputs[0][0] if len(outputs[0].shape) == 2 else outputs[0]
|
||||
# 输出是 [1, num_tags] 的 logits 数组,需要 sigmoid 转换
|
||||
logits = outputs[0][0] if len(outputs[0].shape) == 2 else outputs[0]
|
||||
scores = 1.0 / (1.0 + np.exp(-logits)) # sigmoid 激活
|
||||
|
||||
# 提取分级分数
|
||||
ratings = {}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user