load_trained_model.py

#!/usr/bin/env python3
"""
Properly load the trained ultimate model by recreating the class structure
"""

import pickle
import numpy as np
import json
from pathlib import Path
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.linear_model import Ridge, ElasticNet

# Recreate the exact UltimateEnsemble class structure so pickle can load it
class UltimateEnsemble:
    """Recreated UltimateEnsemble class for loading"""
    
    def __init__(self):
        self.models = {}
        self.weights = {}
        self.fitted = False
    
    def predict(self, X):
        """Make ensemble prediction"""
        predictions = []
        total_weight = 0
        
        for name, model in self.models.items():
            try:
                weight = self.weights.get(name, 1.0)
                pred = model.predict(X)
                predictions.append(pred * weight)
                total_weight += weight
            except Exception as e:
                print(f"Model {name} failed: {e}")
                continue
        
        if predictions and total_weight > 0:
            return np.sum(predictions, axis=0)
        else:
            return np.array([12.0])
    
    def predict_with_uncertainty(self, X):
        """Predict with uncertainty estimation"""
        # Get individual predictions
        individual_preds = []
        
        for name, model in self.models.items():
            try:
                pred = model.predict(X)
                individual_preds.append(pred)
            except:
                continue
        
        if individual_preds:
            # Ensemble prediction
            ensemble_pred = self.predict(X)
            
            # Uncertainty as standard deviation
            if len(individual_preds) > 1:
                uncertainty = np.std(individual_preds, axis=0)
            else:
                uncertainty = np.array([0.5])
            
            # Confidence based on uncertainty
            confidence = np.clip(100 - uncertainty * 20, 60, 95)
            
            return ensemble_pred, uncertainty, confidence
        else:
            pred = np.array([12.0])
            uncertainty = np.array([1.0])
            confidence = np.array([70.0])
            return pred, uncertainty, confidence

