diabetes / train_diabetes.py
LovnishVerma's picture
Create train_diabetes.py
d9db1ec verified
import pandas as pd
import numpy as np
import os
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report
from imblearn.over_sampling import SMOTE
from joblib import dump
# Load dataset
data = pd.read_csv("dia.csv")
# Features & Target
X = data.drop(columns=["Outcome"])
y = data["Outcome"]
# Replace zero values in certain columns (except Pregnancies)
cols_with_zero = ["Glucose", "BloodPressure",
"SkinThickness", "Insulin", "BMI"]
X[cols_with_zero] = X[cols_with_zero].replace(0, np.nan)
# Save median values used for imputation
medians = X.median().to_dict()
# Fill missing with median
X = X.fillna(medians)
# Train-Test Split
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)
# Scale
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)
# Balance dataset with SMOTE
sm = SMOTE(random_state=42)
X_train_bal, y_train_bal = sm.fit_resample(X_train_scaled, y_train)
# Train Random Forest
model = RandomForestClassifier(n_estimators=200, random_state=42)
model.fit(X_train_bal, y_train_bal)
# Evaluation @ Default Threshold 0.5
y_pred = model.predict(X_test_scaled)
acc = accuracy_score(y_test, y_pred)
print(" Default Threshold Accuracy:", acc)
print("\nClassification Report (Threshold=0.5):\n",
classification_report(y_test, y_pred))
# Threshold Tuning
print("\n Threshold Tuning Results")
y_proba = model.predict_proba(X_test_scaled)[:, 1]
for thresh in [0.3, 0.4, 0.5, 0.6]:
y_pred_thresh = (y_proba >= thresh).astype(int)
acc_thresh = accuracy_score(y_test, y_pred_thresh)
print(f"\nThreshold = {thresh}")
print("Accuracy:", acc_thresh)
print(classification_report(y_test, y_pred_thresh, digits=3))
# Save model, scaler & medians
os.makedirs("models", exist_ok=True)
dump(model, "models/diabetes.sav")
dump(scaler, "models/scaler.sav")
dump(medians, "models/medians.sav")
print("Final Model, Scaler, and Medians saved in 'models/' folder.")