Update tongue_model.py

tongue_model.py

@@ -6,7 +6,6 @@ import numpy as np
 from PIL import Image
 from torchvision import models, transforms
 
-# --- 模型架構定義 ---
 class CBAM(nn.Module):
     def __init__(self, channels, reduction=16):
         super(CBAM, self).__init__()
@@ -51,13 +50,19 @@ class TongueArcResNet(nn.Module):
         features = self.avgpool(x).flatten(1)
         return self.arcface(features)
 
-# --- 2. 定義預處理與推論類別 ---
 class TongueModelWrapper:
     def __init__(self, model_path, num_classes=3):
         self.device = torch.device("cpu")
         self.model = TongueArcResNet(num_classes=num_classes)
         torch.serialization.add_safe_globals([np._core.multiarray.scalar])
-        self.model.load_state_dict(torch.load(model_path, map_location=self.device, weights_only=False))
+        checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
+        
+        if isinstance(checkpoint, dict) and 'model_state' in checkpoint:
+            self.model.load_state_dict(checkpoint['model_state'])
+            print(f"成功載入權重！模型訓練指標：AUC={checkpoint.get('auc', 0):.4f}")
+        else:
+            self.model.load_state_dict(checkpoint)
+            
         self.model.eval()
         self.transform = transforms.Compose([
             transforms.ToTensor(),