File size: 5,721 Bytes
5d0dc03 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor, ExtraTreesRegressor
from sklearn.neural_network import MLPRegressor
from sklearn.linear_model import LinearRegression, Ridge, Lasso, ElasticNet
from sklearn.svm import SVR
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import joblib
import os
class IVIMRegressor:
"""
Machine Learning wrapper for estimating IVIM/DKI parameters from diffusion MRI signals.
This class provides a unified interface for training and applying various regression models
(Random Forest, Extra Trees, MLP, etc.) to map signal attenuation curves directly to
tissue parameters (D, f, D*, K), bypassing iterative non-linear least squares fitting.
Supported architectures:
- 'random_forest': Robust baseline, handles noise well.
- 'extra_trees': Often faster and slightly more accurate than RF. In our experiments, this model showed superior robustness to noise.
- 'mlp': Multi-layer Perceptron for capturing complex non-linear mappings.
- 'xgboost': Gradient boosting (requires xgboost package).
- 'svr': Support Vector Regression.
"""
def __init__(self, model_type='extra_trees', params=None):
self.model_type = model_type
self.params = params if params else {}
self.model = self._build_model()
def _build_model(self):
if self.model_type == 'random_forest':
# Default params from paper/notebook
n_estimators = self.params.get('n_estimators', 100)
return RandomForestRegressor(n_estimators=n_estimators, random_state=42, n_jobs=-1)
elif self.model_type == 'extra_trees':
n_estimators = self.params.get('n_estimators', 100)
return ExtraTreesRegressor(n_estimators=n_estimators, random_state=42, n_jobs=-1)
elif self.model_type == 'mlp':
hidden_layer_sizes = self.params.get('hidden_layer_sizes', (100, 50))
return MLPRegressor(hidden_layer_sizes=hidden_layer_sizes, max_iter=500, random_state=42)
elif self.model_type == 'xgboost':
try:
from xgboost import XGBRegressor
return XGBRegressor(n_estimators=1000, learning_rate=0.01, n_jobs=-1, random_state=42)
except ImportError:
print("XGBoost not installed. Falling back to Random Forest.")
return RandomForestRegressor(n_estimators=100, random_state=42)
elif self.model_type == 'svr':
C = self.params.get('C', 100)
return SVR(C=C)
else:
raise ValueError(f"Unknown model type: {self.model_type}")
def train(self, X, y, test_size=0.2, verbose=True):
"""
Trains the regression model using the provided signal-parameter pairs.
Args:
X: Input feature matrix (Normalized Signal vs b-values). Shape: [n_samples, n_b_values]
y: Target parameter vector (e.g., Diffusion Coefficient D). Shape: [n_samples]
test_size: Fraction of data to reserve for validation (default: 0.2).
verbose: If True, prints training progress and validation metrics.
Returns:
Dictionary containing validation metrics (MAE, MSE, RMSE, R2).
"""
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=42)
if verbose:
print(f"Training {self.model_type} on {len(X_train)} samples...")
self.model.fit(X_train, y_train)
# Evaluate
predictions = self.model.predict(X_test)
metrics = self._evaluate(y_test, predictions)
if verbose:
print("--- Validation Metrics ---")
for k, v in metrics.items():
print(f"{k}: {v:.6f}")
return metrics
def predict(self, X):
"""Predicts parameters for new data."""
return self.model.predict(X)
def _evaluate(self, y_true, y_pred):
return {
'MAE': mean_absolute_error(y_true, y_pred),
'MSE': mean_squared_error(y_true, y_pred),
'RMSE': np.sqrt(mean_squared_error(y_true, y_pred)),
'R2': r2_score(y_true, y_pred)
}
def save(self, filepath):
"""Saves the trained model to disk."""
joblib.dump(self.model, filepath)
print(f"Model saved to {filepath}")
def load(self, filepath):
"""Loads a trained model from disk."""
if os.path.exists(filepath):
self.model = joblib.load(filepath)
print(f"Model loaded from {filepath}")
else:
raise FileNotFoundError(f"Model file not found: {filepath}")
def load_training_data(data_dir, dataset_name='MR701'):
"""
Helper to load X and Y CSV files from the data directory.
Expected format: Data_X2_{dataset}.csv and Data_Y_{dataset}.csv
"""
x_path = os.path.join(data_dir, f'Data_X2_{dataset_name}.csv')
y_path = os.path.join(data_dir, f'Data_Y_{dataset_name}.csv')
if not os.path.exists(x_path) or not os.path.exists(y_path):
raise FileNotFoundError(f"Data files not found for {dataset_name} in {data_dir}")
X = np.loadtxt(x_path)
Y = np.loadtxt(y_path) # Assuming Y contains [D, f, D*, K] columns or similar
return X, Y
|