Update utils/model_loader.py

utils/model_loader.py

@@ -7,11 +7,20 @@ import pickle
 import pandas as pd
 import numpy as np
 import logging
+import os
 from pathlib import Path
 from typing import Dict, Any, Optional
 import warnings
 warnings.filterwarnings('ignore')
 
+# 安全模型加载 - 从私有HF仓库加载
+try:
+    from huggingface_hub import hf_hub_download
+    HF_HUB_AVAILABLE = True
+except ImportError:
+    HF_HUB_AVAILABLE = False
+    print("⚠️ huggingface_hub未安装，将使用本地模型文件")
+
 # 配置日志
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -24,20 +33,37 @@ class ModelManager:
         self.advisory_models = {}
         self.model_configs = {}
         self.thresholds = {}
-        
+
         # 模型路径配置 - 支持本地和云端部署
         self.app_path = Path(__file__).parent.parent
-        
-        # 使用正确的模型路径
-        self.base_path = Path("/Users/ning/Desktop/idea/代码forSarcoAdvisor")
-        self.screening_paths = {
-            'sarcoI': self.base_path / "3.建模/SarcoI_results",
-            'sarcoII': self.base_path / "3.建模/SarcoII_results"
-        }
-        self.advisory_paths = {
-            'sarcoI': self.base_path / "4.DICE建模/SarcoI/individual_models",
-            'sarcoII': self.base_path / "4.DICE建模/SarcoII/individual_models"
-        }
+
+        # 检查是否使用HF模型
+        self.use_hf_models = os.getenv("USE_HF_MODELS", "false").lower() == "true"
+        self.hf_model_repo = os.getenv("HF_MODEL_REPO", "Ning311/sarco-advisor-models")
+        self.hf_token = os.getenv("HF_TOKEN", None)
+
+        if self.use_hf_models and HF_HUB_AVAILABLE:
+            logger.info(f"🔒 使用HF私有仓库模型: {self.hf_model_repo}")
+            # HF模式下的模型路径
+            self.screening_paths = {
+                'sarcoI': "models/screening/sarcoI",
+                'sarcoII': "models/screening/sarcoII"
+            }
+            self.advisory_paths = {
+                'sarcoI': "models/advisory/sarcoI",
+                'sarcoII': "models/advisory/sarcoII"
+            }
+        else:
+            logger.info("📁 使用本地模型文件")
+            # 本地模式下的模型路径
+            self.screening_paths = {
+                'sarcoI': self.app_path / "models/screening/sarcoI",
+                'sarcoII': self.app_path / "models/screening/sarcoII"
+            }
+            self.advisory_paths = {
+                'sarcoI': self.app_path / "models/advisory/sarcoI",
+                'sarcoII': self.app_path / "models/advisory/sarcoII"
+            }
         
     def load_all_models(self):
         """加载所有模型"""
@@ -53,19 +79,40 @@ class ModelManager:
     def _load_screening_models(self):
         """加载筛查类模型"""
         try:
-            # SarcoI筛查模型 - RandomForest (更新后的最佳模型)
-            sarcoI_rf_path = self.screening_paths['sarcoI'] / "randomforest_model.pkl"
+            # SarcoI筛查模型 - RandomForest
+            if self.use_hf_models and HF_HUB_AVAILABLE:
+                # 从HF下载模型
+                sarcoI_rf_path = hf_hub_download(
+                    repo_id=self.hf_model_repo,
+                    filename=f"{self.screening_paths['sarcoI']}/randomforest_model.pkl",
+                    token=self.hf_token
+                )
+            else:
+                # 使用本地模型
+                sarcoI_rf_path = self.screening_paths['sarcoI'] / "randomforest_model.pkl"
+
             with open(sarcoI_rf_path, 'rb') as f:
                 self.screening_models['sarcoI'] = pickle.load(f)
-            
+            logger.info("✅ SarcoI筛查模型加载成功")
+
             # SarcoII筛查模型 - CatBoost (.cbm格式)
