This commit is contained in:
parent
f5e535d16d
commit
41d61a2f76
8
.env
8
.env
|
|
@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue