{
 "metadata": {
  "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.5"
  },
  "orig_nbformat": 4,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.5 64-bit"
  },
  "interpreter": {
   "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnxruntime\n",
    "import torch\n",
    "from PIL import Image\n",
    "import cv2\n",
    "import torchvision.transforms as transforms\n",
    "import numpy as np\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "handwrite_mean=(0.877, 0.877, 0.877)\n",
    "handwrite_std=(0.200, 0.200, 0.200)\n",
    "\n",
    "# 减少图像白边\n",
    "def cutImg(img):\n",
    "    dst = 255- img\n",
    "    gray = cv2.cvtColor(dst,cv2.COLOR_BGR2GRAY)  \n",
    "    ret, binary = cv2.threshold(gray,127,255,cv2.THRESH_BINARY)  \n",
    "    \n",
    "    contours, hierarchy = cv2.findContours(binary,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)\n",
    "\n",
    "    minx = 65523\n",
    "    miny = 65523\n",
    "    maxx = 0\n",
    "    maxy = 0\n",
    "\n",
    "    for contour in contours:\n",
    "        boundRect = cv2.boundingRect(contour)\n",
    "        (x,y,w,h)=boundRect\n",
    "        if x<minx:\n",
    "            minx=x\n",
    "        if y<miny:\n",
    "            miny=y\n",
    "        if x+w>maxx:\n",
    "            maxx=x+w\n",
    "        if y+h>maxy:\n",
    "            maxy=y+h\n",
    "    cutw=maxx-minx\n",
    "    cuth=maxy-miny\n",
    "    # w>h\n",
    "    # if cutw>cuth:\n",
    "    #     cuty1=int((maxy+miny)/2-cutw/2)\n",
    "    #     cuty2=int((maxy+miny)/2+cutw/2)\n",
    "    #     img = img[cuty1:cuty2,minx:maxx]\n",
    "    # if cutw<cuth:\n",
    "    #     cutx1=int((maxx+minx)/2-cuth/2)\n",
    "    #     cutx2=int((maxx+minx)/2+cuth/2)\n",
    "    #     img = img[miny:maxy,cutx1:cutx2] \n",
    "    img = img[miny:maxy,minx:maxx]\n",
    "    return img\n",
    "\n",
    "def to_numpy(tensor):\n",
    "    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()\n",
    "    \n",
    "test = 'data/test/receive.png'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = cv2.imread(test)\n",
    "img = cutImg(img)\n",
    "image = Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))\n",
    "# cv2.imwrite('dealcut.png',img)\n",
    "# image = Image.open(test).convert('RGB')\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean = handwrite_mean, std = handwrite_std)\n",
    "])\n",
    "img_tensor = transform(image).unsqueeze(0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "ort_session = onnxruntime.InferenceSession('onnx/handwriting.onnx')\n",
    "ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(img_tensor)}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[array([[-63.502773, -50.50067 , -58.770657, ..., -60.615524, -62.879475,\n        -58.971222]], dtype=float32)]\n"
     ]
    }
   ],
   "source": [
    "ort_outs = ort_session.run(None, ort_inputs)\n",
    "print(ort_outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[3594 2715 1533 ... 3279 1428 1398]\n"
     ]
    }
   ],
   "source": [
    "ret = np.array(ort_outs[0][0])\n",
    "result = np.argsort(-ret)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ]
}