orpose_color_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YYBGy3NmE6Yb"
   },
   "source": [
    "Copyright (c) University of Strasbourg. All Rights Reserved.\n",
    "\n",
    "# ORPose-Color\n",
    "-------------------\n",
    "**Self-supervision on Unlabelled OR Data for Multi-person 2D/3D Human Pose Estimation (MICCAI-2020)**\n",
    "\n",
    "_Vinkle Srivastav, Afshin Gangi, Nicolas Padoy_\n",
    "\n",
    "This repository contains the inference demo and evaluation scripts.\n",
    "\n",
    "[![arXiv](https://img.shields.io/badge/arxiv-2007.08354-red)](https://arxiv.org/abs/2007.08354) \n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/CAMMA-public/ORPose-Color/blob/master/orpose_color_demo.ipynb)\n",
    "\n",
    "**This demo notebook contains the inference scripts for the following models: _ORPose_fixed_8x_, _ORPose_fixed_10x_, and _ORPose_fixed_12x_**\n",
    "\n",
    "<br/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EbEWVSn5E6Yc"
   },
   "source": [
    "## Install dependencies and download the trained models\n",
    "**Needed only for the colab demo. Please make sure to enable \"GPU\" using EDIT->Notebook settings.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 93816,
     "status": "ok",
     "timestamp": 1597833723825,
     "user": {
      "displayName": "teaching camma",
      "photoUrl": "",
      "userId": "15642510192522395139"
     },
     "user_tz": -120
    },
    "id": "IPCzU-3cE6Yd",
    "outputId": "9e24efde-c938-45b4-83fe-afa9f520a296"
   },
   "outputs": [],
   "source": [
    "!pip install pyyaml==5.1 pycocotools>=2.0.1\n",
    "import torch, torchvision\n",
    "print(torch.__version__, torch.cuda.is_available())\n",
    "!gcc --version\n",
    "\n",
    "# Install pytorch and detectron2\n",
    "!python -m pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 -f https://download.pytorch.org/whl/torch_stable.html\n",
    "!python -m pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu111/torch1.10/index.html\n",
    "# Restart the runtime after installing pytorch and detectron2 (Runtime -> Restart runtime)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "zEJlgPGBkIvi"
   },
   "outputs": [],
   "source": [
    "# download the trained models\n",
    "!git clone https://github.com/CAMMA-public/ORPose-Color.git\n",
    "%cd ORPose-Color\n",
    "\n",
    "!wget https://s3.unistra.fr/camma_public/github/ORPose/models.zip\n",
    "!unzip -q models.zip\n",
    "!rm models.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "K-T0n3_dE6Yg"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "y1DpKoDyE6Yh"
   },
   "outputs": [],
   "source": [
    "import torch, torchvision\n",
    "import detectron2\n",
    "from detectron2.utils.logger import setup_logger\n",
    "setup_logger()\n",
    "\n",
    "# import some common libraries\n",
    "import os\n",
    "import numpy as np\n",
    "import cv2\n",
    "import random\n",
    "from base64 import b64encode\n",
    "from IPython.display import HTML\n",
    "from PIL import Image\n",
    "import glob\n",
    "from tqdm import tqdm \n",
    "# import some common detectron2 utilities\n",
    "from detectron2.engine import DefaultPredictor\n",
    "from detectron2.config import get_cfg\n",
    "from detectron2.data import detection_utils as utils\n",
    "from detectron2.data import transforms as T\n",
    "from keypoints_3d import add_keypoints3d_config, draw_2d_keypoints, bgr2rgb, coco_to_h36_kps, images_to_video, progress_bar\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "46PYhTPNE6Yj"
   },
   "source": [
    "## Choose the model type and set the paths to the data and model directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "D24BHKrjE6Yk"
   },
   "outputs": [],
   "source": [
    "# Paths and params\n",
    "SCALE = 12 # 10, 12\n",
    "\n",
    "MODEL_WEIGHTS = \"models/orpose_fixed_{}x_model_final.pth\".format(SCALE)\n",
    "CONFIG_FILE = \"configs/orpose_fixed_{}x.yaml\".format(SCALE)\n",
    "VIDEO_DIR = \"data/mvor_seq1_x{}_color\".format(SCALE)\n",
    "    \n",
    "H36_STATS_FILE = \"data/h36stats.json\"\n",
    "OUTPUT_DIR = \"output\"\n",
    "OUTPUT_VID_NAME = os.path.join(OUTPUT_DIR, \"output.mp4\")\n",
    "IMG_FORMAT = \"BGR\"\n",
    "IMG_RESIZE_SHAPE = (480, 640)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FIcyM01JE6Ym"
   },
   "source": [
    "## Setup the config file and load the model weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5xw91za9E6Ym"
   },
   "outputs": [],
   "source": [
    "cfg = get_cfg()\n",
    "add_keypoints3d_config(cfg)\n",
    "cfg.merge_from_file(CONFIG_FILE)\n",
    "\n",
    "cfg.MODEL.WEIGHTS = MODEL_WEIGHTS\n",
    "cfg.DATASETS.H36_STATS_PATH = H36_STATS_FILE\n",
    "predictor = DefaultPredictor(cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "T09jWw7VE6Yp"
   },
   "source": [
    "##  Inference demo on the video frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 70
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 192419,
     "status": "ok",
     "timestamp": 1597836761185,
     "user": {
      "displayName": "teaching camma",
      "photoUrl": "",
      "userId": "15642510192522395139"
     },
     "user_tz": -120
    },
    "id": "dhIci7IZE6Yp",
    "outputId": "6bf08379-ea8d-4b9b-f54f-a3c7494689a7"
   },
   "outputs": [],
   "source": [
    "# read the files\n",
    "files = sorted(glob.glob(VIDEO_DIR+\"/*.png\"))\n",
    "\n",
    "# create the output directory\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "# perform inference on video\n",
    "print(\"inference on video started\")\n",
    "# Initialize the progress-bar\n",
    "bar = display(progress_bar(1, len(files)), display_id=True)\n",
    "for index, (img_path) in enumerate(files):\n",
    "    img_lr = utils.read_image(img_path, format=\"BGR\")\n",
    "    # Resize using bilinear interpolation\n",
    "    img_sr, _ = T.apply_transform_gens([T.Resize(shape= IMG_RESIZE_SHAPE, interp=Image.BILINEAR)], img_lr)\n",
    "    predictions = predictor(img_sr)\n",
    "    predictions = predictions[\"instances\"].to(\"cpu\")\n",
    "    \n",
    "    # get boxes, keypoints2d, scores\n",
    "    boxes = predictions.pred_boxes if predictions.has(\"pred_boxes\") else None\n",
    "    scores = predictions.scores if predictions.has(\"scores\") else None\n",
    "    keypoints2d_coco =  np.asarray(predictions.pred_keypoints) if predictions.has(\"pred_keypoints\") else None\n",
    "    \n",
    "    # convert pytorch tensors to numpy\n",
    "    boxes = boxes.tensor.numpy()\n",
    "    for sc, box, kpts_2d in zip(scores, boxes, keypoints2d_coco):\n",
    "        if sc > 0.5:\n",
    "            img_sr = draw_2d_keypoints(img_sr, coco_to_h36_kps(kpts_2d), style=\"h36\", box=box)  \n",
    "    cv2.imwrite(os.path.join(OUTPUT_DIR, f\"{index:06d}\" + \".png\"), img_sr)\n",
    "    bar.update(progress_bar(index + 1, len(files)))\n",
    "print(\"inference on video finished\")\n",
    "# Convert the rendered images to video\n",
    "images_to_video(OUTPUT_DIR, OUTPUT_VID_NAME)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Cffd6vGcE6Yr"
   },
   "source": [
    "## Show the output video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 320
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 5416,
     "status": "ok",
     "timestamp": 1597836819059,
     "user": {
      "displayName": "teaching camma",
      "photoUrl": "",
      "userId": "15642510192522395139"
     },
     "user_tz": -120
    },
    "id": "Qyh6-xFmE6Ys",
    "outputId": "ad6c6bef-6e6d-4553-8334-2b34029997c0"
   },
   "outputs": [],
   "source": [
    "mp4 = open(OUTPUT_VID_NAME, \"rb\").read()\n",
    "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
    "HTML(\"\"\" <video width=400 controls>\n",
    "         <source src=\"%s\" type=\"video/mp4\">\n",
    "         </video> \n",
    "     \"\"\" % data_url\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "demo_orpose.ipynb",
   "provenance": []
  },
  "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
}