-            sarcoII_cat_path = self.screening_paths['sarcoII'] / "catboost_model.cbm"
-            
+            if self.use_hf_models and HF_HUB_AVAILABLE:
+                # 从HF下载模型
+                sarcoII_cat_path = hf_hub_download(
+                    repo_id=self.hf_model_repo,
+                    filename=f"{self.screening_paths['sarcoII']}/catboost_model.cbm",
+                    token=self.hf_token
+                )
+            else:
+                # 使用本地模型
+                sarcoII_cat_path = self.screening_paths['sarcoII'] / "catboost_model.cbm"
+
             # 需要特殊处理CatBoost模型加载
             try:
                 import catboost as cb
                 self.screening_models['sarcoII'] = cb.CatBoostClassifier()
                 self.screening_models['sarcoII'].load_model(str(sarcoII_cat_path))
+                logger.info("✅ SarcoII筛查模型加载成功")
             except ImportError:
                 logger.error("CatBoost未安装，无法加载SarcoII筛查模型")
                 raise
@@ -80,17 +127,35 @@ class ModelManager:
         """加载建议类模型(高精确率)"""
         try:
             # SarcoI建议模型 (CatBoost)
-            sarcoI_cat_path = self.advisory_paths['sarcoI'] / "CatBoost_model.pkl"
+            if self.use_hf_models and HF_HUB_AVAILABLE:
+                sarcoI_cat_path = hf_hub_download(
+                    repo_id=self.hf_model_repo,
+                    filename=f"{self.advisory_paths['sarcoI']}/CatBoost_model.pkl",
+                    token=self.hf_token
+                )
+            else:
+                sarcoI_cat_path = self.advisory_paths['sarcoI'] / "CatBoost_model.pkl"
+
             with open(sarcoI_cat_path, 'rb') as f:
                 self.advisory_models['sarcoI'] = pickle.load(f)
-            
+            logger.info("✅ SarcoI建议模型加载成功")
+
             # SarcoII建议模型 (RandomForest)
-            sarcoII_rf_path = self.advisory_paths['sarcoII'] / "RandomForest_model.pkl"
+            if self.use_hf_models and HF_HUB_AVAILABLE:
+                sarcoII_rf_path = hf_hub_download(
+                    repo_id=self.hf_model_repo,
+                    filename=f"{self.advisory_paths['sarcoII']}/RandomForest_model.pkl",
+                    token=self.hf_token
+                )
+            else:
+                sarcoII_rf_path = self.advisory_paths['sarcoII'] / "RandomForest_model.pkl"
+
             with open(sarcoII_rf_path, 'rb') as f:
                 self.advisory_models['sarcoII'] = pickle.load(f)
-                
+            logger.info("✅ SarcoII建议模型加载成功")
+
             logger.info("建议模型加载成功")
-            
+
         except Exception as e:
             logger.error(f"建议模型加载失败: {str(e)}")
             raise
@@ -234,10 +299,10 @@ class ModelManager:
         """准备模型特征 - 基于实际训练数据的特征顺序"""
         if model_type == 'sarcoI':
             if mode == 'screening':
-                # SarcoI筛查模型特征 - 必须与训练时的特征顺序完全一致
-                # 从调试结果得知模型期望的顺序：body_mass_index, age_years, WWI
+                # SarcoI筛查模型特征 - 基于实际模型期望的特征顺序
+                # 模型期望：['age_years', 'WWI', 'body_mass_index']
                 features = [
-                    'body_mass_index', 'age_years', 'WWI'
+                    'age_years', 'WWI', 'body_mass_index'
                 ]
             else:  # advisory
                 # SarcoI建议模型特征 (来自/Users/ning/Desktop/idea/代码forSarcoAdvisor/4.DICE建模/预筛选/SarcoI_train_final.csv)
@@ -247,10 +312,10 @@ class ModelManager:
                 ]
         else:  # sarcoII
             if mode == 'screening':
-                # SarcoII筛查模型特征 (来自/Users/ning/Desktop/idea/模型forSarcoAdvisor/SarcoII/SarcoII_train_selected.csv)
-                # 必须与训练时的特征顺序完全一致：age_years, race_ethnicity, body_mass_index, WWI
+                # SarcoII筛查模型特征 - 基于实际模型期望的特征顺序
+                # 模型期望：['age_years', 'race_ethnicity', 'WWI', 'body_mass_index']
                 features = [
-                    'age_years', 'race_ethnicity', 'body_mass_index', 'WWI'
+                    'age_years', 'race_ethnicity', 'WWI', 'body_mass_index'
                 ]
             else:  # advisory
                 # SarcoII建议模型特征 (来自/Users/ning/Desktop/idea/代码forSarcoAdvisor/4.DICE建模/预筛选/SarcoII_train_final.csv)
