train_cardd_resnet50.py

print("SCRIPT STARTED")

import os
print("Current working directory:", os.getcwd())

RESULTS_DIR = "results_resnet50"
os.makedirs(RESULTS_DIR, exist_ok=True)

BEST_MODEL_PATH = os.path.join(RESULTS_DIR, "best_cardd_resnet50.keras")
FINAL_MODEL_PATH = os.path.join(RESULTS_DIR, "final_cardd_resnet50.keras")

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from tensorflow.keras.applications import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input

from tensorflow.keras.models import Model, load_model
from tensorflow.keras.layers import Dense, Dropout, GlobalAveragePooling2D
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
from pathlib import Path

BASE_DIR = Path(__file__).resolve().parent

from sklearn.metrics import (
    classification_report,
    multilabel_confusion_matrix,
    ConfusionMatrixDisplay,
    f1_score,
    precision_score,
    recall_score
)
# =========================
# SETTINGS
# =========================

CSV_PATH = "cardd_multilabel.csv"

LABELS = [
    "dent",
    "scratch",
    "crack",
    "glass_shatter",
    "lamp_broken",
    "tire_flat"
]

IMG_SIZE = 224
BATCH_SIZE = 16
EPOCHS = 30


# =========================
# LOAD CSV
# =========================

CSV_PATH = BASE_DIR / "cardd_multilabel.csv"

df = pd.read_csv(CSV_PATH)

print("Dataset shape:", df.shape)
print("\nFirst rows:")
print(df.head())

print("\nSplit distribution:")
print(df["split"].value_counts())

print("\nClass distribution:")
print(df[LABELS].sum())

# =========================
# CLASS DISTRIBUTION GRAPH
# =========================

class_counts = df[LABELS].sum()

plt.figure(figsize=(8, 5))
class_counts.plot(kind="bar")
plt.title("Class Distribution in Dataset")
plt.xlabel("Damage Class")
plt.ylabel("Number of Images")
plt.xticks(rotation=45)
plt.grid(axis="y")
plt.tight_layout()
plt.savefig(os.path.join(RESULTS_DIR, "class_distribution.png"), dpi=300, bbox_inches="tight")
plt.show()

# Separate train / validation / test
train_df = df[df["split"] == "train"].copy()
val_df = df[df["split"] == "val"].copy()
test_df = df[df["split"] == "test"].copy()

print("\nTrain:", train_df.shape)
print("Val:", val_df.shape)
print("Test:", test_df.shape)


# =========================
# IMAGE GENERATORS
# =========================

train_datagen = ImageDataGenerator(
    preprocessing_function=preprocess_input,
    rotation_range=15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    zoom_range=0.1,
    horizontal_flip=True
)

val_test_datagen = ImageDataGenerator(
    preprocessing_function=preprocess_input
)

train_generator = train_datagen.flow_from_dataframe(
    dataframe=train_df,
    x_col="image_path",
    y_col=LABELS,
    target_size=(IMG_SIZE, IMG_SIZE),
    batch_size=BATCH_SIZE,
    class_mode="raw",
    shuffle=True
)

val_generator = val_test_datagen.flow_from_dataframe(
    dataframe=val_df,
    x_col="image_path",
    y_col=LABELS,
    target_size=(IMG_SIZE, IMG_SIZE),
    batch_size=BATCH_SIZE,
    class_mode="raw",
    shuffle=False
)

test_generator = val_test_datagen.flow_from_dataframe(
    dataframe=test_df,
    x_col="image_path",
    y_col=LABELS,
    target_size=(IMG_SIZE, IMG_SIZE),
    batch_size=BATCH_SIZE,
    class_mode="raw",
    shuffle=False
)


# =========================
# BUILD RESNET50 MODEL
# =========================

# Load pretrained ResNet50 without its final ImageNet classification layer
base_model = ResNet50(
    weights="imagenet",
    include_top=False,
    input_shape=(IMG_SIZE, IMG_SIZE, 3)
)

# Freeze ResNet50 first
# This means we use it as a feature extractor
for layer in base_model.layers:
    layer.trainable = False

x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(256, activation="relu")(x)
x = Dropout(0.4)(x)

# Multi-label output:
# sigmoid because each damage type is independent yes/no
outputs = Dense(len(LABELS), activation="sigmoid")(x)

model = Model(inputs=base_model.input, outputs=outputs)

model.compile(
    optimizer=Adam(learning_rate=0.0001),
    loss="binary_crossentropy",
    metrics=["binary_accuracy"]
)

model.summary()


# =========================
# TRAIN MODEL
# =========================

checkpoint = ModelCheckpoint(
    BEST_MODEL_PATH,
    monitor="val_loss",
    save_best_only=True,
    verbose=1
)

early_stop = EarlyStopping(
    monitor="val_loss",
    patience=3,
    restore_best_weights=True
)

history = model.fit(
    train_generator,
    validation_data=val_generator,
    epochs=EPOCHS,
    callbacks=[checkpoint, early_stop]
)

# Load the best model saved according to validation loss
model = load_model(BEST_MODEL_PATH)
print("\nLoaded best model from:", BEST_MODEL_PATH)

