vehicle is masked as 255
Re implementation of U-Net in Tensorflow
- to check how image segmentations can be used for detection problems
Original Paper
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://arxiv.org/abs/1505.04597)
## Summary
Vehicle Detection using U-Net
Objective: detect vehicles
Find a function f such that y = f(X)
| Input | Shape | Explanation | Example |
|---|---|---|---|
| X: 3-D Tensor | (640, 960, 3) | RGB image in an array | |
| y: 3-D Tensor | (640, 960, 1) | Binarized image. Bacground is 0 vehicle is masked as 255 |
|
## Get Started
### Download dataset
- the annotated driving dataset is provided by [Udacity](https://github.com/udacity/self-driving-car/tree/master/annotations)
- In total, 9,423 frames with 65,000 labels at 1920x1200 resolution.
```bash
make download
```
### Resize image and generate mask images
- [utils/data.py](./utils/data.py) is used to resize images and generate masks
```bash
make generate
```
### Train Test Split
Make sure masks and bounding boxes
```bash
jupyter notebook "Visualization & Train Test Split.ipynb"
```
### Train
```bash
# Train for 1 epoch
python train.py
```
or
```bash
$ python train.py --help
usage: train.py [-h] [--epochs EPOCHS] [--batch-size BATCH_SIZE]
[--logdir LOGDIR] [--reg REG] [--ckdir CKDIR]
optional arguments:
-h, --help show this help message and exit
--epochs EPOCHS Number of epochs (default: 1)
--batch-size BATCH_SIZE
Batch size (default: 4)
--logdir LOGDIR Tensorboard log directory (default: logdir)
--reg REG L2 Regularizer Term (default: 0.1)
--ckdir CKDIR Checkpoint directory (default: models)
```
### Test
- Open the Jupyter notebook file to run against test data
```bash
jupyter notebook "./Test Run After Training.ipynb"
```
================================================
FILE: Test Run After Training.ipynb
================================================
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# U-Net test"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import tensorflow as tf\n",
"import cv2\n",
"import pandas as pd\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"\n",
"from scipy.ndimage.measurements import label\n",
"%matplotlib inline"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"saver = tf.train.import_meta_graph(\"./models/model.ckpt.meta\")"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"sess = tf.InteractiveSession()"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO:tensorflow:Restoring parameters from models/model.ckpt\n"
]
}
],
"source": [
"saver.restore(sess, \"models/model.ckpt\")"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"X, mode = tf.get_collection(\"inputs\")\n",
"pred = tf.get_collection(\"outputs\")[0]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Helper functions"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"def plot_image(image, title=None, **kwargs):\n",
" \"\"\"Plots a single image\n",
"\n",
" Args:\n",
" image (2-D or 3-D array): image as a numpy array (H, W) or (H, W, C)\n",
" title (str, optional): title for a plot\n",
" **kwargs: keyword arguemtns for `plt.imshow`\n",
" \"\"\"\n",
" shape = image.shape\n",
"\n",
" if len(shape) == 3:\n",
" plt.imshow(image, **kwargs)\n",
" elif len(shape) == 2:\n",
" plt.imshow(image, **kwargs)\n",
" else:\n",
" raise TypeError(\n",
" \"2-D array or 3-D array should be given but {} was given\".format(shape))\n",
"\n",
" if title:\n",
" plt.title(title)\n",
"\n",
"\n",
"def plot_two_images(image_A, title_A, image_B, title_B, figsize=(15, 15), kwargs_1={}, kwargs_2={}):\n",
" \"\"\"Plots two images side by side\"\"\"\n",
" plt.figure(figsize=figsize)\n",
" plt.subplot(1, 2, 1)\n",
" plot_image(image_A, title=title_A, **kwargs_1)\n",
"\n",
" plt.subplot(1, 2, 2)\n",
" plot_image(image_B, title=title_B, **kwargs_2)\n",
" \n",
"\n",
"def plot_three_images(image_A, image_B, image_C, figsize=(15, 15)):\n",
" \"\"\"Plots three images side by side\"\"\"\n",
" plt.figure(figsize=figsize)\n",
" \n",
" plt.subplot(1, 3, 1)\n",
" plot_image(image_A)\n",
" \n",
" plt.subplot(1, 3, 2)\n",
" plot_image(image_B)\n",
" \n",
" plt.subplot(1, 3, 3)\n",
" plot_image(image_C)\n",
" \n",
" \n",
"def read_image(image_path, gray=False):\n",
" \"\"\"Returns an image array\n",
"\n",
" Args:\n",
" image_path (str): Path to image.jpg\n",
"\n",
" Returns:\n",
" 3-D array: RGB numpy image array\n",
" \"\"\"\n",
" if gray:\n",
" return cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)\n",
" \n",
" image = cv2.imread(image_path) \n",
" return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Pipeline function"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"def already_drawn_bbox(bbox, left_top, right_bot):\n",
" \n",
" for (a, b), (c, d) in bbox:\n",
" \n",
" if a <= left_top[0] <= c:\n",
" if a <= right_bot[0] <= c:\n",
" if b <= left_top[1] <= d:\n",
" if b <= left_top[1] <= d:\n",
" return True\n",
" \n",
" return False"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"def pipeline(image, threshold=0.9999, image_WH=(960, 640)):\n",
" image = np.copy(image)\n",
" H, W, C = image.shape\n",
" \n",
" if (W, H) != image_WH:\n",
" image = cv2.resize(image, image_WH)\n",
" \n",
" mask_pred = sess.run(pred, feed_dict={X: np.expand_dims(image, 0),\n",
" mode: False})\n",
" \n",
" mask_pred = np.squeeze(mask_pred)\n",
" mask_pred = mask_pred > threshold\n",
" \n",
" labeled_heatmap, n_labels = label(mask_pred)\n",
" \n",
" bbox = []\n",
" \n",
" for i in range(n_labels):\n",
" mask_i = labeled_heatmap == (i + 1)\n",
" \n",
" nonzero = np.nonzero(mask_i)\n",
" \n",
" nonzero_row = nonzero[0]\n",
" nonzero_col = nonzero[1]\n",
" \n",
" left_top = min(nonzero_col), min(nonzero_row)\n",
" right_bot = max(nonzero_col), max(nonzero_row)\n",
" \n",
" if not already_drawn_bbox(bbox, left_top, right_bot):\n",
" image = cv2.rectangle(image, left_top, right_bot, color=(0, 255, 0), thickness=3)\n",
" \n",
" bbox.append((left_top, right_bot))\n",
" \n",
" return image"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"| \n", " | Frame | \n", "Mask | \n", "
|---|---|---|
| 0 | \n", "data_resize/1479498371963069978.jpg | \n", "mask/1479498371963069978.jpg | \n", "
| 1 | \n", "data_resize/1479498371963069978.jpg | \n", "mask/1479498371963069978.jpg | \n", "
| 2 | \n", "data_resize/1479498371963069978.jpg | \n", "mask/1479498371963069978.jpg | \n", "
| 3 | \n", "data_resize/1479498371963069978.jpg | \n", "mask/1479498371963069978.jpg | \n", "
| 4 | \n", "data_resize/1479498371963069978.jpg | \n", "mask/1479498371963069978.jpg | \n", "
| \n", " | xmin | \n", "ymin | \n", "xmax | \n", "ymax | \n", "Frame | \n", "Label | \n", "Preview URL | \n", "Mask | \n", "
|---|---|---|---|---|---|---|---|---|
| 0 | \n", "392 | \n", "284 | \n", "452 | \n", "343 | \n", "data_resize/1479498371963069978.jpg | \n", "Car | \n", "http://crowdai.com/images/Wwj-gorOCisE7uxA/vis... | \n", "mask/1479498371963069978.jpg | \n", "
| 1 | \n", "44 | \n", "293 | \n", "145 | \n", "362 | \n", "data_resize/1479498371963069978.jpg | \n", "Car | \n", "http://crowdai.com/images/Wwj-gorOCisE7uxA/vis... | \n", "mask/1479498371963069978.jpg | \n", "
| 2 | \n", "134 | \n", "291 | \n", "191 | \n", "346 | \n", "data_resize/1479498371963069978.jpg | \n", "Car | \n", "http://crowdai.com/images/Wwj-gorOCisE7uxA/vis... | \n", "mask/1479498371963069978.jpg | \n", "
| 3 | \n", "227 | \n", "278 | \n", "274 | \n", "328 | \n", "data_resize/1479498371963069978.jpg | \n", "Truck | \n", "http://crowdai.com/images/Wwj-gorOCisE7uxA/vis... | \n", "mask/1479498371963069978.jpg | \n", "
| 4 | \n", "274 | \n", "278 | \n", "312 | \n", "322 | \n", "data_resize/1479498371963069978.jpg | \n", "Truck | \n", "http://crowdai.com/images/Wwj-gorOCisE7uxA/vis... | \n", "mask/1479498371963069978.jpg | \n", "