Untitled.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "88833101",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2, os, torch, onnxruntime\n",
    "import imgaug.augmenters as iaa\n",
    "import imgaug as ia\n",
    "from imgaug.augmentables.bbs import BoundingBox\n",
    "from imgaug import Polygon\n",
    "from yolox.exp import get_exp\n",
    "from yolox.models.network_blocks import SiLU\n",
    "from yolox.utils import replace_module\n",
    "from torch import nn\n",
    "\n",
    "from pycocotools.coco import COCO\n",
    "from pycocotools.mask import encode,decode,area,toBbox\n",
    "from pycocotools import mask as maskUtils\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ed7371d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.randn(1,20,2)\n",
    "b = a.repeat(1,1,4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bc432cd2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([-1.1475,  0.1468]),\n",
       " tensor([-1.1475,  0.1468, -1.1475,  0.1468, -1.1475,  0.1468, -1.1475,  0.1468]))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a[0, 0], b[0, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "5e285180",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rle_encode(binary_mask):\n",
    "    pixels = binary_mask.flatten()\n",
    "    pixels = np.concatenate([[0], pixels, [0]])\n",
    "    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n",
    "    runs[1::2] -= runs[::2]\n",
    "    return ' '.join(str(x) for x in runs)\n",
    "def rle_decode(mask_rle, shape):\n",
    "    s = mask_rle.split()\n",
    "    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n",
    "    starts -= 1\n",
    "    ends = starts + lengths\n",
    "    binary_mask = np.zeros(shape[0] * shape[1], dtype=np.uint8)\n",
    "    for lo, hi in zip(starts, ends):\n",
    "        binary_mask[lo:hi] = 1\n",
    "    return binary_mask.reshape(shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "11f77bb3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'100201 200 100701 200 101201 200 101701 200 102201 200 102701 200 103201 200 103701 200 104201 200 104701 200 105201 200 105701 200 106201 200 106701 200 107201 200 107701 200 108201 200 108701 200 109201 200 109701 200 110201 200 110701 200 111201 200 111701 200 112201 200 112701 200 113201 200 113701 200 114201 200 114701 200 115201 200 115701 200 116201 200 116701 200 117201 200 117701 200 118201 200 118701 200 119201 200 119701 200 120201 200 120701 200 121201 200 121701 200 122201 200 122701 200 123201 200 123701 200 124201 200 124701 200 125201 200 125701 200 126201 200 126701 200 127201 200 127701 200 128201 200 128701 200 129201 200 129701 200 130201 200 130701 200 131201 200 131701 200 132201 200 132701 200 133201 200 133701 200 134201 200 134701 200 135201 200 135701 200 136201 200 136701 200 137201 200 137701 200 138201 200 138701 200 139201 200 139701 200 140201 200 140701 200 141201 200 141701 200 142201 200 142701 200 143201 200 143701 200 144201 200 144701 200 145201 200 145701 200 146201 200 146701 200 147201 200 147701 200 148201 200 148701 200 149201 200 149701 200 150201 100 150351 50 150701 100 150851 50 151201 100 151351 50 151701 100 151851 50 152201 100 152351 50 152701 100 152851 50 153201 100 153351 50 153701 100 153851 50 154201 100 154351 50 154701 100 154851 50 155201 100 155351 50 155701 100 155851 50 156201 100 156351 50 156701 100 156851 50 157201 100 157351 50 157701 100 157851 50 158201 100 158351 50 158701 100 158851 50 159201 100 159351 50 159701 100 159851 50 160201 100 160351 50 160701 100 160851 50 161201 100 161351 50 161701 100 161851 50 162201 100 162351 50 162701 100 162851 50 163201 100 163351 50 163701 100 163851 50 164201 100 164351 50 164701 100 164851 50 165201 100 165351 50 165701 100 165851 50 166201 100 166351 50 166701 100 166851 50 167201 100 167351 50 167701 100 167851 50 168201 100 168351 50 168701 100 168851 50 169201 100 169351 50 169701 100 169851 50 170201 100 170351 50 170701 100 170851 50 171201 100 171351 50 171701 100 171851 50 172201 100 172351 50 172701 100 172851 50 173201 100 173351 50 173701 100 173851 50 174201 100 174351 50 174701 100 174851 50 175201 200 175701 200 176201 200 176701 200 177201 200 177701 200 178201 200 178701 200 179201 200 179701 200 180201 200 180701 200 181201 200 181701 200 182201 200 182701 200 183201 200 183701 200 184201 200 184701 200 185201 200 185701 200 186201 200 186701 200 187201 200 187701 200 188201 200 188701 200 189201 200 189701 200 190201 200 190701 200 191201 200 191701 200 192201 200 192701 200 193201 200 193701 200 194201 200 194701 200 195201 200 195701 200 196201 200 196701 200 197201 200 197701 200 198201 200 198701 200 199201 200 199701 200'"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rle_encode(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "cad50221",
   "metadata": {},
   "outputs": [],
   "source": [
    "contours, hierarchy = cv2.findContours((img*255).astype(np.uint8), cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)\n",
    "# mask_new, contours, hierarchy = cv2.findContours((mask).astype(np.uint8), cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)\n",
    "segmentation = []\n",
    "for contour in contours:\n",
    "    contour_list = contour.flatten().tolist()\n",
    "    if len(contour_list) > 4:# and cv2.contourArea(contour)>10000\n",
    "        segmentation.append(contour_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "930ac1cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "rles = maskUtils.frPyObjects(segmentation, 500, 500)\n",
    "rle = maskUtils.merge(rles)\n",
    "m = maskUtils.decode(rle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "bd1993e7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "c3d67140",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = np.zeros((500,500), np.uint8)\n",
    "img[200:400, 200:400] = 1\n",
    "img[300:350, 300:350] = 0\n",
    "cv2.imshow('1', img*255)\n",
    "\n",
    "mask_merge = np.asfortranarray(img)    \n",
    "mask_str = encode(mask_merge)\n",
    "mask_str['counts'] = str(mask_str['counts'], 'utf-8')\n",
    "mask=decode(mask_str)\n",
    "\n",
    "cv2.imshow('3', mask*255)\n",
    "cv2.waitKey(0)\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "7e179cb3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "b'XkQ3X6\\\\9000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000lLVH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7f400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000Xd`1'"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = str(mask_str['counts'], 'utf-8')\n",
    "bytes(a, encoding='utf8')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "7221b6f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'size': [500, 500],\n",
       " 'counts': b'XkQ3X6\\\\9000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000lLVH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7b1VH^Nj7f400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000Xd`1'}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mask_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "d2260f5b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "**** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ******** for onnx ****\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "outputs shape:  torch.Size([1, 2100, 6]) torch.Size([1, 2, 320, 320])\n"
     ]
    }
   ],
   "source": [
    "# ###   for onnx\n",
    "# exp = get_exp(None, 'yolox-plaster')\n",
    "# model = exp.get_model()\n",
    "# ckpt = torch.load('model_weights/plaster_s.pth', map_location=\"cpu\")\n",
    "# model.eval()\n",
    "# if \"model\" in ckpt:\n",
    "#     ckpt = ckpt[\"model\"]\n",
    "# model.load_state_dict(ckpt)\n",
    "# model = replace_module(model, nn.SiLU, SiLU)\n",
    "# model.head.decode_in_inference = False\n",
    "# img_channel = exp.img_channel\n",
    "# dummy_input = torch.randn(1, img_channel, exp.test_size[0], exp.test_size[1])\n",
    "# outputs, seg_output = model(dummy_input)\n",
    "# print('outputs shape: ', outputs.shape, seg_output.shape)\n",
    "# torch.onnx.export(model, dummy_input, 'out.onnx', opset_version=12, input_names=['images'],\n",
    "#         output_names=['output'], dynamic_axes={'images': {0: 'batch'}, 'output': {0: 'batch'}})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "36a4dbc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "### 计算算子核大小\n",
    "# for outsize in [1,2,3,6]:\n",
    "#     stridesz = np.floor(80/outsize).astype(np.int32)\n",
    "#     kernelsz = 80-(outsize-1)*stridesz\n",
    "#     print(kernelsz, stridesz)\n",
    "# avg1 = nn.AvgPool2d(kernel_size=80, stride=80)\n",
    "# avg2 = nn.AvgPool2d(kernel_size=40, stride=40)\n",
    "# avg3 = nn.AvgPool2d(kernel_size=28, stride=26)\n",
    "# avg4 = nn.AvgPool2d(kernel_size=15, stride=13)\n",
    "# pool1 = nn.AdaptiveAvgPool2d(1)\n",
    "# pool2 = nn.AdaptiveAvgPool2d(2)\n",
    "# pool3 = nn.AdaptiveAvgPool2d(3)\n",
    "# pool4 = nn.AdaptiveAvgPool2d(6)\n",
    "# img = torch.randn(1, 64, 80, 80)\n",
    "# aa = avg4(img)\n",
    "# aa.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4a859e87",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Anaconda3\\lib\\site-packages\\onnxruntime\\capi\\onnxruntime_inference_collection.py:350: UserWarning: Deprecation warning. This ORT build has ['CUDAExecutionProvider', 'CPUExecutionProvider'] enabled. The next release (ORT 1.10) will require explicitly setting the providers parameter (as opposed to the current behavior of providers getting set/registered by default based on the build flags) when instantiating InferenceSession.For example, onnxruntime.InferenceSession(..., providers=[\"CUDAExecutionProvider\"], ...)\n",
      "  warnings.warn(\"Deprecation warning. This ORT build has {} enabled. \".format(available_providers) +\n"
     ]
    }
   ],
   "source": [
    "model = onnxruntime.InferenceSession('out.onnx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "8d9a48d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def py_cpu_nms(dets, scores, thresh):\n",
    "    x1 = dets[:, 0]\n",
    "    y1 = dets[:, 1]\n",
    "    x2 = dets[:, 2]\n",
    "    y2 = dets[:, 3]\n",
    "    areas = (x2 - x1 + 1) * (y2 - y1 + 1)\n",
    "    order = scores.argsort()[::-1]\n",
    "    keep = []\n",
    "    while order.size > 0:\n",
    "        i = order[0]\n",
    "        keep.append(i)\n",
    "        xx1 = np.maximum(x1[i], x1[order[1:]])\n",
    "        yy1 = np.maximum(y1[i], y1[order[1:]])\n",
    "        xx2 = np.minimum(x2[i], x2[order[1:]])\n",
    "        yy2 = np.minimum(y2[i], y2[order[1:]])\n",
    "        w = np.maximum(0.0, xx2 - xx1 + 1)\n",
    "        h = np.maximum(0.0, yy2 - yy1 + 1)\n",
    "        inter = w * h\n",
    "        ovr = inter / (areas[i] + areas[order[1:]] - inter)\n",
    "        inds = np.where(ovr <= thresh)[0]\n",
    "        order = order[inds + 1]\n",
    "    return np.array(keep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "b1576303",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_pred = torch.from_numpy(image_pred)\n",
    "class_conf, class_pred = torch.max(torch_pred[:, 5: 5 + num_classes], 1, keepdim=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "403fc008",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2100])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class_conf.squeeze().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "05093ac3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2100,)"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class_conf = np.max(image_pred[:, 5: 5 + num_classes], axis=1)\n",
    "class_pred = np.argmax(image_pred[:, 5: 5 + num_classes], axis=1)#.reshape(-1, 1)\n",
    "(image_pred[:, 4]* class_conf>= conf_thre).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "85275a32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[100.77322388,  28.80500793, 304.89938354, 319.77435303,\n",
       "           0.99929237,   0.97557169,   0.        ],\n",
       "        [  0.52454185,  30.3742218 ,  24.95177269,  66.66628265,\n",
       "           0.9509362 ,   0.90179068,   0.        ]])]"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def postprocess(prediction, num_classes, conf_thre=0.7, nms_thre=0.45, class_agnostic=False, keypoints=False):\n",
    "\n",
    "    box_corner = np.zeros_like(prediction)\n",
    "    box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2\n",
    "    box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2\n",
    "    box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2\n",
    "    box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2\n",
    "    prediction[:, :, :4] = box_corner[:, :, :4]\n",
    "\n",
    "    output = [None for _ in range(len(prediction))]\n",
    "    for i, image_pred in enumerate(prediction):\n",
    "        if not image_pred.shape[0]:\n",
    "            continue\n",
    "\n",
    "        class_conf = np.max(image_pred[:, 5: 5 + num_classes], axis=1).reshape(-1, 1)\n",
    "        class_pred = np.argmax(image_pred[:, 5: 5 + num_classes], axis=1).reshape(-1, 1)\n",
    "        conf_mask = (image_pred[:, 4] * class_conf.squeeze() >= conf_thre).squeeze()\n",
    "\n",
    "        if keypoints>0:\n",
    "            detections = torch.cat((image_pred[:, :5],\n",
    "                                    class_conf,\n",
    "                                    class_pred.float(),\n",
    "                                    image_pred[:, 5+num_classes:]), 1)\n",
    "        else:\n",
    "            detections = np.concatenate((image_pred[:, :5], class_conf, class_pred), 1)\n",
    "        detections = detections[conf_mask]\n",
    "\n",
    "        if not detections.shape[0]:\n",
    "            continue\n",
    "        nms_out_index = py_cpu_nms(\n",
    "            detections[:, :4],\n",
    "            detections[:, 4] * detections[:, 5],nms_thre)\n",
    "\n",
    "        detections = detections[nms_out_index]\n",
    "        if output[i] is None:\n",
    "            output[i] = detections\n",
    "        else:\n",
    "            output[i] = np.concatenate((output[i], detections))\n",
    "\n",
    "    return output\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "351614d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def yolox_decode_outputs(outputs, seg_output, yolox_hw, yolox_strides, keypoints=0):\n",
    "    grids = []\n",
    "    strides = []\n",
    "    for (hsize, wsize), stride in zip(yolox_hw, yolox_strides):\n",
    "        yv, xv = np.meshgrid(np.arange(hsize), np.arange(wsize))\n",
    "        yv = np.transpose(yv, (1, 0))\n",
    "        xv = np.transpose(xv, (1, 0))\n",
    "        grid = np.stack((xv, yv), 2).reshape(1, -1, 2)\n",
    "        grids.append(grid)\n",
    "        shape = grid.shape[:2]\n",
    "        strides.append(np.full((*shape, 1), stride))\n",
    "\n",
    "    grids = np.concatenate(grids, axis=1)\n",
    "    strides = np.concatenate(strides, axis=1)\n",
    "    outputs[..., :2] = (outputs[..., :2] + grids) * strides  # xy\n",
    "    outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * strides  # wh\n",
    "    if keypoints > 0:\n",
    "        for i in [i for i in range(2, self.keypoints * 2 + 1, 2)][::-1]:\n",
    "            if i != 2:\n",
    "                outputs[..., -1 * i:-1 * (i - 2)] = outputs[..., -1 * i:-1 * (i - 2)] * strides + outputs[..., :2]\n",
    "            else:\n",
    "                outputs[..., -2:] = outputs[..., -2:] * strides + outputs[..., :2]\n",
    "    seg_ls = []\n",
    "    if seg_output is not None:\n",
    "        for seg in seg_output:\n",
    "            seg = np.argmax(seg, axis=0)\n",
    "            seg_ls.append(seg)\n",
    "    return outputs, seg_ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "355c5e63",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 4) (2, 1) (2, 1)\n"
     ]
    }
   ],
   "source": [
    "def preproc(img, input_size, swap=(2, 0, 1)):\n",
    "    h, w, c = img.shape\n",
    "    if len(img.shape) == 3:\n",
    "        padded_img = np.ones((input_size[0], input_size[1], c), dtype=np.uint8) * 114\n",
    "    else:\n",
    "        padded_img = np.ones(input_size, dtype=np.uint8) * 114\n",
    "\n",
    "    r = min(input_size[0] / h, input_size[1] / w)\n",
    "    resized_img = cv2.resize(img, (int(w * r), int(h * r)), \n",
    "                             interpolation=cv2.INTER_LINEAR,).astype(np.uint8)\n",
    "    padded_img[: int(h * r), : int(w * r)] = resized_img\n",
    "    padded_img = padded_img.transpose(swap)\n",
    "    padded_img = np.expand_dims(padded_img, axis=0)\n",
    "    return padded_img\n",
    "\n",
    "img = cv2.imread('imgs/plaster2/3.png', -1)\n",
    "height, width = img.shape[:2]\n",
    "in_shapes = [[height, width]]\n",
    "test_size = [320, 320]\n",
    "\n",
    "ratio = min(test_size[0] / height, test_size[1] / width)\n",
    "ratio_ls = [ratio]\n",
    "\n",
    "image = preproc(img, test_size)\n",
    "\n",
    "outputs, seg_output = model.run(['output', '1280'], {\"images\": np.array(image, np.float32)})\n",
    "\n",
    "yolox_hw = []\n",
    "input_h, input_w = test_size\n",
    "yolox_strides = [8, 16, 32]\n",
    "for stride in yolox_strides:\n",
    "    yolox_hw.append([input_h//stride, input_w//stride])\n",
    "outputs, seg_outputs = yolox_decode_outputs(outputs, seg_output, yolox_hw, yolox_strides, keypoints=0)\n",
    "\n",
    "num_classes, nms_thre, conf_thre = 1, 0.45, 0.25\n",
    "outputs = postprocess(outputs, num_classes, confthre, nmsthre)\n",
    "new_outputs, new_seg_outputs = [], []\n",
    "for in_shape, ratio, output, seg in zip(in_shapes, ratio_ls, outputs, seg_outputs):\n",
    "    box, cls = output[:, 0:4], output[:, 6]\n",
    "    scores = output[:, 4] * output[:, 5]\n",
    "    box /= ratio\n",
    "    cls, scores = np.expand_dims(cls, axis=1), np.expand_dims(scores, axis=1)\n",
    "    output = np.concatenate([box, scores, cls], axis=1)\n",
    "    h, w = in_shape\n",
    "    sh, sw = seg.shape\n",
    "    seg = cv2.resize(seg, (int(sw/ratio), int(sh/ratio)), interpolation=cv2.INTER_NEAREST)[:h, :w]\n",
    "    new_outputs.append(output)\n",
    "    new_seg_outputs.append(seg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "c8f99058",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 7)"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_outputs[0].shape"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}