# =========================
# FIND BETTER THRESHOLDS USING VALIDATION SET
# =========================
# Default threshold is 0.5 for all classes.
# For weaker classes, we search for a better threshold using validation F1-score.

BAD_LABELS_TO_TUNE = ["dent", "crack", "lamp_broken"]

thresholds = np.array([0.5] * len(LABELS))

val_generator.reset()
y_val_prob = model.predict(val_generator)
y_val_true = val_df[LABELS].values.astype(int)

if len(y_val_true) != len(y_val_prob):
    raise ValueError("Mismatch between validation labels and predictions.")

threshold_results = []

print("\nThreshold tuning on validation set:")

for i, label in enumerate(LABELS):
    best_threshold = 0.5
    best_f1 = 0.0
    best_precision = 0.0
    best_recall = 0.0

    if label in BAD_LABELS_TO_TUNE:
        for threshold in np.arange(0.10, 0.91, 0.05):
            y_val_pred_label = (y_val_prob[:, i] >= threshold).astype(int)

            f1 = f1_score(y_val_true[:, i], y_val_pred_label, zero_division=0)
            precision = precision_score(y_val_true[:, i], y_val_pred_label, zero_division=0)
            recall = recall_score(y_val_true[:, i], y_val_pred_label, zero_division=0)

            if f1 > best_f1:
                best_f1 = f1
                best_threshold = threshold
                best_precision = precision
                best_recall = recall

        thresholds[i] = best_threshold
    else:
        y_val_pred_label = (y_val_prob[:, i] >= 0.5).astype(int)

        best_f1 = f1_score(y_val_true[:, i], y_val_pred_label, zero_division=0)
        best_precision = precision_score(y_val_true[:, i], y_val_pred_label, zero_division=0)
        best_recall = recall_score(y_val_true[:, i], y_val_pred_label, zero_division=0)

    threshold_results.append({
        "label": label,
        "threshold": thresholds[i],
        "validation_precision": best_precision,
        "validation_recall": best_recall,
        "validation_f1": best_f1
    })

    print(
        f"{label}: threshold={thresholds[i]:.2f}, "
        f"precision={best_precision:.3f}, "
        f"recall={best_recall:.3f}, "
        f"f1={best_f1:.3f}"
    )

threshold_df = pd.DataFrame(threshold_results)
threshold_df.to_csv(os.path.join(RESULTS_DIR, "threshold_tuning_results.csv"), index=False)

print("\nFinal thresholds used:")
for label, threshold in zip(LABELS, thresholds):
    print(f"{label}: {threshold:.2f}")


# =========================
# PLOT TRAINING GRAPHS
# =========================

epochs_range = range(1, len(history.history["loss"]) + 1)

