Pose_classification_(extended).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Pose classification (extended).ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "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.8.3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/aimanyounises1/Cadence_Counter/blob/main/Pose_classification_(extended).ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vL2VrLjtHZCs"
      },
      "source": [
        "# Overview\n",
        "\n",
        "This Colab helps to create and validate a training set for the k-NN classifier described in the MediaPipe [Pose Classification](https://google.github.io/mediapipe/solutions/pose_classification.html) soultion, test it on an arbitrary video, export to a CSV and then use it in the [ML Kit sample app](https://developers.google.com/ml-kit/vision/pose-detection/classifying-poses#4_integrate_with_the_ml_kit_quickstart_app)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j8OxqytxxV-e"
      },
      "source": [
        "# Step 0: Start Colab"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OGJFajURxZii"
      },
      "source": [
        "Connect the Colab to hosted Python3 runtime (check top-right corner) and then install required dependencies."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UsA8WJi60PaX",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "3765ca88-6806-4051-b8bc-72f30eac4763"
      },
      "source": [
        "!pip install pillow==8.1.0\n",
        "!pip install matplotlib==3.3.4\n",
        "!pip install numpy==1.19.3\n",
        "!pip install opencv-python==4.5.1.48\n",
        "!pip install tqdm==4.56.0\n",
        "!pip install requests==2.25.1\n",
        "\n",
        "!pip install mediapipe==0.8.3"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting pillow==8.1.0\n",
            "  Downloading Pillow-8.1.0-cp37-cp37m-manylinux1_x86_64.whl (2.2 MB)\n",
            "\u001b[K     |████████████████████████████████| 2.2 MB 5.0 MB/s \n",
            "\u001b[?25hInstalling collected packages: pillow\n",
            "  Attempting uninstall: pillow\n",
            "    Found existing installation: Pillow 7.1.2\n",
            "    Uninstalling Pillow-7.1.2:\n",
            "      Successfully uninstalled Pillow-7.1.2\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible.\u001b[0m\n",
            "Successfully installed pillow-8.1.0\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "PIL"
                ]
              }
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting matplotlib==3.3.4\n",
            "  Downloading matplotlib-3.3.4-cp37-cp37m-manylinux1_x86_64.whl (11.5 MB)\n",
            "\u001b[K     |████████████████████████████████| 11.5 MB 12 kB/s \n",
            "\u001b[?25hRequirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.4) (2.8.2)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.4) (1.3.1)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.4) (8.1.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.4) (0.10.0)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.3 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.4) (2.4.7)\n",
            "Requirement already satisfied: numpy>=1.15 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.4) (1.19.5)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from cycler>=0.10->matplotlib==3.3.4) (1.15.0)\n",
            "Installing collected packages: matplotlib\n",
            "  Attempting uninstall: matplotlib\n",
            "    Found existing installation: matplotlib 3.2.2\n",
            "    Uninstalling matplotlib-3.2.2:\n",
            "      Successfully uninstalled matplotlib-3.2.2\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible.\u001b[0m\n",
            "Successfully installed matplotlib-3.3.4\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "matplotlib",
                  "mpl_toolkits"
                ]
              }
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting numpy==1.19.3\n",
            "  Downloading numpy-1.19.3-cp37-cp37m-manylinux2010_x86_64.whl (14.9 MB)\n",
            "\u001b[K     |████████████████████████████████| 14.9 MB 100 kB/s \n",
            "\u001b[?25hInstalling collected packages: numpy\n",
            "  Attempting uninstall: numpy\n",
            "    Found existing installation: numpy 1.19.5\n",
            "    Uninstalling numpy-1.19.5:\n",
            "      Successfully uninstalled numpy-1.19.5\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.\n",
            "albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible.\u001b[0m\n",
            "Successfully installed numpy-1.19.3\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "numpy"
                ]
              }
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting opencv-python==4.5.1.48\n",
            "  Downloading opencv_python-4.5.1.48-cp37-cp37m-manylinux2014_x86_64.whl (50.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 50.4 MB 14 kB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy>=1.14.5 in /usr/local/lib/python3.7/dist-packages (from opencv-python==4.5.1.48) (1.19.3)\n",
            "Installing collected packages: opencv-python\n",
            "  Attempting uninstall: opencv-python\n",
            "    Found existing installation: opencv-python 4.1.2.30\n",
            "    Uninstalling opencv-python-4.1.2.30:\n",
            "      Successfully uninstalled opencv-python-4.1.2.30\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible.\u001b[0m\n",
            "Successfully installed opencv-python-4.5.1.48\n",
            "Collecting tqdm==4.56.0\n",
            "  Downloading tqdm-4.56.0-py2.py3-none-any.whl (72 kB)\n",
            "\u001b[K     |████████████████████████████████| 72 kB 878 kB/s \n",
            "\u001b[?25hInstalling collected packages: tqdm\n",
            "  Attempting uninstall: tqdm\n",
            "    Found existing installation: tqdm 4.62.0\n",
            "    Uninstalling tqdm-4.62.0:\n",
            "      Successfully uninstalled tqdm-4.62.0\n",
            "Successfully installed tqdm-4.56.0\n",
            "Collecting requests==2.25.1\n",
            "  Downloading requests-2.25.1-py2.py3-none-any.whl (61 kB)\n",
            "\u001b[K     |████████████████████████████████| 61 kB 3.5 MB/s \n",
            "\u001b[?25hRequirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests==2.25.1) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests==2.25.1) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests==2.25.1) (2021.5.30)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests==2.25.1) (1.24.3)\n",
            "Installing collected packages: requests\n",
            "  Attempting uninstall: requests\n",
            "    Found existing installation: requests 2.23.0\n",
            "    Uninstalling requests-2.23.0:\n",
            "      Successfully uninstalled requests-2.23.0\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "google-colab 1.0.0 requires requests~=2.23.0, but you have requests 2.25.1 which is incompatible.\n",
            "datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.\u001b[0m\n",
            "Successfully installed requests-2.25.1\n",
            "Collecting mediapipe==0.8.3\n",
            "  Downloading mediapipe-0.8.3-cp37-cp37m-manylinux2014_x86_64.whl (67.0 MB)\n",
            "\u001b[K     |████████████████████████████████| 67.0 MB 31 kB/s \n",
            "\u001b[?25hRequirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from mediapipe==0.8.3) (1.15.0)\n",
            "Requirement already satisfied: opencv-python in /usr/local/lib/python3.7/dist-packages (from mediapipe==0.8.3) (4.5.1.48)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.7/dist-packages (from mediapipe==0.8.3) (0.12.0)\n",
            "Requirement already satisfied: numpy==1.19.3 in /usr/local/lib/python3.7/dist-packages (from mediapipe==0.8.3) (1.19.3)\n",
            "Requirement already satisfied: protobuf>=3.11.4 in /usr/local/lib/python3.7/dist-packages (from mediapipe==0.8.3) (3.17.3)\n",
            "Requirement already satisfied: wheel in /usr/local/lib/python3.7/dist-packages (from mediapipe==0.8.3) (0.37.0)\n",
            "Requirement already satisfied: attrs in /usr/local/lib/python3.7/dist-packages (from mediapipe==0.8.3) (21.2.0)\n",
            "Collecting dataclasses\n",
            "  Downloading dataclasses-0.6-py3-none-any.whl (14 kB)\n",
            "Installing collected packages: dataclasses, mediapipe\n",
            "Successfully installed dataclasses-0.6 mediapipe-0.8.3\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l7S1Dl8Dhfa2"
      },
      "source": [
        "# Codebase"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N5Ay8WCLqosN"
      },
      "source": [
        "## Commons"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "swiAP0RYqqVM"
      },
      "source": [
        "from matplotlib import pyplot as plt\n",
        "\n",
        "\n",
        "def show_image(img, figsize=(10, 10)):\n",
        "  \"\"\"Shows output PIL image.\"\"\"\n",
        "  plt.figure(figsize=figsize)\n",
        "  plt.imshow(img)\n",
        "  plt.show()"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0L6HLFd9AXmh"
      },
      "source": [
        "## Pose embedding"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QBrKOeP30RAx"
      },
      "source": [
        "class FullBodyPoseEmbedder(object):\n",
        "  \"\"\"Converts 3D pose landmarks into 3D embedding.\"\"\"\n",
        "\n",
        "  def __init__(self, torso_size_multiplier=2.5):\n",
        "    # Multiplier to apply to the torso to get minimal body size.\n",
        "    self._torso_size_multiplier = torso_size_multiplier\n",
        "\n",
        "    # Names of the landmarks as they appear in the prediction.\n",
        "    self._landmark_names = [\n",
        "        'nose',\n",
        "        'left_eye_inner', 'left_eye', 'left_eye_outer',\n",
        "        'right_eye_inner', 'right_eye', 'right_eye_outer',\n",
        "        'left_ear', 'right_ear',\n",
        "        'mouth_left', 'mouth_right',\n",
        "        'left_shoulder', 'right_shoulder',\n",
        "        'left_elbow', 'right_elbow',\n",
        "        'left_wrist', 'right_wrist',\n",
        "        'left_pinky_1', 'right_pinky_1',\n",
        "        'left_index_1', 'right_index_1',\n",
        "        'left_thumb_2', 'right_thumb_2',\n",
        "        'left_hip', 'right_hip',\n",
        "        'left_knee', 'right_knee',\n",
        "        'left_ankle', 'right_ankle',\n",
        "        'left_heel', 'right_heel',\n",
        "        'left_foot_index', 'right_foot_index',\n",
        "    ]\n",
        "\n",
        "  def __call__(self, landmarks):\n",
        "    \"\"\"Normalizes pose landmarks and converts to embedding\n",
        "    \n",
        "    Args:\n",
        "      landmarks - NumPy array with 3D landmarks of shape (N, 3).\n",
        "\n",
        "    Result:\n",
        "      Numpy array with pose embedding of shape (M, 3) where `M` is the number of\n",
        "      pairwise distances defined in `_get_pose_distance_embedding`.\n",
        "    \"\"\"\n",
        "    assert landmarks.shape[0] == len(self._landmark_names), 'Unexpected number of landmarks: {}'.format(landmarks.shape[0])\n",
        "\n",
        "    # Get pose landmarks.\n",
        "    landmarks = np.copy(landmarks)\n",
        "\n",
        "    # Normalize landmarks.\n",
        "    landmarks = self._normalize_pose_landmarks(landmarks)\n",
        "\n",
        "    # Get embedding.\n",
        "    embedding = self._get_pose_distance_embedding(landmarks)\n",
        "\n",
        "    return embedding\n",
        "\n",
        "  def _normalize_pose_landmarks(self, landmarks):\n",
        "    \"\"\"Normalizes landmarks translation and scale.\"\"\"\n",
        "    landmarks = np.copy(landmarks)\n",
        "\n",
        "    # Normalize translation.\n",
        "    pose_center = self._get_pose_center(landmarks)\n",
        "    landmarks -= pose_center\n",
        "\n",
        "    # Normalize scale.\n",
        "    pose_size = self._get_pose_size(landmarks, self._torso_size_multiplier)\n",
        "    landmarks /= pose_size\n",
        "    # Multiplication by 100 is not required, but makes it eaasier to debug.\n",
        "    landmarks *= 100\n",
        "\n",
        "    return landmarks\n",
        "\n",
        "  def _get_pose_center(self, landmarks):\n",
        "    \"\"\"Calculates pose center as point between hips.\"\"\"\n",
        "    left_hip = landmarks[self._landmark_names.index('left_hip')]\n",
        "    right_hip = landmarks[self._landmark_names.index('right_hip')]\n",
        "    center = (left_hip + right_hip) * 0.5\n",
        "    return center\n",
        "\n",
        "  def _get_pose_size(self, landmarks, torso_size_multiplier):\n",
        "    \"\"\"Calculates pose size.\n",
        "    \n",
        "    It is the maximum of two values:\n",
        "      * Torso size multiplied by `torso_size_multiplier`\n",
        "      * Maximum distance from pose center to any pose landmark\n",
        "    \"\"\"\n",
        "    # This approach uses only 2D landmarks to compute pose size.\n",
        "    landmarks = landmarks[:, :2]\n",
        "\n",
        "    # Hips center.\n",
        "    left_hip = landmarks[self._landmark_names.index('left_hip')]\n",
        "    right_hip = landmarks[self._landmark_names.index('right_hip')]\n",
        "    hips = (left_hip + right_hip) * 0.5\n",
        "\n",
        "    # Shoulders center.\n",
        "    left_shoulder = landmarks[self._landmark_names.index('left_shoulder')]\n",
        "    right_shoulder = landmarks[self._landmark_names.index('right_shoulder')]\n",
        "    shoulders = (left_shoulder + right_shoulder) * 0.5\n",
        "\n",
        "    # Torso size as the minimum body size.\n",
        "    torso_size = np.linalg.norm(shoulders - hips)\n",
        "\n",
        "    # Max dist to pose center.\n",
        "    pose_center = self._get_pose_center(landmarks)\n",
        "    max_dist = np.max(np.linalg.norm(landmarks - pose_center, axis=1))\n",
        "\n",
        "    return max(torso_size * torso_size_multiplier, max_dist)\n",
        "\n",
        "  def _get_pose_distance_embedding(self, landmarks):\n",
        "    \"\"\"Converts pose landmarks into 3D embedding.\n",
        "\n",
        "    We use several pairwise 3D distances to form pose embedding. All distances\n",
        "    include X and Y components with sign. We differnt types of pairs to cover\n",
        "    different pose classes. Feel free to remove some or add new.\n",
        "    \n",
        "    Args:\n",
        "      landmarks - NumPy array with 3D landmarks of shape (N, 3).\n",
        "\n",
        "    Result:\n",
        "      Numpy array with pose embedding of shape (M, 3) where `M` is the number of\n",
        "      pairwise distances.\n",
        "    \"\"\"\n",
        "    embedding = np.array([\n",
        "        # One joint.\n",
        "\n",
        "        self._get_distance(\n",
        "            self._get_average_by_names(landmarks, 'left_hip', 'right_hip'),\n",
        "            self._get_average_by_names(landmarks, 'left_shoulder', 'right_shoulder')),\n",
        "\n",
        "        self._get_distance_by_names(landmarks, 'left_shoulder', 'left_elbow'),\n",
        "        self._get_distance_by_names(landmarks, 'right_shoulder', 'right_elbow'),\n",
        "\n",
        "        self._get_distance_by_names(landmarks, 'left_elbow', 'left_wrist'),\n",
        "        self._get_distance_by_names(landmarks, 'right_elbow', 'right_wrist'),\n",
        "\n",
        "        self._get_distance_by_names(landmarks, 'left_hip', 'left_knee'),\n",
        "        self._get_distance_by_names(landmarks, 'right_hip', 'right_knee'),\n",
        "\n",
        "        self._get_distance_by_names(landmarks, 'left_knee', 'left_ankle'),\n",
        "        self._get_distance_by_names(landmarks, 'right_knee', 'right_ankle'),\n",
        "\n",
        "        # Two joints.\n",
        "\n",
        "        self._get_distance_by_names(landmarks, 'left_shoulder', 'left_wrist'),\n",
        "        self._get_distance_by_names(landmarks, 'right_shoulder', 'right_wrist'),\n",
        "\n",
        "        self._get_distance_by_names(landmarks, 'left_hip', 'left_ankle'),\n",
        "        self._get_distance_by_names(landmarks, 'right_hip', 'right_ankle'),\n",
        "\n",
        "        # Four joints.\n",
        "\n",
        "        self._get_distance_by_names(landmarks, 'left_hip', 'left_wrist'),\n",
        "        self._get_distance_by_names(landmarks, 'right_hip', 'right_wrist'),\n",
        "\n",
        "        # Five joints.\n",
        "\n",
        "        self._get_distance_by_names(landmarks, 'left_shoulder', 'left_ankle'),\n",
        "        self._get_distance_by_names(landmarks, 'right_shoulder', 'right_ankle'),\n",
        "        \n",
        "        self._get_distance_by_names(landmarks, 'left_hip', 'left_wrist'),\n",
        "        self._get_distance_by_names(landmarks, 'right_hip', 'right_wrist'),\n",
        "\n",
        "        # Cross body.\n",
        "\n",
        "        self._get_distance_by_names(landmarks, 'left_elbow', 'right_elbow'),\n",
        "        self._get_distance_by_names(landmarks, 'left_knee', 'right_knee'),\n",
        "\n",
        "        self._get_distance_by_names(landmarks, 'left_wrist', 'right_wrist'),\n",
        "        self._get_distance_by_names(landmarks, 'left_ankle', 'right_ankle'),\n",
        "\n",
        "        # Body bent direction.\n",
        "\n",
        "        # self._get_distance(\n",
        "        #     self._get_average_by_names(landmarks, 'left_wrist', 'left_ankle'),\n",
        "        #     landmarks[self._landmark_names.index('left_hip')]),\n",
        "        # self._get_distance(\n",
        "        #     self._get_average_by_names(landmarks, 'right_wrist', 'right_ankle'),\n",
        "        #     landmarks[self._landmark_names.index('right_hip')]),\n",
        "    ])\n",
        "\n",
        "    return embedding\n",
        "\n",
        "  def _get_average_by_names(self, landmarks, name_from, name_to):\n",
        "    lmk_from = landmarks[self._landmark_names.index(name_from)]\n",
        "    lmk_to = landmarks[self._landmark_names.index(name_to)]\n",
        "    return (lmk_from + lmk_to) * 0.5\n",
        "\n",
        "  def _get_distance_by_names(self, landmarks, name_from, name_to):\n",
        "    lmk_from = landmarks[self._landmark_names.index(name_from)]\n",
        "    lmk_to = landmarks[self._landmark_names.index(name_to)]\n",
        "    return self._get_distance(lmk_from, lmk_to)\n",
        "\n",
        "  def _get_distance(self, lmk_from, lmk_to):\n",
        "    return lmk_to - lmk_from"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hH-efWS61Tfy"
      },
      "source": [
        "## Pose classification"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TR9M0uAK1WzX"
      },
      "source": [
        "class PoseSample(object):\n",
        "\n",
        "  def __init__(self, name, landmarks, class_name, embedding):\n",
        "    self.name = name\n",
        "    self.landmarks = landmarks\n",
        "    self.class_name = class_name\n",
        "    \n",
        "    self.embedding = embedding\n",
        "\n",
        "\n",
        "class PoseSampleOutlier(object):\n",
        "\n",
        "  def __init__(self, sample, detected_class, all_classes):\n",
        "    self.sample = sample\n",
        "    self.detected_class = detected_class\n",
        "    self.all_classes = all_classes"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y230jVvP1u33"
      },
      "source": [
        "import csv\n",
        "import numpy as np\n",
        "import os\n",
        "\n",
        "class PoseClassifier(object):\n",
        "  \"\"\"Classifies pose landmarks.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               pose_samples_folder,\n",
        "               pose_embedder,\n",
        "               file_extension='csv',\n",
        "               file_separator=',',\n",
        "               n_landmarks=33,\n",
        "               n_dimensions=3,\n",
        "               top_n_by_max_distance=30,\n",
        "               top_n_by_mean_distance=10,\n",
        "               axes_weights=(1., 1., 0.2)):\n",
        "    self._pose_embedder = pose_embedder\n",
        "    self._n_landmarks = n_landmarks\n",
        "    self._n_dimensions = n_dimensions\n",
        "    self._top_n_by_max_distance = top_n_by_max_distance\n",
        "    self._top_n_by_mean_distance = top_n_by_mean_distance\n",
        "    self._axes_weights = axes_weights\n",
        "\n",
        "    self._pose_samples = self._load_pose_samples(pose_samples_folder,\n",
        "                                                 file_extension,\n",
        "                                                 file_separator,\n",
        "                                                 n_landmarks,\n",
        "                                                 n_dimensions,\n",
        "                                                 pose_embedder)\n",
        "\n",
        "  def _load_pose_samples(self,\n",
        "                         pose_samples_folder,\n",
        "                         file_extension,\n",
        "                         file_separator,\n",
        "                         n_landmarks,\n",
        "                         n_dimensions,\n",
        "                         pose_embedder):\n",
        "    \"\"\"Loads pose samples from a given folder.\n",
        "    \n",
        "    Required folder structure:\n",
        "      neutral_standing.csv\n",
        "      pushups_down.csv\n",
        "      pushups_up.csv\n",
        "      squats_down.csv\n",
        "      ...\n",
        "\n",
        "    Required CSV structure:\n",
        "      sample_00001,x1,y1,z1,x2,y2,z2,....\n",
        "      sample_00002,x1,y1,z1,x2,y2,z2,....\n",
        "      ...\n",
        "    \"\"\"\n",
        "    # Each file in the folder represents one pose class.\n",
        "    file_names = [name for name in os.listdir(pose_samples_folder) if name.endswith(file_extension)]\n",
        "\n",
        "    pose_samples = []\n",
        "    for file_name in file_names:\n",
        "      # Use file name as pose class name.\n",
        "      class_name = file_name[:-(len(file_extension) + 1)]\n",
        "      \n",
        "      # Parse CSV.\n",
        "      with open(os.path.join(pose_samples_folder, file_name)) as csv_file:\n",
        "        csv_reader = csv.reader(csv_file, delimiter=file_separator)\n",
        "        for row in csv_reader:\n",
        "          assert len(row) == n_landmarks * n_dimensions + 1, 'Wrong number of values: {}'.format(len(row))\n",
        "          landmarks = np.array(row[1:], np.float32).reshape([n_landmarks, n_dimensions])\n",
        "          pose_samples.append(PoseSample(\n",
        "              name=row[0],\n",
        "              landmarks=landmarks,\n",
        "              class_name=class_name,\n",
        "              embedding=pose_embedder(landmarks),\n",
        "          ))\n",
        "\n",
        "    return pose_samples\n",
        "\n",
        "  def find_pose_sample_outliers(self):\n",
        "    \"\"\"Classifies each sample against the entire database.\"\"\"\n",
        "    # Find outliers in target poses\n",
        "    outliers = []\n",
        "    for sample in self._pose_samples:\n",
        "      # Find nearest poses for the target one.\n",
        "      pose_landmarks = sample.landmarks.copy()\n",
        "      pose_classification = self.__call__(pose_landmarks)\n",
        "      class_names = [class_name for class_name, count in pose_classification.items() if count == max(pose_classification.values())]\n",
        "\n",
        "      # Sample is an outlier if nearest poses have different class or more than\n",
        "      # one pose class is detected as nearest.\n",
        "      if sample.class_name not in class_names or len(class_names) != 1:\n",
        "        outliers.append(PoseSampleOutlier(sample, class_names, pose_classification))\n",
        "\n",
        "    return outliers\n",
        "\n",
        "  def __call__(self, pose_landmarks):\n",
        "    \"\"\"Classifies given pose.\n",
        "\n",
        "    Classification is done in two stages:\n",
        "      * First we pick top-N samples by MAX distance. It allows to remove samples\n",
        "        that are almost the same as given pose, but has few joints bent in the\n",
        "        other direction.\n",
        "      * Then we pick top-N samples by MEAN distance. After outliers are removed\n",
        "        on a previous step, we can pick samples that are closes on average.\n",
        "    \n",
        "    Args:\n",
        "      pose_landmarks: NumPy array with 3D landmarks of shape (N, 3).\n",
        "\n",
        "    Returns:\n",
        "      Dictionary with count of nearest pose samples from the database. Sample:\n",
        "        {\n",
        "          'pushups_down': 8,\n",
        "          'pushups_up': 2,\n",
        "        }\n",
        "    \"\"\"\n",
        "    # Check that provided and target poses have the same shape.\n",
        "    assert pose_landmarks.shape == (self._n_landmarks, self._n_dimensions), 'Unexpected shape: {}'.format(pose_landmarks.shape)\n",
        "\n",
        "    # Get given pose embedding.\n",
        "    pose_embedding = self._pose_embedder(pose_landmarks)\n",
        "    flipped_pose_embedding = self._pose_embedder(pose_landmarks * np.array([-1, 1, 1]))\n",
        "\n",
        "    # Filter by max distance.\n",
        "    #\n",
        "    # That helps to remove outliers - poses that are almost the same as the\n",
        "    # given one, but has one joint bent into another direction and actually\n",
        "    # represnt a different pose class.\n",
        "    max_dist_heap = []\n",
        "    for sample_idx, sample in enumerate(self._pose_samples):\n",
        "      max_dist = min(\n",
        "          np.max(np.abs(sample.embedding - pose_embedding) * self._axes_weights),\n",
        "          np.max(np.abs(sample.embedding - flipped_pose_embedding) * self._axes_weights),\n",
        "      )\n",
        "      max_dist_heap.append([max_dist, sample_idx])\n",
        "\n",
        "    max_dist_heap = sorted(max_dist_heap, key=lambda x: x[0])\n",
        "    max_dist_heap = max_dist_heap[:self._top_n_by_max_distance]\n",
        "\n",
        "    # Filter by mean distance.\n",
        "    #\n",
        "    # After removing outliers we can find the nearest pose by mean distance.\n",
        "    mean_dist_heap = []\n",
        "    for _, sample_idx in max_dist_heap:\n",
        "      sample = self._pose_samples[sample_idx]\n",
        "      mean_dist = min(\n",
        "          np.mean(np.abs(sample.embedding - pose_embedding) * self._axes_weights),\n",
        "          np.mean(np.abs(sample.embedding - flipped_pose_embedding) * self._axes_weights),\n",
        "      )\n",
        "      mean_dist_heap.append([mean_dist, sample_idx])\n",
        "\n",
        "    mean_dist_heap = sorted(mean_dist_heap, key=lambda x: x[0])\n",
        "    mean_dist_heap = mean_dist_heap[:self._top_n_by_mean_distance]\n",
        "\n",
        "    # Collect results into map: (class_name -> n_samples)\n",
        "    class_names = [self._pose_samples[sample_idx].class_name for _, sample_idx in mean_dist_heap]\n",
        "    result = {class_name: class_names.count(class_name) for class_name in set(class_names)}\n",
        "\n",
        "    return result"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g-VRo98tE1JH"
      },
      "source": [
        "## Classification smoothing"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "POC4_eQsE3VO"
      },
      "source": [
        "class EMADictSmoothing(object):\n",
        "  \"\"\"Smoothes pose classification.\"\"\"\n",
        "\n",
        "  def __init__(self, window_size=10, alpha=0.2):\n",
        "    self._window_size = window_size\n",
        "    self._alpha = alpha\n",
        "\n",
        "    self._data_in_window = []\n",
        "\n",
        "  def __call__(self, data):\n",
        "    \"\"\"Smoothes given pose classification.\n",
        "\n",
        "    Smoothing is done by computing Exponential Moving Average for every pose\n",
        "    class observed in the given time window. Missed pose classes arre replaced\n",
        "    with 0.\n",
        "    \n",
        "    Args:\n",
        "      data: Dictionary with pose classification. Sample:\n",
        "          {\n",
        "            'pushups_down': 8,\n",
        "            'pushups_up': 2,\n",
        "          }\n",
        "\n",
        "    Result:\n",
        "      Dictionary in the same format but with smoothed and float instead of\n",
        "      integer values. Sample:\n",
        "        {\n",
        "          'pushups_down': 8.3,\n",
        "          'pushups_up': 1.7,\n",
        "        }\n",
        "    \"\"\"\n",
        "    # Add new data to the beginning of the window for simpler code.\n",
        "    self._data_in_window.insert(0, data)\n",
        "    self._data_in_window = self._data_in_window[:self._window_size]\n",
        "\n",
        "    # Get all keys.\n",
        "    keys = set([key for data in self._data_in_window for key, _ in data.items()])\n",
        "\n",
        "    # Get smoothed values.\n",
        "    smoothed_data = dict()\n",
        "    for key in keys:\n",
        "      factor = 1.0\n",
        "      top_sum = 0.0\n",
        "      bottom_sum = 0.0\n",
        "      for data in self._data_in_window:\n",
        "        value = data[key] if key in data else 0.0\n",
        "\n",
        "        top_sum += factor * value\n",
        "        bottom_sum += factor\n",
        "\n",
        "        # Update factor.\n",
        "        factor *= (1.0 - self._alpha)\n",
        "\n",
        "      smoothed_data[key] = top_sum / bottom_sum\n",
        "\n",
        "    return smoothed_data"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eWuA2OYgGtZn"
      },
      "source": [
        "## Repetition counter"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jpZcekQ-mbDS",
        "outputId": "f0e2620d-a0c2-4593-c480-4cc026c1d0e0"
      },
      "source": [
        "!\n",
        "!git clone https://github.com/aimanyounises1/Cadence_Counter.git\n"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'Cadence_Counter'...\n",
            "fatal: could not read Username for 'https://github.com': No such device or address\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TEs_lgNiGv-j"
      },
      "source": [
        "class RepetitionCounter(object):\n",
        "  \"\"\"Counts number of repetitions of given target pose class.\"\"\"\n",
        "\n",
        "  def __init__(self, class_name, enter_threshold=6, exit_threshold=4):\n",
        "    self._class_name = class_name\n",
        "\n",
        "    # If pose counter passes given threshold, then we enter the pose.\n",
        "    self._enter_threshold = enter_threshold\n",
        "    self._exit_threshold = exit_threshold\n",
        "\n",
        "    # Either we are in given pose or not.\n",
        "    self._pose_entered = False\n",
        "\n",
        "    # Number of times we exited the pose.\n",
        "    self._n_repeats = 0\n",
        "\n",
        "  @property\n",
        "  def n_repeats(self):\n",
        "    return self._n_repeats\n",
        "\n",
        "  def __call__(self, pose_classification):\n",
        "    \"\"\"Counts number of repetitions happend until given frame.\n",
        "\n",
        "    We use two thresholds. First you need to go above the higher one to enter\n",
        "    the pose, and then you need to go below the lower one to exit it. Difference\n",
        "    between the thresholds makes it stable to prediction jittering (which will\n",
        "    cause wrong counts in case of having only one threshold).\n",
        "    \n",
        "    Args:\n",
        "      pose_classification: Pose classification dictionary on current frame.\n",
        "        Sample:\n",
        "          {\n",
        "            'pushups_down': 8.3,\n",
        "            'pushups_up': 1.7,\n",
        "          }\n",
        "\n",
        "    Returns:\n",
        "      Integer counter of repetitions.\n",
        "    \"\"\"\n",
        "    # Get pose confidence.\n",
        "    pose_confidence = 0.0\n",
        "    if self._class_name in pose_classification:\n",
        "      pose_confidence = pose_classification[self._class_name]\n",
        "\n",
        "    # On the very first frame or if we were out of the pose, just check if we\n",
        "    # entered it on this frame and update the state.\n",
        "    if not self._pose_entered:\n",
        "      self._pose_entered = pose_confidence > self._enter_threshold\n",
        "      return self._n_repeats\n",
        "\n",
        "    # If we were in the pose and are exiting it, then increase the counter and\n",
        "    # update the state.\n",
        "    if pose_confidence < self._exit_threshold:\n",
        "      self._n_repeats += 1\n",
        "      self._pose_entered = False\n",
        "\n",
        "    return self._n_repeats"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QBVyanN2Ic4W"
      },
      "source": [
        "## Classification visualizer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hgFLe1oTIgJH"
      },
      "source": [
        "import io\n",
        "from PIL import Image\n",
        "from PIL import ImageFont\n",
        "from PIL import ImageDraw\n",
        "import requests\n",
        "\n",
        "class PoseClassificationVisualizer(object):\n",
        "  \"\"\"Keeps track of claassifcations for every frame and renders them.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               class_name,\n",
        "               plot_location_x=0.05,\n",
        "               plot_location_y=0.05,\n",
        "               plot_max_width=0.4,\n",
        "               plot_max_height=0.4,\n",
        "               plot_figsize=(9, 4),\n",
        "               plot_x_max=None,\n",
        "               plot_y_max=None,\n",
        "               counter_location_x=0.85,\n",
        "               counter_location_y=0.05,\n",
        "               counter_font_path='https://github.com/googlefonts/roboto/blob/main/src/hinted/Roboto-Regular.ttf?raw=true',\n",
        "               counter_font_color='red',\n",
        "               counter_font_size=0.15):\n",
        "    self._class_name = class_name\n",
        "    self._plot_location_x = plot_location_x\n",
        "    self._plot_location_y = plot_location_y\n",
        "    self._plot_max_width = plot_max_width\n",
        "    self._plot_max_height = plot_max_height\n",
        "    self._plot_figsize = plot_figsize\n",
        "    self._plot_x_max = plot_x_max\n",
        "    self._plot_y_max = plot_y_max\n",
        "    self._counter_location_x = counter_location_x\n",
        "    self._counter_location_y = counter_location_y\n",
        "    self._counter_font_path = counter_font_path\n",
        "    self._counter_font_color = counter_font_color\n",
        "    self._counter_font_size = counter_font_size\n",
        "\n",
        "    self._counter_font = None\n",
        "\n",
        "    self._pose_classification_history = []\n",
        "    self._pose_classification_filtered_history = []\n",
        "\n",
        "  def __call__(self,\n",
        "               frame,\n",
        "               pose_classification,\n",
        "               pose_classification_filtered,\n",
        "               repetitions_count):\n",
        "    \"\"\"Renders pose classifcation and counter until given frame.\"\"\"\n",
        "    # Extend classification history.\n",
        "    self._pose_classification_history.append(pose_classification)\n",
        "    self._pose_classification_filtered_history.append(pose_classification_filtered)\n",
        "\n",
        "    # Output frame with classification plot and counter.\n",
        "    output_img = Image.fromarray(frame)\n",
        "\n",
        "    output_width = output_img.size[0]\n",
        "    output_height = output_img.size[1]\n",
        "\n",
        "    # Draw the plot.\n",
        "    img = self._plot_classification_history(output_width, output_height)\n",
        "    img.thumbnail((int(output_width * self._plot_max_width),\n",
        "                   int(output_height * self._plot_max_height)),\n",
        "                  Image.ANTIALIAS)\n",
        "    output_img.paste(img,\n",
        "                     (int(output_width * self._plot_location_x),\n",
        "                      int(output_height * self._plot_location_y)))\n",
        "\n",
        "    # Draw the count.\n",
        "    output_img_draw = ImageDraw.Draw(output_img)\n",
        "    if self._counter_font is None:\n",
        "      font_size = int(output_height * self._counter_font_size)\n",
        "      font_request = requests.get(self._counter_font_path, allow_redirects=True)\n",
        "      self._counter_font = ImageFont.truetype(io.BytesIO(font_request.content), size=font_size)\n",
        "    output_img_draw.text((output_width * self._counter_location_x,\n",
        "                          output_height * self._counter_location_y),\n",
        "                         str(repetitions_count),\n",
        "                         font=self._counter_font,\n",
        "                         fill=self._counter_font_color)\n",
        "\n",
        "    return output_img\n",
        "\n",
        "  def _plot_classification_history(self, output_width, output_height):\n",
        "    fig = plt.figure(figsize=self._plot_figsize)\n",
        "\n",
        "    for classification_history in [self._pose_classification_history,\n",
        "                                   self._pose_classification_filtered_history]:\n",
        "      y = []\n",
        "      for classification in classification_history:\n",
        "        if classification is None:\n",
        "          y.append(None)\n",
        "        elif self._class_name in classification:\n",
        "          y.append(classification[self._class_name])\n",
        "        else:\n",
        "          y.append(0)\n",
        "      plt.plot(y, linewidth=7)\n",
        "\n",
        "    plt.grid(axis='y', alpha=0.75)\n",
        "    plt.xlabel('Frame')\n",
        "    plt.ylabel('Confidence')\n",
        "    plt.title('Classification history for `{}`'.format(self._class_name))\n",
        "    plt.legend(loc='upper right')\n",
        "\n",
        "    if self._plot_y_max is not None:\n",
        "      plt.ylim(top=self._plot_y_max)\n",
        "    if self._plot_x_max is not None:\n",
        "      plt.xlim(right=self._plot_x_max)\n",
        "\n",
        "    # Convert plot to image.\n",
        "    buf = io.BytesIO()\n",
        "    dpi = min(\n",
        "        output_width * self._plot_max_width / float(self._plot_figsize[0]),\n",
        "        output_height * self._plot_max_height / float(self._plot_figsize[1]))\n",
        "    fig.savefig(buf, dpi=dpi)\n",
        "    buf.seek(0)\n",
        "    img = Image.open(buf)\n",
        "    plt.close()\n",
        "\n",
        "    return img"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KlUbZ_c-Aq4B"
      },
      "source": [
        "## Bootstrap helper"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tw2xYlGmAt3q"
      },
      "source": [
        "import cv2\n",
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "import os\n",
        "from PIL import Image\n",
        "import sys\n",
        "import tqdm\n",
        "\n",
        "from mediapipe.python.solutions import drawing_utils as mp_drawing\n",
        "from mediapipe.python.solutions import pose as mp_pose\n",
        "\n",
        "\n",
        "class BootstrapHelper(object):\n",
        "  \"\"\"Helps to bootstrap images and filter pose samples for classification.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               images_in_folder,\n",
        "               images_out_folder,\n",
        "               csvs_out_folder):\n",
        "    self._images_in_folder = images_in_folder\n",
        "    self._images_out_folder = images_out_folder\n",
        "    self._csvs_out_folder = csvs_out_folder\n",
        "\n",
        "    # Get list of pose classes and print image statistics.\n",
        "    self._pose_class_names = sorted([n for n in os.listdir(self._images_in_folder) if not n.startswith('.')])\n",
        "    \n",
        "  def bootstrap(self, per_pose_class_limit=None):\n",
        "    \"\"\"Bootstraps images in a given folder.\n",
        "    \n",
        "    Required image in folder (same use for image out folder):\n",
        "      pushups_up/\n",
        "        image_001.jpg\n",
        "        image_002.jpg\n",
        "        ...\n",
        "      pushups_down/\n",
        "        image_001.jpg\n",
        "        image_002.jpg\n",
        "        ...\n",
        "      ...\n",
        "\n",
        "    Produced CSVs out folder:\n",
        "      pushups_up.csv\n",
        "      pushups_down.csv\n",
        "\n",
        "    Produced CSV structure with pose 3D landmarks:\n",
        "      sample_00001,x1,y1,z1,x2,y2,z2,....\n",
        "      sample_00002,x1,y1,z1,x2,y2,z2,....\n",
        "    \"\"\"\n",
        "    # Create output folder for CVSs.\n",
        "    if not os.path.exists(self._csvs_out_folder):\n",
        "      os.makedirs(self._csvs_out_folder)\n",
        "\n",
        "    for pose_class_name in self._pose_class_names:\n",
        "      print('Bootstrapping ', pose_class_name, file=sys.stderr)\n",
        "\n",
        "      # Paths for the pose class.\n",
        "      images_in_folder = os.path.join(self._images_in_folder, pose_class_name)\n",
        "      images_out_folder = os.path.join(self._images_out_folder, pose_class_name)\n",
        "      csv_out_path = os.path.join(self._csvs_out_folder, pose_class_name + '.csv')\n",
        "      if not os.path.exists(images_out_folder):\n",
        "        os.makedirs(images_out_folder)\n",
        "\n",
        "      with open(csv_out_path, 'w') as csv_out_file:\n",
        "        csv_out_writer = csv.writer(csv_out_file, delimiter=',', quoting=csv.QUOTE_MINIMAL)\n",
        "        # Get list of images.\n",
        "        image_names = sorted([n for n in os.listdir(images_in_folder) if not n.startswith('.')])\n",
        "        if per_pose_class_limit is not None:\n",
        "          image_names = image_names[:per_pose_class_limit]\n",
        "\n",
        "        # Bootstrap every image.\n",
        "        for image_name in tqdm.tqdm(image_names):\n",
        "          # Load image.\n",
        "          input_frame = cv2.imread(os.path.join(images_in_folder, image_name))\n",
        "          input_frame = cv2.cvtColor(input_frame, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "          # Initialize fresh pose tracker and run it.\n",
        "          with mp_pose.Pose(upper_body_only=False) as pose_tracker:\n",
        "            result = pose_tracker.process(image=input_frame)\n",
        "            pose_landmarks = result.pose_landmarks\n",
        "\n",
        "          # Save image with pose prediction (if pose was detected).\n",
        "          output_frame = input_frame.copy()\n",
        "          if pose_landmarks is not None:\n",
        "            mp_drawing.draw_landmarks(\n",
        "                image=output_frame,\n",
        "                landmark_list=pose_landmarks,\n",
        "                connections=mp_pose.POSE_CONNECTIONS)\n",
        "          output_frame = cv2.cvtColor(output_frame, cv2.COLOR_RGB2BGR)\n",
        "          cv2.imwrite(os.path.join(images_out_folder, image_name), output_frame)\n",
        "          \n",
        "          # Save landmarks if pose was detected.\n",
        "          if pose_landmarks is not None:\n",
        "            # Get landmarks.\n",
        "            frame_height, frame_width = output_frame.shape[0], output_frame.shape[1]\n",
        "            pose_landmarks = np.array(\n",
        "                [[lmk.x * frame_width, lmk.y * frame_height, lmk.z * frame_width]\n",
        "                 for lmk in pose_landmarks.landmark],\n",
        "                dtype=np.float32)\n",
        "            assert pose_landmarks.shape == (33, 3), 'Unexpected landmarks shape: {}'.format(pose_landmarks.shape)\n",
        "            csv_out_writer.writerow([image_name] + pose_landmarks.flatten().astype(np.str).tolist())\n",
        "\n",
        "          # Draw XZ projection and concatenate with the image.\n",
        "          projection_xz = self._draw_xz_projection(\n",
        "              output_frame=output_frame, pose_landmarks=pose_landmarks)\n",
        "          output_frame = np.concatenate((output_frame, projection_xz), axis=1)\n",
        "\n",
        "  def _draw_xz_projection(self, output_frame, pose_landmarks, r=0.5, color='red'):\n",
        "    frame_height, frame_width = output_frame.shape[0], output_frame.shape[1]\n",
        "    img = Image.new('RGB', (frame_width, frame_height), color='white')\n",
        "\n",
        "    if pose_landmarks is None:\n",
        "      return np.asarray(img)\n",
        "\n",
        "    # Scale radius according to the image width.\n",
        "    r *= frame_width * 0.01\n",
        "\n",
        "    draw = ImageDraw.Draw(img)\n",
        "    for idx_1, idx_2 in mp_pose.POSE_CONNECTIONS:\n",
        "      # Flip Z and move hips center to the center of the image.\n",
        "      x1, y1, z1 = pose_landmarks[idx_1] * [1, 1, -1] + [0, 0, frame_height * 0.5]\n",
        "      x2, y2, z2 = pose_landmarks[idx_2] * [1, 1, -1] + [0, 0, frame_height * 0.5]\n",
        "\n",
        "      draw.ellipse([x1 - r, z1 - r, x1 + r, z1 + r], fill=color)\n",
        "      draw.ellipse([x2 - r, z2 - r, x2 + r, z2 + r], fill=color)\n",
        "      draw.line([x1, z1, x2, z2], width=int(r), fill=color)\n",
        "\n",
        "    return np.asarray(img)\n",
        "\n",
        "  def align_images_and_csvs(self, print_removed_items=False):\n",
        "    \"\"\"Makes sure that image folders and CSVs have the same sample.\n",
        "    \n",
        "    Leaves only intersetion of samples in both image folders and CSVs.\n",
        "    \"\"\"\n",
        "    for pose_class_name in self._pose_class_names:\n",
        "      # Paths for the pose class.\n",
        "      images_out_folder = os.path.join(self._images_out_folder, pose_class_name)\n",
        "      csv_out_path = os.path.join(self._csvs_out_folder, pose_class_name + '.csv')\n",
        "\n",
        "      # Read CSV into memory.\n",
        "      rows = []\n",
        "      with open(csv_out_path) as csv_out_file:\n",
        "        csv_out_reader = csv.reader(csv_out_file, delimiter=',')\n",
        "        for row in csv_out_reader:\n",
        "          rows.append(row)\n",
        "\n",
        "      # Image names left in CSV.\n",
        "      image_names_in_csv = []\n",
        "\n",
        "      # Re-write the CSV removing lines without corresponding images.\n",
        "      with open(csv_out_path, 'w') as csv_out_file:\n",
        "        csv_out_writer = csv.writer(csv_out_file, delimiter=',', quoting=csv.QUOTE_MINIMAL)\n",
        "        for row in rows:\n",
        "          image_name = row[0]\n",
        "          image_path = os.path.join(images_out_folder, image_name)\n",
        "          if os.path.exists(image_path):\n",
        "            image_names_in_csv.append(image_name)\n",
        "            csv_out_writer.writerow(row)\n",
        "          elif print_removed_items:\n",
        "            print('Removed image from CSV: ', image_path)\n",
        "\n",
        "      # Remove images without corresponding line in CSV.\n",
        "      for image_name in os.listdir(images_out_folder):\n",
        "        if image_name not in image_names_in_csv:\n",
        "          image_path = os.path.join(images_out_folder, image_name)\n",
        "          os.remove(image_path)\n",
        "          if print_removed_items:\n",
        "            print('Removed image from folder: ', image_path)\n",
        "\n",
        "  def analyze_outliers(self, outliers):\n",
        "    \"\"\"Classifies each sample agains all other to find outliers.\n",
        "    \n",
        "    If sample is classified differrrently than the original class - it sould\n",
        "    either be deleted or more similar samples should be aadded.\n",
        "    \"\"\"\n",
        "    for outlier in outliers:\n",
        "      image_path = os.path.join(self._images_out_folder, outlier.sample.class_name, outlier.sample.name)\n",
        "\n",
        "      print('Outlier')\n",
        "      print('  sample path =    ', image_path)\n",
        "      print('  sample class =   ', outlier.sample.class_name)\n",
        "      print('  detected class = ', outlier.detected_class)\n",
        "      print('  all classes =    ', outlier.all_classes)\n",
        "\n",
        "      img = cv2.imread(image_path)\n",
        "      img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "      show_image(img, figsize=(20, 20))\n",
        "\n",
        "  def remove_outliers(self, outliers):\n",
        "    \"\"\"Removes outliers from the image folders.\"\"\"\n",
        "    for outlier in outliers:\n",
        "      image_path = os.path.join(self._images_out_folder, outlier.sample.class_name, outlier.sample.name)\n",
        "      os.remove(image_path)\n",
        "\n",
        "  def print_images_in_statistics(self):\n",
        "    \"\"\"Prints statistics from the input image folder.\"\"\"\n",
        "    self._print_images_statistics(self._images_in_folder, self._pose_class_names)\n",
        "\n",
        "  def print_images_out_statistics(self):\n",
        "    \"\"\"Prints statistics from the output image folder.\"\"\"\n",
        "    self._print_images_statistics(self._images_out_folder, self._pose_class_names)\n",
        "\n",
        "  def _print_images_statistics(self, images_folder, pose_class_names):\n",
        "    print('Number of images per pose class:')\n",
        "    for pose_class_name in pose_class_names:\n",
        "      n_images = len([\n",
        "          n for n in os.listdir(os.path.join(images_folder, pose_class_name))\n",
        "          if not n.startswith('.')])\n",
        "      print('  {}: {}'.format(pose_class_name, n_images))"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QIiEj8Tx_x-q"
      },
      "source": [
        "# Step 1: Build classifier"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KpRszzilECFz"
      },
      "source": [
        "## Upload image samples\n",
        "\n",
        "Locally create a folder named `fitness_poses_images_in` with image samples.\n",
        "\n",
        "Images should repesent terminal states of desired pose classes. I.e. if you want to classify push-up provide iamges for two classes: when person is up, and when person is down.\n",
        "\n",
        "There should be about a few hundred samples per class covering different camera angles, environment conditions, body shapes, and exercise variations to build a good classifier.\n",
        "\n",
        "Required structure of the images_in_folder:\n",
        "```\n",
        "fitness_poses_images_in/\n",
        "  pushups_up/\n",
        "    image_001.jpg\n",
        "    image_002.jpg\n",
        "    ...\n",
        "  pushups_down/\n",
        "    image_001.jpg\n",
        "    image_002.jpg\n",
        "    ...\n",
        "  ...\n",
        "```\n",
        "\n",
        "Zip the `fitness_poses_images_in` folder:\n",
        "```\n",
        "zip -r fitness_poses_images_in.zip fitness_poses_images_in\n",
        "```\n",
        "\n",
        "And run the code below to upload it to the Colab runtime"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DErZEmKUEDK0"
      },
      "source": [
        "from google.colab import files\n",
        "import os\n",
        "\n",
        "uploaded = files.upload()\n",
        "os.listdir('.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EtTCrVF9EQZd"
      },
      "source": [
        "Unzip the archive:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FwMkvJ1EEQ6b"
      },
      "source": [
        "import zipfile\n",
        "import io\n",
        "\n",
        "zf = zipfile.ZipFile(io.BytesIO(uploaded['fitness_poses_images_in.zip']), \"r\")\n",
        "zf.extractall()\n",
        "os.listdir('.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1QBS_P_Y_2mg"
      },
      "source": [
        "## Bootstrap images"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bERVPO8Ja6j7"
      },
      "source": [
        "# Required structure of the images_in_folder:\n",
        "#\n",
        "#   fitness_poses_images_in/\n",
        "#     pushups_up/\n",
        "#       image_001.jpg\n",
        "#       image_002.jpg\n",
        "#       ...\n",
        "#     pushups_down/\n",
        "#       image_001.jpg\n",
        "#       image_002.jpg\n",
        "#       ...\n",
        "#     ...\n",
        "bootstrap_images_in_folder = 'fitness_poses_images_in'\n",
        "\n",
        "# Output folders for bootstrapped images and CSVs.\n",
        "bootstrap_images_out_folder = 'fitness_poses_images_out'\n",
        "bootstrap_csvs_out_folder = 'fitness_poses_csvs_out'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PVYsbbJbOW7W"
      },
      "source": [
        "# Initialize helper.\n",
        "bootstrap_helper = BootstrapHelper(\n",
        "    images_in_folder=bootstrap_images_in_folder,\n",
        "    images_out_folder=bootstrap_images_out_folder,\n",
        "    csvs_out_folder=bootstrap_csvs_out_folder,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e832H-X6b-v7"
      },
      "source": [
        "# Check how many pose classes and images for them are available.\n",
        "bootstrap_helper.print_images_in_statistics()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EAWtcZSHcQHc"
      },
      "source": [
        "# Bootstrap all images.\n",
        "# Set limit to some small number for debug.\n",
        "bootstrap_helper.bootstrap(per_pose_class_limit=None)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xRdqXeUScko9"
      },
      "source": [
        "# Check how many images were bootstrapped.\n",
        "bootstrap_helper.print_images_out_statistics()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rTc3dlvFgg50"
      },
      "source": [
        "# After initial bootstrapping images without detected poses were still saved in\n",
        "# the folderd (but not in the CSVs) for debug purpose. Let's remove them.\n",
        "bootstrap_helper.align_images_and_csvs(print_removed_items=False)\n",
        "bootstrap_helper.print_images_out_statistics()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F3Sjl9JcJQzN"
      },
      "source": [
        "## Manual filtration\n",
        "\n",
        "Please manually verify predictions and remove samples (images) that has wrong pose prediction. Check as if you were asked to classify pose just from predicted landmarks. If you can't - remove it.\n",
        "\n",
        "Align CSVs and image folders once you are done."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fI7OytDZiHmg"
      },
      "source": [
        "# Align CSVs with filtered images.\n",
        "bootstrap_helper.align_images_and_csvs(print_removed_items=False)\n",
        "bootstrap_helper.print_images_out_statistics()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uGxcMoDLK8L0"
      },
      "source": [
        "## Automatic filtration\n",
        "\n",
        "Classify each sample against database of all other samples and check if it gets in the same class as annotated after classification.\n",
        "\n",
        "There can be two reasons for the outliers:\n",
        "\n",
        "  * **Wrong pose prediction**: In this case remove such outliers.\n",
        "\n",
        "  * **Wrong classification** (i.e. pose is predicted correctly and you aggree with original pose class assigned to the sample): In this case sample is from the underrepresented group (e.g. unusual angle or just very few samples). Add more similar samples and run bootstrapping from the very beginning.\n",
        "\n",
        "Even if you just removed some samples it makes sence to re-run automatic filtration one more time as database of poses has changed.\n",
        "\n",
        "**Important!!** Check that you are using the same parameters when classifying whole videos later."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "txnkReCiJr4Y"
      },
      "source": [
        "# Find outliers.\n",
        "\n",
        "# Transforms pose landmarks into embedding.\n",
        "pose_embedder = FullBodyPoseEmbedder()\n",
        "\n",
        "# Classifies give pose against database of poses.\n",
        "pose_classifier = PoseClassifier(\n",
        "    pose_samples_folder=bootstrap_csvs_out_folder,\n",
        "    pose_embedder=pose_embedder,\n",
        "    top_n_by_max_distance=30,\n",
        "    top_n_by_mean_distance=10)\n",
        "\n",
        "outliers = pose_classifier.find_pose_sample_outliers()\n",
        "print('Number of outliers: ', len(outliers))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HS6TkprSklIw"
      },
      "source": [
        "# Analyze outliers.\n",
        "bootstrap_helper.analyze_outliers(outliers)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZTWgo8-rlDFv"
      },
      "source": [
        "# Remove all outliers (if you don't want to manually pick).\n",
        "bootstrap_helper.remove_outliers(outliers)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LR9xOhPdlrkN"
      },
      "source": [
        "# Align CSVs with images after removing outliers.\n",
        "bootstrap_helper.align_images_and_csvs(print_removed_items=False)\n",
        "bootstrap_helper.print_images_out_statistics()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nd0OSLWXMJFC"
      },
      "source": [
        "## Dump for the App\n",
        "\n",
        "Dump filtered poses to CSV and download it.\n",
        "\n",
        "Please check this [guide](https://developers.google.com/ml-kit/vision/pose-detection/classifying-poses#4_integrate_with_the_ml_kit_quickstart_app) on how to use this CSV in the ML Kit sample app."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-2VfIiLPML8A"
      },
      "source": [
        "import csv\n",
        "import os\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "def dump_for_the_app():\n",
        "  pose_samples_folder = 'fitness_poses_csvs_out'\n",
        "  pose_samples_csv_path = 'fitness_poses_csvs_out.csv'\n",
        "  file_extension = 'csv'\n",
        "  file_separator = ','\n",
        "\n",
        "  # Each file in the folder represents one pose class.\n",
        "  file_names = [name for name in os.listdir(pose_samples_folder) if name.endswith(file_extension)]\n",
        "\n",
        "  with open(pose_samples_csv_path, 'w') as csv_out:\n",
        "    csv_out_writer = csv.writer(csv_out, delimiter=file_separator, quoting=csv.QUOTE_MINIMAL)\n",
        "    for file_name in file_names:\n",
        "      # Use file name as pose class name.\n",
        "      class_name = file_name[:-(len(file_extension) + 1)]\n",
        "\n",
        "      # One file line: `sample_00001,x1,y1,x2,y2,....`.\n",
        "      with open(os.path.join(pose_samples_folder, file_name)) as csv_in:\n",
        "        csv_in_reader = csv.reader(csv_in, delimiter=file_separator)\n",
        "        for row in csv_in_reader:\n",
        "          row.insert(1, class_name)\n",
        "          csv_out_writer.writerow(row)\n",
        "\n",
        "  files.download(pose_samples_csv_path)\n",
        "\n",
        "\n",
        "dump_for_the_app()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cfE7C9wHel_7"
      },
      "source": [
        "# Step 2: Classification\n",
        "\n",
        "**Important!!** Check that you are using the same classification parameters as while building classifier."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k4eaOJpoEvFl"
      },
      "source": [
        "# Upload your video.\n",
        "uploaded = files.upload()\n",
        "os.listdir('.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nYkZJ3a_2MW_"
      },
      "source": [
        "# Specify your video name and target pose class to count the repetitions.\n",
        "video_path = 'pushups-sample.mov'\n",
        "class_name='pushups_down'\n",
        "out_video_path = 'pushups-sample-out.mov'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A3gGMDE0R2pe"
      },
      "source": [
        "# Open the video.\n",
        "import cv2\n",
        "\n",
        "video_cap = cv2.VideoCapture(video_path)\n",
        "\n",
        "# Get some video parameters to generate output video with classificaiton.\n",
        "video_n_frames = video_cap.get(cv2.CAP_PROP_FRAME_COUNT)\n",
        "video_fps = video_cap.get(cv2.CAP_PROP_FPS)\n",
        "video_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
        "video_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7t_ACEmTSOhr"
      },
      "source": [
        "# Initilize tracker, classifier and counter.\n",
        "# Do that before every video as all of them have state.\n",
        "from mediapipe.python.solutions import pose as mp_pose\n",
        "\n",
        "\n",
        "# Folder with pose class CSVs. That should be the same folder you using while\n",
        "# building classifier to output CSVs.\n",
        "pose_samples_folder = 'fitness_poses_csvs_out'\n",
        "\n",
        "# Initialize tracker.\n",
        "pose_tracker = mp_pose.Pose(upper_body_only=False)\n",
        "\n",
        "# Initialize embedder.\n",
        "pose_embedder = FullBodyPoseEmbedder()\n",
        "\n",
        "# Initialize classifier.\n",
        "# Ceck that you are using the same parameters as during bootstrapping.\n",
        "pose_classifier = PoseClassifier(\n",
        "    pose_samples_folder=pose_samples_folder,\n",
        "    pose_embedder=pose_embedder,\n",
        "    top_n_by_max_distance=30,\n",
        "    top_n_by_mean_distance=10)\n",
        "\n",
        "# # Uncomment to validate target poses used by classifier and find outliers.\n",
        "# outliers = pose_classifier.find_pose_sample_outliers()\n",
        "# print('Number of pose sample outliers (consider removing them): ', len(outliers))\n",
        "\n",
        "# Initialize EMA smoothing.\n",
        "pose_classification_filter = EMADictSmoothing(\n",
        "    window_size=10,\n",
        "    alpha=0.2)\n",
        "\n",
        "# Initialize counter.\n",
        "repetition_counter = RepetitionCounter(\n",
        "    class_name=class_name,\n",
        "    enter_threshold=6,\n",
        "    exit_threshold=4)\n",
        "\n",
        "# Initialize renderer.\n",
        "pose_classification_visualizer = PoseClassificationVisualizer(\n",
        "    class_name=class_name,\n",
        "    plot_x_max=video_n_frames,\n",
        "    # Graphic looks nicer if it's the same as `top_n_by_mean_distance`.\n",
        "    plot_y_max=10)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4lXymkneOjgZ"
      },
      "source": [
        "# Run classification on a video.\n",
        "import os\n",
        "import tqdm\n",
        "\n",
        "from mediapipe.python.solutions import drawing_utils as mp_drawing\n",
        "\n",
        "\n",
        "# Open output video.\n",
        "out_video = cv2.VideoWriter(out_video_path, cv2.VideoWriter_fourcc(*'mp4v'), video_fps, (video_width, video_height))\n",
        "\n",
        "frame_idx = 0\n",
        "output_frame = None\n",
        "with tqdm.tqdm(total=video_n_frames, position=0, leave=True) as pbar:\n",
        "  while True:\n",
        "    # Get next frame of the video.\n",
        "    success, input_frame = video_cap.read()\n",
        "    if not success:\n",
        "      break\n",
        "\n",
        "    # Run pose tracker.\n",
        "    input_frame = cv2.cvtColor(input_frame, cv2.COLOR_BGR2RGB)\n",
        "    result = pose_tracker.process(image=input_frame)\n",
        "    pose_landmarks = result.pose_landmarks\n",
        "\n",
        "    # Draw pose prediction.\n",
        "    output_frame = input_frame.copy()\n",
        "    if pose_landmarks is not None:\n",
        "      mp_drawing.draw_landmarks(\n",
        "          image=output_frame,\n",
        "          landmark_list=pose_landmarks,\n",
        "          connections=mp_pose.POSE_CONNECTIONS)\n",
        "    \n",
        "    if pose_landmarks is not None:\n",
        "      # Get landmarks.\n",
        "      frame_height, frame_width = output_frame.shape[0], output_frame.shape[1]\n",
        "      pose_landmarks = np.array([[lmk.x * frame_width, lmk.y * frame_height, lmk.z * frame_width]\n",
        "                                 for lmk in pose_landmarks.landmark], dtype=np.float32)\n",
        "      assert pose_landmarks.shape == (33, 3), 'Unexpected landmarks shape: {}'.format(pose_landmarks.shape)\n",
        "\n",
        "      # Classify the pose on the current frame.\n",
        "      pose_classification = pose_classifier(pose_landmarks)\n",
        "\n",
        "      # Smooth classification using EMA.\n",
        "      pose_classification_filtered = pose_classification_filter(pose_classification)\n",
        "\n",
        "      # Count repetitions.\n",
        "      repetitions_count = repetition_counter(pose_classification_filtered)\n",
        "    else:\n",
        "      # No pose => no classification on current frame.\n",
        "      pose_classification = None\n",
        "\n",
        "      # Still add empty classification to the filter to maintaing correct\n",
        "      # smoothing for future frames.\n",
        "      pose_classification_filtered = pose_classification_filter(dict())\n",
        "      pose_classification_filtered = None\n",
        "\n",
        "      # Don't update the counter presuming that person is 'frozen'. Just\n",
        "      # take the latest repetitions count.\n",
        "      repetitions_count = repetition_counter.n_repeats\n",
        "\n",
        "    # Draw classification plot and repetition counter.\n",
        "    output_frame = pose_classification_visualizer(\n",
        "        frame=output_frame,\n",
        "        pose_classification=pose_classification,\n",
        "        pose_classification_filtered=pose_classification_filtered,\n",
        "        repetitions_count=repetitions_count)\n",
        "\n",
        "    # Save the output frame.\n",
        "    out_video.write(cv2.cvtColor(np.array(output_frame), cv2.COLOR_RGB2BGR))\n",
        "\n",
        "    # Show intermediate frames of the video to track progress.\n",
        "    if frame_idx % 50 == 0:\n",
        "      show_image(output_frame)\n",
        "\n",
        "    frame_idx += 1\n",
        "    pbar.update()\n",
        "\n",
        "# Close output video.\n",
        "out_video.release()\n",
        "\n",
        "# Release MediaPipe resources.\n",
        "pose_tracker.close()\n",
        "\n",
        "# Show the last frame of the video.\n",
        "if output_frame is not None:\n",
        "  show_image(output_frame)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xJFftPiQE56E"
      },
      "source": [
        "# Download generated video\n",
        "files.download(out_video_path)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}