{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ZeaVis Edu — Corn Leaf Disease Classifier\n",
    "\n",
    "Mengklasifikasikan penyakit daun jagung (Bercak Daun, Hawar Daun, Karat Daun, Daun Sehat) menggunakan EfficientNetV2B0 dengan transfer learning.\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Persiapan Lingkungan\n",
    "\n",
    "Mengimpor pustaka, mengatur seed, dan mengoptimalkan konfigurasi.\n",
    "**Presisi float32** digunakan untuk komputasi.\n",
    "Resolusi gambar: **224x224** (EfficientNetV2B0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94627f52",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bad56dcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import zipfile\n",
    "import random\n",
    "from collections import Counter\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import splitfolders\n",
    "from PIL import Image\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import layers, models, callbacks\n",
    "from tensorflow.keras.applications import EfficientNetV2B0\n",
    "from tensorflow.keras.applications.efficientnet_v2 import preprocess_input\n",
    "from tensorflow.keras.optimizers import AdamW\n",
    "from sklearn.metrics import classification_report, confusion_matrix\n",
    "from sklearn.utils.class_weight import compute_class_weight\n",
    "\n",
    "# Detect environment\n",
    "try:\n",
    "    from google.colab import drive\n",
    "    IS_COLAB = True\n",
    "    print(\"Running on Google Colab\")\n",
    "except ModuleNotFoundError:\n",
    "    IS_COLAB = False\n",
    "    print(f\"Running locally (TF {tf.__version__}, GPU: {tf.config.list_physical_devices('GPU')})\")\n",
    "\n",
    "tf.keras.mixed_precision.set_global_policy('float32')\n",
    "\n",
    "# Hyperparams\n",
    "IMG_SIZE = (224, 224)\n",
    "BATCH_SIZE = 32\n",
    "SEED = 42\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "tf.random.set_seed(SEED)\n",
    "\n",
    "AUTOTUNE = tf.data.AUTOTUNE\n",
    "print(f\"Setup OK. IMG={IMG_SIZE}, BATCH={BATCH_SIZE}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83235964",
   "metadata": {},
   "source": [
    "## 2. Download dan Ekstraksi Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c18354b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "if IS_COLAB:\n",
    "    drive.mount('/content/drive')\n",
    "    archive_path = '/content/drive/MyDrive/jagung/dataset_jagung.zip'\n",
    "    destination_path = '/content/dataset_jagung.zip'\n",
    "    extract_path = '/content/dataset'\n",
    "else:\n",
    "    import gdown\n",
    "    base = os.getcwd()\n",
    "    archive_path = os.path.join(base, 'dataset_jagung.zip')\n",
    "    destination_path = archive_path\n",
    "    extract_path = os.path.join(base, 'dataset')\n",
    "    DRIVE_FILE_ID = \"1s0H2lDOQVCixywk5eZXJz2i9jj4JihxJ\"\n",
    "    if not os.path.exists(archive_path):\n",
    "        print(\"Downloading from Google Drive...\")\n",
    "        try:\n",
    "            gdown.download(f\"https://drive.google.com/uc?id={DRIVE_FILE_ID}\", archive_path, quiet=False)\n",
    "        except Exception as e:\n",
    "            print(f\"Download failed: {e}\")\n",
    "\n",
    "if os.path.exists(destination_path):\n",
    "    if not os.path.exists(extract_path) or len(os.listdir(extract_path)) == 0:\n",
    "        os.makedirs(extract_path, exist_ok=True)\n",
    "        print(\"Extracting dataset...\")\n",
    "        try:\n",
    "            with zipfile.ZipFile(destination_path, 'r') as zip_ref:\n",
    "                zip_ref.extractall(path=extract_path)\n",
    "            print(\"Extraction completed!\")\n",
    "        except Exception as e:\n",
    "            print(f\"Extraction failed: {e}\")\n",
    "    else:\n",
    "        print(\"Dataset ready.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e19a255",
   "metadata": {},
   "source": [
    "## 3. Data Cleaning dan Validasi Gambar (RGB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63ee82cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Determine dataset path ---\n",
    "# Colab: zip extracts to dataset/ with dataset_jagung_v1/ subfolder\n",
    "# Local: zip extracts classes directly to dataset/\n",
    "if IS_COLAB:\n",
    "    dataset_path = \"/content/dataset/dataset_jagung_v1\"\n",
    "else:\n",
    "    candidate = os.path.join(extract_path, \"dataset_jagung_v1\")\n",
    "    if os.path.isdir(candidate):\n",
    "        dataset_path = candidate\n",
    "    else:\n",
    "        dataset_path = extract_path  # classes are directly in dataset/\n",
    "\n",
    "print(f\"Dataset path: {dataset_path}\")\n",
    "\n",
    "def clean_image_data(directory):\n",
    "    removed_count = 0\n",
    "    for root, dirs, files in os.walk(directory):\n",
    "        for file in files:\n",
    "            file_path = os.path.join(root, file)\n",
    "            try:\n",
    "                img = Image.open(file_path)\n",
    "                img.verify()\n",
    "                img = Image.open(file_path)\n",
    "                if img.mode != 'RGB':\n",
    "                    img = img.convert('RGB')\n",
    "                    img.save(file_path)\n",
    "            except Exception:\n",
    "                print(f\"Removing: {file_path}\")\n",
    "                os.remove(file_path)\n",
    "                removed_count += 1\n",
    "    return removed_count\n",
    "\n",
    "print(\"Cleaning data...\")\n",
    "removed = clean_image_data(dataset_path)\n",
    "print(f\"Done. {removed} problematic files removed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c93f0a16",
   "metadata": {},
   "source": [
    "## 4. Data Splitting (Train:Validation:Test = 70:15:15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d34ff6",
   "metadata": {},
   "outputs": [],
   "source": [
    "if IS_COLAB:\n",
    "    output_dir = \"/content/dataset_split\"\n",
    "else:\n",
    "    output_dir = os.path.join(os.getcwd(), \"dataset_split\")\n",
    "\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "\n",
    "print(\"Splitting dataset 70:15:15...\")\n",
    "splitfolders.ratio(dataset_path, output=output_dir,\n",
    "                   seed=SEED, ratio=(0.7, 0.15, 0.15),\n",
    "                   group_prefix=None, move=False)\n",
    "print(\"Done.\")\n",
    "\n",
    "train_dir = os.path.join(output_dir, 'train')\n",
    "val_dir = os.path.join(output_dir, 'val')\n",
    "test_dir = os.path.join(output_dir, 'test')\n",
    "\n",
    "# Count\n",
    "def count_images(path):\n",
    "    return sum(len(files) for _, _, files in os.walk(path))\n",
    "\n",
    "print(f'Train: {count_images(train_dir)} | Val: {count_images(val_dir)} | Test: {count_images(test_dir)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Data Loader dengan Augmentasi Lanjutan\n",
    "\n",
    "Menggunakan pipeline augmentasi berlapis:\n",
    "- **Geometric**: RandomFlip, Rotation, Zoom, Translation, Contrast, Brightness\n",
    "- **CutMix & MixUp**: mencampur dua gambar dan label (probabilitas 50% masing-masing)\n",
    "- **RandomErasing**: menghapus area acak pada gambar\n",
    "- **GaussianNoise** di dalam model (layer terpisah)\n",
    "- **Class weights**: menangani ketidakseimbangan kelas (capped max 2.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Load datasets ---\n",
    "train_ds = tf.keras.utils.image_dataset_from_directory(\n",
    "    train_dir, shuffle=True, batch_size=BATCH_SIZE, image_size=IMG_SIZE)\n",
    "val_ds = tf.keras.utils.image_dataset_from_directory(\n",
    "    val_dir, shuffle=True, batch_size=BATCH_SIZE, image_size=IMG_SIZE)\n",
    "test_ds = tf.keras.utils.image_dataset_from_directory(\n",
    "    test_dir, shuffle=False, batch_size=BATCH_SIZE, image_size=IMG_SIZE)\n",
    "\n",
    "class_names = train_ds.class_names\n",
    "NUM_CLASSES = len(class_names)\n",
    "print(f\"Classes ({NUM_CLASSES}): {class_names}\")\n",
    "\n",
    "# --- Helper: one-hot ---\n",
    "def to_one_hot(image, label):\n",
    "    return image, tf.one_hot(label, NUM_CLASSES)\n",
    "\n",
    "# --- Beta distribution for MixUp/CutMix ---\n",
    "def sample_beta_distribution(size, concentration_0=0.2, concentration_1=0.2):\n",
    "    gamma_1 = tf.random.gamma(shape=[size], alpha=concentration_1, dtype=tf.float32)\n",
    "    gamma_2 = tf.random.gamma(shape=[size], alpha=concentration_0, dtype=tf.float32)\n",
    "    return gamma_2 / (gamma_1 + gamma_2 + 1e-8)\n",
    "\n",
    "# --- MixUp ---\n",
    "def mix_up(images, labels, alpha=0.2):\n",
    "    batch_size = tf.shape(images)[0]\n",
    "    lambda_param = sample_beta_distribution(batch_size, alpha, alpha)\n",
    "    lambda_param = tf.reshape(lambda_param, [batch_size, 1, 1, 1])\n",
    "    random_indices = tf.random.shuffle(tf.range(batch_size))\n",
    "    mixed_images = lambda_param * images + (1 - lambda_param) * tf.gather(images, random_indices)\n",
    "    labels = tf.cast(labels, tf.float32)\n",
    "    mixed_labels = lambda_param[...,0,0] * labels + (1 - lambda_param[...,0,0]) * tf.gather(labels, random_indices)\n",
    "    return mixed_images, mixed_labels\n",
    "\n",
    "# --- CutMix ---\n",
    "def cut_mix(images, labels, alpha=0.2):\n",
    "    batch_size = tf.shape(images)[0]\n",
    "    lambda_param = sample_beta_distribution(batch_size, alpha, alpha)\n",
    "    random_indices = tf.random.shuffle(tf.range(batch_size))\n",
    "    h, w = tf.shape(images)[1], tf.shape(images)[2]\n",
    "    cut_ratio = tf.sqrt(1.0 - lambda_param)\n",
    "    r_h = tf.cast(cut_ratio * tf.cast(h, tf.float32), tf.int32)\n",
    "    r_w = tf.cast(cut_ratio * tf.cast(w, tf.float32), tf.int32)\n",
    "    cx = tf.random.uniform([], 0, w, tf.int32)\n",
    "    cy = tf.random.uniform([], 0, h, tf.int32)\n",
    "    x1 = tf.clip_by_value(cx - r_w // 2, 0, w)\n",
    "    x2 = tf.clip_by_value(cx + r_w // 2, 0, w)\n",
    "    y1 = tf.clip_by_value(cy - r_h // 2, 0, h)\n",
    "    y2 = tf.clip_by_value(cy + r_h // 2, 0, h)\n",
    "    # Create binary mask for the cut region\n",
    "    mask = tf.ones((h, w, 3), dtype=tf.float32)\n",
    "    ones = tf.ones((y2 - y1, x2 - x1, 3), dtype=tf.float32)\n",
    "    mask = tf.tensor_scatter_nd_update(mask, tf.constant([[y1, x1, 0]]), tf.expand_dims(ones[0,0], 0))\n",
    "    # Simpler approach: pad the patch\n",
    "    patch = tf.image.crop_to_bounding_box(tf.gather(images, random_indices), y1, x1, y2 - y1, x2 - x1)\n",
    "    paddings = [[0, 0], [y1, h - y2], [x1, w - x2], [0, 0]]\n",
    "    patch = tf.pad(patch, paddings, constant_values=0)\n",
    "    mixed_images = images * (1 - tf.cast(tf.cast(patch, tf.bool), tf.float32)) + patch\n",
    "    mixed_images = tf.where(tf.cast(patch, tf.bool), patch, images)\n",
    "    labels = tf.cast(labels, tf.float32)\n",
    "    lambda_reshaped = tf.reshape(lambda_param, [-1, 1])\n",
    "    mixed_labels = lambda_reshaped * labels + (1 - lambda_reshaped) * tf.gather(labels, random_indices)\n",
    "    return mixed_images, mixed_labels\n",
    "\n",
    "# --- RandomErasing ---\n",
    "def random_erasing(images, probability=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3)):\n",
    "    batch_size = tf.shape(images)[0]\n",
    "    h, w = tf.shape(images)[1], tf.shape(images)[2]\n",
    "    target_area = tf.random.uniform([], scale[0], scale[1]) * tf.cast(h * w, tf.float32)\n",
    "    aspect_ratio = tf.random.uniform([], ratio[0], ratio[1])\n",
    "    erasing_h = tf.cast(tf.sqrt(target_area / aspect_ratio), tf.int32)\n",
    "    erasing_w = tf.cast(tf.sqrt(target_area * aspect_ratio), tf.int32)\n",
    "    erasing_h = tf.clip_by_value(erasing_h, 1, h - 1)\n",
    "    erasing_w = tf.clip_by_value(erasing_w, 1, w - 1)\n",
    "    cx = tf.random.uniform([], 0, w - erasing_w, tf.int32)\n",
    "    cy = tf.random.uniform([], 0, h - erasing_h, tf.int32)\n",
    "    noise = tf.random.uniform([batch_size, erasing_h, erasing_w, 3], 0, 255, dtype=tf.float32)\n",
    "    mask = tf.ones([batch_size, h, w, 3], dtype=tf.float32)\n",
    "    ones_patch = tf.zeros([batch_size, erasing_h, erasing_w, 3], dtype=tf.float32)\n",
    "    # Scatter the erase region\n",
    "    updates = tf.ones([batch_size, erasing_h * erasing_w * 3], dtype=tf.float32)\n",
    "    # Use simple approach with slicing\n",
    "    result = tf.identity(images)\n",
    "    # Apply noise patch via scatter_nd\n",
    "    batch_indices = tf.reshape(tf.repeat(tf.range(batch_size), erasing_h * erasing_w * 3), [-1, 1])\n",
    "    y_indices = tf.reshape(tf.repeat(tf.range(cy, cy + erasing_h), erasing_w * 3), [-1, 1])\n",
    "    x_indices = tf.reshape(tf.tile(tf.repeat(tf.range(cx, cx + erasing_w), 3), [erasing_h]), [-1, 1])\n",
    "    c_indices = tf.reshape(tf.tile(tf.range(3), [erasing_h * erasing_w]), [-1, 1])\n",
    "    indices = tf.concat([batch_indices, y_indices, x_indices, c_indices], axis=1)\n",
    "    noise_flat = tf.reshape(noise, [-1])\n",
    "    noise_flat = noise_flat[:tf.shape(indices)[0]]\n",
    "    result = tf.tensor_scatter_nd_update(result, indices, noise_flat)\n",
    "    return tf.cond(tf.random.uniform([]) < probability, lambda: result, lambda: images)\n",
    "\n",
    "# --- Geometric augmentation (kept for TTA compatibility) ---\n",
    "data_augmentation = tf.keras.Sequential([\n",
    "    layers.RandomFlip(\"horizontal_and_vertical\"),\n",
    "    layers.RandomRotation(0.2),\n",
    "    layers.RandomZoom(0.2),\n",
    "    layers.RandomTranslation(0.1, 0.1),\n",
    "    layers.RandomContrast(0.2),\n",
    "    layers.RandomBrightness(0.2)\n",
    "], name=\"data_augmentation\")\n",
    "\n",
    "# --- Augment + MixUp/CutMix + RandomErasing pipeline ---\n",
    "def augment_and_mix(images, labels):\n",
    "    # images are uint8 [0,255], labels are int\n",
    "    # 1. Geometric augmentation\n",
    "    images = data_augmentation(images, training=True)\n",
    "    # 2. Convert to float32 for MixUp/CutMix\n",
    "    images = tf.cast(images, tf.float32)\n",
    "    # 3. Apply CutMix or MixUp with 50% probability\n",
    "    choice = tf.random.uniform([])\n",
    "    labels_onehot = tf.one_hot(labels, NUM_CLASSES)\n",
    "    images, labels_onehot = tf.cond(\n",
    "        choice < 0.3,  # 30% MixUp\n",
    "        lambda: mix_up(images, labels_onehot),\n",
    "        lambda: tf.cond(\n",
    "            choice < 0.6,  # 30% CutMix\n",
    "            lambda: cut_mix(images, labels_onehot),\n",
    "            lambda: (images, labels_onehot)  # 40% no mix\n",
    "        )\n",
    "    )\n",
    "    # 4. Random Erasing\n",
    "    images = random_erasing(images, probability=0.25)\n",
    "    return images, labels_onehot\n",
    "\n",
    "# --- Preprocessing (EfficientNetV2: [-1,1]) ---\n",
    "def preprocess_fn(image, label):\n",
    "    return preprocess_input(image), label\n",
    "\n",
    "# --- Class Weights ---\n",
    "train_class_counts = Counter()\n",
    "for cn in class_names:\n",
    "    p = os.path.join(train_dir, cn)\n",
    "    if os.path.exists(p):\n",
    "        train_class_counts[cn] = len(os.listdir(p))\n",
    "\n",
    "y_integer = []\n",
    "for i, cn in enumerate(class_names):\n",
    "    y_integer.extend([i] * train_class_counts[cn])\n",
    "\n",
    "class_weights_array = compute_class_weight('balanced', classes=np.unique(y_integer), y=y_integer)\n",
    "class_weights_capped = [min(w, 2.5) for w in class_weights_array]\n",
    "class_weights_tensor = tf.constant(class_weights_capped, dtype=tf.float32)\n",
    "\n",
    "def add_sample_weight(image, label):\n",
    "    class_indices = tf.argmax(label, axis=-1)\n",
    "    sample_weights = tf.gather(class_weights_tensor, class_indices)\n",
    "    return image, label, sample_weights\n",
    "\n",
    "print(\"Class weights (capped at 2.5):\")\n",
    "for i, cn in enumerate(class_names):\n",
    "    print(f\"  {cn}: {class_weights_capped[i]:.4f}\")\n",
    "\n",
    "# --- Build pipelines ---\n",
    "# Train: augment → mix → erase → preprocess → onehot → sample_weight → prefetch\n",
    "train_ds = (train_ds\n",
    "    .map(augment_and_mix, num_parallel_calls=AUTOTUNE)\n",
    "    .map(preprocess_fn, num_parallel_calls=AUTOTUNE)\n",
    "    .map(add_sample_weight, num_parallel_calls=AUTOTUNE)\n",
    "    .prefetch(AUTOTUNE))\n",
    "\n",
    "# Val/Test: preprocess → onehot → prefetch\n",
    "val_ds = (val_ds\n",
    "    .map(lambda img, lbl: (preprocess_input(tf.cast(img, tf.float32)), lbl), num_parallel_calls=AUTOTUNE)\n",
    "    .map(to_one_hot, num_parallel_calls=AUTOTUNE)\n",
    "    .prefetch(AUTOTUNE))\n",
    "\n",
    "test_ds = (test_ds\n",
    "    .map(lambda img, lbl: (preprocess_input(tf.cast(img, tf.float32)), lbl), num_parallel_calls=AUTOTUNE)\n",
    "    .map(to_one_hot, num_parallel_calls=AUTOTUNE)\n",
    "    .prefetch(AUTOTUNE))\n",
    "\n",
    "print(\"Data pipelines ready.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Visualisasi Sampel Data (dengan Augmentasi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualisasi augmentasi (non-preprocessed)\n",
    "vis_ds = tf.keras.utils.image_dataset_from_directory(\n",
    "    train_dir, shuffle=True, batch_size=BATCH_SIZE, image_size=IMG_SIZE)\n",
    "\n",
    "plt.figure(figsize=(16, 10))\n",
    "for images, labels in vis_ds.take(1):\n",
    "    # Original (4)\n",
    "    for i in range(4):\n",
    "        plt.subplot(3, 4, i + 1)\n",
    "        plt.imshow(images[i].numpy().astype(\"uint8\"))\n",
    "        plt.title(f\"Asli: {class_names[labels[i].numpy()]}\", fontsize=11)\n",
    "        plt.axis(\"off\")\n",
    "    # Augmented (4) — geometric only\n",
    "    aug = data_augmentation(images, training=True)\n",
    "    for i in range(4):\n",
    "        plt.subplot(3, 4, i + 5)\n",
    "        plt.imshow(aug[i].numpy().astype(\"uint8\"))\n",
    "        plt.title(f\"Aug: {class_names[labels[i].numpy()]}\", fontsize=11)\n",
    "        plt.axis(\"off\")\n",
    "    # MixUp/CutMix result (4)\n",
    "    aug_f = tf.cast(images, tf.float32)\n",
    "    mixed, _ = mix_up(aug_f, tf.one_hot(labels, NUM_CLASSES))\n",
    "    for i in range(4):\n",
    "        plt.subplot(3, 4, i + 9)\n",
    "        plt.imshow(tf.clip_by_value(mixed[i], 0, 255).numpy().astype(\"uint8\"))\n",
    "        plt.title(\"MixUp/CutMix\", fontsize=11)\n",
    "        plt.axis(\"off\")\n",
    "\n",
    "plt.suptitle(\"Contoh Augmentasi Berlapis\", fontsize=16)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Arsitektur Model CNN (EfficientNetV2B0)\n",
    "\n",
    "Base model **EfficientNetV2B0** (imagenet, frozen) + Conv2D tambahan + BatchNorm + Dropout + Dense.\n",
    "Ditambah **GaussianNoise(0.1)** untuk regularisasi tambahan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(num_classes):\n",
    "    base_model = EfficientNetV2B0(\n",
    "        input_shape=IMG_SIZE + (3,),\n",
    "        include_top=False,\n",
    "        weights='imagenet',\n",
    "    )\n",
    "    base_model.trainable = False\n",
    "\n",
    "    inputs = tf.keras.Input(shape=IMG_SIZE + (3,))\n",
    "    x = layers.GaussianNoise(0.1)(inputs)\n",
    "    x = base_model(x, training=False)\n",
    "    x = layers.Conv2D(512, (3, 3), padding='same', activation='swish')(x)\n",
    "    x = layers.BatchNormalization()(x)\n",
    "    x = layers.MaxPooling2D((2, 2))(x)\n",
    "    x = layers.Dropout(0.2)(x)\n",
    "    x = layers.Conv2D(256, (3, 3), padding='same', activation='swish')(x)\n",
    "    x = layers.BatchNormalization()(x)\n",
    "    x = layers.GlobalAveragePooling2D()(x)\n",
    "    x = layers.Dropout(0.3)(x)\n",
    "    x = layers.Dense(1024, activation='swish')(x)\n",
    "    x = layers.BatchNormalization()(x)\n",
    "    x = layers.Dropout(0.4)(x)\n",
    "    outputs = layers.Dense(num_classes, activation='softmax', dtype='float32')(x)\n",
    "    return models.Model(inputs, outputs), base_model\n",
    "\n",
    "if IS_COLAB:\n",
    "    ckpt_dir = '/content/best_model'\n",
    "else:\n",
    "    ckpt_dir = os.path.join(os.getcwd(), 'best_model')\n",
    "checkpoint_path = os.path.join(ckpt_dir, 'best_model.keras')\n",
    "\n",
    "if os.path.exists(checkpoint_path):\n",
    "    print(f\"Loading checkpoint: {checkpoint_path}\")\n",
    "    try:\n",
    "        model = models.load_model(checkpoint_path, compile=False)\n",
    "    except Exception as e:\n",
    "        print(f\"Load failed: {e}. Building fresh.\")\n",
    "        model, base_model = build_model(NUM_CLASSES)\n",
    "else:\n",
    "    print(\"No checkpoint. Building fresh model.\")\n",
    "    model, base_model = build_model(NUM_CLASSES)\n",
    "\n",
    "os.makedirs(ckpt_dir, exist_ok=True)\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Callbacks\n",
    "\n",
    "- **ModelCheckpoint**: simpan yang terbaik (val_accuracy)\n",
    "- **EarlyStopping**: patience 15 (lebih panjang untuk fine-tuning)\n",
    "- **ReduceLROnPlateau**: faktor 0.2, patience 5, min 1e-8\n",
    "- **CSVLogger**: riwayat training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_dir = \"logs/fit/\" + time.strftime(\"%Y%m%d-%H%M%S\")\n",
    "import time\n",
    "\n",
    "checkpoint_cb = callbacks.ModelCheckpoint(\n",
    "    checkpoint_path,\n",
    "    save_best_only=True,\n",
    "    monitor=\"val_accuracy\",\n",
    "    mode=\"max\",\n",
    "    verbose=1\n",
    ")\n",
    "\n",
    "early_stopping_cb = callbacks.EarlyStopping(\n",
    "    monitor=\"val_accuracy\",\n",
    "    patience=15,\n",
    "    restore_best_weights=True,\n",
    "    mode=\"max\",\n",
    "    verbose=1\n",
    ")\n",
    "\n",
    "reduce_lr_cb = callbacks.ReduceLROnPlateau(\n",
    "    monitor='val_loss',\n",
    "    factor=0.2,\n",
    "    patience=5,\n",
    "    min_lr=1e-8,\n",
    "    verbose=1,\n",
    "    mode=\"min\"\n",
    ")\n",
    "\n",
    "csv_logger = callbacks.CSVLogger(os.path.join(ckpt_dir, 'training_log.csv'))\n",
    "\n",
    "callbacks_list = [checkpoint_cb, early_stopping_cb, reduce_lr_cb, csv_logger]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Pelatihan Fase 1 — Head Only\n",
    "\n",
    "Optimizer **AdamW** dengan EMA (Exponential Moving Average) dan **Label Smoothing** 0.2.\n",
    "Base model beku, hanya head yang dilatih."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(\n",
    "    optimizer=AdamW(use_ema=True, ema_momentum=0.999,\n",
    "                    learning_rate=1e-3, weight_decay=1e-4),\n",
    "    loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.2),\n",
    "    metrics=['accuracy']\n",
    ")\n",
    "\n",
    "print(\"Fase 1: Head training (base frozen)...\")\n",
    "history_1 = model.fit(\n",
    "    train_ds,\n",
    "    validation_data=val_ds,\n",
    "    epochs=30,\n",
    "    callbacks=callbacks_list\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Fine-tuning Fase 2 — Unfreeze Layer Atas\n",
    "\n",
    "Membuka **100 layer teratas** base model. Learning rate diturunkan ke 1e-4."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_model = model.layers[1]  # layer[0]=GaussianNoise, layer[1]=EfficientNetV2B0\n",
    "base_model.trainable = True\n",
    "for layer in base_model.layers[:-100]:\n",
    "    layer.trainable = False\n",
    "\n",
    "model.compile(\n",
    "    optimizer=AdamW(use_ema=True, ema_momentum=0.999,\n",
    "                    learning_rate=1e-4, weight_decay=1e-4),\n",
    "    loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.2),\n",
    "    metrics=['accuracy']\n",
    ")\n",
    "\n",
    "print(\"Fase 2: Fine-tuning top 100 layers...\")\n",
    "history_2 = model.fit(\n",
    "    train_ds,\n",
    "    validation_data=val_ds,\n",
    "    epochs=60,\n",
    "    initial_epoch=history_1.epoch[-1] + 1,\n",
    "    callbacks=callbacks_list\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Fine-tuning Fase 3 — Full Unfreeze\n",
    "\n",
    "Membuka **semua layer** base model. Learning rate 5e-5 (sangat kecil agar tidak merusak bobot pretrained)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_model.trainable = True\n",
    "\n",
    "model.compile(\n",
    "    optimizer=AdamW(use_ema=True, ema_momentum=0.999,\n",
    "                    learning_rate=5e-5, weight_decay=1e-4),\n",
    "    loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.2),\n",
    "    metrics=['accuracy']\n",
    ")\n",
    "\n",
    "print(\"Fase 3: Full fine-tuning...\")\n",
    "history_3 = model.fit(\n",
    "    train_ds,\n",
    "    validation_data=val_ds,\n",
    "    epochs=90,\n",
    "    initial_epoch=history_2.epoch[-1] + 1,\n",
    "    callbacks=callbacks_list\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 12. Plot Akurasi dan Loss (Gabungan Semua Fase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gabungkan history\n",
    "acc = history_1.history['accuracy'] + history_2.history['accuracy'] + history_3.history['accuracy']\n",
    "val_acc = history_1.history['val_accuracy'] + history_2.history['val_accuracy'] + history_3.history['val_accuracy']\n",
    "loss = history_1.history['loss'] + history_2.history['loss'] + history_3.history['loss']\n",
    "val_loss = history_1.history['val_loss'] + history_2.history['val_loss'] + history_3.history['val_loss']\n",
    "\n",
    "# Batas antar fase\n",
    "boundary_1 = len(history_1.history['accuracy']) - 1\n",
    "boundary_2 = boundary_1 + len(history_2.history['accuracy'])\n",
    "\n",
    "plt.figure(figsize=(16, 6))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(acc, label='Training Accuracy', linewidth=2)\n",
    "plt.plot(val_acc, label='Validation Accuracy', linewidth=2)\n",
    "plt.axvline(x=boundary_1, color='gray', linestyle='--', alpha=0.7, label='Fase 2 start')\n",
    "plt.axvline(x=boundary_2, color='black', linestyle='--', alpha=0.7, label='Fase 3 start')\n",
    "plt.legend(fontsize=12)\n",
    "plt.title('Training & Validation Accuracy', fontsize=14)\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.grid(alpha=0.3)\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(loss, label='Training Loss', linewidth=2)\n",
    "plt.plot(val_loss, label='Validation Loss', linewidth=2)\n",
    "plt.axvline(x=boundary_1, color='gray', linestyle='--', alpha=0.7, label='Fase 2 start')\n",
    "plt.axvline(x=boundary_2, color='black', linestyle='--', alpha=0.7, label='Fase 3 start')\n",
    "plt.legend(fontsize=12)\n",
    "plt.title('Training & Validation Loss', fontsize=14)\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 13. Test Time Augmentation (TTA)\n",
    "\n",
    "Menggunakan model terbaik (EMA) dan menerapkan augmentasi geometrik saat inferensi untuk meningkatkan akurasi."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load model terbaik\n",
    "print(f\"Loading best model from {checkpoint_path}...\")\n",
    "best_model = tf.keras.models.load_model(checkpoint_path, compile=False)\n",
    "\n",
    "# Kumpulkan gambar test asli\n",
    "raw_test_ds = tf.keras.utils.image_dataset_from_directory(\n",
    "    test_dir, shuffle=False, batch_size=BATCH_SIZE, image_size=IMG_SIZE)\n",
    "\n",
    "test_images = []\n",
    "test_labels_true_raw = []\n",
    "for images, labels in raw_test_ds.unbatch():\n",
    "    test_images.append(images.numpy())\n",
    "    test_labels_true_raw.append(labels.numpy())\n",
    "\n",
    "test_images = np.array(test_images)\n",
    "test_labels_true = tf.one_hot(np.array(test_labels_true_raw), NUM_CLASSES).numpy()\n",
    "\n",
    "TTA_STEPS = 5\n",
    "tta_predictions = []\n",
    "\n",
    "for i in range(TTA_STEPS):\n",
    "    aug_images = data_augmentation(test_images, training=True)\n",
    "    aug_images = preprocess_input(aug_images)\n",
    "    preds = best_model.predict(aug_images, batch_size=BATCH_SIZE, verbose=0)\n",
    "    tta_predictions.append(preds)\n",
    "    print(f\"  TTA step {i+1}/{TTA_STEPS}\")\n",
    "\n",
    "mean_tta = np.mean(tta_predictions, axis=0)\n",
    "\n",
    "test_preds = np.argmax(mean_tta, axis=1)\n",
    "test_true = np.argmax(test_labels_true, axis=1)\n",
    "tta_acc = np.mean(test_preds == test_true)\n",
    "\n",
    "print(f\"\\nTest Accuracy (TTA {TTA_STEPS}x): {tta_acc*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 14. Evaluasi — Classification Report & Confusion Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\nClassification Report:\\n\")\n",
    "print(classification_report(test_true, test_preds, target_names=class_names))\n",
    "\n",
    "cm = confusion_matrix(test_true, test_preds)\n",
    "plt.figure(figsize=(10, 8))\n",
    "sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',\n",
    "            xticklabels=class_names, yticklabels=class_names)\n",
    "plt.title('Confusion Matrix', fontsize=14)\n",
    "plt.xlabel('Predicted', fontsize=12)\n",
    "plt.ylabel('True', fontsize=12)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 15. Simpan Model Akhir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simpan model training final (dengan EMA weights)\n",
    "final_path = os.path.join(ckpt_dir, 'best_model.keras')\n",
    "model.save(final_path)\n",
    "print(f\"Model saved to {final_path}\")\n",
    "\n",
    "# Simpan juga versi tanpa EMA untuk fallback\n",
    "model.save(os.path.join(ckpt_dir, 'final_model.keras'))\n",
    "print(\"Final model saved.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 16. Export Model untuk Produksi\n",
    "\n",
    "Setelah training selesai, jalankan pipeline ekspor secara berurutan:\n",
    "\n",
    "### 1. SavedModel + TFLite\n",
    "```bash\n",
    "python save_model.py\n",
    "```\n",
    "Memuat `best_model/best_model.keras`, membangun arsitektur bersih (tanpa augmentasi), dan mengekspor ke:\n",
    "- `model/saved_model/` — format produksi\n",
    "- `model/model.tflite` — untuk perangkat mobile/edge\n",
    "\n",
    "### 2. ONNX (Rust ML Service)\n",
    "```bash\n",
    "python convert_onnx.py\n",
    "```\n",
    "Mengonversi SavedModel ke `model/model.onnx` untuk digunakan oleh Rust/Axum/ONNX Runtime.\n",
    "\n",
    "### 3. TensorFlow.js (Web)\n",
    "```bash\n",
    "export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python\n",
    "tensorflowjs_converter   --input_format=tf_saved_model   --output_format=tfjs_graph_model   --signature_name=serving_default   --saved_model_tags=serve   model/saved_model   model/tfjs_model\n",
    "```\n",
    "\n",
    "> **Catatan**: Pipeline ekspor terpisah dari notebook karena `save_model.py` membangun ulang arsitektur tanpa layer training (GaussianNoise, augmentasi) untuk produksi."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 17. Model Card — ZeaVis Edu Corn Disease Classifier\n",
    "\n",
    "| Atribut | Detail |\n",
    "|---|---|\n",
    "| **Nama Model** | ZeaVis Edu — EfficientNetV2B0 Classifier |\n",
    "| **Versi** | 2.0 |\n",
    "| **Arsitektur** | EfficientNetV2B0 (Transfer Learning) + GaussianNoise + Conv2D(512) + Conv2D(256) + Dense(1024) |\n",
    "| **Framework** | TensorFlow 2.x / Keras (float32) |\n",
    "| **Dataset** | ~6000-8000 gambar daun jagung (4 kelas) |\n",
    "| **Kelas** | Bercak Daun, Hawar Daun, Karat Daun, Daun Sehat |\n",
    "| **Input** | Gambar RGB 224×224 piksel |\n",
    "| **Output** | Probabilitas per kelas (softmax) |\n",
    "| **Augmentasi** | Flip, Rotation, Zoom, Translation, Contrast, Brightness, MixUp, CutMix, RandomErasing, GaussianNoise |\n",
    "| **Optimizer** | AdamW + EMA + Label Smoothing 0.2 |\n",
    "| **Training** | 3 fase: Head (lr=1e-3) → Partial FT (lr=1e-4) → Full FT (lr=5e-5) |\n",
    "| **Target Akurasi** | ≥95% test accuracy |\n",
    "| **Cara Pakai** | Upload gambar daun jagung → model memprediksi kelas penyakit |\n",
    "| **Etika** | Model ini hanya untuk tujuan edukasi/penelitian pertanian. Jangan gunakan sebagai satu-satunya alat diagnosis. Konsultasi dengan ahli pertanian tetap diperlukan. |\n",
    "\n",
    "---\n",
    "\n",
    "© ZeaVis Edu — Proyek klasifikasi penyakit daun jagung"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}