train_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration for the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "cwd = os.getcwd()\n",
    "\n",
    "if not os.path.exists('model'):\n",
    "    os.mkdir('model')\n",
    "\n",
    "if not os.path.exists(f'{cwd}/result'):\n",
    "    os.mkdir(f'{cwd}/result')\n",
    "\n",
    "TRAIN_DATASET_PATH = f'{cwd}/dataset/train'\n",
    "VALID_DATASET_PATH = f'{cwd}/dataset/valid'\n",
    "TEST_DATASET_PATH = f'{cwd}/dataset/test'\n",
    "MODEL_PATH = f'{cwd}/model'\n",
    "\n",
    "MODEL = 'efficientdet_lite0'\n",
    "MODEL_NAME = 'fish.tflite'\n",
    "CLASSES = ['fish', 'jellyfish', 'penguin', 'shark', 'puffin', 'stingray', 'starfish']\n",
    "EPOCHS = 20\n",
    "BATCH_SIZE = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tflite_model_maker import model_spec\n",
    "from tflite_model_maker import object_detector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create dataset for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = object_detector.DataLoader.from_pascal_voc(\n",
    "    TRAIN_DATASET_PATH,\n",
    "    TRAIN_DATASET_PATH,\n",
    "    CLASSES\n",
    ")\n",
    "\n",
    "val_data = object_detector.DataLoader.from_pascal_voc(\n",
    "    VALID_DATASET_PATH,\n",
    "    VALID_DATASET_PATH,\n",
    "    CLASSES\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec = model_spec.get(MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = object_detector.create(\n",
    "    train_data,\n",
    "    model_spec=spec,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    train_whole_model=True,\n",
    "    epochs=EPOCHS,\n",
    "    validation_data=val_data\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.evaluate(val_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.export(export_dir=MODEL_PATH, tflite_filename=MODEL_NAME)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.15 ('envtf39')",
   "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.9.15"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "4ec9bab4747f9b0862eeffc942c248941fd079eff6e0ea072e5b4a42a65614b4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}