This commit is contained in:
Rzy 2026-09-08 16:43:07 +08:00
parent f5e535d16d
commit 41d61a2f76
2 changed files with 18 additions and 2 deletions

8
.env
View File

@ -25,3 +25,11 @@ BRAND_MODEL_DIR=/app/services/brand/model
BRAND_SERVICE_NAME=brand_classifier BRAND_SERVICE_NAME=brand_classifier
BRAND_TYPE=brand BRAND_TYPE=brand
# ==================== 分类模型 ====================
CATE_MODEL_ID=circles1/cate
CATE_MODEL_DIR=/app/services/cate/model
CATE_SERVICE_NAME=cate_classifier
CATE_TYPE=cate

View File

@ -195,10 +195,19 @@ class BasePredictor:
def _format_output(self, label: str, confidence: float) -> Dict: def _format_output(self, label: str, confidence: float) -> Dict:
"""根据模型类型格式化输出""" """根据模型类型格式化输出"""
# 确保所有类型都是 Python 原生类型
confidence = float(confidence)
# 确保 label 是字符串
if not isinstance(label, str):
if hasattr(label, 'item'):
label = str(label.item())
else:
label = str(label)
result = {"confidence": confidence} result = {"confidence": confidence}
if self.model_type == 'tax': if self.model_type == 'tax':
# tax 模型type_tax 格式
parts = label.split('_') parts = label.split('_')
if len(parts) >= 2: if len(parts) >= 2:
result['type'] = parts[0] result['type'] = parts[0]
@ -211,7 +220,6 @@ class BasePredictor:
elif self.model_type == 'brand': elif self.model_type == 'brand':
result['brand'] = label result['brand'] = label
else: else:
# 通用模型:直接输出 label
result['label'] = label result['label'] = label
return result return result