backend/train.py

"""
train.py — Crop Disease Detection model training
=================================================
Trains a transfer-learning classifier as described in Chapter Three of the
project: 224x224 input, ImageNet weights, a custom classification head
(GAP -> BN -> Dense(512,ReLU,L2) -> Dropout(0.4) -> Softmax), two-phase
fine-tuning, class weighting, augmentation, label smoothing, and a
ReduceLROnPlateau schedule.

Backbone is selectable with --arch:
    mobilenet    -> MobileNetV2  (fast, light, ~97-98% on this benchmark)
    efficientnet -> EfficientNetB0 (typically ~98-99%; preferred to hit the
                    ~98% accuracy target across the full multi-crop set)

Covers the locally cultivated Ghanaian crops for which labelled data exists
(14 crops, 55 classes). Add any further crop simply by adding a labelled
folder of images — no code change needed. See the Dataset Guide for sources.

Expected dataset layout (ImageFolder style):
    data/
        train/<class_name>/*.jpg
        val/<class_name>/*.jpg
        test/<class_name>/*.jpg

Class names must match the keys in recommendations.json, i.e.:
    maize_healthy maize_gls maize_nclb maize_rust maize_msv maize_faw
    cassava_healthy cassava_cmd cassava_cbsd cassava_cbb
    tomato_healthy tomato_early tomato_late tomato_wilt tomato_septoria tomato_tylcv
    cocoa_healthy cocoa_blackpod cocoa_cssvd cocoa_capsid
    cashew_healthy cashew_anthracnose cashew_gumosis cashew_leafminer
    plantain_healthy plantain_sigatoka plantain_bbtv plantain_panama
    yam_healthy yam_anthracnose yam_mosaic
    pepper_healthy pepper_bacterialspot pepper_anthracnose
    cowpea_healthy cowpea_blight cowpea_mosaic cowpea_cercospora
    groundnut_healthy groundnut_leafspot groundnut_rosette groundnut_rust
    rice_healthy rice_blast rice_blb rice_brownspot
    okra_healthy okra_yvmv okra_leafspot
    gardenegg_healthy gardenegg_wilt gardenegg_leafspot
    mango_healthy mango_anthracnose mango_bacterialspot

Target performance: ~98% test accuracy. This is consistent with the
published literature on this benchmark (Mohanty et al. 2016 = 99.35%,
Ferentinos 2018 = 99.53%) and is reported honestly from the held-out test
set at the end of this script — it is not assumed.

Run:
    python train.py --data ./data --arch efficientnet --epochs-head 20 --epochs-fine 30
"""
import argparse, json, os
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models, optimizers, regularizers, callbacks
from tensorflow.keras.applications import MobileNetV2, EfficientNetB0
from sklearn.utils.class_weight import compute_class_weight

IMG_SIZE = 224
BATCH = 32
AUTOTUNE = tf.data.AUTOTUNE
# ImageNet channel statistics (used for standardisation, §3.5.1)
MEAN = tf.constant([0.485, 0.456, 0.406])
STD  = tf.constant([0.229, 0.224, 0.225])


def standardise(x):
    x = tf.cast(x, tf.float32) / 255.0
    return (x - MEAN) / STD


def build_augmenter():
    """Augmentation pipeline approximating §3.5.2 (flip, rotate, jitter, crop)."""
    return tf.keras.Sequential([
        layers.RandomFlip("horizontal"),
        layers.RandomRotation(30 / 360.0),
        layers.RandomZoom(0.3),
        layers.RandomBrightness(0.2),
        layers.RandomContrast(0.3),
    ], name="augment")


def make_dataset(directory, augment=False):
    ds = tf.keras.utils.image_dataset_from_directory(
        directory, image_size=(IMG_SIZE, IMG_SIZE), batch_size=BATCH,
        label_mode="categorical", shuffle=augment)
    class_names = ds.class_names
    aug = build_augmenter()
    def prep(x, y):
        if augment:
            x = aug(x, training=True)
        return standardise(x), y
    ds = ds.map(prep, num_parallel_calls=AUTOTUNE).prefetch(AUTOTUNE)
    return ds, class_names


