Vehicle Engine Health Predictor

Train ML models to predict engine maintenance needs from sensor data.

Introduction

30 Second Summary

Every fleet manager dreads the unexpected breakdown. A truck stalls on the highway, a delivery van misses its route, and suddenly one failed engine costs thousands in repairs and lost revenue. What if you could predict those failures before they happen, using the same sensor data the engine already generates?

In this project, you will train two machine learning models on real diesel engine sensor data and deploy the best one into an interactive dashboard that predicts whether an engine needs maintenance in real time.

What You'll Build

A predictive maintenance system that uses Random Forest and XGBoost classifiers trained on commercial diesel engine sensor data, deployed as an interactive Streamlit dashboard.

By the end of this project, you'll have:

  • A trained Random Forest and XGBoost classifier that predicts engine maintenance needs from sensor data like RPM, oil pressure, and coolant temperature.
  • A Streamlit dashboard where you can adjust sensor inputs via sliders and see real-time maintenance predictions with confidence scores.
  • A complete model evaluation workflow with confusion matrix, classification report, and feature importance visualizations.
  • Secret Mission: Add SHAP explainability to your dashboard so you can see exactly which sensors are driving each prediction.

Are there any prerequisites?

You should have basic Python familiarity (variables, functions, loops) and some comfort with pandas for data manipulation. No prior machine learning experience is required.

Set Up Your Environment

Before you can train machine learning models on engine sensor data, you need a clean workspace with the right tools installed.

In this step, you will create a project folder, set up an isolated Python environment, and install all the packages needed for data analysis, model training, and dashboard building.

In this step, get ready to:
  • Create a project directory and virtual environment.
  • Install machine learning and visualization packages.
  • Verify that all packages import correctly.
Create your project directory and virtual environment

A virtual environment is an isolated Python workspace. It keeps your project's packages separate from your system Python so nothing conflicts.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Visual Studio Code and press Enter to open it.
  • Open the integrated terminal by pressing Ctrl+` (backtick key, top-left of your keyboard).
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder by running this command:
mkdir engine-health-predictor
  • Open the project folder in VS Code:
  • Click File in the top menu bar.
  • Click Open Folder.
  • Navigate to your Desktop, select the engine-health-predictor folder, and click Open (macOS) or Select Folder (Windows).

VS Code reloads with your project folder as the workspace. You should see engine-health-predictor in the Explorer sidebar on the left.

  • Open a new integrated terminal by pressing Ctrl+` again.

The terminal now opens inside your project folder. Next, create and activate your virtual environment.

macOS

  • Create your virtual environment by running:
python3 -m venv venv
  • Activate it by running:
source venv/bin/activate

Windows

  • Create your virtual environment by running:
python -m venv venv
  • Activate it by running:
venv\Scripts\activate

You should see (venv) appear at the beginning of your terminal prompt. This confirms you are working inside the virtual environment.

Don't see (venv) in your prompt?

Make sure you ran the activation command (not just the creation command). If you see a permissions error on Windows, try running Set-ExecutionPolicy -Scope CurrentUser RemoteSigned in PowerShell first, then retry the activation.

help me activate my Python virtual environment

Install the required packages

Your project needs packages for machine learning, data handling, and visualization. You will install them all in a single command.

  • Install all required packages by running:
pip install scikit-learn xgboost streamlit datasets pandas matplotlib seaborn

What are these packages?

  • scikit-learn provides classical machine learning algorithms and evaluation tools.
  • xgboost is a high-performance gradient boosting library optimized for structured data.
  • streamlit turns Python scripts into interactive web dashboards with no frontend code.
  • datasets is Hugging Face's library for loading and processing ML datasets.
  • pandas, matplotlib, and seaborn handle data manipulation and visualization.

This may take a minute or two as pip downloads and installs each package along with its dependencies.

You should see a series of "Successfully installed" messages once the process completes.

Seeing installation errors?

Make sure your virtual environment is active (you see (venv) in your terminal). If you see a network timeout, check your internet connection and try the command again.

help me fix pip install errors in my virtual environment

Verify the installation

A quick import check confirms that all packages installed correctly and can be loaded by Python.

  • Verify your installation by running:
python -c "import sklearn, xgboost, streamlit, datasets; print('All packages installed successfully')"

You should see All packages installed successfully printed in your terminal with no errors.

Seeing a ModuleNotFoundError?

This usually means the package did not install properly. Double check that your virtual environment is active, then re-run the pip install command from the previous substep.

help me fix ModuleNotFoundError after pip install

Your environment is ready with all the tools you need. Next up, you will load real engine sensor data and start exploring it.

Load and Explore the Dataset

Your development environment is set up and ready to go. Now it's time to get your hands on the data that will power your engine health predictor.

The dataset you'll use contains real sensor telemetry from commercial diesel engines. Before training any models, you need to understand what the data looks like, spot any quality issues, and identify which sensors differ most between healthy and failing engines.

In this step, get ready to:
  • Load the Predictive Maintenance Engine Sensor Dataset from Hugging Face.
  • Check for missing values, class imbalance, and basic statistics.
  • Visualize sensor correlations and feature distributions.
Create your training script and load the dataset

Every machine learning project starts with loading and inspecting the data. You'll create a single Python script called train_model.py that will grow throughout this project.

  • In VS Code, create a new file called train_model.py inside your engine-health-predictor folder.
  • Add the following imports and directory setup to the top of the file:
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datasets import load_dataset
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix
import xgboost as xgb
import joblib

# Create output directories
os.makedirs("plots", exist_ok=True)
os.makedirs("models", exist_ok=True)

What does this code do?

The imports pull in all the libraries you installed in Step 1. You're importing more than you need right now because this file will grow as you add model training in later steps.

The os.makedirs lines create two folders for storing your outputs. The exist_ok=True flag means the script won't crash if the folders already exist.

  • Below the directory setup, add the code to load the dataset from Hugging Face and print its basic structure:
