train.py

import numpy as np
import pandas as pd
from src.models.cnn_model import build_model
from sklearn.utils.class_weight import compute_class_weight
from tensorflow.keras.preprocessing.image import ImageDataGenerator



# 1. load the pre-split data
data = np.load("data/dataset.npz")
X_train, X_test = data["X_train"], data["X_test"]
y_train, y_test = data["y_train"], data["y_test"]

# 2. compute class weights
#    this will give higher weight to the minority class
classes = np.unique(y_train)
weights = compute_class_weight(
    class_weight="balanced",
    classes=classes,
    y=y_train
)
class_weights = dict(zip(classes, weights))
print("Class weights:", class_weights)

# 3. build and train the model
model = build_model(input_shape=(64, 64, 1))
datagen = ImageDataGenerator(
    rotation_range=10,
    width_shift_range=0.1,
    height_shift_range=0.1,
    zoom_range=0.1,
    horizontal_flip=True
)
datagen.fit(X_train)

history = model.fit(
    datagen.flow(X_train, y_train, batch_size=32),
    epochs=10,
    validation_data=(X_test, y_test),
    class_weight=class_weights
)

# 4. save the trained model
model.save("src/models/best_model.h5")

# 5. save training history for plotting
pd.DataFrame(history.history).to_csv("results/history.csv", index=False)
print("Model and history saved in results/")

# 6. evaluate on test set
loss, acc = model.evaluate(X_test, y_test)
print(f"Test Loss: {loss:.4f}, Test Accuracy: {acc:.4f}")