notebooks/train_mobilenet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import numpy as np\n",
    "from os.path import join\n",
    "import tensorflow as tf\n",
    "\n",
    "from modules.datasets import ImageTargetDataset, RandomConcatDataset, ConcatDataset\n",
    "from modules.segm_transforms import train_transforms, test_transforms, ToTensorColor\n",
    "from modules.metrics import FbSegm\n",
    "from train.train import Model\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.config.threading.set_intra_op_parallelism_threads(16)\n",
    "tf.config.threading.set_inter_op_parallelism_threads(16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_batch_size = 32\n",
    "val_batch_size = 32\n",
    "INPUT_SIZE = (224, 224)\n",
    "AUG_PARAMS = [0.75, 1.25, 0.75, 1.25, 0.6, 1.4]\n",
    "ANG_RANGE = 15\n",
    "\n",
    "device = 'GPU:0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_trns = train_transforms(dataset='picsart', scale_size=INPUT_SIZE, ang_range=ANG_RANGE,\n",
    "                                      augment_params=AUG_PARAMS, add_background=False,\n",
    "                                      crop_scale=0.02)\n",
    "val_trns = test_transforms(dataset='picsart', scale_size=INPUT_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dirs_hq = [\n",
    "    '/workdir/data/datasets/picsart/',\n",
    "    '/workdir/data/datasets/supervisely_person/',\n",
    "]\n",
    "\n",
    "data_dirs_coco = [\n",
    "    '/workdir/data/datasets/coco_person/'\n",
    "#   '/workdir/data/datasets/cityscapes_person/',\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dirs_hq = [join(d, 'train') for d in data_dirs_hq]\n",
    "val_dirs_hq = [join(d, 'val') for d in data_dirs_hq]\n",
    "train_dirs_coco = [join(d, 'train') for d in data_dirs_coco]\n",
    "val_dirs_coco = [join(d, 'val') for d in data_dirs_coco]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset_hq = ImageTargetDataset(train_dirs_hq,\n",
    "                                           train_batch_size,\n",
    "                                           shuffle=True,\n",
    "                                           device=device,\n",
    "                                           **train_trns,\n",
    "                                           IMG_EXTN='.jpg',\n",
    "                                           TRGT_EXTN='.png')\n",
    "val_dataset_hq = ImageTargetDataset(val_dirs_hq,\n",
    "                                           val_batch_size,\n",
    "                                           shuffle=False,\n",
    "                                           device=device,\n",
    "                                           **val_trns,\n",
    "                                           IMG_EXTN='.jpg',\n",
    "                                           TRGT_EXTN='.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset_coco = ImageTargetDataset(train_dirs_coco,\n",
    "                                           train_batch_size,\n",
    "                                           shuffle=True,\n",
    "                                           device=device,\n",
    "                                           **train_trns,\n",
    "                                           IMG_EXTN='.jpg',\n",
    "                                           TRGT_EXTN='.png')\n",
    "val_dataset_coco = ImageTargetDataset(val_dirs_coco,\n",
    "                                           val_batch_size,\n",
    "                                           shuffle=False,\n",
    "                                           device=device,\n",
    "                                           **val_trns,\n",
    "                                           IMG_EXTN='.jpg',\n",
    "                                           TRGT_EXTN='.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = RandomConcatDataset([train_dataset_hq, train_dataset_coco],\n",
    "                                    [0.95, 0.05], size=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Train dataset len:\", len(train_dataset))\n",
    "print(\"Val dataset len:\", len(val_dataset_hq))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vis_dataset(dataset, num_samples=train_batch_size):\n",
    "    for x in dataset:\n",
    "        img, target = x[0], x[1]\n",
    "        for i in range(num_samples):\n",
    "            print(\"Image shape: {}, target shape: {}\".format(img[i].shape, target[i].shape))\n",
    "            plt.imshow(img[i])\n",
    "            plt.imshow(np.squeeze(target[i]), alpha=0.4)\n",
    "            plt.show()\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vis_dataset(train_dataset, 8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vis_dataset(val_dataset_hq, 8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize model params\n",
    "model_name = 'mobilenet_small'\n",
    "n_class=1\n",
    "old_model_path = None  # Or path to the previous saved model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train params\n",
    "n_train = len(train_dataset)\n",
    "n_val = len(val_dataset_hq)\n",
    "\n",
    "loss_name = 'fb_combined'\n",
    "optimizer = 'Adam'\n",
    "lr = 0.00005\n",
    "batch_size = train_batch_size\n",
    "max_epoches = 1000\n",
    "save_directory = '/workdir/data/experiments/mobilenetv3_test'\n",
    "reduce_factor = 0.75\n",
    "epoches_limit = 5\n",
    "early_stoping = 100\n",
    "metrics = [FbSegm(channel_axis=-1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mobilenet_model = Model(device=device,\n",
    "                        model_name=model_name,\n",
    "                        n_class=n_class,\n",
    "                        input_shape=(train_batch_size, INPUT_SIZE[0],INPUT_SIZE[1],3),\n",
    "                        old_model_path=old_model_path, shape=INPUT_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mobilenet_model.prepare_train(train_loader=train_dataset,\n",
    "                              val_loader=val_dataset_hq,\n",
    "                              n_train=n_train,\n",
    "                              n_val=n_val,\n",
    "                              loss_name=loss_name,\n",
    "                              optimizer=optimizer,\n",
    "                              lr = lr,\n",
    "                              batch_size = batch_size,\n",
    "                              max_epoches = max_epoches,\n",
    "                              save_directory = save_directory,\n",
    "                              reduce_factor=reduce_factor,\n",
    "                              epoches_limit=epoches_limit,\n",
    "                              early_stoping=early_stoping,\n",
    "                              metrics=metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "mobilenet_model.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mobilenet_model.validate(val_dataset_hq, n_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test on some images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_imgs = os.listdir('/workdir/data/test_examples/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for img_path in test_imgs:\n",
    "    img_path = os.path.join('/workdir/data/test_examples/', img_path)\n",
    "    test_img = cv2.imread(img_path)\n",
    "    test_img = test_img[:,:,::-1]\n",
    "    test_img = cv2.resize(test_img, INPUT_SIZE)\n",
    "    test_tensor = ToTensorColor()(test_img)\n",
    "    test_tensor = tf.expand_dims(test_tensor, 0)\n",
    "    out = mobilenet_model.predict(test_tensor)\n",
    "    out_img = np.squeeze(out)\n",
    "    print(\"Prediction shape:\", out_img.shape)\n",
    "    plt.imshow(test_img)\n",
    "    plt.imshow((out_img>0.5)*255, alpha=0.4)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": 2
}