train.py

from config import TRAIN_DIR, TEST_DIR, CLASSIFIER_TYPE
from data import prepare_dataset
from model import train_classifier, evaluate_classifier
import joblib

def main():
    print("🔍 Loading data...")
    X_train, y_train, X_test, y_test, label_encoder = prepare_dataset(TRAIN_DIR, TEST_DIR)

    print("🎯 Training model...")
    model = train_classifier(X_train, y_train, CLASSIFIER_TYPE)

    print("📊 Evaluating...")
    evaluate_classifier(model, X_test, y_test, label_encoder)

    # Save label encoder
    joblib.dump(label_encoder, "label_encoder.joblib")

if __name__ == "__main__":
    main()