Tracking_using_feature_maps.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tracking using feature maps\n",
    "\n",
    "### Basics\n",
    "**Author**: Paul Willot  \n",
    "**Creation date**: 2019/01/15  \n",
    "**Tested on**: Ubuntu 18.04, and OSX 10.14 (Mojave).\n",
    "\n",
    "\n",
    "### Idea\n",
    "\n",
    "The goal is to **track individual objects through time** in a video.\n",
    "\n",
    "My attempt to solve this rely on a [pre-trained MaskRCNN model](https://github.com/facebookresearch/maskrcnn-benchmark), infering the bounding boxes image by image independently.\n",
    "\n",
    "Obviously as expected from such a common task, solutions have already been proposed, like simply measuring the IOU of bounding boxes over frames ([Simple Online and Realtime Tracking](https://arxiv.org/abs/1602.00763) paper), or a bit more advanced ones like adding a Kalman filter and smarter heuristics ([DeepSORT](https://arxiv.org/abs/1703.07402)).\n",
    "\n",
    "But code for this already exist so it's not fun to just reproduce it! I wanted to try something a bit different.  \n",
    "So below I **extract the feature maps associated with each bounding boxes**, track feature maps/objects relation in time and **compute the distance between objects across frames**.  \n",
    "Objects across frames are then **paired by increasing distance** (objects can only be paired once). Leftover objects are considered new, and objects above a threshold are considered too distant and not paired, and considered new as well.\n",
    "\n",
    "### Data\n",
    "\n",
    "I used a pre-trained model so I don't need to train anything.  \n",
    "So instead of using a nice dataset, I just searched youtube for \"dashcam\", copied the first interesting links then downloaded with `youtube-dl`, and extracted some frames with `ffmpeg`.\n",
    "\n",
    "If you ran `setup.sh` you should now have some data under `local_data`, already separated into images.\n",
    "\n",
    "Of course, that shoud work for any video (and most content, not driving related only), you could do something like this to create a new set:\n",
    "```sh\n",
    "youtube-dl -f 135 \"URL\" -o \"set42.%(ext)s\"\n",
    "mkdir -p set42\n",
    "ffmpeg -i \"set42.mp4\" -ss 00:00:10 -t 00:00:16 -vf fps=4 set42/img_%d.png\n",
    "```\n",
    "\n",
    "### Notes\n",
    "\n",
    "See the markdown blocks for some general comments per section.\n",
    "\n",
    "The code itself is documented and should be pretty straightforward.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bunch of imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "from IPython.display import display\n",
    "\n",
    "import joblib\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "%matplotlib inline\n",
    "\n",
    "from maskrcnn_benchmark.config import cfg\n",
    "from predictor import COCODemo, to_image_list\n",
    "import torch\n",
    "\n",
    "import re\n",
    "from pathlib import Path\n",
    "from collections import OrderedDict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare the pytorch model.\n",
    "First time this is called, it will download about 300Mb of data for the weights.\n",
    "\n",
    "I'm using a pretrained MaskRCNN implementation by Facebook research, despite the pace of DL research it's still close to the state of the art (as of January 2019), so that will do.  \n",
    "Also, it's nicely packaged :D\n",
    "\n",
    "I setup a color palette so that tracked objects keep their color though time.  \n",
    "It's a global for this notebook so not very clean, with a bit more time I would put most of the code below in a class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file = \"./maskrcnn-benchmark/configs/caffe2/e2e_mask_rcnn_R_50_FPN_1x_caffe2.yaml\"\n",
    "\n",
    "cfg.merge_from_file(config_file)\n",
    "cfg.merge_from_list([\"MODEL.DEVICE\", \"cpu\"])\n",
    "\n",
    "coco_demo = COCODemo(\n",
    "    cfg,\n",
    "    min_image_size=800,\n",
    "    confidence_threshold=0.7,\n",
    ")\n",
    "\n",
    "palette = sns.color_palette(\"deep\", 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extract feature maps\n",
    "Get the feature maps associated with each bounding boxes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Caching predictions to disk because they take ~12s each to generate,\n",
    "# feel free to remove it if using GPU or beefier machine.\n",
    "mem = joblib.Memory(\"./joblib_cache\", verbose=0)\n",
    "\n",
    "@mem.cache\n",
    "def get_pred_feat(image):\n",
    "    \"\"\"Make predictions on an image and keep the feature maps.\"\"\"\n",
    "    original_image = image.copy()\n",
    "    \n",
    "    # Setup the image\n",
    "    img = coco_demo.transforms(original_image)\n",
    "    image_list = to_image_list(img, coco_demo.cfg.DATALOADER.SIZE_DIVISIBILITY)\n",
    "    image_list = image_list.to(coco_demo.device)\n",
    "    \n",
    "    # Predict\n",
    "    with torch.no_grad():\n",
    "        targets = None\n",
    "        images = to_image_list(image_list)\n",
    "        features = coco_demo.model.backbone(images.tensors)\n",
    "        proposals, proposal_losses = coco_demo.model.rpn(images, features, targets)\n",
    "        x, result, detector_losses = coco_demo.model.roi_heads(features, proposals, targets)\n",
    "        \n",
    "    # Filter all predictions\n",
    "    scores = result[0].get_field(\"scores\")\n",
    "    keep = torch.nonzero(scores > coco_demo.confidence_threshold).squeeze(1)\n",
    "    pred = result[0][keep]\n",
    "    scores = pred.get_field(\"scores\")\n",
    "    _, idx = scores.sort(0, descending=True)\n",
    "    pred = pred[idx]\n",
    "    \n",
    "    # Resize to fit original image shape\n",
    "    pred = pred.resize(original_image.shape[:-1][::-1])\n",
    "\n",
    "    # Filter the feature maps as well\n",
    "    feat_map = x[keep]\n",
    "    feat_map = feat_map[idx].data.numpy()\n",
    "    \n",
    "    return pred, feat_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Two ways to compute distance\n",
    "\n",
    "`compute_dist` is just the norm of the difference between two sets of feature maps.\n",
    "\n",
    "`compute_bbox_dist` simply measure the distance between the center of two bounding boxes.\n",
    "\n",
    "Turns out that the difference between feature maps was enough for it to work, but I didn't know."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_dist(dict_objs: dict, features: np.ndarray):\n",
    "    \"\"\"Compute the distance between previous features and current ones\"\"\"\n",
    "    dist_mat = np.zeros((len(dict_objs), features.shape[0]))\n",
    "    \n",
    "    for idxi, (k, v) in enumerate(dict_objs.items()):\n",
    "        for idxj, j in enumerate(features):\n",
    "            i = v[\"features\"]\n",
    "            dist = np.linalg.norm(i - j)\n",
    "            dist_mat[idxi, idxj] = dist\n",
    "    \n",
    "    return dist_mat.round(2)\n",
    "\n",
    "def compute_bbox_dist(dict_objs: dict, pred: torch.tensor):\n",
    "    \"\"\"Compute the distance between two bounding boxes\"\"\"\n",
    "    bboxes = pred.bbox.data.numpy()\n",
    "    dist_mat = np.zeros((len(dict_objs), bboxes.shape[0]))\n",
    "    \n",
    "    for idxi, (k, v) in enumerate(dict_objs.items()):\n",
    "        for idxj, j in enumerate(bboxes):\n",
    "            i = v[\"center\"]\n",
    "            assert j[2] > j[0]\n",
    "            center = np.array([\n",
    "                (j[2] - j[0]) / 2,\n",
    "                (j[3] - j[1]) / 2,\n",
    "            ])\n",
    "            dist = np.abs(np.sum((i - center).flatten()))\n",
    "            dist_mat[idxi, idxj] = dist\n",
    "    \n",
    "    return dist_mat.round(2)\n",
    "\n",
    "def merge_dist(\n",
    "    dist_features: np.ndarray, dist_bbox: np.ndarray, alpha: float):\n",
    "    \"\"\"Rescale and merge of distances.\n",
    "    \n",
    "    Alpha control the ratio of each, 1 for distance from features only.\n",
    "    Hardcoded for quick prototyping.\n",
    "    \"\"\"\n",
    "    a = dist_features / 200\n",
    "    b = dist_bbox / 50\n",
    "    return a * alpha + b * (1 - alpha)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keeping track of objects\n",
    "\n",
    "Main method to keep track of objects in time.\n",
    "\n",
    "It's just a big dictionary, with IDs for objects being simply the order in which they appear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_obj_dict(pred: torch.tensor, features: np.ndarray):\n",
    "    \"\"\"Initialize a dict with all objects from current frame.\"\"\"\n",
    "    labels = pred.get_field(\"labels\")\n",
    "    # bbox are each: [x_min, y_min, x_max, y_max]\n",
    "    bboxes = pred.bbox.data.numpy()\n",
    "    current_objs = {\n",
    "        k: {\n",
    "            # age is to keep track of how long the object \n",
    "            # has not been seen, and discard old ones\n",
    "            \"age\": 0,\n",
    "            \"features\": v,\n",
    "            \"color\": palette[k % len(palette)],\n",
    "            \"label\": coco_demo.CATEGORIES[labels[k]],\n",
    "            # center of bbox\n",
    "            \"center\": np.array([\n",
    "                (bboxes[k][2] - bboxes[k][0]) / 2,\n",
    "                (bboxes[k][3] - bboxes[k][1]) / 2,\n",
    "            ])\n",
    "        } for k, v in enumerate(features)}\n",
    "    return OrderedDict(sorted(current_objs.items()))\n",
    "\n",
    "def update_obj_dict(\n",
    "    dict_objs: dict,\n",
    "    features: np.ndarray,\n",
    "    dist_mat: np.ndarray,\n",
    "    pred: torch.tensor,\n",
    "    age_limit: int=2):\n",
    "    \"\"\"Update the dict of objects. Also return objects to plot at this frame.\n",
    "    \n",
    "    This method is a bit big and should be divided but well...\n",
    "    \"\"\"\n",
    "    \n",
    "    # Threshold for distance with previous objects, above is considered new\n",
    "    threshold = 1.\n",
    "    \n",
    "    # index to keep track of objects from previous/current frame\n",
    "    x = np.tile(np.arange(len(features)), len(dict_objs))\n",
    "    y = np.repeat(np.arange(len(dict_objs.keys())), len(features))\n",
    "    z = np.array(sorted(zip(dist_mat.flatten(), x, y), key=lambda x: x[0]))\n",
    "\n",
    "    previous_obj = {i: None for i in range(len(dict_objs.keys()))}\n",
    "    new_obj = {i: None for i in range(len(features))}\n",
    "    \n",
    "    # List of paired objects\n",
    "    match_list = []\n",
    "    try:\n",
    "        to_check = z[z[:, 0] < threshold]\n",
    "        for dist, a, b in to_check:\n",
    "            try:\n",
    "                previous_obj.pop(int(b))\n",
    "                try:\n",
    "                    new_obj.pop(int(a))\n",
    "                except KeyError:\n",
    "                    # Put back the key\n",
    "                    previous_obj[int(b)] = None\n",
    "                    continue\n",
    "                match_list.append((int(a), int(b)))\n",
    "            except KeyError:\n",
    "                continue\n",
    "    except IndexError:\n",
    "        # Nothing above threshold\n",
    "        pass\n",
    "            \n",
    "    new_labels = pred.get_field(\"labels\")\n",
    "    new_bboxes = pred.bbox.data.numpy()\n",
    "\n",
    "    # Prepare new dictionary of updated objects in current frame\n",
    "    new_dict = OrderedDict()\n",
    "    _keys = list(dict_objs.keys())\n",
    "    to_plot = [(a, _keys[b]) for a, b in match_list]\n",
    "    # Objects still here in current frame\n",
    "    still_here = {_keys[x[1]]: x[0] for x in match_list}\n",
    "\n",
    "    for k, v in dict_objs.items():\n",
    "        age = v[\"age\"]\n",
    "        \n",
    "        # If the object is still in the frame, reset it's age\n",
    "        if k in still_here.keys():\n",
    "            age = -1\n",
    "            new_id = still_here[k]\n",
    "            # Update feature map\n",
    "            v[\"features\"] = features[new_id]\n",
    "            # Update bbox center\n",
    "            _bbox = new_bboxes[new_id]\n",
    "            v[\"center\"] = np.array([\n",
    "                (_bbox[2] - _bbox[0]) / 2,\n",
    "                (_bbox[3] - _bbox[1]) / 2,\n",
    "            ])\n",
    "            # Update label\n",
    "            v[\"label\"] = coco_demo.CATEGORIES[new_labels[new_id]]\n",
    "            \n",
    "        # Don't keep it if the object is too old\n",
    "        if age > age_limit:\n",
    "            continue\n",
    "            \n",
    "        # Icrement all ages\n",
    "        v[\"age\"] = age + 1\n",
    "        new_dict[k] = v\n",
    "\n",
    "    # Add new objects\n",
    "    biggest_key = max(new_dict.keys(), default=0)\n",
    "    for k, _ in new_obj.items():\n",
    "        biggest_key += 1\n",
    "        new_dict[biggest_key] = {\n",
    "            \"age\": 0,\n",
    "            \"features\": features[k],\n",
    "            \"color\": palette[biggest_key % len(palette)],\n",
    "            \"label\": coco_demo.CATEGORIES[new_labels[k]],\n",
    "            \"center\": np.array([\n",
    "                (new_bboxes[k][2] - new_bboxes[k][0]) / 2,\n",
    "                (new_bboxes[k][3] - new_bboxes[k][1]) / 2,\n",
    "            ])\n",
    "        }\n",
    "        to_plot.append((k, biggest_key))\n",
    "        \n",
    "    to_plot = [x[1] for x in sorted(to_plot, key=lambda x: x[0])]\n",
    "    return new_dict, to_plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting, loading and saving\n",
    "\n",
    "Convenience methods to plot and save figures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def natsort(l, key=lambda x: x):\n",
    "    \"\"\"Alphanumeric sort, for convenient listing of files\"\"\"\n",
    "    r = re.compile(\"([0-9]+)\")\n",
    "\n",
    "    def alphanum_key(s):\n",
    "        return [int(c) if c.isdigit() else c.lower() for c in r.split(key(s))]\n",
    "\n",
    "    return sorted(l, key=alphanum_key)\n",
    "\n",
    "def fill_bbox(bbox, ax, img, color, crop=False, alpha=0.6, label=None):\n",
    "    t = bbox.data.numpy().round().astype(int)\n",
    "    x0, y0 = t[0], t[1]\n",
    "    x1, y1 = t[2], t[3]\n",
    "    if crop:\n",
    "        x1 = min(img.shape[1], x1)\n",
    "        y1 = min(img.shape[0], y1)\n",
    "    ax.fill([x0, x0, x1, x1], [y0, y1, y1, y0], alpha=alpha, color=color, label=label)\n",
    "    \n",
    "def plot_bbox(img, pred, meta, ids):\n",
    "    fig, ax = plt.subplots(figsize=(16, 6))\n",
    "    colors = []\n",
    "    labels = []\n",
    "    for i in ids:\n",
    "        colors.append(meta[i][\"color\"])\n",
    "        labels.append(\"{:>3}, {:<12}\".format(i, meta[i][\"label\"]))\n",
    "    ax.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))\n",
    "    for idx, bbox in enumerate(pred.bbox):\n",
    "        fill_bbox(bbox, ax, img, colors[idx], label=labels[idx])\n",
    "    \n",
    "    # Sort the legend for readability\n",
    "    hand, lab = ax.get_legend_handles_labels()\n",
    "    if len(hand) > 0:\n",
    "        # sort both labels and handles by labels\n",
    "        lab, hand = zip(*sorted(zip(lab, hand), key=lambda t: int(t[0].split(\",\")[0])))\n",
    "        # Crop legends too long\n",
    "        hand, lab = hand[:20], lab[:20]\n",
    "        ax.legend(hand, lab, loc=6, bbox_to_anchor=(1., 0.5), frameon=False)\n",
    "    \n",
    "    return fig, ax\n",
    "\n",
    "def save_fig(\n",
    "    fig, ax, filename, save_dir=\"./results/\"\n",
    "):\n",
    "    \"\"\"Save a figure and create directory if necessary.\"\"\"\n",
    "    root = Path(save_dir)\n",
    "    root.mkdir(parents=True, exist_ok=True)\n",
    "    path = root.joinpath(filename)\n",
    "    plt.savefig(str(path), bbox_inches=\"tight\", pad_inches=0.1)\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Predicting\n",
    "\n",
    "Now we can just put it together and make our predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_and_update(obj_di: dict, img_path: str, alpha=0.5):\n",
    "    \"\"\"Convenience method to init or update objects dict, and make predictions.\"\"\"\n",
    "    \n",
    "    # Get the image and predict\n",
    "    img = cv2.imread(str(img_path))\n",
    "    pred, fm = get_pred_feat(img)\n",
    "    \n",
    "    if obj_di is None:\n",
    "        # Initialize with objects from current frame\n",
    "        obj_di = init_obj_dict(pred, fm)\n",
    "        to_plot = list(range(len(obj_di)))\n",
    "    else:\n",
    "        # Compare with previous frames\n",
    "        dist_features = compute_dist(obj_di, fm)\n",
    "        dist_bbox = compute_bbox_dist(obj_di, pred)\n",
    "        dist_mat = merge_dist(dist_features, dist_bbox, alpha=alpha)\n",
    "        obj_di, to_plot = update_obj_dict(obj_di, fm, dist_mat, pred)\n",
    "    return img, pred, obj_di, to_plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# No objects to start with\n",
    "obj_di = None\n",
    "\n",
    "# Location of directory containing the images.\n",
    "# Replace with \n",
    "img_dir = Path(\"./local_data/set1/\")\n",
    "# img_dir = Path(\"./local_data/set2/\")\n",
    "# img_dir = Path(\"./local_data/set3/\")\n",
    "\n",
    "# Iter over all in order, and save the predictions\n",
    "for img_path in natsort(img_dir.glob(\"*.png\"), key=lambda x: x.stem):\n",
    "    img, pred, obj_di, to_plot = predict_and_update(obj_di, img_path, alpha=0.8)\n",
    "    \n",
    "    fig, ax = plot_bbox(img, pred, obj_di, to_plot)\n",
    "\n",
    "    save_fig(fig, ax, img_path.stem, save_dir=Path(\"./results\").joinpath(\n",
    "        img_path.parent.name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Result\n",
    "\n",
    "Clearly, results are far from perfect, but it does work in most simple cases.  \n",
    "Work in the sense that objects that don't change orientation drastically are matched from frame to frame and keep their ID. \n",
    "\n",
    "I played with the `alpha` parameter a bit, and not using the distance from bounding boxes (setting `alpha=1`) doesn't make much difference, so the feature map seems to be enough for a simple tracking.\n",
    "\n",
    "Note, I generated the gif using this command:\n",
    "```sh\n",
    "ffmpeg -f image2 -framerate 4 -y -i ./results/set1/img_%d.png set1.gif\n",
    "```\n",
    "\n",
    "Again, if you ran `setup.sh` you should already have some results under `./results`"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}