Fix: extract feature names from XGBoost model to match training feature selection (30 vs 63)

src/pwat_estimator.py

@@ -23,7 +23,6 @@ ITEM_NAMES = {
 }
 
 # Debiasing correction factors (calibrated from 61 DFU images)
-# Applied as: adjusted = clip(raw + factor, 0, 4)
 CORRECTION_FACTORS = {
     "I":   {3: 0.0,  4: 0.0,  5: 0.0,  6: 0.0,  7: 0.0,  8: 0.0},
     "II":  {3: 0.0,  4: 0.0,  5: 0.0,  6: 0.0,  7: 0.0,  8: 0.0},
@@ -36,8 +35,8 @@ CORRECTION_FACTORS = {
 
 @dataclass
 class PWATResult:
-    scores_raw: dict = field(default_factory=dict)       # {item: int}
-    scores_adjusted: dict = field(default_factory=dict)  # {item: float} (debiased)
+    scores_raw: dict = field(default_factory=dict)
+    scores_adjusted: dict = field(default_factory=dict)
     total_raw: int = 0
     total_adjusted: float = 0.0
     fitzpatrick_type: str = ""
@@ -45,10 +44,7 @@ class PWATResult:
 
 
 def extract_features(img_bgr: np.ndarray, ulcer_mask: np.ndarray) -> Optional[dict]:
-    """Extract 63 features from the wound region for PWAT prediction.
-
-    Features: color (RGB/HSV/Lab), tissue composition, morphology, texture.
-    """
+    """Extract features from the wound region for PWAT prediction."""
     b = ulcer_mask > 0 if ulcer_mask.dtype == bool else ulcer_mask > 127
     npx = int(np.sum(b))
     if npx < 50:
@@ -56,7 +52,7 @@ def extract_features(img_bgr: np.ndarray, ulcer_mask: np.ndarray) -> Optional[di
 
     feats = {}
 
-    # --- Color features (45) ---
+    # Color features (45)
     hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
     lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2Lab).astype(np.float32)
     rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB).astype(np.float32)
@@ -70,7 +66,7 @@ def extract_features(img_bgr: np.ndarray, ulcer_mask: np.ndarray) -> Optional[di
             feats[f"{cs}_{cn}_p25"] = float(np.percentile(vals, 25))
             feats[f"{cs}_{cn}_p75"] = float(np.percentile(vals, 75))
 
-    # --- Tissue composition (5) ---
+    # Tissue composition (5)
     h, s, v = hsv[b, 0], hsv[b, 1], hsv[b, 2]
     l_ch = lab[b, 0] * (100 / 255)
     a_ch = lab[b, 1] - 128
@@ -86,14 +82,14 @@ def extract_features(img_bgr: np.ndarray, ulcer_mask: np.ndarray) -> Optional[di
     feats["tissue_necro_pct"] = float(np.sum(necro) / npx * 100)
     feats["tissue_necro_total"] = feats["tissue_eschar_pct"] + feats["tissue_slough_pct"] + feats["tissue_necro_pct"]
 
-    # --- Morphological features (7) ---
+    # Morphological features (7)
     mask_u8 = b.astype(np.uint8) if b.dtype == bool else (ulcer_mask > 127).astype(np.uint8)
     cnts, _ = cv2.findContours(mask_u8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
     if cnts:
         cnt = max(cnts, key=cv2.contourArea)
         area = cv2.contourArea(cnt)
         perim = cv2.arcLength(cnt, True)
-        circ = 4 * np.pi * area / (perim**2) if perim > 0 else 0
+        circ = 4 * np.pi * area / (perim ** 2) if perim > 0 else 0
         feats["morph_area"] = float(area)
         feats["morph_perimeter"] = float(perim)
         feats["morph_circularity"] = float(circ)
@@ -104,7 +100,7 @@ def extract_features(img_bgr: np.ndarray, ulcer_mask: np.ndarray) -> Optional[di
         hull = cv2.convexHull(cnt)
         feats["morph_solidity"] = float(area / (cv2.contourArea(hull) + 1e-8))
 
-    # --- Texture features (4) ---
+    # Texture features (4)
     gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
     wound_gray = gray[b]
     feats["texture_mean"] = float(np.mean(wound_gray))
@@ -118,7 +114,7 @@ def extract_features(img_bgr: np.ndarray, ulcer_mask: np.ndarray) -> Optional[di
     if np.any(edge_zone):
         feats["edge_gradient"] = float(np.mean(np.abs(cv2.Sobel(gray.astype(np.float32), cv2.CV_32F, 1, 0)[edge_zone])))
 
-    # --- ROI features (2) ---
+    # ROI features (2)
     feats["wound_npx"] = float(npx)
     feats["wound_ratio"] = float(npx / (img_bgr.shape[0] * img_bgr.shape[1]))
 
@@ -126,52 +122,56 @@ def extract_features(img_bgr: np.ndarray, ulcer_mask: np.ndarray) -> Optional[di
 
 
 class PWATPredictor:
-    """Predicts PWAT items 3-8 from wound features using trained XGBoost models."""
+    """Predicts PWAT items 3-8 using trained XGBoost models."""
 
     def __init__(self, models_dir: str):
         self.models = {}
+        self.feature_names = {}
         models_path = Path(models_dir)
         for item in ITEMS:
             pkl = models_path / f"xgb_pwat{item}.pkl"
             if pkl.exists():
-                self.models[item] = joblib.load(pkl)
+                model = joblib.load(pkl)
+                self.models[item] = model
+                # Extract expected feature names from the trained model
+                try:
+                    names = model.get_booster().feature_names
+                    if names:
+                        self.feature_names[item] = names
+                except Exception:
+                    pass
 
     def predict(
         self,
         img_bgr: np.ndarray,
         ulcer_mask: np.ndarray,
         fitzpatrick_type: str = "III",
-        feature_cols: Optional[list] = None,
     ) -> PWATResult:
-        """Predict PWAT scores for a single image.
-
-        Args:
-            img_bgr: BGR image
-            ulcer_mask: Binary ulcer mask (H, W)
-            fitzpatrick_type: Fitzpatrick type for debiasing ("I" .. "VI")
-            feature_cols: Ordered feature column names (must match training order).
-                          If None, uses all extracted features sorted alphabetically.
-        """
+        """Predict PWAT scores for a single image."""
         feats = extract_features(img_bgr, ulcer_mask)
         if feats is None:
             return PWATResult(fitzpatrick_type=fitzpatrick_type)
 
-        # Build feature vector
-        if feature_cols is None:
-            feature_cols = sorted(feats.keys())
-        X = np.array([[feats.get(c, 0.0) for c in feature_cols]])
-
         scores_raw = {}
         scores_adj = {}
         for item in ITEMS:
-            if item in self.models:
-                pred = int(self.models[item].predict(X)[0])
-                scores_raw[item] = pred
-                factor = CORRECTION_FACTORS.get(fitzpatrick_type, {}).get(item, 0.0)
-                scores_adj[item] = float(np.clip(pred + factor, 0, 4))
-            else:
+            if item not in self.models:
                 scores_raw[item] = 0
                 scores_adj[item] = 0.0
+                continue
+
+            # Use model's expected feature names if available
+            if item in self.feature_names:
+                cols = self.feature_names[item]
+            else:
+                cols = sorted(feats.keys())
+
+            X = np.array([[feats.get(c, 0.0) for c in cols]])
+
+            pred = int(self.models[item].predict(X)[0])
+            scores_raw[item] = pred
+            factor = CORRECTION_FACTORS.get(fitzpatrick_type, {}).get(item, 0.0)
+            scores_adj[item] = float(np.clip(pred + factor, 0, 4))
 
         return PWATResult(
             scores_raw=scores_raw,