retinanet-chess-dataset-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "import os\n",
    "import tempfile\n",
    "import zipfile\n",
    "import tqdm\n",
    "import glob\n",
    "\n",
    "from cral.pipeline import ObjectDetectionPipe\n",
    "from cral.models.object_detection import RetinanetConfig\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Download Chess dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_chess_dataset(dataset_path=tempfile.gettempdir()):\n",
    "    zip_url = 'https://public.roboflow.ai/ds/uBYkFHtqpy?key=HZljsh2sXY'\n",
    "    path_to_zip_file = tf.keras.utils.get_file(\n",
    "        'chess_pieces.zip',\n",
    "        zip_url,\n",
    "        cache_dir=dataset_path, \n",
    "        cache_subdir='',\n",
    "        extract=False)\n",
    "    directory_to_extract_to = os.path.join(dataset_path,'chess_pieces')\n",
    "    with zipfile.ZipFile(path_to_zip_file, 'r') as zip_ref:\n",
    "        zip_ref.extractall(directory_to_extract_to)\n",
    "\n",
    "    images_dir = os.path.join(dataset_path, 'chess_pieces','train')\n",
    "    annotation_dir = os.path.join(dataset_path, 'chess_pieces','train')\n",
    "\n",
    "    for image in tqdm.tqdm(glob.glob(os.path.join(images_dir, '*.jpg'))):\n",
    "        new_name = image.replace('_jpg.rf.', '')\n",
    "        os.rename(image, new_name)\n",
    "\n",
    "        annotation = image.replace('.jpg', '.xml')\n",
    "        new_name = annotation.replace('_jpg.rf.', '')\n",
    "        os.rename(annotation, new_name)\n",
    "\n",
    "    return images_dir, annotation_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_dir, annotation_dir = download_chess_dataset('./')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SetUp pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = ObjectDetectionPipe()\n",
    "\n",
    "pipe.add_data(\n",
    "    train_images_dir=images_dir,\n",
    "    train_anno_dir=annotation_dir,\n",
    "    annotation_format='pascal_voc',\n",
    "    split=0.2)\n",
    "\n",
    "pipe.lock_data()\n",
    "\n",
    "pipe.set_algo(feature_extractor='resnet50', config=RetinanetConfig())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from segmind.keras import CheckpointCallback, KerasCallback\n",
    "\n",
    "keras_cb = KerasCallback()\n",
    "ckpt_cb = CheckpointCallback(\n",
    "            snapshot_interval=10,\n",
    "            snapshot_path='/tmp',\n",
    "            checkpoint_prefix='retinanet_test')\n",
    "pipe.train(\n",
    "    num_epochs=10,\n",
    "    snapshot_prefix='test_retinanet',\n",
    "    snapshot_path='/tmp',\n",
    "    snapshot_every_n=10,\n",
    "    batch_size=4,\n",
    "    #steps_per_epoch=100,\n",
    "    callbacks=[keras_cb, ckpt_cb])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cral.models.object_detection import annotate_image\n",
    "from PIL import Image\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load checkpoint files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_file= '/tmp/test_retinanet_final'\n",
    "prediction_func= pipe.prediction_model(checkpoint_file= checkpoint_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict and annotate Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path= 'chess_pieces/test/5a35ba2ec3e0d0b2b12b1758a8ac29aa_jpg.rf.280f9940defacbb5d840aef65a9257e5.jpg'\n",
    "\n",
    "# make prediction\n",
    "bboxes, scores, labels= prediction_func(img_path)\n",
    "\n",
    "# draw predictions on image and save\n",
    "image_array= np.array(Image.open(img_path))\n",
    "drawn_image= annotate_image(image_array, bboxes, scores, labels, threshold=0.5)\n",
    "drawn_image.save('image_retinanet.jpeg')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}