# Load dataset from Hugging Face
print("Loading dataset...")
dataset = load_dataset("jskswamy/predictive-maintenance-data")
train_df = dataset["train"].to_pandas()
test_df = dataset["test"].to_pandas()

print(f"Training data shape: {train_df.shape}")
print(f"Test data shape: {test_df.shape}")
print(f"\nFirst few rows:\n{train_df.head()}")

What does this code do?

The load_dataset function downloads the dataset directly from Hugging Face using the ID jskswamy/predictive-maintenance-data. This dataset contains engine sensor readings from commercial diesel vehicles.

The .to_pandas() call converts each split into a pandas DataFrame so you can work with familiar tools like .head() and .describe(). The dataset comes pre-split into training and test sets.

  • Save train_model.py.
  • Run the script in your terminal to confirm the dataset loads by running:
python train_model.py

You should see output showing the training data shape as (14651, 7) and a preview of the first few rows with columns like Engine RPM, Lub Oil Pressure, Fuel Pressure, Coolant Pressure, Lub Oil Temp, Coolant Temp, and Engine Condition.

Seeing a ModuleNotFoundError?

Make sure your virtual environment is activated. You should see (venv) at the start of your terminal prompt. If not, activate it again with source venv/bin/activate (macOS/Linux) or venv\Scripts\activate (Windows).

If you see a network error, check your internet connection. The dataset is downloaded from Hugging Face on the first run. Help me debug this ModuleNotFoundError when running train_model.py

Perform exploratory data analysis

Before building models, you need to check the data quality and understand what you're working with. EDA (exploratory data analysis) reveals missing values, class imbalance, and relationships between features.

  • In train_model.py, add the following code below the print statements to check data quality:
# EDA
print(f"\nMissing values:\n{train_df.isnull().sum()}")
print(f"\nTarget distribution:\n{train_df['Engine Condition'].value_counts()}")
print(f"\nDescriptive statistics:\n{train_df.describe()}")

What does this code do?

  • .isnull().sum() counts missing values per column. Missing data can break model training if not handled.
  • .value_counts() on the target column Engine Condition shows how many engines are Normal (0) versus Needs Maintenance (1). A heavily imbalanced dataset can bias your model toward the majority class.
  • .describe() gives you min, max, mean, and standard deviation for each sensor. This helps you spot outliers and understand typical sensor ranges.
  • Next, add the correlation heatmap below the EDA section to visualize relationships between sensor features:
# Correlation heatmap
plt.figure(figsize=(10, 8))
sns.heatmap(train_df.corr(), annot=True, cmap="coolwarm", center=0)
plt.title("Sensor Feature Correlation Heatmap")
plt.tight_layout()
plt.savefig("plots/correlation_heatmap.png")
plt.close()

Why create a correlation heatmap?

A correlation heatmap shows how strongly each pair of sensors moves together. If two features are highly correlated (close to 1.0 or -1.0), they carry redundant information.

The seaborn library's heatmap function turns the correlation matrix into a color-coded grid. The plt.close() call frees memory after saving the figure to disk.

  • Save train_model.py and run it again:
python train_model.py

You should see the EDA output in your terminal (missing values, target distribution, descriptive statistics) plus a new file at plots/correlation_heatmap.png.

  • Confirm the heatmap was saved by running:
ls plots/

You should see correlation_heatmap.png listed in the output.

Don't see the heatmap file?

Check that you're running the script from inside the engine-health-predictor folder. The plots/ directory is created relative to where you run the command.

If you see a matplotlib backend warning, that's OK. The file still saves correctly even if no GUI window opens. Help me figure out why my heatmap isn't saving

Visualize feature distributions by engine condition

The correlation heatmap shows relationships between sensors. Now you need to see how each sensor's readings differ between healthy engines and engines that need maintenance.

  • In train_model.py, add the distribution plots below the heatmap code:
# Distribution plots by Engine Condition
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
features = ["Engine RPM", "Lub Oil Pressure", "Fuel Pressure",
            "Coolant Pressure", "Lub Oil Temp", "Coolant Temp"]

