mindspore_yolo/tools/visualization.ipynb

{
 "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": {