{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0ab662ce",
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cfd4b844",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "weigths = torch.load('./weights/yolov7-e6e.pt')\n",
    "model = weigths['model']\n",
    "model = model.half().to(device)\n",
    "_ = model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1ee054f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = cv2.imread('./images/person.jpg')  # 504x378 image\n",
    "image = letterbox(image, 1280, 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": 4,
   "id": "7ec9e6ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "obj1 = output[1][0][0, 0, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj2 = output[1][0][0, 1, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj3 = output[1][0][0, 2, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj4 = output[1][1][0, 0, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj5 = output[1][1][0, 1, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj6 = output[1][1][0, 2, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj7 = output[1][2][0, 0, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj8 = output[1][2][0, 1, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj9 = output[1][2][0, 2, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj10 = output[1][3][0, 0, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj11 = output[1][3][0, 1, :, :, 4].sigmoid().cpu().numpy()\n",
    "obj12 = output[1][3][0, 2, :, :, 4].sigmoid().cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "38878c81",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 456064 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 33572 characters]\n",
      "text/plain": [
       "<Figure size 1152x864 with 12 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(image_[:,:,[2,1,0]])\n",
    "plt.show()\n",
    "fig, ax = plt.subplots(4,3,figsize=(16,12))\n",
    "#[ax_.axis('off') for ax_ in ax.ravel()]\n",
    "[ax_.set_xticklabels([]) for ax_ in ax.ravel()]\n",
    "[ax_.set_yticklabels([]) for ax_ in ax.ravel()]\n",
    "ax.ravel()[0].imshow(obj1)\n",
    "ax.ravel()[1].imshow(obj2)\n",
    "ax.ravel()[2].imshow(obj3)\n",
    "ax.ravel()[3].imshow(obj4)\n",
    "ax.ravel()[4].imshow(obj5)\n",
    "ax.ravel()[5].imshow(obj6)\n",
    "ax.ravel()[6].imshow(obj7)\n",
    "ax.ravel()[7].imshow(obj8)\n",
    "ax.ravel()[8].imshow(obj9)\n",
    "ax.ravel()[9].imshow(obj10)\n",
    "ax.ravel()[10].imshow(obj11)\n",
    "ax.ravel()[11].imshow(obj12)\n",
    "plt.subplots_adjust(wspace=-0.52, hspace=0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8536ecc8",
   "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
}
