mindspore_yolo/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",
    "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": {