class UltimateFeatureExtractor:
    """Recreated UltimateFeatureExtractor class for loading"""
    
    def __init__(self):
        from sklearn.preprocessing import StandardScaler
        self.scaler = StandardScaler()
        self.fitted = False
    
    def extract_hemoglobin_features(self, image_path):
        """Extract features for hemoglobin estimation"""
        try:
            from PIL import Image
            
            img = Image.open(image_path).convert('RGB')
            img = img.resize((150, 150), Image.Resampling.LANCZOS)
            img_array = np.array(img).astype(np.float32) / 255.0
            
            features = []
            
            # Core RGB Statistics (21 features)
            for channel in range(3):
                data = img_array[:, :, channel].flatten()
                features.extend([
                    np.mean(data), np.std(data), np.median(data),
                    np.percentile(data, 25), np.percentile(data, 75),
                    np.min(data), np.max(data)
                ])
            
            # Color Ratios (15 features)
            mean_r, mean_g, mean_b = [np.mean(img_array[:, :, i]) for i in range(3)]
            eps = 1e-8
            
            features.extend([
                mean_r / (mean_g + eps), mean_r / (mean_b + eps), mean_g / (mean_b + eps),
                mean_r / (mean_g + mean_b + eps), (mean_r - mean_g) / (mean_r + mean_g + eps),
                (mean_r - mean_b) / (mean_r + mean_b + eps), mean_r * mean_g / (mean_b + eps),
                (mean_r + mean_g) / (mean_b + eps), mean_r / (mean_r + mean_g + mean_b + eps),
                mean_g / (mean_r + mean_g + mean_b + eps), mean_b / (mean_r + mean_g + mean_b + eps),
                abs(mean_r - mean_g), abs(mean_r - mean_b), abs(mean_g - mean_b),
                max(mean_r, mean_g, mean_b) - min(mean_r, mean_g, mean_b)
            ])
            
            # Histograms (24 features)
            for channel in range(3):
                hist, _ = np.histogram(img_array[:, :, channel], bins=8, range=(0, 1))
                hist = hist.astype(float) / (hist.sum() + eps)
                features.extend(hist)
            
            # Regional Analysis (20 features)
            h, w = img_array.shape[:2]
            regions = [
                img_array[:h//2, :w//2], img_array[:h//2, w//2:],
                img_array[h//2:, :w//2], img_array[h//2:, w//2:]
            ]
            
            for region in regions:
                if region.size > 0:
                    features.extend([
                        np.mean(region[:, :, 0]), np.mean(region[:, :, 1]),
                        np.mean(region[:, :, 2]), np.std(region), np.mean(region)
                    ])
                else:
                    features.extend([0, 0, 0, 0, 0])
            
            # Texture (10 features)
            gray = 0.299 * img_array[:, :, 0] + 0.587 * img_array[:, :, 1] + 0.114 * img_array[:, :, 2]
            features.extend([
                np.mean(gray), np.std(gray), np.median(gray), np.min(gray), np.max(gray),
                np.percentile(gray, 90), np.percentile(gray, 10), np.var(gray),
                np.mean(gray > np.mean(gray)), np.mean(gray < np.mean(gray))
            ])
            
            # Enhanced Color (10 features)
            red_channel = img_array[:, :, 0]
            features.extend([
                np.mean(red_channel > 0.5), np.mean(red_channel < 0.3),
                np.std(red_channel) / (np.mean(red_channel) + eps),
                np.mean(red_channel > np.percentile(red_channel, 75)),
                np.mean(red_channel < np.percentile(red_channel, 25)),
                np.max(red_channel) - np.min(red_channel),
                np.mean(np.abs(red_channel - np.mean(red_channel))),
                np.corrcoef(red_channel.flatten(), img_array[:, :, 1].flatten())[0, 1],
                np.corrcoef(red_channel.flatten(), img_array[:, :, 2].flatten())[0, 1],
                np.mean(red_channel) / (np.mean(img_array) + eps)
            ])
            
            return np.array(features, dtype=np.float32)
            
        except Exception as e:
            print(f"Feature extraction error: {e}")
            return np.zeros(100, dtype=np.float32)
    
    def transform(self, image_paths):
        """Transform multiple images"""
        if isinstance(image_paths, str):
            return self.extract_hemoglobin_features(image_paths)
        
        features = []
        for path in image_paths:
            feat = self.extract_hemoglobin_features(path)
            features.append(feat)
        return np.array(features)

def load_actual_trained_model():
    """Load the actual trained ultimate model"""
    try:
        ensemble_path = Path('weights/ultimate_ensemble.pkl')
        extractor_path = Path('weights/ultimate_extractor.pkl')
        metadata_path = Path('weights/ultimate_metadata.json')
        
        if not all(p.exists() for p in [ensemble_path, extractor_path, metadata_path]):
            print("❌ Model files not found")
            return None, None, None
        
        # Load metadata
        with open(metadata_path, 'r') as f:
            metadata = json.load(f)
        
        print(f"📊 Loading trained model with MAE: {metadata['mae']:.3f} g/dL")
        
        # Load the trained ensemble
        with open(ensemble_path, 'rb') as f:
            ensemble = pickle.load(f)
        
        # Load the trained feature extractor
        with open(extractor_path, 'rb') as f:
            extractor = pickle.load(f)
        
        print(f"✅ Loaded ACTUAL trained model!")
        print(f"   - {len(ensemble.models)} trained algorithms")
        print(f"   - Expected MAE: {metadata['mae']:.3f} g/dL")
        print(f"   - Trained on {metadata['training_samples']} samples")
        
        return ensemble, extractor, metadata
        
    except Exception as e:
        print(f"❌ Failed to load trained model: {e}")
        return None, None, None

if __name__ == "__main__":
    ensemble, extractor, metadata = load_actual_trained_model()
    if ensemble:
        print("✅ SUCCESS: Loaded the actual trained model!")
        
        # Test with dummy data
        dummy_features = np.random.randn(1, 100)
        pred, unc, conf = ensemble.predict_with_uncertainty(dummy_features)
        print(f"Test prediction: {pred[0]:.2f} g/dL (confidence: {conf[0]:.1f}%)")
    else:
        print("❌ FAILED to load trained model")