{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import cv2\n",
    "from glob import glob\n",
    "from os.path import join as pjoin\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def resize_label(bboxes, d_height, gt_height, bias=0):\n",
    "    bboxes_new = []\n",
    "    scale = gt_height / d_height\n",
    "    for bbox in bboxes:\n",
    "        bbox = [int(b * scale + bias) for b in bbox]\n",
    "        bboxes_new.append(bbox)\n",
    "    return bboxes_new\n",
    "\n",
    "\n",
    "def draw_bounding_box(org, corners, color=(0, 255, 0), line=2, show=False):\n",
    "    board = org.copy()\n",
    "    for i in range(len(corners)):\n",
    "        board = cv2.rectangle(board, (corners[i][0], corners[i][1]), (corners[i][2], corners[i][3]), color, line)\n",
    "    if show:\n",
    "        cv2.imshow('a', cv2.resize(board, (500, 1000)))\n",
    "        cv2.waitKey(0)\n",
    "    return board\n",
    "\n",
    "\n",
    "def load_detect_result_json(reslut_file_root, shrink=0):\n",
    "    def is_bottom_or_top(corner):\n",
    "        column_min, row_min, column_max, row_max = corner\n",
    "        if row_max < 36 or row_min > 725:\n",
    "            return True\n",
    "        return False\n",
    "\n",
    "    result_files = glob(pjoin(reslut_file_root, '*.json'))\n",
    "    compos_reform = {}\n",
    "    print('Loading %d detection results' % len(result_files))\n",
    "    for reslut_file in tqdm(result_files):\n",
    "        img_name = reslut_file.split('\\\\')[-1].split('.')[0]\n",
    "        compos = json.load(open(reslut_file, 'r'))['compos']\n",
    "        for compo in compos:\n",
    "            if is_bottom_or_top((compo['column_min'], compo['row_min'], compo['column_max'], compo['row_max'])):\n",
    "                continue\n",
    "            if img_name not in compos_reform:\n",
    "                compos_reform[img_name] = {'bboxes': [[compo['column_min'] + shrink, compo['row_min'] + shrink, compo['column_max'] - shrink, compo['row_max'] - shrink]],\n",
    "                                           'categories': [compo['category']]}\n",
    "            else:\n",
    "                compos_reform[img_name]['bboxes'].append([compo['column_min'] + shrink, compo['row_min'] + shrink, compo['column_max'] - shrink, compo['row_max'] - shrink])\n",
    "                compos_reform[img_name]['categories'].append(compo['category'])\n",
    "    return compos_reform\n",
    "\n",
    "\n",
    "def load_ground_truth_json(gt_file):\n",
    "    def get_img_by_id(img_id):\n",
    "        for image in images:\n",
    "            if image['id'] == img_id:\n",
    "                return image['file_name'].split('/')[-1][:-4], (image['height'], image['width'])\n",
    "\n",
    "    def cvt_bbox(bbox):\n",
    "        '''\n",
    "        :param bbox: [x,y,width,height]\n",
    "        :return: [col_min, row_min, col_max, row_max]\n",
    "        '''\n",
    "        bbox = [int(b) for b in bbox]\n",
    "        return [bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]]\n",
    "\n",
    "    data = json.load(open(gt_file, 'r'))\n",
    "    images = data['images']\n",
    "    annots = data['annotations']\n",
    "    compos = {}\n",
    "    print('Loading %d ground truth' % len(annots))\n",
    "    for annot in tqdm(annots):\n",
    "        img_name, size = get_img_by_id(annot['image_id'])\n",
    "        if img_name not in compos:\n",
    "            compos[img_name] = {'bboxes': [cvt_bbox(annot['bbox'])], 'categories': [annot['category_id']], 'size': size}\n",
    "        else:\n",
    "            compos[img_name]['bboxes'].append(cvt_bbox(annot['bbox']))\n",
    "            compos[img_name]['categories'].append(annot['category_id'])\n",
    "    return compos\n",
    "\n",
    "\n",
    "def eval(detection, ground_truth, img_root, show=True, no_text=False, only_text=False):\n",
    "    def compo_filter(compos, flag):\n",
    "        if not no_text and not only_text:\n",
    "            return compos\n",
    "        compos_new = {'bboxes': [], 'categories': []}\n",
    "        for k, category in enumerate(compos['categories']):\n",
    "            if only_text:\n",
    "                if flag == 'det' and category != 'TextView':\n",
    "                    continue\n",
    "                if flag == 'gt' and int(category) != 14:\n",
    "                    continue\n",
    "            elif no_text:\n",
    "                if flag == 'det' and category == 'TextView':\n",
    "                    continue\n",
    "                if flag == 'gt' and int(category) == 14:\n",
    "                    continue\n",
    "\n",
    "            compos_new['bboxes'].append(compos['bboxes'][k])\n",
    "            compos_new['categories'].append(category)\n",
    "        return compos_new\n",
    "\n",
    "    def match(org, d_bbox, gt_bboxes, matched):\n",
    "        '''\n",
    "        :param matched: mark if the ground truth component is matched\n",
    "        :param d_bbox: [col_min, row_min, col_max, row_max]\n",
    "        :param gt_bboxes: list of ground truth [[col_min, row_min, col_max, row_max]]\n",
    "        :return: Boolean: if IOU large enough or detected box is contained by ground truth\n",
    "        '''\n",
    "        area_d = (d_bbox[2] - d_bbox[0]) * (d_bbox[3] - d_bbox[1])\n",
    "        for i, gt_bbox in enumerate(gt_bboxes):\n",
    "            if matched[i] == 0:\n",
    "                continue\n",
    "            area_gt = (gt_bbox[2] - gt_bbox[0]) * (gt_bbox[3] - gt_bbox[1])\n",
    "            col_min = max(d_bbox[0], gt_bbox[0])\n",
    "            row_min = max(d_bbox[1], gt_bbox[1])\n",
    "            col_max = min(d_bbox[2], gt_bbox[2])\n",
    "            row_max = min(d_bbox[3], gt_bbox[3])\n",
    "            # if not intersected, area intersection should be 0\n",
    "            w = max(0, col_max - col_min)\n",
    "            h = max(0, row_max - row_min)\n",
    "            area_inter = w * h\n",
    "            if area_inter == 0:\n",
    "                continue\n",
    "            iod = area_inter / area_d\n",
    "            iou = area_inter / (area_d + area_gt - area_inter)\n",
    "            # if show:\n",
    "            #     cv2.putText(org, (str(round(iou, 2)) + ',' + str(round(iod, 2))), (d_bbox[0], d_bbox[1]),\n",
    "            #                 cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)\n",
    "\n",
    "            if iou > 0.9 or iod == 1:\n",
    "                matched[i] = 0\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "    amount = len(detection)\n",
    "    TP, FP, FN = 0, 0, 0\n",
    "    pres, recalls, f1s = [], [], []\n",
    "    for i, image_id in enumerate(detection):\n",
    "        TP_this, FP_this, FN_this = 0, 0, 0\n",
    "        img = cv2.imread(pjoin(img_root, image_id + '.jpg'))\n",
    "        d_compos = detection[image_id]\n",
    "        gt_compos = ground_truth[image_id]\n",
    "\n",
    "        org_height = gt_compos['size'][0]\n",
    "\n",
    "        d_compos = compo_filter(d_compos, 'det')\n",
    "        gt_compos = compo_filter(gt_compos, 'gt')\n",
    "\n",
    "        d_compos['bboxes'] = resize_label(d_compos['bboxes'], 800, org_height)\n",
    "        matched = np.ones(len(gt_compos['bboxes']), dtype=int)\n",
    "        for d_bbox in d_compos['bboxes']:\n",
    "            if match(img, d_bbox, gt_compos['bboxes'], matched):\n",
    "                TP += 1\n",
    "                TP_this += 1\n",
    "            else:\n",
    "                FP += 1\n",
    "                FP_this += 1\n",
    "        FN += sum(matched)\n",
    "        FN_this = sum(matched)\n",
    "\n",
    "        try:\n",
    "            pre_this = TP_this / (TP_this + FP_this)\n",
    "            recall_this = TP_this / (TP_this + FN_this)\n",
    "            f1_this = 2 * (pre_this * recall_this) / (pre_this + recall_this)\n",
    "        except:\n",
    "            print('empty')\n",
    "            continue\n",
    "\n",
    "        pres.append(pre_this)\n",
    "        recalls.append(recall_this)\n",
    "        f1s.append(f1_this)\n",
    "        if show:\n",
    "            print(image_id + '.jpg')\n",
    "            print('[%d/%d] TP:%d, FP:%d, FN:%d, Precesion:%.3f, Recall:%.3f' % (\n",
    "                i, amount, TP_this, FP_this, FN_this, pre_this, recall_this))\n",
    "            cv2.imshow('org', cv2.resize(img, (500, 1000)))\n",
    "            broad = draw_bounding_box(img, d_compos['bboxes'], color=(255, 0, 0), line=3)\n",
    "            draw_bounding_box(broad, gt_compos['bboxes'], color=(0, 0, 255), show=True, line=2)\n",
    "\n",
    "        if i % 200 == 0:\n",
    "            precision = TP / (TP + FP)\n",
    "            recall = TP / (TP + FN)\n",
    "            f1 = 2 * (precision * recall) / (precision + recall)\n",
    "            print(\n",
    "                '[%d/%d] TP:%d, FP:%d, FN:%d, Precesion:%.3f, Recall:%.3f, F1:%.3f' % (i, amount, TP, FP, FN, precision, recall, f1))\n",
    "\n",
    "    precision = TP / (TP + FP)\n",
    "    recall = TP / (TP + FN)\n",
    "    print('[%d/%d] TP:%d, FP:%d, FN:%d, Precesion:%.3f, Recall:%.3f, F1:%.3f' % (i, amount, TP, FP, FN, precision, recall, f1))\n",
    "    # print(\"Average precision:%.4f; Average recall:%.3f\" % (sum(pres)/len(pres), sum(recalls)/len(recalls)))\n",
    "\n",
    "    return pres, recalls, f1s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "def draw_plot(data, title='Score for our approach'):\n",
    "    for i in range(len(data)):\n",
    "        data[i] = [d for d in data[i] if not math.isnan(d)]\n",
    "#     plt.title(title)\n",
    "    labels = ['Precision', 'Recall', 'F1']\n",
    "    bplot = plt.boxplot(data, patch_artist=True, labels=labels)  # 设置箱型图可填充\n",
    "    colors = ['pink', 'lightblue', 'lightgreen']\n",
    "    for patch, color in zip(bplot['boxes'], colors):\n",
    "        patch.set_facecolor(color) \n",
    "    plt.grid(axis='y')\n",
    "    plt.xticks(fontsize=16)\n",
    "    plt.yticks(fontsize=16)\n",
    "    plt.savefig(title + '.png')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "detect = load_detect_result_json('E:\\\\Mulong\\\\Result\\\\rico\\\\rico_uied\\\\rico_new_uied_cls\\\\merge')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gt = load_ground_truth_json('E:\\\\Mulong\\\\Datasets\\\\rico\\\\instances_test.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "no_text = False\n",
    "only_text = False\n",
    "pres_all, recalls_all, f1_all = eval(detect, gt, 'E:\\\\Mulong\\\\Datasets\\\\rico\\\\combined', show=False, no_text=no_text, only_text=only_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "no_text = True\n",
    "only_text = False\n",
    "pres_non_text, recalls_non_text, f1_non_text = eval(detect, gt, 'E:\\\\Mulong\\\\Datasets\\\\rico\\\\combined', show=False, no_text=no_text, only_text=only_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "no_text = False\n",
    "only_text = True\n",
    "pres_text, recalls_text, f1_text = eval(detect, gt, 'E:\\\\Mulong\\\\Datasets\\\\rico\\\\combined', show=False, no_text=no_text, only_text=only_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "draw_plot([pres_all, recalls_all])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "draw_plot([pres_non_text, recalls_non_text])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "\n",
    "pres1 = pd.DataFrame({'score_type':'Precision', 'score': pres_non_text, 'class':'Non_text'})\n",
    "pres2 = pd.DataFrame({'score_type':'Precision', 'score': pres_all, 'class':'All_element'})\n",
    "\n",
    "recalls1 = pd.DataFrame({'score_type':'Recall', 'score':recalls_non_text, 'class':'Non_text'})\n",
    "recalls2 = pd.DataFrame({'score_type':'Recall', 'score':recalls_all, 'class':'All_element'})\n",
    "\n",
    "f1s1 = pd.DataFrame({'score_type':'F1', 'score':f1_non_text, 'class':'Non_text'})\n",
    "f1s2 = pd.DataFrame({'score_type':'F1', 'score':f1_all, 'class':'All_element'})\n",
    "\n",
    "data=pd.concat([pres1, pres2, recalls1, recalls2, f1s1, f1s2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.boxenplot(x='score_type', y='score', hue='class', data=data, width=0.5, linewidth=1.0, palette=\"Set3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "draw_plot([pres_all, recalls_all, f1_all], title='Scores for All Elements')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "draw_plot([pres_non_text, recalls_non_text, f1_non_text], title='Score for Non-text Elements')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}