@@ -334,9 +399,79 @@ class ModelManager:
         else:
             return 'low'
     
+    def get_comprehensive_risk(self, sarcoI_screening_result: Dict, sarcoI_advisory_result: Dict = None,
+                              sarcoII_screening_result: Dict = None, sarcoII_advisory_result: Dict = None) -> Dict:
+        """
+        计算新的综合风险等级 - 基于建议模型优先的融合方案
+
+        Args:
+            sarcoI_screening_result: SarcoI筛查模型结果
+            sarcoI_advisory_result: SarcoI建议模型结果 (可选)
+            sarcoII_screening_result: SarcoII筛查模型结果 (可选)
+            sarcoII_advisory_result: SarcoII建议模型结果 (可选)
+
+        Returns:
+            Dict: 包含SarcoI和SarcoII综合风险的字典
+        """
+        results = {}
+
+        # SarcoI 综合风险判定
+        if sarcoI_screening_result:
+            P_recall_I = sarcoI_screening_result['probability']
+            P_precision_I = sarcoI_advisory_result['probability'] if sarcoI_advisory_result else 0.0
+
+            # 使用实际的模型阈值
+            sarcoI_advisory_threshold = self.thresholds['sarcoI']['advisory']
+            sarcoI_screening_threshold = self.thresholds['sarcoI']['screening']
+
+            if P_precision_I >= sarcoI_advisory_threshold:  # 建议模型高风险阈值
+                sarcoI_comprehensive_risk = "high"
+                sarcoI_risk_reason = "advisory_model_high_risk"
+            elif P_recall_I >= sarcoI_screening_threshold:   # 筛查模型高风险阈值
+                sarcoI_comprehensive_risk = "medium"
+                sarcoI_risk_reason = "screening_model_risk"
+            else:
+                sarcoI_comprehensive_risk = "low"
+                sarcoI_risk_reason = "both_models_low_risk"
+
+            results['sarcoI'] = {
+                'comprehensive_risk': sarcoI_comprehensive_risk,
+                'screening_probability': P_recall_I,
+                'advisory_probability': P_precision_I,
+                'risk_reason': sarcoI_risk_reason
+            }
+
+        # SarcoII 综合风险判定
+        if sarcoII_screening_result:
+            P_recall_II = sarcoII_screening_result['probability']
+            P_precision_II = sarcoII_advisory_result['probability'] if sarcoII_advisory_result else 0.0
+
+            # 使用实际的模型阈值
+            sarcoII_advisory_threshold = self.thresholds['sarcoII']['advisory']
+            sarcoII_screening_threshold = self.thresholds['sarcoII']['screening']
+
+            if P_precision_II >= sarcoII_advisory_threshold:  # 建议模型高风险阈值
+                sarcoII_comprehensive_risk = "high"
+                sarcoII_risk_reason = "advisory_model_high_risk"
+            elif P_recall_II >= sarcoII_screening_threshold:   # 筛查模型高风险阈值
+                sarcoII_comprehensive_risk = "medium"
+                sarcoII_risk_reason = "screening_model_risk"
+            else:
+                sarcoII_comprehensive_risk = "low"
+                sarcoII_risk_reason = "both_models_low_risk"
+
+            results['sarcoII'] = {
+                'comprehensive_risk': sarcoII_comprehensive_risk,
+                'screening_probability': P_recall_II,
+                'advisory_probability': P_precision_II,
+                'risk_reason': sarcoII_risk_reason
+            }
+
+        return results
+
     def get_overall_risk(self, sarcoI_result: Dict, sarcoII_result: Dict) -> str:
         """
-        计算综合风险等级
+        计算综合风险等级 (保持向后兼容)
 
         基于两个模型的预测结果，使用更科学的综合评估方法
         """