Spaces:
Sleeping
Sleeping
| from sklearn.neighbors import KNeighborsClassifier | |
| from sklearn.model_selection import train_test_split | |
| from sklearn.metrics import accuracy_score | |
| from .base_model import BaseModel | |
| from ..config import HEART_DISEASE_MODEL_PATH, RANDOM_STATE, TEST_SIZE | |
| import numpy as np | |
| import pandas as pd | |
| class HeartDiseaseModel(BaseModel): | |
| def __init__(self): | |
| super().__init__(HEART_DISEASE_MODEL_PATH) | |
| self.model = KNeighborsClassifier( | |
| n_neighbors=5, | |
| weights='distance', # Weight by distance for better local sensitivity | |
| metric='manhattan' # Manhattan distance for better feature importance | |
| ) | |
| self.feature_names = [ | |
| 'age', 'sex', 'cp', 'trestbps', 'chol', 'fbs', 'restecg', | |
| 'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal' | |
| ] | |
| self.X_train = None | |
| self.y_train = None | |
| # Define risk thresholds | |
| self.high_risk_threshold = 0.5 | |
| # Feature importance weights | |
| self.feature_weights = { | |
| 'age': 1.5, # Age is important | |
| 'cp': 2.0, # Chest pain type is very important | |
| 'trestbps': 1.2, # Blood pressure | |
| 'chol': 1.2, # Cholesterol | |
| 'thalach': 1.5, # Max heart rate | |
| 'oldpeak': 1.8, # ST depression | |
| 'ca': 2.0, # Number of vessels | |
| 'thal': 1.5 # Thalassemia | |
| } | |
| def train(self, X, y): | |
| X = X[self.feature_names] | |
| # Apply feature weights | |
| for feature, weight in self.feature_weights.items(): | |
| if feature in X.columns: | |
| X[feature] = X[feature] * weight | |
| X_train, X_test, y_train, y_test = train_test_split( | |
| X, y, test_size=TEST_SIZE, random_state=RANDOM_STATE, | |
| stratify=y # Ensure balanced split | |
| ) | |
| self.X_train = X_train | |
| self.y_train = y_train | |
| self.model.fit(X_train, y_train) | |
| return self.evaluate(X_train, X_test, y_train, y_test) | |
| def predict(self, X): | |
| if self.scaler: | |
| X = self.scaler.transform(X) | |
| X = pd.DataFrame(X, columns=self.feature_names) | |
| # Apply feature weights | |
| for feature, weight in self.feature_weights.items(): | |
| if feature in X.columns: | |
| X[feature] = X[feature] * weight | |
| # Get nearest neighbors | |
| distances, indices = self.model.kneighbors(X) | |
| # Get similar cases | |
| similar_cases = self.X_train.iloc[indices[0]] | |
| similar_outcomes = self.y_train.iloc[indices[0]] | |
| # Calculate risk score based on weighted voting | |
| weights = 1 / (distances[0] + 1e-6) | |
| weighted_prob = np.sum(similar_outcomes * weights) / np.sum(weights) | |
| # Calculate additional risk factors | |
| risk_factors = [] | |
| # Convert X back to original scale if scaler exists | |
| if self.scaler: | |
| X_orig = pd.DataFrame(self.scaler.inverse_transform(X), columns=self.feature_names) | |
| else: | |
| X_orig = X | |
| # Check various risk factors | |
| if X_orig['age'].iloc[0] > 60: | |
| weighted_prob += 0.1 | |
| if X_orig['cp'].iloc[0] >= 2: # Non-typical chest pain | |
| weighted_prob += 0.1 | |
| if X_orig['trestbps'].iloc[0] > 140: # High blood pressure | |
| weighted_prob += 0.1 | |
| if X_orig['chol'].iloc[0] > 240: # High cholesterol | |
| weighted_prob += 0.1 | |
| if X_orig['thalach'].iloc[0] < 120: # Low max heart rate | |
| weighted_prob += 0.1 | |
| if X_orig['oldpeak'].iloc[0] > 2: # High ST depression | |
| weighted_prob += 0.15 | |
| if X_orig['ca'].iloc[0] > 0: # Presence of vessels colored by fluoroscopy | |
| weighted_prob += 0.15 * X_orig['ca'].iloc[0] | |
| # Make final prediction based on threshold | |
| prediction = np.array([1 if weighted_prob >= self.high_risk_threshold else 0]) | |
| return prediction, similar_cases, similar_outcomes, distances[0] | |
| def evaluate(self, X_train, X_test, y_train, y_test): | |
| train_accuracy = accuracy_score(y_train, self.model.predict(X_train)) | |
| test_accuracy = accuracy_score(y_test, self.model.predict(X_test)) | |
| return train_accuracy, test_accuracy |