for idx, feature in enumerate(features):
    ax = axes[idx // 3, idx % 3]
    for condition in [0, 1]:
        subset = train_df[train_df["Engine Condition"] == condition]
        ax.hist(subset[feature], bins=30, alpha=0.5,
                label=f"{'Normal' if condition == 0 else 'Maintenance'}")
    ax.set_title(feature)
    ax.legend()

plt.tight_layout()
plt.savefig("plots/feature_distributions.png")
plt.close()

What does this code do?

This creates a 2x3 grid of histograms, one for each sensor feature. Each histogram overlays two distributions: Normal engines (condition 0) in one color and Maintenance-needed engines (condition 1) in another.

Where the two distributions overlap heavily, that sensor doesn't help distinguish healthy from failing engines. Where they separate, that sensor is a strong predictor. This visual insight guides your feature engineering in the next step.

  • Save train_model.py and run the script one more time:
python train_model.py
  • Confirm both plots were saved by running:
ls plots/

You should see both correlation_heatmap.png and feature_distributions.png listed.

Seeing a KeyError for Engine Condition?

Make sure the column name in your code matches exactly: Engine Condition (with a space and capital letters). Column names are case-sensitive in pandas.

Double check that the dataset loaded correctly by looking at the earlier terminal output showing the column names. Help me fix this KeyError in my distribution plots

✔️ Awesome, I've got everything!

Great. Double check you've saved train_model.py and that both plot files exist in your plots/ directory.

ⓧ I'd like to double check the full code

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datasets import load_dataset
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix
import xgboost as xgb
import joblib

# Create output directories
os.makedirs("plots", exist_ok=True)
os.makedirs("models", exist_ok=True)

# Load dataset from Hugging Face
print("Loading dataset...")
dataset = load_dataset("jskswamy/predictive-maintenance-data")
train_df = dataset["train"].to_pandas()
test_df = dataset["test"].to_pandas()

print(f"Training data shape: {train_df.shape}")
print(f"Test data shape: {test_df.shape}")
print(f"\nFirst few rows:\n{train_df.head()}")

# EDA
print(f"\nMissing values:\n{train_df.isnull().sum()}")
print(f"\nTarget distribution:\n{train_df['Engine Condition'].value_counts()}")
print(f"\nDescriptive statistics:\n{train_df.describe()}")

# Correlation heatmap
plt.figure(figsize=(10, 8))
sns.heatmap(train_df.corr(), annot=True, cmap="coolwarm", center=0)
plt.title("Sensor Feature Correlation Heatmap")
plt.tight_layout()
plt.savefig("plots/correlation_heatmap.png")
plt.close()

# Distribution plots by Engine Condition
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
features = ["Engine RPM", "Lub Oil Pressure", "Fuel Pressure",
            "Coolant Pressure", "Lub Oil Temp", "Coolant Temp"]

for idx, feature in enumerate(features):
    ax = axes[idx // 3, idx % 3]
    for condition in [0, 1]:
        subset = train_df[train_df["Engine Condition"] == condition]
        ax.hist(subset[feature], bins=30, alpha=0.5,
                label=f"{'Normal' if condition == 0 else 'Maintenance'}")
    ax.set_title(feature)
    ax.legend()

plt.tight_layout()
plt.savefig("plots/feature_distributions.png")
plt.close()

Your data is loaded, inspected, and visualized. Next up, you'll engineer new features from these raw sensor readings and prepare the data for model training.

Engineer Features and Prepare Data

Your EDA revealed relationships between sensor readings and engine health. Raw sensor values alone tell part of the story, but combining them unlocks patterns that models can exploit more effectively.

In this step, you will create new features that capture domain-relevant signals (thermal differentials, pressure efficiency, stress indicators) and then prepare all features for training by normalizing them to a common scale.

In this step, get ready to:
  • Engineer three new features from the raw sensor data.
  • Separate features from the target variable and split into train/test sets.
  • Normalize all features using StandardScaler.
Engineer new features from sensor data

Raw sensor readings are useful, but combining them captures patterns that individual values miss. A temperature differential between oil and coolant reveals cooling system stress. A pressure ratio shows fuel delivery efficiency relative to lubrication.

  • In train_model.py, scroll to the bottom of the file (after the distribution plots code).
  • Add the feature engineering function by copying this code at the end of the file:
# Feature Engineering
def engineer_features(df):
    df = df.copy()
    df["Temp_Difference"] = df["Lub Oil Temp"] - df["Coolant Temp"]
    df["Pressure_Ratio"] = df["Fuel Pressure"] / (df["Lub Oil Pressure"] + 1e-6)
    rpm_75th = train_df["Engine RPM"].quantile(0.75)
    df["High_RPM"] = (df["Engine RPM"] > rpm_75th).astype(int)
    return df

train_df = engineer_features(train_df)
test_df = engineer_features(test_df)

What does this code do?

The engineer_features() function creates three new columns from the existing sensor data:

  • Temp_Difference subtracts coolant temperature from oil temperature. A large gap signals the cooling system is struggling to keep up.
  • Pressure_Ratio divides fuel pressure by oil pressure. This captures fuel delivery efficiency relative to lubrication. The + 1e-6 prevents division by zero.
  • High_RPM flags engines running above the 75th percentile of RPM in the training data. Sustained high RPM accelerates wear.

The function is applied to both train_df and test_df so both datasets get the same engineered columns.

Why use the training set's 75th percentile for both?

The rpm_75th threshold is computed from train_df and applied to both train and test. This prevents data leakage. If you computed the threshold on the test set, your model would be using information it shouldn't have during evaluation.

  • Save train_model.py.
  • Run the script in your terminal to confirm the new features are created without errors:
python train_model.py

You should see the same EDA output as before (dataset shape, missing values, target distribution) with no new errors. The feature engineering runs silently since it doesn't print anything yet.

Seeing a NameError or KeyError?

  • If you see NameError: name 'train_df' is not defined, make sure the feature engineering code is placed after the distribution plots section where train_df is already loaded.
  • If you see a KeyError, double check that the column names match exactly: Lub Oil Temp, Coolant Temp, Fuel Pressure, Lub Oil Pressure, Engine RPM.

Help me debug my feature engineering function in train_model.py

Prepare the data for modeling

Machine learning models need their inputs in a specific format. You need to separate the features (sensor readings) from the target (engine condition), then scale everything so no single feature dominates due to its unit of measurement.

  • In train_model.py, add the following code directly below test_df = engineer_features(test_df):
# Prepare features and target
feature_columns = ["Engine RPM", "Lub Oil Pressure", "Fuel Pressure",
                   "Coolant Pressure", "Lub Oil Temp", "Coolant Temp",
                   "Temp_Difference", "Pressure_Ratio", "High_RPM"]

X_train = train_df[feature_columns]
y_train = train_df["Engine Condition"]
X_test = test_df[feature_columns]
y_test = test_df["Engine Condition"]

# Scale features
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

What does this code do?

  • feature_columns lists all 9 features (6 original sensors + 3 engineered) the model will learn from.
  • X_train and X_test hold the input features. y_train and y_test hold the target labels (0 = Normal, 1 = Needs Maintenance).
  • StandardScaler normalizes each feature to have a mean of 0 and standard deviation of 1. This prevents high-magnitude features like Engine RPM (0-2500) from overshadowing low-magnitude ones like pressures (0-10).
  • fit_transform learns the scaling parameters from training data only. transform applies those same parameters to the test data, preventing data leakage.
  • Save train_model.py.
  • Run the script again to verify everything works:
python train_model.py

The script should complete without errors. You will see the same EDA output as before. The data preparation happens silently in memory, ready for the models you will train in the next step.

Script failing at the scaler step?

  • If you see an error about StandardScaler not being defined, check that your import at the top of the file includes from sklearn.preprocessing import StandardScaler. This import already exists from the previous step.
  • If a KeyError mentions one of the engineered feature names, make sure the feature engineering code runs before the data preparation code. The order matters.

Help me fix my data preparation code

✔️ Awesome, I've got everything!

Great. Double check you have saved train_model.py and the script runs without errors.

ⓧ I'd like to double check the full code

Here is the complete train_model.py file as it should look at the end of this step:

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datasets import load_dataset
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix
import xgboost as xgb
import joblib

# Create output directories
os.makedirs("plots", exist_ok=True)
os.makedirs("models", exist_ok=True)

# Load dataset from Hugging Face
print("Loading dataset...")
dataset = load_dataset("jskswamy/predictive-maintenance-data")
train_df = dataset["train"].to_pandas()
test_df = dataset["test"].to_pandas()

print(f"Training data shape: {train_df.shape}")
print(f"Test data shape: {test_df.shape}")
print(f"\nFirst few rows:\n{train_df.head()}")

# EDA
print(f"\nMissing values:\n{train_df.isnull().sum()}")
print(f"\nTarget distribution:\n{train_df['Engine Condition'].value_counts()}")
print(f"\nDescriptive statistics:\n{train_df.describe()}")

# Correlation heatmap
plt.figure(figsize=(10, 8))
sns.heatmap(train_df.corr(), annot=True, cmap="coolwarm", center=0)
plt.title("Sensor Feature Correlation Heatmap")
plt.tight_layout()
plt.savefig("plots/correlation_heatmap.png")
plt.close()

# Distribution plots by Engine Condition
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
features = ["Engine RPM", "Lub Oil Pressure", "Fuel Pressure",
            "Coolant Pressure", "Lub Oil Temp", "Coolant Temp"]

for idx, feature in enumerate(features):
    ax = axes[idx // 3, idx % 3]
    for condition in [0, 1]:
        subset = train_df[train_df["Engine Condition"] == condition]
        ax.hist(subset[feature], bins=30, alpha=0.5,
                label=f"{'Normal' if condition == 0 else 'Maintenance'}")
    ax.set_title(feature)
    ax.legend()

plt.tight_layout()
plt.savefig("plots/feature_distributions.png")
plt.close()

# Feature Engineering
def engineer_features(df):
    df = df.copy()
    df["Temp_Difference"] = df["Lub Oil Temp"] - df["Coolant Temp"]
    df["Pressure_Ratio"] = df["Fuel Pressure"] / (df["Lub Oil Pressure"] + 1e-6)
    rpm_75th = train_df["Engine RPM"].quantile(0.75)
    df["High_RPM"] = (df["Engine RPM"] > rpm_75th).astype(int)
    return df

train_df = engineer_features(train_df)
test_df = engineer_features(test_df)

# Prepare features and target
feature_columns = ["Engine RPM", "Lub Oil Pressure", "Fuel Pressure",
                   "Coolant Pressure", "Lub Oil Temp", "Coolant Temp",
                   "Temp_Difference", "Pressure_Ratio", "High_RPM"]

X_train = train_df[feature_columns]
y_train = train_df["Engine Condition"]
X_test = test_df[feature_columns]
y_test = test_df["Engine Condition"]

# Scale features
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

Your data is now fully prepared. You have 9 carefully chosen features, properly scaled and split into training and test sets. Next up, you will train two competing models and see which one better predicts engine health.

Train and Compare Models

Your data is prepped and scaled. Now it's time for the core of any predictive maintenance system: training classifiers that can tell you whether an engine needs attention before it breaks down.

In this step, you'll train two different machine learning models on the same data and compare their performance. Random Forest uses many decision trees voting together, while XGBoost builds trees sequentially, with each one correcting the mistakes of the last.

In this step, get ready to:
  • Train a Random Forest classifier on your scaled sensor data.
  • Train an XGBoost classifier on the same data for comparison.
  • Evaluate both models with classification reports, confusion matrices, and feature importance plots.
Train the Random Forest classifier

A Random Forest trains many decision trees in parallel, each on a random subset of data. The final prediction comes from a majority vote across all trees.

  • In train_model.py, find the line that says X_test_scaled = scaler.transform(X_test) at the bottom of the file.
  • Add the following code below it to train your Random Forest model:
# Train Random Forest
print("\nTraining Random Forest...")
rf_model = RandomForestClassifier(n_estimators=100, random_state=42)
rf_model.fit(X_train_scaled, y_train)
rf_preds = rf_model.predict(X_test_scaled)

print("\nRandom Forest Classification Report:")
print(classification_report(y_test, rf_preds, target_names=["Normal", "Maintenance"]))

What does this code do?

RandomForestClassifier(n_estimators=100, random_state=42) creates a forest of 100 decision trees. The random_state ensures you get reproducible results every run.

classification_report prints precision, recall, and F1-score for each class. These metrics tell you how well the model catches engines that need maintenance (recall) and how often its maintenance alerts are correct (precision).

You won't see output yet. The next chunk adds XGBoost for comparison before you run the full script.

Train the XGBoost classifier

XGBoost (Extreme Gradient Boosting) takes a different approach. Instead of training trees in parallel, it trains them one after another, with each new tree focusing on the mistakes the previous ones made.

  • Add this code directly below the Random Forest classification report in train_model.py:
# Train XGBoost
print("Training XGBoost...")
xgb_model = xgb.XGBClassifier(
    n_estimators=100,
    learning_rate=0.1,
    max_depth=5,
    random_state=42,
    eval_metric="logloss"
)
xgb_model.fit(X_train_scaled, y_train)
xgb_preds = xgb_model.predict(X_test_scaled)

print("\nXGBoost Classification Report:")
print(classification_report(y_test, xgb_preds, target_names=["Normal", "Maintenance"]))

What does this code do?

learning_rate=0.1 controls how aggressively each new tree corrects errors. A lower rate means more conservative learning, which often improves generalization.

max_depth=5 limits tree complexity to prevent overfitting. eval_metric="logloss" tells XGBoost to optimize for binary classification accuracy using logarithmic loss.

Why compare two models?

Random Forest and XGBoost represent two fundamental ensemble strategies: bagging (parallel trees) and boosting (sequential trees). Comparing them on the same data shows which strategy suits engine sensor data better.

In industry, you'd always compare multiple approaches before selecting one for production deployment.

You won't see output yet. The next substep adds evaluation plots and runs the full script.

Evaluate and save the best model

Classification reports give you the numbers, but confusion matrices and feature importance plots make those numbers visual. They tell you exactly where each model gets confused and which sensors matter most for predictions.

  • Add the following code below the XGBoost classification report in train_model.py:
# Confusion matrices
fig, axes = plt.subplots(1, 2, figsize=(12, 5))

for ax, preds, title in zip(axes, [rf_preds, xgb_preds],
                             ["Random Forest", "XGBoost"]):
    cm = confusion_matrix(y_test, preds)
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", ax=ax,
                xticklabels=["Normal", "Maintenance"],
                yticklabels=["Normal", "Maintenance"])
    ax.set_title(f"{title} Confusion Matrix")
    ax.set_xlabel("Predicted")
    ax.set_ylabel("Actual")

plt.tight_layout()
plt.savefig("plots/confusion_matrices.png")
plt.close()

# Feature importance comparison
fig, axes = plt.subplots(1, 2, figsize=(14, 6))

rf_importance = pd.Series(rf_model.feature_importances_, index=feature_columns)
rf_importance.sort_values().plot(kind="barh", ax=axes[0], color="steelblue")
axes[0].set_title("Random Forest Feature Importance")

xgb_importance = pd.Series(xgb_model.feature_importances_, index=feature_columns)
xgb_importance.sort_values().plot(kind="barh", ax=axes[1], color="darkorange")
axes[1].set_title("XGBoost Feature Importance")

plt.tight_layout()
plt.savefig("plots/feature_importance.png")
plt.close()

# Save the better model (XGBoost typically performs better on this data)
rf_f1 = classification_report(y_test, rf_preds, output_dict=True)["weighted avg"]["f1-score"]
xgb_f1 = classification_report(y_test, xgb_preds, output_dict=True)["weighted avg"]["f1-score"]

if xgb_f1 >= rf_f1:
    best_model = xgb_model
    print(f"\nXGBoost selected as best model (F1: {xgb_f1:.4f})")
else:
    best_model = rf_model
    print(f"\nRandom Forest selected as best model (F1: {rf_f1:.4f})")

joblib.dump(best_model, "models/best_model.joblib")
joblib.dump(scaler, "models/scaler.joblib")
joblib.dump(feature_columns, "models/feature_columns.joblib")

print("\nModel and scaler saved to models/ directory.")
print("Plots saved to plots/ directory.")
print("Done!")

What does this code do?

  • The confusion matrix shows where each model gets it right and wrong. The top-left cell counts correctly identified normal engines. The bottom-right counts correctly caught maintenance cases.
  • Feature importance bars reveal which sensors each model relies on most. If both models agree on the top features, those sensors are genuinely predictive.
  • The script compares weighted F1-scores to pick the best model automatically. It saves the winner, the scaler, and the feature column names to the models/ directory using joblib so the Streamlit app can load them later.
  • Save train_model.py.
  • Run the full training script by executing this command in your terminal:
python train_model.py
  • You should see the dataset loading, both classification reports printing with precision/recall/F1 scores, and a final message saying which model was selected and that files were saved.

What should I see?

Your terminal will print two classification reports (one for Random Forest, one for XGBoost) showing precision, recall, and F1-score columns. At the bottom, you'll see a message like XGBoost selected as best model followed by its F1 score.

Check your plots/ folder. You should see confusion_matrices.png and feature_importance.png. Your models/ folder should contain best_model.joblib, scaler.joblib, and feature_columns.joblib.

Seeing an error when running the script?

  • If you see a NameError for rf_model or xgb_model, make sure the training code comes before the evaluation code in your file.
  • If the dataset download is slow, this is normal on first run. Hugging Face caches it locally for future runs.
  • If you see an indentation error, double-check that the for loop inside the confusion matrix block is properly indented.

Help me debug my train_model.py script

  • Confirm the saved files exist by running:
ls models/
  • You should see best_model.joblib, scaler.joblib, and feature_columns.joblib in the output.

✔️ Awesome, I've got everything!

Great. Double check you've saved your file.

ⓧ I'd like to double check the full code

Here is the complete train_model.py file for comparison:

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datasets import load_dataset
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix
import xgboost as xgb
import joblib

# Create output directories
os.makedirs("plots", exist_ok=True)
os.makedirs("models", exist_ok=True)

# Load dataset from Hugging Face
print("Loading dataset...")
dataset = load_dataset("jskswamy/predictive-maintenance-data")
train_df = dataset["train"].to_pandas()
test_df = dataset["test"].to_pandas()

print(f"Training data shape: {train_df.shape}")
print(f"Test data shape: {test_df.shape}")
print(f"\nFirst few rows:\n{train_df.head()}")

# EDA
print(f"\nMissing values:\n{train_df.isnull().sum()}")
print(f"\nTarget distribution:\n{train_df['Engine Condition'].value_counts()}")
print(f"\nDescriptive statistics:\n{train_df.describe()}")

# Correlation heatmap
plt.figure(figsize=(10, 8))
sns.heatmap(train_df.corr(), annot=True, cmap="coolwarm", center=0)
plt.title("Sensor Feature Correlation Heatmap")
plt.tight_layout()
plt.savefig("plots/correlation_heatmap.png")
plt.close()

# Distribution plots by Engine Condition
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
features = ["Engine RPM", "Lub Oil Pressure", "Fuel Pressure",
            "Coolant Pressure", "Lub Oil Temp", "Coolant Temp"]

for idx, feature in enumerate(features):
    ax = axes[idx // 3, idx % 3]
    for condition in [0, 1]:
        subset = train_df[train_df["Engine Condition"] == condition]
        ax.hist(subset[feature], bins=30, alpha=0.5,
                label=f"{'Normal' if condition == 0 else 'Maintenance'}")
    ax.set_title(feature)
    ax.legend()

plt.tight_layout()
plt.savefig("plots/feature_distributions.png")
plt.close()

# Feature Engineering
def engineer_features(df):
    df = df.copy()
    df["Temp_Difference"] = df["Lub Oil Temp"] - df["Coolant Temp"]
    df["Pressure_Ratio"] = df["Fuel Pressure"] / (df["Lub Oil Pressure"] + 1e-6)
    rpm_75th = train_df["Engine RPM"].quantile(0.75)
    df["High_RPM"] = (df["Engine RPM"] > rpm_75th).astype(int)
    return df

train_df = engineer_features(train_df)
test_df = engineer_features(test_df)

# Prepare features and target
feature_columns = ["Engine RPM", "Lub Oil Pressure", "Fuel Pressure",
                   "Coolant Pressure", "Lub Oil Temp", "Coolant Temp",
                   "Temp_Difference", "Pressure_Ratio", "High_RPM"]

X_train = train_df[feature_columns]
y_train = train_df["Engine Condition"]
X_test = test_df[feature_columns]
y_test = test_df["Engine Condition"]

# Scale features
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# Train Random Forest
print("\nTraining Random Forest...")
rf_model = RandomForestClassifier(n_estimators=100, random_state=42)
rf_model.fit(X_train_scaled, y_train)
rf_preds = rf_model.predict(X_test_scaled)

print("\nRandom Forest Classification Report:")
print(classification_report(y_test, rf_preds, target_names=["Normal", "Maintenance"]))

# Train XGBoost
print("Training XGBoost...")
xgb_model = xgb.XGBClassifier(
    n_estimators=100,
    learning_rate=0.1,
    max_depth=5,
    random_state=42,
    eval_metric="logloss"
)
xgb_model.fit(X_train_scaled, y_train)
xgb_preds = xgb_model.predict(X_test_scaled)

print("\nXGBoost Classification Report:")
print(classification_report(y_test, xgb_preds, target_names=["Normal", "Maintenance"]))

# Confusion matrices
fig, axes = plt.subplots(1, 2, figsize=(12, 5))

for ax, preds, title in zip(axes, [rf_preds, xgb_preds],
                             ["Random Forest", "XGBoost"]):
    cm = confusion_matrix(y_test, preds)
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", ax=ax,
                xticklabels=["Normal", "Maintenance"],
                yticklabels=["Normal", "Maintenance"])
    ax.set_title(f"{title} Confusion Matrix")
    ax.set_xlabel("Predicted")
    ax.set_ylabel("Actual")

plt.tight_layout()
plt.savefig("plots/confusion_matrices.png")
plt.close()

# Feature importance comparison
fig, axes = plt.subplots(1, 2, figsize=(14, 6))

rf_importance = pd.Series(rf_model.feature_importances_, index=feature_columns)
rf_importance.sort_values().plot(kind="barh", ax=axes[0], color="steelblue")
axes[0].set_title("Random Forest Feature Importance")

xgb_importance = pd.Series(xgb_model.feature_importances_, index=feature_columns)
xgb_importance.sort_values().plot(kind="barh", ax=axes[1], color="darkorange")
axes[1].set_title("XGBoost Feature Importance")

plt.tight_layout()
plt.savefig("plots/feature_importance.png")
plt.close()

# Save the better model (XGBoost typically performs better on this data)
rf_f1 = classification_report(y_test, rf_preds, output_dict=True)["weighted avg"]["f1-score"]
xgb_f1 = classification_report(y_test, xgb_preds, output_dict=True)["weighted avg"]["f1-score"]

if xgb_f1 >= rf_f1:
    best_model = xgb_model
    print(f"\nXGBoost selected as best model (F1: {xgb_f1:.4f})")
else:
    best_model = rf_model
    print(f"\nRandom Forest selected as best model (F1: {rf_f1:.4f})")

joblib.dump(best_model, "models/best_model.joblib")
joblib.dump(scaler, "models/scaler.joblib")
joblib.dump(feature_columns, "models/feature_columns.joblib")

print("\nModel and scaler saved to models/ directory.")
print("Plots saved to plots/ directory.")
print("Done!")

What is a confusion matrix?

A confusion matrix is a 2x2 grid that breaks down predictions into four categories: true positives (correctly flagged for maintenance), true negatives (correctly identified as normal), false positives (flagged but actually fine), and false negatives (missed maintenance needs).

For engine health, false negatives are the most dangerous. A missed maintenance case means an engine could fail unexpectedly.

Your models are trained, evaluated, and saved. Next up, you'll build an interactive Streamlit dashboard where anyone can adjust sensor sliders and see real-time maintenance predictions from your best model.

Build the Streamlit Dashboard

You have a trained model that can predict engine health from sensor data. But a saved .joblib file sitting in a folder is not something a fleet manager can use.

In this step, you will build an interactive Streamlit dashboard that loads your best model and lets users adjust sensor readings via sliders. The app will display a real-time maintenance prediction with confidence scores and risk alerts.

In this step, get ready to:
  • Create streamlit_app.py with a page layout and model loading.
  • Build a sidebar with sensor input sliders and computed engineered features.
  • Display predictions, confidence metrics, a risk bar, and conditional alerts.
Create the Streamlit app file and load your model

Streamlit apps are plain Python scripts. When you run one, Streamlit re-executes the entire script from top to bottom every time a user interacts with a widget.

  • Create a new file called streamlit_app.py in your engine-health-predictor folder.
  • Add the following code to set up the page layout and load your saved model:
import streamlit as st
import pandas as pd
import numpy as np
import joblib

st.set_page_config(page_title="Engine Health Predictor", layout="wide")

st.title("Vehicle Engine Health Predictor")
st.markdown(
    "This tool predicts whether a commercial diesel engine needs maintenance "
    "based on real-time sensor readings. Adjust the sensor values in the "
    "sidebar to simulate different engine conditions."
)

# Load model and scaler
model = joblib.load("models/best_model.joblib")
scaler = joblib.load("models/scaler.joblib")
feature_columns = joblib.load("models/feature_columns.joblib")

What does this code do?

st.set_page_config sets the browser tab title and uses the full page width. st.title and st.markdown display the heading and description at the top of the app.

The three joblib.load calls load the model, scaler, and feature column names you saved in the previous step. These are everything the app needs to make predictions.

  • Save streamlit_app.py.

Seeing a file not found error later?

Make sure streamlit_app.py is in the root of your engine-health-predictor folder, not inside models/ or plots/. The joblib.load paths are relative to where you run the app from.

help me fix my Streamlit file path error

Create the sidebar input interface

Streamlit's sidebar is a collapsible panel on the left side of the app. It keeps input controls separate from the main display area so the dashboard stays clean.

  • Add the following code below the model loading section in streamlit_app.py:
# Sidebar inputs
st.sidebar.header("Sensor Inputs")

engine_rpm = st.sidebar.slider("Engine RPM", 0, 2500, 800)
lub_oil_pressure = st.sidebar.slider("Lub Oil Pressure (bar)", 0.0, 10.0, 4.0, 0.1)
fuel_pressure = st.sidebar.slider("Fuel Pressure (bar)", 0.0, 15.0, 6.0, 0.1)
coolant_pressure = st.sidebar.slider("Coolant Pressure (bar)", 0.0, 8.0, 3.0, 0.1)
lub_oil_temp = st.sidebar.slider("Lub Oil Temp (C)", 50.0, 120.0, 75.0, 0.5)
coolant_temp = st.sidebar.slider("Coolant Temp (C)", 50.0, 120.0, 75.0, 0.5)

# Compute engineered features
temp_difference = lub_oil_temp - coolant_temp
pressure_ratio = fuel_pressure / (lub_oil_pressure + 1e-6)
high_rpm = 1 if engine_rpm > 1200 else 0

# Create input dataframe
input_data = pd.DataFrame([[
    engine_rpm, lub_oil_pressure, fuel_pressure, coolant_pressure,
    lub_oil_temp, coolant_temp, temp_difference, pressure_ratio, high_rpm
]], columns=feature_columns)

What does this code do?

Each st.sidebar.slider creates a draggable slider with a label, min, max, default, and step size. The ranges match the sensor data from your training set.

The engineered features section computes the same three derived values you created in train_model.py: temperature difference, pressure ratio, and a high RPM indicator. The model expects these as input alongside the raw sensor values.

Finally, the code packs all 9 features into a single-row DataFrame that matches the column order the scaler and model were trained on.

  • Save streamlit_app.py.

Slider not showing in the sidebar?

Make sure every slider call uses st.sidebar.slider, not st.slider. Without the sidebar prefix, the slider renders in the main content area instead.

help me fix my Streamlit sidebar layout

Display predictions and confidence

The final section of the app takes the slider input, runs it through the model, and displays the result. You will show three key metrics, a visual risk bar, and conditional alerts.

  • Add the following code below the input dataframe section in streamlit_app.py:
# Scale and predict
input_scaled = scaler.transform(input_data)
prediction = model.predict(input_scaled)[0]
probabilities = model.predict_proba(input_scaled)[0]

# Display results
col1, col2, col3 = st.columns(3)

with col1:
    status = "Normal" if prediction == 0 else "Needs Maintenance"
    status_color = "green" if prediction == 0 else "red"
    st.metric("Engine Status", status)

with col2:
    confidence = probabilities[prediction] * 100
    st.metric("Confidence", f"{confidence:.1f}%")

with col3:
    risk = probabilities[1] * 100
    st.metric("Maintenance Risk", f"{risk:.1f}%")

# Risk progress bar
st.subheader("Maintenance Risk Level")
st.progress(probabilities[1])

if probabilities[1] > 0.7:
    st.error("HIGH RISK: Immediate maintenance recommended!")
elif probabilities[1] > 0.4:
    st.warning("MODERATE RISK: Schedule maintenance soon.")
else:
    st.success("LOW RISK: Engine operating normally.")

# Display input values
st.subheader("Current Sensor Readings")
st.dataframe(input_data.T.rename(columns={0: "Value"}), use_container_width=True)

What does this code do?

scaler.transform normalizes the input using the same scaling learned during training. predict_proba returns the probability for each class (Normal vs Needs Maintenance), which powers both the confidence score and the risk bar.

st.columns(3) creates a three-column layout. Each column shows a metric card: the predicted status, the model's confidence in that prediction, and the raw maintenance risk percentage.

st.progress renders a horizontal bar from 0 to 1. The conditional st.error, st.warning, and st.success blocks show color-coded alerts based on the risk level thresholds.

  • Save streamlit_app.py.
  • In your terminal, run the app:
streamlit run streamlit_app.py
  • Your browser should open automatically to http://localhost:8501. You should see the title, three metric cards showing "Normal" status, a green progress bar, and a "LOW RISK" success message.
  • Drag the Engine RPM slider to the far right (2500) and the Lub Oil Pressure slider down to 1.0. Watch the prediction change to "Needs Maintenance" with a red risk bar.

App not opening in the browser?

If the browser does not open automatically, copy the Local URL printed in your terminal (http://localhost:8501) and paste it into your browser manually.

If you see an import error, confirm your virtual environment is still active. You should see (venv) at the start of your terminal prompt.

help me get my Streamlit app running

✔️ Awesome, I've got everything!

Great. Double check you have saved streamlit_app.py and that the app is running in your browser.

ⓧ I'd like to double check the full code

import streamlit as st
import pandas as pd
import numpy as np
import joblib

st.set_page_config(page_title="Engine Health Predictor", layout="wide")

st.title("Vehicle Engine Health Predictor")
st.markdown(
    "This tool predicts whether a commercial diesel engine needs maintenance "
    "based on real-time sensor readings. Adjust the sensor values in the "
    "sidebar to simulate different engine conditions."
)

# Load model and scaler
model = joblib.load("models/best_model.joblib")
scaler = joblib.load("models/scaler.joblib")
feature_columns = joblib.load("models/feature_columns.joblib")

# Sidebar inputs
st.sidebar.header("Sensor Inputs")

engine_rpm = st.sidebar.slider("Engine RPM", 0, 2500, 800)
lub_oil_pressure = st.sidebar.slider("Lub Oil Pressure (bar)", 0.0, 10.0, 4.0, 0.1)
fuel_pressure = st.sidebar.slider("Fuel Pressure (bar)", 0.0, 15.0, 6.0, 0.1)
coolant_pressure = st.sidebar.slider("Coolant Pressure (bar)", 0.0, 8.0, 3.0, 0.1)
lub_oil_temp = st.sidebar.slider("Lub Oil Temp (C)", 50.0, 120.0, 75.0, 0.5)
coolant_temp = st.sidebar.slider("Coolant Temp (C)", 50.0, 120.0, 75.0, 0.5)

# Compute engineered features
temp_difference = lub_oil_temp - coolant_temp
pressure_ratio = fuel_pressure / (lub_oil_pressure + 1e-6)
high_rpm = 1 if engine_rpm > 1200 else 0

# Create input dataframe
input_data = pd.DataFrame([[
    engine_rpm, lub_oil_pressure, fuel_pressure, coolant_pressure,
    lub_oil_temp, coolant_temp, temp_difference, pressure_ratio, high_rpm
]], columns=feature_columns)

# Scale and predict
input_scaled = scaler.transform(input_data)
prediction = model.predict(input_scaled)[0]
probabilities = model.predict_proba(input_scaled)[0]

# Display results
col1, col2, col3 = st.columns(3)

with col1:
    status = "Normal" if prediction == 0 else "Needs Maintenance"
    status_color = "green" if prediction == 0 else "red"
    st.metric("Engine Status", status)

with col2:
    confidence = probabilities[prediction] * 100
    st.metric("Confidence", f"{confidence:.1f}%")

with col3:
    risk = probabilities[1] * 100
    st.metric("Maintenance Risk", f"{risk:.1f}%")

# Risk progress bar
st.subheader("Maintenance Risk Level")
st.progress(probabilities[1])

if probabilities[1] > 0.7:
    st.error("HIGH RISK: Immediate maintenance recommended!")
elif probabilities[1] > 0.4:
    st.warning("MODERATE RISK: Schedule maintenance soon.")
else:
    st.success("LOW RISK: Engine operating normally.")

# Display input values
st.subheader("Current Sensor Readings")
st.dataframe(input_data.T.rename(columns={0: "Value"}), use_container_width=True)

Your engine health predictor is live. You now have a complete end-to-end pipeline: raw sensor data goes in, a trained model processes it, and an interactive dashboard presents the results. In the Secret Mission, you will add SHAP explainability to show exactly which sensors are driving each prediction.

Secret mission

Add SHAP Explainability

Your dashboard predicts maintenance needs, but can you explain WHY? Add SHAP explainability so a fleet manager can see exactly which sensors are driving each prediction. This is a highly valued skill in data science interviews.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs or cloud resources to worry about. Decide whether to keep your project files for future reference, pause by stopping the running server, or delete everything entirely.

Resources you used:

  • The engine-health-predictor project folder (contains code, virtual environment, trained models, and saved plots).
  • The running Streamlit server process in your terminal.

Keep everything running

No action needed. Choose this if you want to keep experimenting with the dashboard or plan to extend the project further.

  • Your Streamlit app continues running at the local URL shown in your terminal.
  • All trained models, plots, and code remain in the engine-health-predictor folder on your Desktop.
  • Your virtual environment stays activated until you close the terminal or run deactivate.

Pause - I'll come back to this later

Stop the running Streamlit server to free up your terminal, but keep all project files for later.

  • In the terminal where Streamlit is running, press Ctrl+C to stop the server.
  • Deactivate the virtual environment by running:
deactivate

When you want to resume, open the engine-health-predictor folder in VS Code, activate the virtual environment, and run streamlit run streamlit_app.py again.

Delete - I don't want to use this again

Remove all project files, the virtual environment, and trained models. This frees up disk space completely.

  • In the terminal where Streamlit is running, press Ctrl+C to stop the server.
  • Deactivate the virtual environment by running:
deactivate
  • Navigate to your Desktop and delete the entire project folder.

macOS

cd ~/Desktop
rm -rf engine-health-predictor

Windows

cd ~\Desktop
rmdir /s /q engine-health-predictor

This removes all Python scripts, the virtual environment (including all installed packages), trained model files, and saved plots.

Nice Work!

Nice Work!

You just built a complete predictive maintenance system for commercial diesel vehicle engines. From raw sensor telemetry to a live interactive dashboard, you took real automotive data through a full machine learning pipeline and delivered a deployable tool.

You learned how to:

  • Load, explore, and feature-engineer a real automotive sensor dataset (19,535 records) to prepare it for binary classification.
  • Train and compare Random Forest and XGBoost classifiers using precision, recall, F1-score, and confusion matrices to select the best-performing model.
  • Build an interactive Streamlit dashboard with sensor input sliders, real-time predictions, and confidence metrics for engine health monitoring.
  • Secret Mission: Added SHAP explainability to make predictions interpretable for non-technical stakeholders using waterfall and summary plots.

Ready to quiz yourself?