notebooks/train_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook for model training\n",
    "\n",
    "### Training Flow\n",
    "\n",
    "Load Training Data:\n",
    "- Load set of input data and label files\n",
    "- Prefilter positive set only to rule out heavily-masked patches\n",
    "- Create train/test set\n",
    "- Define augmentation parameters\n",
    "\n",
    "Train Model:\n",
    "- Define model architecture\n",
    "- Compile model\n",
    "- Train and evaluate model\n",
    "- Save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import date\n",
    "import os\n",
    "import pickle\n",
    "import sys\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from sklearn.metrics import classification_report\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.utils import shuffle\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "parent_dir = os.path.split(os.getcwd())[0]\n",
    "if parent_dir not in sys.path:\n",
    "    sys.path.insert(0, parent_dir)\n",
    "\n",
    "from scripts import dl_utils\n",
    "from scripts import viz_tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resolution = 48\n",
    "\n",
    "data_file_names = [\n",
    "                    \"bolivar_2020_thresh_0.8_1_negatives\",\n",
    "                    \"amazonas_2020_thresh_0.5_2_negatives\",\n",
    "                    \"riverbank_negatives\",\n",
    "                    \"amazonas_2020_thresh_0.8_sumbsample_3_positives\",\n",
    "                    \"MinesPos2018-2020Sentinel_points\",\n",
    "                    \"bolivar_2020_thresh_0.8_sumbsample_5_positives\",\n",
    "                    \"full_amazon_v9_negatives\",\n",
    "                    \"v2.0_bolivar_negatives\",\n",
    "                    \"v2.1.1_bolivar_negatives\",\n",
    "\n",
    "              ]\n",
    "\n",
    "data_file_names = [\n",
    "                    'bolivar_2020_thresh_0.8_1_negatives',\n",
    "                    'amazonas_2020_thresh_0.5_2_negatives',\n",
    "                    'riverbank_negatives',\n",
    "                    'amazonas_2020_thresh_0.8_sumbsample_3_positives',\n",
    "                    'MinesPos2018-2020Sentinel_points',\n",
    "                    #'bolivar_2020_thresh_0.8_sumbsample_5_positives',\n",
    "                    'full_amazon_v9_negatives',\n",
    "                    'v2.0_bolivar_negatives',\n",
    "                    'v2.1.1_bolivar_negatives',\n",
    "                    'v2.4_amazonas_negatives',\n",
    "                    'v2.4_amazon_negatives',\n",
    "                    'v2.4_amazon_positives',\n",
    "                    'v2.6_amazon_thresh_0.8_negatives',\n",
    "                    'v2.6_amazon_negatives',\n",
    "                    'v2.6_amazon_negatives_v2'\n",
    "              ]\n",
    "\n",
    "\n",
    "\n",
    "#start_dates = ['2020-01-01', '2021-01-01', '2022-01-01']\n",
    "#end_dates = ['2021-01-01', '2022-01-01', '2023-01-01']\n",
    "start_dates = ['2019-01-01']\n",
    "end_dates = ['2020-01-01']\n",
    "dates = [f'{sd}_{ed}' for sd,ed in zip(start_dates, end_dates)]\n",
    "data_files = [f'{fn}_{sd}_{ed}_patch_arrays.pkl' for fn in data_file_names for sd, ed in zip(start_dates, end_dates)]\n",
    "label_files = [f.split('s.pkl')[0] + '_labels.pkl' for f in data_files]\n",
    "\n",
    "\n",
    "data_files += [\n",
    "    'v3.1_2023_negatives_2023-01-01_2024-01-01_patch_arrays.pkl', \n",
    "    'amazon_all_48px_v3.1_2023_positives_0.999_2023-01-01_2024-01-01_patch_arrays.pkl', \n",
    "    'v3.2_2023_negatives_2023-01-01_2024-01-01_patch_arrays.pkl',\n",
    "    'v3.3_2023_positives_2023-01-01_2024-01-01_patch_arrays.pkl',\n",
    "    'v3.4_2023_negatives_2023-01-01_2024-01-01_patch_arrays.pkl',\n",
    "    'v3.5_2023_negatives_2023-01-01_2024-01-01_patch_arrays.pkl',\n",
    "    'v3.6_2023_positives_2023-01-01_2024-01-01_patch_arrays.pkl'\n",
    "    ]\n",
    "label_files += [\n",
    "    'v3.1_2023_negatives_2023-01-01_2024-01-01_patch_array_labels.pkl', \n",
    "    'amazon_all_48px_v3.1_2023_positives_0.999_2023-01-01_2024-01-01_patch_array_labels.pkl', \n",
    "    'v3.2_2023_negatives_2023-01-01_2024-01-01_patch_array_labels.pkl',\n",
    "    'v3.3_2023_positives_2023-01-01_2024-01-01_patch_array_labels.pkl',\n",
    "    'v3.4_2023_negatives_2023-01-01_2024-01-01_patch_array_labels.pkl',\n",
    "    'v3.5_2023_negatives_2023-01-01_2024-01-01_patch_array_labels.pkl',\n",
    "    'v3.6_2023_positives_2023-01-01_2024-01-01_patch_array_labels.pkl'\n",
    "    ]\n",
    "\n",
    "patches = []\n",
    "labels = []\n",
    "\n",
    "data_dir = os.path.join('..', 'data', 'training_data', f\"{resolution}_px\")\n",
    "\n",
    "for data, label in tqdm(zip(data_files, label_files), total=len(data_files)):\n",
    "    with open(os.path.join(data_dir, data), 'rb') as f:\n",
    "        data = pickle.load(f)\n",
    "        for elem in data:\n",
    "            #patch = dl_utils.pad_patch(elem, resolution)\n",
    "            patches.append(elem)\n",
    "    with open(os.path.join(data_dir, label), 'rb') as f:\n",
    "        label = pickle.load(f)\n",
    "        labels = np.concatenate((labels, label))\n",
    "\n",
    "patches = np.array(patches)\n",
    "positive_patches = patches[labels == 1]\n",
    "negative_patches = patches[labels == 0]\n",
    "\n",
    "print(len(patches), \"samples loaded\")\n",
    "print(sum(labels == 1), \"positive samples\")\n",
    "print(sum(labels == 0), \"negative samples\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 64\n",
    "indices = np.random.randint(0, len(patches), num_samples)\n",
    "viz_tools.plot_image_grid(patches[indices], labels=[int(label) for label in labels[indices]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_black(data, mask_limit=0.1, return_rejects=False):\n",
    "    masked_fraction = np.array([np.sum(np.mean(patch, axis=-1) < 10) / np.size(np.mean(patch, axis=-1)) for patch in data])\n",
    "    filtered_data = data[masked_fraction < mask_limit]\n",
    "    print(f\"{len(filtered_data) / len(data) :.1%} of data below brightness limit\")\n",
    "    if return_rejects:\n",
    "        rejected_data = data[masked_fraction >= mask_limit]\n",
    "        return filtered_data, rejected_data\n",
    "    else:\n",
    "        return filtered_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter positive pixels that are masked beyond a threshold. Don't want to give positive examples of cloud-masked patches\n",
    "filtered_positives, positive_rejects = filter_black(positive_patches, mask_limit = 0.1, return_rejects=True)\n",
    "if len(positive_rejects) > 0:\n",
    "    num_samples = 16**2\n",
    "    if len(positive_rejects) < num_samples:\n",
    "        num_samples = len(positive_rejects)\n",
    "    indices = np.random.randint(0, len(positive_rejects), num_samples)\n",
    "    viz_tools.plot_numpy_grid(positive_rejects[indices,:,:,3:0:-1] / 3000)\n",
    "    plt.title(\"Positive Masked Rejects\")\n",
    "    plt.show()\n",
    "\n",
    "num_samples = 25 ** 2\n",
    "indices = np.random.randint(0, len(filtered_positives), num_samples)\n",
    "fig = viz_tools.plot_numpy_grid(filtered_positives[indices,:,:,3:0:-1] / 3000)\n",
    "plt.title('Positive Filtered Samples')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_negatives, negative_rejects = filter_black(negative_patches, mask_limit = 0.6, return_rejects=True)\n",
    "if len(negative_rejects) > 0:\n",
    "    num_samples = 16**2\n",
    "    if len(positive_rejects) < num_samples:\n",
    "        num_samples = len(positive_rejects)\n",
    "    indices = np.random.randint(0, len(negative_rejects), num_samples)\n",
    "    viz_tools.plot_numpy_grid(negative_rejects[indices,:,:,3:0:-1] / 3000)\n",
    "    plt.title(\"Negative Mask Rejects\")\n",
    "    plt.show()\n",
    "\n",
    "num_samples = 25 ** 2\n",
    "indices = np.random.randint(0, len(filtered_negatives), num_samples)\n",
    "fig = viz_tools.plot_numpy_grid(filtered_negatives[indices,:,:,3:0:-1] / 3000)\n",
    "plt.title('Negative Filtered Samples')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Data for Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for RGBIR, x = normalize(np.copy(images[:,:,:,[1,2,3,8]]))\n",
    "x = np.concatenate((filtered_negatives, filtered_positives))\n",
    "y = np.concatenate((np.zeros(len(filtered_negatives)), np.ones(len(filtered_positives))))\n",
    "x, y = shuffle(x, y, random_state=33)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Augmentation\n",
    "Augmentation is important to help the model train. These are parameters that generally worked, but there is certainly room for improvement.\n",
    "\n",
    "Plot a set of images to give an example of augmented images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "augmentation_parameters = {\n",
    "    'featurewise_center': False,\n",
    "    'rotation_range': 360,\n",
    "    'width_shift_range': [0.9, 1.1],\n",
    "    'height_shift_range': [0.9, 1.1],\n",
    "    'shear_range': 10,\n",
    "    'zoom_range': [0.9, 1.1],\n",
    "    'vertical_flip': True,\n",
    "    'horizontal_flip': True,\n",
    "    # Fill options: \"constant\", \"nearest\", \"reflect\" or \"wrap\"\n",
    "    'fill_mode': 'reflect'\n",
    "}\n",
    "\n",
    "datagen = ImageDataGenerator(**augmentation_parameters)\n",
    "\n",
    "\n",
    "plt.figure(figsize=(12,12), facecolor=(1,1,1), dpi=150)\n",
    "aug_img, aug_labels = datagen.flow(x, y, batch_size=64).next()\n",
    "viz_tools.plot_image_grid(aug_img, labels=[int(l) for l in aug_labels], norm=True);\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_norm = np.clip(np.array(x.astype(\"float32\") / 10000), 0, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Ensemble"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "for seed, name in zip([1, 2, 3, 4, 5, 6, 7, 8, 9], ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i']):\n",
    "        keras.utils.set_random_seed(seed)\n",
    "        x_train, x_test, y_train, y_test = train_test_split(x_norm, y, test_size=0.30, random_state=seed)\n",
    "        print(f\"Min Value:\\t\\t {np.min(x_train):.3f}\")\n",
    "        print(f\"Max Value:\\t\\t {np.max(x_train):.3f}\")\n",
    "        print(f\"Median Value:\\t\\t {np.median(x_train):.3f}\")\n",
    "        print(\"Num Train:\\t\\t\", len(x_train))\n",
    "        print(\"Num Test:\\t\\t\", len(x_test))\n",
    "        print(f\"Percent Negative Train:\\t {100 * sum(y_train == 0.0) / len(y_train):.1f}\")\n",
    "        print(f\"Percent Negative Test:\\t {100 * sum(y_test == 0.0) / len(y_test):.1f}\")\n",
    "\n",
    "        num_classes = 2\n",
    "        input_shape = x_train.shape[1:]\n",
    "        print(\"Input Shape:\", input_shape)\n",
    "\n",
    "        model = keras.Sequential([\n",
    "                keras.Input(shape=input_shape),\n",
    "                layers.Conv2D(32, kernel_size=(3), padding='same', activation=\"relu\"),\n",
    "                layers.Conv2D(32, kernel_size=(3), padding='same', activation=\"relu\"),\n",
    "                layers.Conv2D(32, kernel_size=(3), padding='same', activation=\"relu\"),\n",
    "                layers.MaxPooling2D(pool_size=(2)),\n",
    "                layers.Conv2D(32, kernel_size=(3), padding='same', activation=\"relu\"),\n",
    "                layers.Conv2D(32, kernel_size=(3), padding='same', activation=\"relu\"),\n",
    "                layers.Conv2D(32, kernel_size=(3), padding='same', activation=\"relu\"),\n",
    "                layers.MaxPooling2D(pool_size=(2)),\n",
    "                layers.Conv2D(32, kernel_size=(3), padding='same', activation=\"relu\"),\n",
    "                layers.Conv2D(32, kernel_size=(3), padding='same', activation=\"relu\"),\n",
    "                layers.Conv2D(32, kernel_size=(3), padding='same', activation=\"relu\"),\n",
    "                layers.MaxPooling2D(pool_size=(3)),\n",
    "                layers.Flatten(),\n",
    "                layers.Dense(64, activation=\"relu\"),\n",
    "                layers.Dropout(0.3),\n",
    "                layers.Dense(64, activation=\"relu\"),\n",
    "                layers.Dropout(0.3),\n",
    "                layers.Dense(32, activation=\"relu\"),\n",
    "                layers.Dropout(0.3),\n",
    "                layers.Dense(1, activation='sigmoid')])\n",
    "        model.summary()\n",
    "        model.compile(\n",
    "        optimizer=keras.optimizers.Adam(3e-4), \n",
    "        loss=keras.losses.BinaryCrossentropy(from_logits=False), \n",
    "        metrics=[keras.metrics.BinaryAccuracy(name=\"acc\")])\n",
    "\n",
    "        train_accuracy = []\n",
    "        test_accuracy = []\n",
    "\n",
    "        batch_size = 32\n",
    "        epochs = 160\n",
    "        # add a weighted loss\n",
    "        class_weight = {0: 1, 1: 1}\n",
    "\n",
    "        model.fit(datagen.flow(x_train, y_train),\n",
    "                batch_size=batch_size, \n",
    "                epochs=epochs, \n",
    "                validation_data = (x_test, y_test),\n",
    "                verbose = 1,\n",
    "                shuffle=True,\n",
    "                class_weight = class_weight\n",
    "                )\n",
    "        train_accuracy += model.history.history['acc']\n",
    "        test_accuracy += model.history.history['val_acc']\n",
    "\n",
    "\n",
    "        plt.figure(figsize=(8,5), dpi=100, facecolor=(1,1,1))\n",
    "        plt.plot([i for i in range(1, len(train_accuracy) + 1)], train_accuracy, label='Train Acc')\n",
    "        plt.plot([i for i in range(1, len(train_accuracy) + 1)], test_accuracy, c='r', label='Val Acc')\n",
    "        percent_negative = (sum(y_train == 0.0) / len(y_train))\n",
    "        plt.plot([1, len(train_accuracy) + 1], [percent_negative, percent_negative], '--', c='gray', label='Baseline')\n",
    "        plt.grid()\n",
    "        plt.xlabel('Epoch')\n",
    "        plt.ylabel('Accuracy')\n",
    "        plt.legend()\n",
    "        plt.title('Network Train and Val Accuracy')\n",
    "        plt.show()\n",
    "\n",
    "        threshold = 0.8\n",
    "        print(\"Test Set Metrics:\")\n",
    "        preds = model.predict(x_test)\n",
    "        print(classification_report(y_test, preds > threshold, \n",
    "                                target_names=['No Mine', 'Mine']))\n",
    "\n",
    "        version_number = f'3.7-{name}'\n",
    "        current_date = date.today()\n",
    "        model_name = f\"{resolution}px_v{version_number}_{current_date.isoformat()}\"\n",
    "\n",
    "        #assert not os.path.exists('../models/' + model_name + '.h5'), f\"Model of name {model_name} already exists\"\n",
    "\n",
    "        with open('../models/' + model_name + '_config.txt', 'w') as f:\n",
    "                f.write('Input Data:\\n')\n",
    "                [f.write('\\t' + file + '\\n') for file in data_files]\n",
    "                f.write('\\n\\nAugmentation Parameters:\\n')\n",
    "                for k, v in zip(augmentation_parameters.keys(), augmentation_parameters.values()):\n",
    "                        f.write(f\"\\t{k}: {v}\\n\")\n",
    "                f.write(f\"\\nBatch Size: {batch_size}\")\n",
    "                f.write(f\"\\nTraining Epochs: {len(train_accuracy)}\")\n",
    "                f.write(f'\\n\\nClassification Report at {threshold}\\n')\n",
    "                f.write(classification_report(y_test, model.predict(x_test) > threshold, \n",
    "                                        target_names=['No Mine', 'Mine']))\n",
    "                        \n",
    "\n",
    "        model.save(f'../models/{model_name}.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate Model Performance Characteristics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find the threshold that maximizes performance on the test set. Note that while this may be the optimum performance on the test set, it does not account for the fact that false positives are functionally worse than false negatives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_model = model\n",
    "\n",
    "val_images = x_test\n",
    "val_labels = y_test\n",
    "\n",
    "test_preds = test_model.predict(val_images)\n",
    "thresh = []\n",
    "score = []\n",
    "for threshold in range(2, 100, 2):\n",
    "    threshold /= 100\n",
    "    thresh.append(threshold)\n",
    "    test_labels = [np.argmax(y) for y in val_labels]\n",
    "    test_out = [pred > threshold for pred in test_preds]\n",
    "    score.append(1 - (np.sum(np.array(test_labels) != np.array(test_out)[:,0]) / len(test_labels)))\n",
    "    #print(np.sum(np.array(test_labels) != np.array(test_preds)), \"of\", len(test_labels), \"test set predictions incorrect\")\n",
    "plt.plot(thresh, score)\n",
    "plt.ylabel('Success Rate')\n",
    "plt.xlabel('Threshold')\n",
    "plt.title(f\"Optimal Threshold: {thresh[np.argmax(score)]:.2f}\")\n",
    "plt.show()\n",
    "optimal_threshold = thresh[np.argmax(score)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot images that the model classifies incorrectly. Can be useful to evaluate model bias."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "threshold = 0.8\n",
    "test_model = model\n",
    "val_images = x_test\n",
    "val_labels = y_test\n",
    "test_labels = val_labels\n",
    "test_preds = test_model.predict(val_images)\n",
    "for index, (label, pred, img) in enumerate(zip(test_labels, test_preds, val_images)):\n",
    "    pred = pred[0]\n",
    "    if pred < threshold:\n",
    "        binary_pred = 0\n",
    "    else:\n",
    "        binary_pred = 1\n",
    "    if label != binary_pred:\n",
    "        rgb = (img[:,:,3:0:-1] * 10000 / 3000)\n",
    "        fig = plt.figure(figsize=(2,2), facecolor=(1,1,1), dpi=150)\n",
    "        plt.imshow(np.clip(rgb, 0, 1))\n",
    "        plt.title(f\"label: {label} - pred: {pred:.2f}\")\n",
    "        plt.axis('off')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.7 ('mining-detector')",
   "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.10.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "bec97cbb607180795486aa419a93884fe3d0b55501c3e5098d64200fe61c3ffb"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}