{
 "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', map_location=device)\n",
    "model = weigths['model']\n",
    "_ = model.float().eval()\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    model.half().to(device)"
   ]
  },
  {
   "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",
    "\n",
    "if torch.cuda.is_available():\n",
    "    image = image.half().to(device)   \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": "[encoded data, 473304 characters]\n",
      "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