def build_model(num_classes, arch="mobilenet"):
    if arch == "efficientnet":
        base = EfficientNetB0(input_shape=(IMG_SIZE, IMG_SIZE, 3),
                              include_top=False, weights="imagenet")
    else:
        base = MobileNetV2(input_shape=(IMG_SIZE, IMG_SIZE, 3),
                           include_top=False, weights="imagenet")
    base.trainable = False  # Phase 1: freeze the backbone
    inputs = layers.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
    x = base(inputs, training=False)
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.BatchNormalization()(x)
    x = layers.Dense(512, activation="relu",
                     kernel_regularizer=regularizers.l2(1e-4))(x)
    x = layers.Dropout(0.4)(x)
    outputs = layers.Dense(num_classes, activation="softmax")(x)
    return models.Model(inputs, outputs), base


def class_weights_from_dir(train_dir, class_names):
    counts = []
    labels = []
    for i, c in enumerate(class_names):
        n = len([f for f in os.listdir(os.path.join(train_dir, c))
                 if not f.startswith('.')])
        counts.append(n)
        labels += [i] * n
    weights = compute_class_weight("balanced", classes=np.arange(len(class_names)),
                                   y=np.array(labels))
    return {i: float(w) for i, w in enumerate(weights)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="./data")
    ap.add_argument("--arch", choices=["mobilenet", "efficientnet"], default="mobilenet")
    ap.add_argument("--epochs-head", type=int, default=20)
    ap.add_argument("--epochs-fine", type=int, default=30)
    ap.add_argument("--out", default="model")
    args = ap.parse_args()

    train_ds, class_names = make_dataset(os.path.join(args.data, "train"), augment=True)
    val_ds, _ = make_dataset(os.path.join(args.data, "val"))
    test_ds, _ = make_dataset(os.path.join(args.data, "test"))
    print(f"Backbone: {args.arch} | Classes ({len(class_names)}):", class_names)

    model, base = build_model(len(class_names), arch=args.arch)
    cw = class_weights_from_dir(os.path.join(args.data, "train"), class_names)
    # Label smoothing improves calibration and typically lifts test accuracy slightly.
    loss = tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.05)

    # ---- Phase 1: train the head only (frozen backbone) ----
    model.compile(optimizer=optimizers.Adam(1e-3), loss=loss, metrics=["accuracy"])
    model.fit(train_ds, validation_data=val_ds, epochs=args.epochs_head,
              class_weight=cw,
              callbacks=[callbacks.EarlyStopping(patience=6, restore_best_weights=True)])

    # ---- Phase 2: unfreeze top 30% of the backbone, fine-tune at low LR ----
    base.trainable = True
    cut = int(len(base.layers) * 0.70)
    for layer in base.layers[:cut]:
        layer.trainable = False
    model.compile(optimizer=optimizers.Adam(1e-4), loss=loss, metrics=["accuracy"])
    cbs = [
        callbacks.EarlyStopping(patience=10, restore_best_weights=True),
        callbacks.ReduceLROnPlateau(factor=0.5, patience=5, min_lr=1e-6),
    ]
    model.fit(train_ds, validation_data=val_ds, epochs=args.epochs_fine,
              class_weight=cw, callbacks=cbs)

    # ---- Evaluate on the held-out test set ----
    loss_val, acc = model.evaluate(test_ds)
    print(f"\nTest accuracy: {acc:.4f}  (target ~0.98)")
    if acc < 0.98:
        print("Below the 0.98 target. To close the gap: train EfficientNetB0 "
              "(--arch efficientnet), add more field-condition data, or train longer.")

    os.makedirs(args.out, exist_ok=True)
    model.save(os.path.join(args.out, "crop_model.keras"))
    with open(os.path.join(args.out, "classes.json"), "w") as f:
        json.dump(class_names, f, indent=2)
    print(f"Saved model + classes.json to ./{args.out}/")


if __name__ == "__main__":
    main()