# Loss graph
plt.figure(figsize=(8, 5))
plt.plot(epochs_range, history.history["loss"], label="Train Loss")
plt.plot(epochs_range, history.history["val_loss"], label="Validation Loss")
plt.title("CarDD ResNet50 - Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.grid(True)
plt.savefig(os.path.join(RESULTS_DIR, "loss_curve.png"), dpi=300, bbox_inches="tight")
plt.show()

# Accuracy graph
plt.figure(figsize=(8, 5))
plt.plot(epochs_range, history.history["binary_accuracy"], label="Train Accuracy")
plt.plot(epochs_range, history.history["val_binary_accuracy"], label="Validation Accuracy")
plt.title("CarDD ResNet50 - Binary Accuracy")
plt.xlabel("Epoch")
plt.ylabel("Binary Accuracy")
plt.legend()
plt.grid(True)
plt.savefig(os.path.join(RESULTS_DIR, "accuracy_curve.png"), dpi=300, bbox_inches="tight")
plt.show()


# =========================
# TEST EVALUATION
# =========================

print("\nEvaluating on test set...")

test_generator.reset()
test_loss, test_acc = model.evaluate(test_generator)

print("\nTest loss:", test_loss)
print("Test binary accuracy:", test_acc)

# Predict probabilities
test_generator.reset()
y_prob = model.predict(test_generator)

# Convert probabilities to 0/1 labels
# Apply class-specific thresholds found from the validation set
y_pred = (y_prob >= thresholds).astype(int)

# True labels
y_true = test_df[LABELS].values.astype(int)

if len(y_true) != len(y_prob):
    raise ValueError("Mismatch between true labels and predictions. Check test generator ordering or missing images.")

# Binary accuracy after applying tuned thresholds
test_acc_tuned = np.mean(y_true == y_pred)

print("Test binary accuracy with tuned thresholds:", test_acc_tuned)

# =========================
# MULTI-LABEL 6x6 CO-OCCURRENCE MATRIX
# =========================
# Rows = true labels
# Columns = predicted labels
# Diagonal = correctly detected labels
# Off-diagonal = labels predicted together with other true labels

cm_multilabel_6x6 = np.matmul(y_true.T, y_pred)

cm_multilabel_6x6_df = pd.DataFrame(
    cm_multilabel_6x6,
    index=LABELS,
    columns=LABELS
)

cm_multilabel_6x6_df.to_csv(
    os.path.join(RESULTS_DIR, "confusion_matrix_multilabel_6x6.csv")
)

fig, ax = plt.subplots(figsize=(8, 6))
im = ax.imshow(cm_multilabel_6x6, cmap="Blues")
plt.colorbar(im)

ax.set_xticks(np.arange(len(LABELS)))
ax.set_yticks(np.arange(len(LABELS)))
ax.set_xticklabels(LABELS, rotation=45, ha="right")
ax.set_yticklabels(LABELS)

ax.set_xlabel("Predicted label")
ax.set_ylabel("True label")
ax.set_title("Multi-label 6x6 Prediction Matrix")

for i in range(cm_multilabel_6x6.shape[0]):
    for j in range(cm_multilabel_6x6.shape[1]):
        ax.text(j, i, cm_multilabel_6x6[i, j], ha="center", va="center", color="black")

plt.tight_layout()
plt.savefig(
    os.path.join(RESULTS_DIR, "confusion_matrix_multilabel_6x6.png"),
    dpi=300,
    bbox_inches="tight"
)
plt.show()


print("\nClassification Report:")

report_text = classification_report(
    y_true,
    y_pred,
    target_names=LABELS,
    zero_division=0
)

print(report_text)

# Save classification report as text
with open(os.path.join(RESULTS_DIR, "classification_report.txt"), "w") as f:
    f.write(report_text)

# Save classification report as CSV
report_dict = classification_report(
    y_true,
    y_pred,
    target_names=LABELS,
    output_dict=True,
    zero_division=0
)

report_df = pd.DataFrame(report_dict).transpose()
report_df.to_csv(os.path.join(RESULTS_DIR, "classification_report.csv"))

# =========================
# TRAIN / VALIDATION / TEST ACCURACY SUMMARY
# =========================

best_epoch = int(np.argmin(history.history["val_loss"]))

best_train_acc = history.history["binary_accuracy"][best_epoch]
best_val_acc = history.history["val_binary_accuracy"][best_epoch]
best_val_loss = history.history["val_loss"][best_epoch]

print("\nBest epoch based on validation loss:", best_epoch + 1)
print("Best epoch train accuracy:", best_train_acc)
print("Best epoch validation accuracy:", best_val_acc)
print("Best epoch validation loss:", best_val_loss)

acc_names = ["Best Epoch Train Accuracy", "Best Epoch Validation Accuracy", "Tuned Test Accuracy"]
acc_values = [best_train_acc, best_val_acc, test_acc_tuned]

plt.figure(figsize=(8, 5))
plt.bar(acc_names, acc_values)
plt.ylim(0, 1)
plt.ylabel("Binary Accuracy")
plt.title("Train / Validation / Test Accuracy Summary")
plt.grid(axis="y")
plt.xticks(rotation=15)
plt.savefig(os.path.join(RESULTS_DIR, "accuracy_summary.png"), dpi=300, bbox_inches="tight")
plt.show()

# =========================
# PRECISION / RECALL / F1 PER CLASS
# =========================

metrics_df = report_df.loc[LABELS, ["precision", "recall", "f1-score"]]

ax = metrics_df.plot(kind="bar", figsize=(10, 6))
ax.set_title("Per-Class Precision, Recall, and F1-score")
ax.set_ylabel("Score")
ax.set_ylim(0, 1)
ax.set_xlabel("Damage Class")
ax.grid(axis="y")
plt.xticks(rotation=45)
plt.tight_layout()
plt.savefig(os.path.join(RESULTS_DIR, "precision_recall_f1_per_class.png"), dpi=300, bbox_inches="tight")
plt.show()


# =========================
# MULTI-LABEL CONFUSION MATRICES
# =========================

print("\nMulti-label Confusion Matrices:")

cms = multilabel_confusion_matrix(y_true, y_pred)

for i, label in enumerate(LABELS):
    print("\nLabel:", label)
    print(cms[i])

    fig, ax = plt.subplots(figsize=(5, 5))

    disp = ConfusionMatrixDisplay(
        confusion_matrix=cms[i],
        display_labels=[f"Not {label}", label]
    )

    disp.plot(ax=ax, values_format="d", colorbar=False)
    plt.title(f"Confusion Matrix - {label}")
    plt.tight_layout()

    plt.savefig(
        os.path.join(RESULTS_DIR, f"confusion_matrix_{label}.png"),
        dpi=300,
        bbox_inches="tight"
    )

    plt.show()


# =========================
# SHOW SOME EXAMPLE PREDICTIONS
# =========================

print("\nExample predictions:")

for i in range(10):
    file_name = test_df.iloc[i]["file_name"]

    true_labels = []
    pred_labels = []

    for j, label in enumerate(LABELS):
        if y_true[i][j] == 1:
            true_labels.append(label)

        if y_pred[i][j] == 1:
            pred_labels.append(label)

    print("\nImage:", file_name)
    print("True:", true_labels)
    print("Predicted:", pred_labels)
    print("Probabilities:", np.round(y_prob[i], 3))


model.save(FINAL_MODEL_PATH)
print("\nSaved final model:", FINAL_MODEL_PATH)