yolov7/tools/keypoint.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c691ce2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import cv2\n",
    "from torchvision import transforms\n",
    "import numpy as np\n",
    "from utils.datasets import letterbox\n",
    "from utils.general import non_max_suppression_kpt\n",
    "from utils.plots import output_to_keypoint, plot_skeleton_kpts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d17fb9ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "weigths = torch.load('yolov7-w6-pose.pt')\n",
    "model = weigths['model']\n",
    "model = model.half().to(device)\n",
    "_ = model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27396504",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = cv2.imread('./person.jpg')\n",
    "image = letterbox(image, 960, stride=64, auto=True)[0]\n",
    "image_ = image.copy()\n",
    "image = transforms.ToTensor()(image)\n",
    "image = torch.tensor(np.array([image.numpy()]))\n",
    "image = image.to(device)\n",
    "image = image.half()\n",
    "\n",
    "output, _ = model(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77662e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = non_max_suppression_kpt(output, 0.25, 0.65, nc=model.yaml['nc'], nkpt=model.yaml['nkpt'], kpt_label=True)\n",
    "with torch.no_grad():\n",
    "    output = output_to_keypoint(output)\n",
    "nimg = image[0].permute(1, 2, 0) * 255\n",
    "nimg = nimg.cpu().numpy().astype(np.uint8)\n",
    "nimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)\n",
    "for idx in range(output.shape[0]):\n",
    "    plot_skeleton_kpts(nimg, output[idx, 7:].T, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7b0882c0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAcwAAAFzCAYAAACzRzkmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9R5NtS5qeiT3uS+61tY4d8sTR5+p7U4uq6iqIBgjrBkia0azNOGoMaBzjd6D5EzihkTOQbAKNRHWJzBKprhZHnxNa7NhaLC3cOYhbQGWhKiuRTTMajfEMwiJ2LOlLvPtz/77XhdaaG2644YYbbrjh1yP/v30AN9xwww033PD/C9wI5g033HDDDTf8BtwI5g033HDDDTf8BtwI5g033HDDDTf8BtwI5g033HDDDTf8BtwI5g033HDDDTf8Bpi/7p//6l/9q5uakxtuuOGGG/7/in/9r/+1+Ns+v4kwb7jhhhtuuOE34EYwb7jhhhtuuOE34EYwb7jhhhtuuOE34EYwb7jhhhtuuOE34EYwb7jhhhtuuOE34EYwb7jhhhtuuOE34EYwb7jhhhtuuOE34EYwb7jhhhtuuOE34NcaF4AGDfy1Es6/b/pMcb3Wb4T4j9v9q1/+xs7+S/mVnf/Nbf3q33+177/9fP7Tsn81X6gQ4nph8Z8fn/j6x1/f1q/uTf/KYQkBf9+xXW/r72vJv62t/uqa/bp2/Kvt/rVl/mPbaX61If/Gdv5LLvCv7E/w91/fv75h8dfagV/5/T9f9m85zr+xzt95vX/lvP92BOJveQb+5jn96rH/7X//HdfrP1vnb94bfz9CiF89rr/1vP4XPFv/+Q7//pfB/2L+ZtuIv/E5/Mbn9B9vBP7GI/KbPyd/fe7gv+30f5Mmub5O+m989ret95u/C391/V93n/2XvGd/0zb+u+7Z33Sd3+Ye+rvX+zveWv8f49cK5r3dFD1Z4zUrpGHASii+/PApgYR4sqBWNpE1wTjIcHOTpuPS3WpzdT7DCsBPDC6DNUkBqTSxjIKyY9GuVKn3PRqtEps7NWrduwR+zGo85N6dfbqDJq6hMb0KYbSiCDy0tcY1DUpOhZPzU8KoIJ8I1r5Bud6jvtuiXF/jh1PmE5PZ6Rn7u7s0br/Nj3/2J6yvjnjrwZtMpxnvf7DLoF9gWC4vj16xmiVU7QqtdpUoi1hdXrB95y7DecrTnx+i/QXNzT1eH5+zsb3Nxr7HxmYNVzqYfoZbLZGXDF4cHTJdREz9lOlE8e79b1HvV3nx2f/MaHKJsEzKRsHO9i129zok+RqdmASrJd1Bk0rtIefnp+zttXj54iXPvjrEtl3qlRbTxZIiSWm0HQLho2KBW9/n/TceoJMRy3XEKkxIWTI7mvHWWw/obitqpSZ+lCKNOks9Yz2b0TQLplGIpEOjVicL57S3qtjKYnp5iSp5nA9PKXSZQXvAaD6lkC2absrtuzv4y4xCSzzHZroaEc5XGGLAZJFgG2U2+nsY5QqJghef/wQdXOKVKuRJRrvVYXDnIavVAbZVEPsJhi2xy5LXrw4Ig4hGs06r0ufug99hsXqNP5nRv/U2o9FroqsJ/X2Xq9UMw9AEwRrDcHHLfereLmXb4XJ4QLPmUmuVSPwAw5YYjgnaZrGYkuqUJMyQmLQ2mkyGPugFy2WAtG2WqxWVkok0bazE5cGbH2C7IyKdka01fuozn6xx4y5X4RGdVp3xYkyS5Yjcpt3dwEQRJFNKZYtKZZOL4wW1hkPFLSEsaLgOsXK4OHmF8ExqzhbLq+eM1jGDzQfYVY0uEharE7TKWI5C0kRTadmESYjExBAgMLhz+022dvaYL89Zrlb0B5sIw2IyH6HDhOHBDKvu0N7okyUKYRjMZ7C5+Q2q5ZTlfIZrbzNZntPd6DIeTbhaLdmqbmOIkFXwjHbHxtQN4tRAWgWd7W9ixYpsOWE4GkJlheUIluMReSKptlqkmSIpQtxyFWk/wnE8siDBEgm5ExJHIzY3PsDEYHn5Mc1WG0yFH12S+CuGZzPsSkHZHLCezHBrTYxSjyy8YLYa4lgt6lUPp9bGlQUqESxViU7rLtPzcwxHUi61OHr6c5LZGCkklimYLhOeX53yu7//iFu3dxmNYhqbdXS0Yr1e0Wg42JbDs1fHjP2YN+5tUDHr/D/+/Z9wcnlCTVb4/u+9xccfnTKeHZEX8M6D9+k2emSzBYcnh7jCY53mhMaKVRjRb9fx3AoffPAOH/7sCwJVIB1otZq8cbvLX/zlU0qyQuwveX56joNC2CZF4WCXLHLjWgryOMewMkQgSDLBu9+4x2a3z1/8/BPC5RKFAsPEsEyKPCNNrt/yssgRlkRKA2EKhBK0qxZuySUPNWfrgP2dGvOriGUYkGYKhEQKjZQGUggQFnubfc6GIxzPIA1jsiSlUAKNSbVSpuQZDC9nWCUTKTRFpgGNUhLLMnFQrLIMkWtMw0CjyMXX32C0xhISaUi0LtBAmgnyXGHaYKJRApRponWB1ALDUKAFaaEhFxSFBkMitMCyFFIK0gykUKAMcqkxlUYIjZAgCkCK675WpSkyibQl/8P/8H/6LxfM99/bQPpNnF6X4fmQIBCIYMrHHx4yXi1ZxDGMXbJM4rkKs2yyOovQhs3WluTtvT2uTi55/GxMFESoVBCu1xxNxhhXBo1aieGBi+cdo0uSPNWkuqC/8rizt83DwQapm7LSAevCIKnGWFaLoyAivbrk4f4dmuY2hmgx9OecD4+Y+V9xp3aPuiHY3jGo9KZ8+2GLQyEIswr+/Ix8ahKVDKzuCqnWlCmx2W5gVB3MrGDQuI3tefS2XN699T5Z5HO+kBydTSiKgtEsZNDp0tpy6T6oEePx1ZMntGs2VVfhzANWkwm5f8ZK7iLMLkIkWPmI/Vs9TLegXPKQeco6irAcg3LZoOLF1J2Yq+lz6lWXSllyOfERtRauZ5BTEIQWw1XEZs/AzH0uRiMebDkMNvs8OxgSLuZYGw20imm1y1SqK1qFx5NX5yzDKYbMKPc8TN/EytdUdhyKuEq9bqNMgWFWcEoGpVKfv/jpE8LIxyRiUHfY3R7Q74AuAgb9LfIiIYkzsgJc16Be8RgMHnDn7pusozXz9ZiHd+8zG60hl1QGNczyFjEWjWaFZt0myULII+bBnEZH0unVqNRLJCufWk2gyDGlQ6OSUq5WWTUgFyM2WyZhkpGEIY6lcewAIX1sR7O9VaHZq2CZmrmTUC3X8YMJOTFlL8NOVlhOlfmyYD45oVX18NcC09KkOsIyE2w7pVxy2bq9Tb27JkojWrJMpBeQ5xRLm929h9iXCaabsfIzHLvMbB4iZECj3kEtAirlDuWqZHOrTpT6VNoO1coOzYrDxWzI/TcGFIUFThWhLWzHpFwqWIcjijglWY+pVSwqbYGfhph1g+AyIpcpO9u3eLDzJlpUIEupOSV8NWYVHSDzGsUqpX+7gaViZvMl82FCuXWHUmXAw/YultWhXlvjJzF7tx8hTwx6vU10BOenn6E2JK3mLTrRgPPpC3q1HUoND2WMsOycilsDp2CdTElyQTCbAD7l+gau51Cxu4TJkqvJGPQrmrUeOGVW4Zp8fECzW6fWGVACinUL24iQTsRsPaSQFl7Jw5Aaq+SSGwG79/rYlocpblM176FESBwLpFmlUasRW1skaYdCWDT6GyyjFZaIqBw4aKGRjoGUJuiYIvExRULJDrm/v8lKxViWg2Va5Mpnsrji9PQAA4mRuLx6ecB6dkYRLXnj+x9QxBKnCHHJOV+sePL6C779/tvcfnCLn3z2C7713lssXk8YHx+RigJhJnyzU+Kjjz4itApeXV7gZYI8mXJ5/orhccj9/V0Oz8/wFyO0Y4CyCSKBJVxyJOkiAQpQGSrVVLwGV8Mrvvj0CcEyIE4DlGtii2tRLNKEIgfPa5AEPsopMKVHo1uj12jSb5eZhktyGZNHAWlU4M/XLNYhUZHjugZSQ7lWQmcFfqJ4GkakWULf9JhP5uQqJUs0Uhno0OY4zbEM0GFGphRGJtCmpEDgCBPSFCEFmyWPlYqZpxkogVaCQio8x0FpTV4kaGW
      "text/plain": [
       "<Figure size 576x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "plt.figure(figsize=(8,8))\n",
    "plt.axis('off')\n",
    "plt.imshow(nimg)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "197b9edd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}