{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install --quiet --no-index --find-links=/kaggle/input/download-benetech-packages/offline-paddle/ segmentation_models_pytorch","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:19:48.220679Z","iopub.execute_input":"2023-06-20T01:19:48.221428Z","iopub.status.idle":"2023-06-20T01:20:05.655150Z","shell.execute_reply.started":"2023-06-20T01:19:48.221389Z","shell.execute_reply":"2023-06-20T01:20:05.653835Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install --quiet --no-index --find-links=/kaggle/input/download-benetech-packages/offline-paddle/ paddlepaddle-gpu","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:20:05.658915Z","iopub.execute_input":"2023-06-20T01:20:05.659335Z","iopub.status.idle":"2023-06-20T01:20:56.868324Z","shell.execute_reply.started":"2023-06-20T01:20:05.659292Z","shell.execute_reply":"2023-06-20T01:20:56.867090Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install --quiet --no-index --find-links=/kaggle/input/download-benetech-packages/offline-paddle/ paddleocr","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:20:56.870551Z","iopub.execute_input":"2023-06-20T01:20:56.870968Z","iopub.status.idle":"2023-06-20T01:21:28.971918Z","shell.execute_reply.started":"2023-06-20T01:20:56.870923Z","shell.execute_reply":"2023-06-20T01:21:28.970623Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install --quiet --no-index --find-links=/kaggle/input/download-benetech-packages/offline-paddle/ imutils","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:21:28.975073Z","iopub.execute_input":"2023-06-20T01:21:28.975399Z","iopub.status.idle":"2023-06-20T01:21:42.902180Z","shell.execute_reply.started":"2023-06-20T01:21:28.975367Z","shell.execute_reply":"2023-06-20T01:21:42.900540Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from paddleocr import PaddleOCR, draw_ocr","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:21:42.904250Z","iopub.execute_input":"2023-06-20T01:21:42.905801Z","iopub.status.idle":"2023-06-20T01:21:46.197770Z","shell.execute_reply.started":"2023-06-20T01:21:42.905754Z","shell.execute_reply":"2023-06-20T01:21:46.196587Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cp -r ../input/detectron-05 ./\n\n!pip install --quiet detectron-05/whls/pycocotools-2.0.2/dist/pycocotools-2.0.2.tar --no-index --find-links ../input/detectron-05/whls\n!pip install --quiet detectron-05/whls/fvcore-0.1.5.post20211019/fvcore-0.1.5.post20211019 --no-index --find-links ../input/detectron-05/whls \n!pip install --quiet detectron-05/whls/antlr4-python3-runtime-4.8/antlr4-python3-runtime-4.8 --no-index --find-links ../input/detectron-05/whls \n!pip install --quiet detectron-05/whls/detectron2-0.5/detectron2 --no-index --find-links ../input/detectron-05/whls \n\n!rm -r detectron-05","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:21:46.199295Z","iopub.execute_input":"2023-06-20T01:21:46.201997Z","iopub.status.idle":"2023-06-20T01:26:06.597640Z","shell.execute_reply.started":"2023-06-20T01:21:46.201951Z","shell.execute_reply":"2023-06-20T01:26:06.596159Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !python --version","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:06.602122Z","iopub.execute_input":"2023-06-20T01:26:06.602474Z","iopub.status.idle":"2023-06-20T01:26:06.607808Z","shell.execute_reply.started":"2023-06-20T01:26:06.602441Z","shell.execute_reply":"2023-06-20T01:26:06.606584Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport glob\nimport json\nimport re\nimport pandas as pd\nimport numpy as np\n\nimport torch\nfrom torch.optim import AdamW\n\nimport matplotlib.pyplot as plt\nfrom tqdm.auto import tqdm\n\nimport cv2\n\nimport albumentations as A \nfrom albumentations.pytorch.transforms import ToTensorV2\n\nfrom segmentation_models_pytorch import Unet, UnetPlusPlus\nfrom paddleocr import PaddleOCR,draw_ocr","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:06.609687Z","iopub.execute_input":"2023-06-20T01:26:06.610149Z","iopub.status.idle":"2023-06-20T01:26:11.519690Z","shell.execute_reply.started":"2023-06-20T01:26:06.610107Z","shell.execute_reply":"2023-06-20T01:26:11.518479Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CHART_MAP = {\n    'vertical_bar': 0,\n    'horizontal_bar': 1,\n    'dot': 2,\n    'scatter': 3,\n    'line': 4\n}\n\nCHART_MAP_REV = {v:k for k,v in CHART_MAP.items()}\n\nDEVICE = 'cuda:0'  \n# DEVICE = 'cpu'\n\nFONT_PATH = \"/kaggle/input/benetech-model0/simfang.ttf\"\n\nTEST_DIR = '/kaggle/input/benetech-making-graphs-accessible/test/images'\nTRAIN_DIR = '/kaggle/input/benetech-making-graphs-accessible/train/images'\nTRAIN_ANN_DIR = '/kaggle/input/benetech-making-graphs-accessible/train/annotations'\n\n\nPLOTBB_WEIGHT_PATHS = '/kaggle/input/benetech-plot-bb-v1'\nLINE_SEG_WEIGHT_PATH = '/kaggle/input/benetech-line-seg-v2/line_seg_v2_exp_mask.pth'\n\nTEXT_REC_MODEL_DIR = \"/kaggle/input/benetech-model0/en_PP-OCRv3_rec_infer/root/.paddleocr/whl/rec/en/en_PP-OCRv3_rec_infer/\"\nTEXT_DET_MODEL_DIR = \"/kaggle/input/benetech-ocr-v2/infer_model\"\nTEXT_CLS_MODEL_DIR = \"/kaggle/input/benetech-model0/ch_ppocr_mobile_v2.0_cls_infer/root/.paddleocr/whl/cls/ch_ppocr_mobile_v2.0_cls_infer\"\n\nINST_SEG_WEIGHT_PATH_1 = '/kaggle/input/benetech-inst-seg-v1/inst_seg_v1_0011999.pth'\nINST_SEG_WEIGHT_PATH_2 = '/kaggle/input/benetech-inst-seg-v2/inst_seg_v2_0006499.pth'","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:11.521832Z","iopub.execute_input":"2023-06-20T01:26:11.522283Z","iopub.status.idle":"2023-06-20T01:26:11.531493Z","shell.execute_reply.started":"2023-06-20T01:26:11.522240Z","shell.execute_reply":"2023-06-20T01:26:11.530190Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/benetech-metric')\nfrom metric import benetech_score","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:11.537844Z","iopub.execute_input":"2023-06-20T01:26:11.538261Z","iopub.status.idle":"2023-06-20T01:26:11.621274Z","shell.execute_reply.started":"2023-06-20T01:26:11.538220Z","shell.execute_reply":"2023-06-20T01:26:11.620147Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class PLOTBB_CFG:\n    im_size = 224\n    backbone=\"tu-eca_nfnet_l1\"\n    pretrained_weights = None\n\n\nclass LINE_SEG_CFG:\n    im_size = 512\n    backbone=\"timm-efficientnet-b5\"\n    pretrained_weights = None\n    \n\nPLOTBB_CFG.val_transforms = A.Compose(\n        [\n            A.Resize(height=PLOTBB_CFG.im_size, width=PLOTBB_CFG.im_size, always_apply=True),\n            A.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n            ToTensorV2(always_apply=True),\n        ],\n        p=1.0,\n       \n    )\n\nLINE_SEG_CFG.val_transforms = A.Compose(\n        [\n            A.Resize(height=LINE_SEG_CFG.im_size, width=LINE_SEG_CFG.im_size, always_apply=True),\n            A.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n            ToTensorV2(always_apply=True),\n        ],\n        p=1.0,\n       \n    )\n","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:11.623645Z","iopub.execute_input":"2023-06-20T01:26:11.624771Z","iopub.status.idle":"2023-06-20T01:26:11.635532Z","shell.execute_reply.started":"2023-06-20T01:26:11.624711Z","shell.execute_reply":"2023-06-20T01:26:11.634221Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def modify_key_name(state_dict):\n    new_state_dict = {}\n    for key, value in state_dict.items():\n        new_key = key.replace('model.', '')\n        new_state_dict[new_key] = value\n    return new_state_dict","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:11.637341Z","iopub.execute_input":"2023-06-20T01:26:11.637869Z","iopub.status.idle":"2023-06-20T01:26:11.650969Z","shell.execute_reply.started":"2023-06-20T01:26:11.637827Z","shell.execute_reply":"2023-06-20T01:26:11.649775Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tools.infer.predict_system import TextSystem\n\ndef __call_unsqueezed_quad__(self, img, cls=True):\n    # print('Herre')\n    time_dict = {'det': 0, 'rec': 0, 'csl': 0, 'all': 0}\n    start = time.time()\n    ori_im = img.copy()\n    dt_boxes, elapse = self.text_detector(img)\n    time_dict['det'] = elapse\n    logger.debug(\"dt_boxes num : {}, elapse : {}\".format(\n        len(dt_boxes), elapse))\n    if dt_boxes is None:\n        return None, None\n    img_crop_list = []\n\n    dt_boxes = sorted_boxes(dt_boxes)\n\n#     from shapely import Polygon\n    from shapely.geometry.polygon import Polygon\n    def approximate_by_rect(sbox, thresh=0.8):\n        x0,y0,x1,y1 = sbox.bounds\n        sbox2 = Polygon(np.array([x0,y0,x0,y1,x1,y1,x1,y0]).reshape(-1,2))\n        if sbox.intersection(sbox2).area / sbox2.area > thresh:\n            return sbox2\n        return sbox\n\n    for k in range(len(dt_boxes)):\n        box = dt_boxes[k]\n        sbox = Polygon(box)\n        sbox = sbox.convex_hull\n        box = np.array(sbox.exterior.coords.__array_interface__['data']).reshape(-1, 2)\n        if len(box) < 4: # Convex hull only has 3 points\n            sbox = Polygon(box)\n            sbox = sbox.minimum_rotated_rectangle\n            sbox = sbox.convex_hull\n        sbox = approximate_by_rect(sbox)\n        box = np.array(sbox.exterior.coords.__array_interface__['data']).reshape(-1, 2).astype('float32')\n        box = np.array([box[0], box[3], box[2], box[1]])\n        dt_boxes[k] = box\n\n\n    import matplotlib.pyplot as plt\n\n    for bno in range(len(dt_boxes)):\n        tmp_box = copy.deepcopy(dt_boxes[bno])\n        if self.args.det_box_type == \"quad\":\n            img_crop = get_rotate_crop_image(ori_im, tmp_box)\n        else:\n            img_crop = get_minarea_rect_crop(ori_im, tmp_box)\n        img_crop_list.append(img_crop)\n        \n    if self.use_angle_cls and cls:\n        img_crop_list, angle_list, elapse = self.text_classifier(\n            img_crop_list)\n        time_dict['cls'] = elapse\n        logger.debug(\"cls num  : {}, elapse : {}\".format(\n            len(img_crop_list), elapse))\n\n    rec_res, elapse = self.text_recognizer(img_crop_list)\n    time_dict['rec'] = elapse\n    logger.debug(\"rec_res num  : {}, elapse : {}\".format(\n        len(rec_res), elapse))\n    if self.args.save_crop_res:\n        self.draw_crop_rec_res(self.args.crop_res_save_dir, img_crop_list,\n                                rec_res)\n    filter_boxes, filter_rec_res = [], []\n    for box, rec_result in zip(dt_boxes, rec_res):\n        text, score = rec_result\n        if score >= 0.0:\n            filter_boxes.append(box)\n            filter_rec_res.append(rec_result)\n    end = time.time()\n    time_dict['all'] = end - start\n    return filter_boxes, filter_rec_res, time_dict\n\nTextSystem.__call__.__code__ = __call_unsqueezed_quad__.__code__","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:11.652644Z","iopub.execute_input":"2023-06-20T01:26:11.653921Z","iopub.status.idle":"2023-06-20T01:26:11.676959Z","shell.execute_reply.started":"2023-06-20T01:26:11.653880Z","shell.execute_reply":"2023-06-20T01:26:11.675880Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Paddle OCR overridden","metadata":{}},{"cell_type":"code","source":"import imutils\nfrom PIL import Image\nfrom shapely.geometry.polygon import Polygon\n\nclass TextBoxes:\n    def __init__(self, boxes, texts, scores):\n        self.boxes = boxes\n        self.texts = texts\n        self.scores = scores\n\ndef nms(boxes, texts, scores, thresh=0.2):\n    orders = np.argsort(scores)[::-1]\n    keeps = []\n    \n    while len(orders) > 0:\n        # append the idx of the instance X with highest score to keeps list\n        current_ind = orders[0]\n        keeps.append(current_ind)\n        current = Polygon(boxes[current_ind].reshape(-1,2))\n        \n        # remove the idx of the top instance from the remaining list\n        orders = orders[1:]\n        \n        inds_overlap = []\n        for i, ind in enumerate(orders):\n            check = Polygon(boxes[ind].reshape(-1,2))\n            if current.intersection(check).area /  min(current.area, check.area) > thresh:\n                inds_overlap.append(i)\n        \n        orders = [ind for i, ind in enumerate(orders) if i not in inds_overlap]\n        \n    return boxes[keeps], texts[keeps], scores[keeps]\n\n\nfrom tools.infer.predict_system import TextSystem\n\ndef __call_unsqueezed_quad__(self, img, cls=True):\n    # print('Herre')\n    time_dict = {'det': 0, 'rec': 0, 'csl': 0, 'all': 0}\n    start = time.time()\n    ori_im = img.copy()\n    dt_boxes, elapse = self.text_detector(img)\n    time_dict['det'] = elapse\n    logger.debug(\"dt_boxes num : {}, elapse : {}\".format(\n        len(dt_boxes), elapse))\n    if dt_boxes is None:\n        return None, None\n    img_crop_list = []\n    img_crop_list_90 = []\n    img_crop_list_180 = []\n\n    dt_boxes = sorted_boxes(dt_boxes)\n\n#     from shapely import Polygon\n    from shapely.geometry.polygon import Polygon\n    def approximate_by_rect(sbox, thresh=0.8):\n        x0,y0,x1,y1 = sbox.bounds\n        sbox2 = Polygon(np.array([x0,y0,x0,y1,x1,y1,x1,y0]).reshape(-1,2))\n        if sbox.intersection(sbox2).area / sbox2.area > thresh:\n            return sbox2\n        return sbox\n\n    for k in range(len(dt_boxes)):\n        box = dt_boxes[k]\n        sbox = Polygon(box)\n        sbox = sbox.convex_hull\n        box = np.array(sbox.exterior.coords.__array_interface__['data']).reshape(-1, 2)\n        if len(box) < 4: # Convex hull only has 3 points\n            sbox = Polygon(box)\n            sbox = sbox.minimum_rotated_rectangle\n            sbox = sbox.convex_hull\n        sbox = approximate_by_rect(sbox)\n        box = np.array(sbox.exterior.coords.__array_interface__['data']).reshape(-1, 2).astype('float32')\n        box = np.array([box[0], box[3], box[2], box[1]])\n        dt_boxes[k] = box\n\n    import matplotlib.pyplot as plt\n    \n    def get_rotate_crop_image(img, points):\n        '''\n        img_height, img_width = img.shape[0:2]\n        left = int(np.min(points[:, 0]))\n        right = int(np.max(points[:, 0]))\n        top = int(np.min(points[:, 1]))\n        bottom = int(np.max(points[:, 1]))\n        img_crop = img[top:bottom, left:right, :].copy()\n        points[:, 0] = points[:, 0] - left\n        points[:, 1] = points[:, 1] - top\n        '''\n        assert len(points) == 4, \"shape of points must be 4*2\"\n        img_crop_width = int(\n            max(\n                np.linalg.norm(points[0] - points[1]),\n                np.linalg.norm(points[2] - points[3])))\n        img_crop_height = int(\n            max(\n                np.linalg.norm(points[0] - points[3]),\n                np.linalg.norm(points[1] - points[2])))\n        pts_std = np.float32([[0, 0], [img_crop_width, 0],\n                              [img_crop_width, img_crop_height],\n                              [0, img_crop_height]])\n        M = cv2.getPerspectiveTransform(points, pts_std)\n        dst_img = cv2.warpPerspective(\n            img,\n            M, (img_crop_width, img_crop_height),\n            borderMode=cv2.BORDER_REPLICATE,\n            flags=cv2.INTER_CUBIC)\n        dst_img_height, dst_img_width = dst_img.shape[0:2]\n\n        return dst_img\n    \n    aspect_ratios = []\n    for bno in range(len(dt_boxes)):\n        tmp_box = copy.deepcopy(dt_boxes[bno])\n        if self.args.det_box_type == \"quad\":\n            img_crop = get_rotate_crop_image(ori_im, tmp_box)\n\n        else:\n            img_crop = get_minarea_rect_crop(ori_im, tmp_box)\n        h, w = img_crop.shape[:2]\n        aspect_ratios.append(w / h)\n        img_crop_list.append(img_crop)\n        img_crop_list_90.append(cv2.rotate(img_crop, cv2.ROTATE_90_CLOCKWISE))\n        \n    rec_res_0, elapse = self.text_recognizer(img_crop_list)\n    rec_res_90, elapse = self.text_recognizer(img_crop_list_90)\n    \n    conf_scores = [\n        [it[1] for it in rec_res_0],\n        [it[1] for it in rec_res_90],\n    ]\n    conf_scores = np.array(conf_scores)\n    indices = np.argmax(conf_scores, axis=0)\n    \n    texts = [\n        [it[0] for it in rec_res_0],\n        [it[0] for it in rec_res_90],\n    ]\n    texts = np.array(texts)\n    \n        \n    final_texts = [texts[k, i] for i, k in enumerate(indices)]\n    final_conf_scores = np.max(conf_scores, axis=0)\n    rec_res = list(zip(final_texts, final_conf_scores))\n\n    time_dict['rec'] = elapse\n    logger.debug(\"rec_res num  : {}, elapse : {}\".format(\n        len(rec_res), elapse))\n    if self.args.save_crop_res:\n        self.draw_crop_rec_res(self.args.crop_res_save_dir, img_crop_list,\n                                rec_res)\n    filter_boxes, filter_rec_res = [], []\n    for box, rec_result in zip(dt_boxes, rec_res):\n        text, score = rec_result\n        if score >= 0.0:\n            filter_boxes.append(box)\n            filter_rec_res.append(rec_result)\n    end = time.time()\n    time_dict['all'] = end - start\n    return filter_boxes, filter_rec_res, time_dict\n\nTextSystem.__call__.__code__ = __call_unsqueezed_quad__.__code__\n\nclass OCR:\n    def __init__(self, font_path=FONT_PATH,\n                 det_model_dir=TEXT_DET_MODEL_DIR, rec_model_dir=TEXT_REC_MODEL_DIR,\n                 cls_model_dir=TEXT_CLS_MODEL_DIR):\n        self.ocr = PaddleOCR(use_angle_cls=True, lang='en', det_algorithm='DB',\n                             det_model_dir=det_model_dir, rec_model_dir=rec_model_dir,\n                             cls_model_dir=cls_model_dir,\n                            det_db_thresh=0.001, det_db_box_thresh=0.001, show_log = False,\n                )\n        self.font_path = font_path\n    \n    def __call__(self, img, resized_width=1000, return_im_show=False):\n        inp = imutils.resize(img, width=resized_width)\n\n        scale = resized_width / img.shape[1]\n\n        result = self.ocr.ocr(inp, cls=True)\n        for idx in range(len(result)):\n            res = result[idx]\n                \n        result = result[0]\n\n        boxes = np.array([np.array(line[0]).reshape(-1,) / scale for line in result])\n        txts = np.array([line[1][0] for line in result])\n        scores = np.array([line[1][1] for line in result])\n        \n        # filter overlapping text instances\n        boxes, txts, scores = nms(boxes, txts, scores)\n        \n        draw = inp.copy()\n        for box in result:\n            box = box[0]\n\n            pts = np.array(box).astype(int)\n            cv2.polylines(draw, [pts], 1, (0,255,0), 3)\n\n            for i, v in enumerate(pts):\n                cv2.putText(draw, str(i), (v[0], v[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)\n\n\n        im_show = None\n        if return_im_show:\n            im_show = draw_ocr(img, boxes, txts, scores, font_path=self.font_path, drop_score=0.0)\n        \n        text_boxes = TextBoxes(boxes, txts, scores)\n        return text_boxes, draw","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:11.678811Z","iopub.execute_input":"2023-06-20T01:26:11.679206Z","iopub.status.idle":"2023-06-20T01:26:11.735324Z","shell.execute_reply.started":"2023-06-20T01:26:11.679165Z","shell.execute_reply":"2023-06-20T01:26:11.734091Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Intance segmentation model","metadata":{}},{"cell_type":"code","source":"import detectron2\nfrom pathlib import Path\nimport random, cv2, os\nimport matplotlib.pyplot as plt\n# import some common detectron2 utilities\nfrom detectron2 import model_zoo\nfrom detectron2.engine import DefaultPredictor, DefaultTrainer\nfrom detectron2.config import get_cfg, CfgNode\nfrom detectron2.utils.visualizer import Visualizer, ColorMode\nfrom detectron2.data import MetadataCatalog, DatasetCatalog\nfrom detectron2.data import build_detection_train_loader, build_detection_test_loader, DatasetMapper\nfrom detectron2.data.datasets import register_coco_instances","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:11.737083Z","iopub.execute_input":"2023-06-20T01:26:11.737588Z","iopub.status.idle":"2023-06-20T01:26:11.980690Z","shell.execute_reply.started":"2023-06-20T01:26:11.737545Z","shell.execute_reply":"2023-06-20T01:26:11.979606Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cfg = get_cfg()\ncfg.merge_from_file(model_zoo.get_config_file(\"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\"))\ncfg.DATASETS.TRAIN = (\"train_set\",)\ncfg.DATASETS.TEST = (\"train_set\", \"val_set\")\ncfg.DATALOADER.NUM_WORKERS = 2\ncfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 256   \ncfg.MODEL.ROI_HEADS.NUM_CLASSES = 2\n\ncfg.INPUT.MAX_SIZE_TEST = 2000\ncfg.INPUT.MAX_SIZE_TRAIN = 2000\ncfg.INPUT.MIN_SIZE_TEST = 768\ncfg.INPUT.MIN_SIZE_TRAIN = (768, )\n\ncfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.6  # set threshold for this model\ncfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = 0.9\ncfg.MODEL.DEVICE = DEVICE\n\nINST_SEG_CONF = {'vertical_bar': 0.7, 'horizontal_bar': 0.9, 'scatter': 0.6, 'dot': 0.9}\nINST_SEG_NMS = {'vertical_bar': 0.1, 'horizontal_bar': 0.1, 'scatter': 0.5, 'dot': 0.1}\n\ncfg.MODEL.WEIGHTS = INST_SEG_WEIGHT_PATH_1\ncfg.INPUT.MIN_SIZE_TEST = 1024\ncfg.INPUT.MIN_SIZE_TRAIN = (1024, )\ninst_seg_model1 = DefaultPredictor(cfg)\n\ncfg.MODEL.WEIGHTS = INST_SEG_WEIGHT_PATH_2\ncfg.INPUT.MIN_SIZE_TEST = 768\ncfg.INPUT.MIN_SIZE_TRAIN = (768, )\ninst_seg_model2 = DefaultPredictor(cfg)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:11.982310Z","iopub.execute_input":"2023-06-20T01:26:11.982705Z","iopub.status.idle":"2023-06-20T01:26:24.086128Z","shell.execute_reply.started":"2023-06-20T01:26:11.982664Z","shell.execute_reply":"2023-06-20T01:26:24.084846Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def predict_inst_seg(img, chart_type):\n    if chart_type == 'horizontal_bar':\n        return inst_seg_model1(img)\n    return inst_seg_model2(img)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.087498Z","iopub.execute_input":"2023-06-20T01:26:24.087888Z","iopub.status.idle":"2023-06-20T01:26:24.093358Z","shell.execute_reply.started":"2023-06-20T01:26:24.087850Z","shell.execute_reply":"2023-06-20T01:26:24.092167Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pycocotools.mask as mask_util\n\ndef inst_seg_post_process(pred, conf_thresh=0.5, nms_thresh=0.5):\n    pred_masks = pred['instances'].pred_masks.detach().cpu().numpy()\n    scores = pred['instances'].scores.detach().cpu().numpy()\n    \n    keep = scores>=conf_thresh\n    pred_masks = pred_masks[keep]\n    scores = scores[keep]\n    \n    if(pred_masks.dtype == np.float32):\n        pred_masks = (pred_masks >= 0.5 ).astype('bool')\n        # print(pred_masks.dtype)\n\n    # calculate IOU\n    enc_preds = [mask_util.encode(np.asarray(p, order='F')) for p in pred_masks]\n    ious = mask_util.iou(enc_preds, enc_preds, [0]*len(enc_preds))\n\n    orders = np.arange(len(scores))\n    keeps = []\n\n    while len(orders) > 0:\n\n        # append the idx of the instance X with highest score to keeps list\n        keeps.append(orders[0])\n        # remove the idx of the above instance from the remaining list\n        orders = orders[1:]\n\n        # check IOU of the instance X with all the instance in the remaining list\n        look_up = ious[keeps[-1], orders]\n\n        # filter those having IOU > nms_thresh\n        mask = look_up < nms_thresh\n        orders = orders[mask]\n\n        if(len(orders)==0):\n            break\n    \n    new_pred = pred.copy()\n    new_pred['instances'] = new_pred['instances'][keeps]\n    return new_pred\n","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.095202Z","iopub.execute_input":"2023-06-20T01:26:24.096145Z","iopub.status.idle":"2023-06-20T01:26:24.151537Z","shell.execute_reply.started":"2023-06-20T01:26:24.096088Z","shell.execute_reply":"2023-06-20T01:26:24.150304Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_chart_data_mask(img, model, transforms, device=DEVICE):\n    X = transforms(image=img)['image']\n    X = X.to(device)\n\n    with torch.no_grad():\n        mask_pred = model(X.unsqueeze(0))\n        mask_pred = mask_pred.cpu().numpy()[0][0]\n    mask_pred = cv2.resize(mask_pred, (img.shape[1], img.shape[0]))\n    mask_pred = np.where(mask_pred > 0.5, 255, 0).astype('uint8')\n    return mask_pred","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.153986Z","iopub.execute_input":"2023-06-20T01:26:24.154944Z","iopub.status.idle":"2023-06-20T01:26:24.164671Z","shell.execute_reply.started":"2023-06-20T01:26:24.154901Z","shell.execute_reply":"2023-06-20T01:26:24.163587Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Axes association","metadata":{}},{"cell_type":"code","source":"class PlotAxes:\n    def __init__(self, X, Y, text_boxes, chart_type):\n        self.text_boxes = text_boxes\n        self.X = X\n        self.Y = Y\n        self.chart_type = chart_type\n\n    def is_rotated(self, box):\n        x1,y1, x2,y2, x3,y3, x4,y4 = box\n        if np.abs(y2 - y1) > 10:\n            return True\n        return False\n    \n    def filter_by_area(self, boxes):\n        all_areas = []\n        for x1,y1,x2,y2,x3,y3,x4,y4 in boxes:\n            area = Polygon(np.array([x1,y1,x2,y2,x3,y3,x4,y4]).reshape(-1,2)).area\n            all_areas.append(area)\n        return detect_outliers(all_areas, threshold=4 if self.chart_type=='line' else 3)\n\n    def filter_by_position(self, all_values):\n        if len(all_values) == 0:\n            return []\n        \n        q1 = np.percentile(all_values, 25)\n        q3 = np.percentile(all_values, 75)\n        iqr = q3 - q1\n        lb = q1 - 1.5*iqr\n        ub = q3 + 1.5*iqr\n        \n        inliers_by_stats = np.where((all_values >= lb) & (all_values <= ub))[0]\n        \n        last = all_values[inliers_by_stats][0]\n        k = 0\n        outlier_inds = []\n        for val in all_values:\n            # print(val, last)\n            if np.abs(val - last) > 10: # outlier by checking with the previous\n                outlier_inds.append(k)\n            # else:\n            #     last = val\n            k+=1\n        return outlier_inds\n  \n\n    # point1 ----- point2\n    #   |            |\n    #   |            |\n    #   |            |\n    #   |            |\n    # point4 ----- point3 \n    def get_x_ticks(self):\n        candidate_boxes = []\n        candidate_texts = []\n        candidate_min_ys = []\n        candidate_min_xs = []\n        \n        for box, text, score in zip(self.text_boxes.boxes, self.text_boxes.texts, self.text_boxes.scores):\n            x1,y1, x2,y2, x3,y3, x4,y4 = box\n            if min(y1,y2,y3,y4) + 5 > self.Y and max(x1,x2,x3,x4) + 10 > self.X:\n                if contains_only_whitespace(text) and self.chart_type!='vertical_bar': # textboxes under the x axis and on the right side of y axis\n                    continue\n                candidate_boxes.append(box)\n                candidate_texts.append(text)\n                candidate_min_ys.append(min(y1,y2,y3,y4))\n                candidate_min_xs.append(min(x1,x2,x3,x4))\n\n#         print(candidate_texts)\n        outlier_inds_by_area = self.filter_by_area(candidate_boxes)\n        \n#         print(outlier_inds_by_area, np.array(candidate_texts)[outlier_inds_by_area])\n        \n        sort_inds = np.argsort(candidate_min_xs)\n        sort_inds = [i for i in sort_inds if i not in outlier_inds_by_area]\n        candidate_boxes = np.array(candidate_boxes)[sort_inds]\n        candidate_texts = np.array(candidate_texts)[sort_inds]\n        candidate_min_ys = np.array(candidate_min_ys)[sort_inds]\n        candidate_min_xs = np.array(candidate_min_xs)[sort_inds]\n\n        x_label_inds = self.filter_by_position(candidate_min_ys)\n\n        # print(x_label_inds)\n\n        x_ticks = []\n\n        for ind, (box, text) in enumerate(zip(candidate_boxes, candidate_texts)):\n            if ind in x_label_inds:\n                continue\n            \n            x1,y1, x2,y2, x3,y3, x4,y4 = box\n            s_ind = np.argsort([y1,y2,y3,y4])\n            box = np.array(box).reshape(-1,2)[s_ind]\n            coord = ((box[0,0] + box[1,0])/2, self.Y)\n            x_ticks.append((coord[0], coord[1], text))\n\n        x_ticks = sorted(x_ticks, key=lambda x: x[0])\n\n        return x_ticks \n    \n    def get_y_ticks(self):\n        candidate_boxes = []\n        candidate_texts = []\n        candidate_max_xs = []\n        candidate_min_ys = []\n\n        for box, text, score in zip(self.text_boxes.boxes, self.text_boxes.texts, self.text_boxes.scores):\n            x1,y1, x2,y2, x3,y3, x4,y4 = box\n            if x2 < self.X + 5: # textboxes on the left of the y-axis\n                if contains_only_whitespace(text) and self.chart_type!='vertical_bar': # textboxes under the x axis and on the right side of y axis\n                    continue\n                candidate_boxes.append(box)\n                candidate_texts.append(text)\n                candidate_max_xs.append(max(x1,x2,x3,x4))\n                candidate_min_ys.append(min(y1,y2,y3,y4))\n                \n        outlier_inds_by_area = self.filter_by_area(candidate_boxes)\n        \n        sort_inds = np.argsort(candidate_min_ys)\n        sort_inds = [i for i in sort_inds if i not in outlier_inds_by_area]\n        \n        candidate_boxes = np.array(candidate_boxes)[sort_inds]\n        candidate_texts = np.array(candidate_texts)[sort_inds]\n        candidate_min_ys = np.array(candidate_min_ys)[sort_inds]\n        candidate_max_xs = np.array(candidate_max_xs)[sort_inds]\n\n        # print(candidate_texts)\n\n        # filter y-label label\n        y_label_inds = self.filter_by_position(candidate_max_xs)\n\n        y_ticks = []\n\n        for ind, (box, text) in enumerate(zip(candidate_boxes, candidate_texts)):\n            if ind in y_label_inds:\n                continue\n            x1,y1, x2,y2, x3,y3, x4,y4 = box\n            # the coordinate of the this tick is the projection of the mean \n            # of the midpoint of the point2 and point3 onto the y-axis\n            coord = (self.X, (y2+y3)/2)\n            y_ticks.append((coord[0], coord[1], text))\n\n        y_ticks = sorted(y_ticks, key=lambda x: x[1])\n\n        return y_ticks\n\n    def draw_ticks(img, ticks):\n        draw = img.copy()\n        for x, y, text in ticks:\n            x, y = int(x), int(y)\n            draw = cv2.circle(draw, (x,y), 3, (255,0,0), -1)\n        return draw","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.166538Z","iopub.execute_input":"2023-06-20T01:26:24.167098Z","iopub.status.idle":"2023-06-20T01:26:24.205158Z","shell.execute_reply.started":"2023-06-20T01:26:24.167057Z","shell.execute_reply":"2023-06-20T01:26:24.204030Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def is_numeric_string(s):\n    try:\n        float(s)\n        return True\n    except ValueError:\n        return False\n\ndef is_numeric_array(arr):\n    for element in arr:\n        if element != 'None':\n            if not is_numeric_string(element):\n                return False\n    return True\n\ndef cast_to_numeric_if_possible(arr):\n    new_arr = []\n    if is_numeric_array(arr):\n        for e in arr:\n            if e == 'None':\n                new_arr.append(0)\n            else:\n                new_arr.append(float(e))\n    else:\n          new_arr = arr\n    return new_arr\n\ndef detect_outliers(values, threshold=3):\n    # Calculate the mean and standard deviation of the values\n    mean = np.mean(values)\n    std = np.std(values)\n\n    # Define the lower and upper bounds for outliers\n    lower_bound = mean - threshold * std\n    upper_bound = mean + threshold * std\n\n    # Find the indices of the outliers\n    outliers = []\n    for i, value in enumerate(values):\n        if value < lower_bound or value > upper_bound:\n            outliers.append(i)\n\n    return outliers\n\n    \ndef filter_tick_outliers(numeric_ticks):\n\n    values = np.array([value[-1] for value in numeric_ticks])\n    diff = values[:-1] - values[1:]\n    \n    outlier_diff_inds = np.array( detect_outliers(diff, threshold=1) )\n    \n    outlier_inds = []\n    for i in outlier_diff_inds:\n        if len(outlier_inds) > 0:\n            if outlier_inds[-1] == i:\n                continue\n        if i == 0:\n            outlier_inds.append(i)\n        else:\n            outlier_inds.append(i+1)\n    new_numeric_ticks = [it for i, it in enumerate(numeric_ticks) if i not in outlier_inds]\n    return new_numeric_ticks\n\ndef numeric_clean(x):\n    clean = x.lower().strip()\n    \n    if re.match('\\d+.\\d{3}$', clean):\n        clean = clean.replace('.', '')\n        \n    if re.match('\\d+,\\d{1,2}$', clean):\n        clean = clean.replace(',', '.')\n        \n    clean = re.sub('[Ss]', '5', clean)\n    \n    if clean.endswith('k') :\n        clean = clean.replace('k', '000')\n    if clean.endswith('m'):\n        clean = clean.replace('m', '000000')\n    clean = re.sub( r\"[^\\d\\-.]\", \"\", clean)\n    return clean\n\ndef contains_only_whitespace(input_string):\n    # Use a regular expression to match whitespace characters\n    pattern = r'^\\s*$'\n    match = re.match(pattern, input_string)\n    \n    # If there is a match, the string contains only whitespace\n    if match:\n        return True\n    else:\n        return False","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.207007Z","iopub.execute_input":"2023-06-20T01:26:24.207633Z","iopub.status.idle":"2023-06-20T01:26:24.228051Z","shell.execute_reply.started":"2023-06-20T01:26:24.207583Z","shell.execute_reply":"2023-06-20T01:26:24.226924Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Chart readers by type","metadata":{}},{"cell_type":"code","source":"class DotReader:\n    def is_numerical(x):\n        \"\"\"Test whether input can be parsed as a Python float.\"\"\"\n        try:\n            float(x)\n            return True\n        except ValueError:\n            return False\n\n    def read(xticks, yticks, X, Y, img, instances, debug=False):\n        bounding_boxes = instances['instances'].pred_boxes.tensor.cpu().numpy()\n\n        bounding_boxes = np.array( sorted(bounding_boxes, key=lambda x: x[0]) )\n\n        # extract the x-coordinates of the bounding boxes\n        x_coords = bounding_boxes[:, 0]\n\n        # cluster the bounding boxes based on their x-coordinates using DBSCAN\n        eps = 10\n        \n        curr_group = 0\n        groups = [curr_group]\n        last = x_coords[0]\n        for x in x_coords[1:]:\n            if np.abs(x-last) > 10:\n                curr_group += 1\n            # print(x, last, curr_group)\n            groups.append(curr_group)\n            last = x\n\n        groups = np.array(groups)\n\n        data_xs = []\n        data_ys = []\n        for g in range(0, max(groups)+1):\n            # print('Group:', g)\n            inds = np.where(groups==g)\n            # print(inds)\n            bin_boxes = bounding_boxes[inds]\n\n            if debug:\n                draw = img.copy()\n                plt.figure()\n                for x1,y1,x2,y2 in bin_boxes:\n                    cv2.rectangle(draw, (x1,y1), (x2,y2), (255,0,0), 2)\n                plt.imshow(draw)\n\n            # find the corresponding x-label\n            mean_x = np.mean( (bin_boxes[:,0] + bin_boxes[:,2])/2 )\n            chosen_tick_text = 'None'\n            for tick_coord_x, tick_coord_y, text in xticks:\n                if np.abs(tick_coord_x - mean_x) < 10:\n                    chosen_tick_text = text\n                    break\n            \n            dot_count = len(bin_boxes)\n\n            data_xs.append(chosen_tick_text)\n            data_ys.append(dot_count)\n\n        data_xs = cast_to_numeric_if_possible(data_xs)\n\n        return data_xs, data_ys","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.229763Z","iopub.execute_input":"2023-06-20T01:26:24.230319Z","iopub.status.idle":"2023-06-20T01:26:24.246961Z","shell.execute_reply.started":"2023-06-20T01:26:24.230275Z","shell.execute_reply":"2023-06-20T01:26:24.245766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_chart_y_value(pixel_yval, pixel_min, pixel_max, tick_min, tick_max ):\n    denominator = 1 if (pixel_min - pixel_max) == 0 else (pixel_min - pixel_max)\n    if pixel_min >= pixel_yval:\n        return (pixel_min - pixel_yval) / denominator * (tick_max - tick_min) + tick_min\n    else:\n        return (pixel_yval - pixel_min) / denominator * (tick_max - tick_min) + tick_min\n    \n\n\nclass VBarReader:\n    \n    def read(xticks, yticks, X, Y, img, instances, debug=False):\n\n        bounding_boxes = instances['instances'].pred_boxes.tensor.cpu().numpy()\n\n        bounding_boxes = np.array( sorted(bounding_boxes, key=lambda x: x[0]) )\n        \n        outlier_indices = detect_outliers([box[3] for box in bounding_boxes], threshold=3) # filter outliers by y2\n\n        bounding_boxes = [box for i, box in enumerate(bounding_boxes) if i not in outlier_indices]\n        \n        numeric_yticks = []\n        for tick in yticks:\n            clean_tick =  numeric_clean(tick[2])\n            if is_numeric_string(clean_tick):\n                numeric_yticks.append((tick[0], tick[1], float(clean_tick)))\n                \n        numeric_yticks = filter_tick_outliers(numeric_yticks)\n        \n#         print(numeric_yticks)\n                \n        if len(numeric_yticks) > 0:\n            min_y_ticks = min(numeric_yticks, key=lambda x: x[-1])\n            max_y_ticks = max(numeric_yticks, key=lambda x: x[-1])\n        else:\n            max_y_ticks = (0,0,0)\n            min_y_ticks = (0,0,0)\n            \n        draw = img.copy()\n\n        data_xs = []\n        data_ys = []\n        assinged_tick_ids = []\n        for i, box in enumerate(bounding_boxes):\n            x1,y1,x2,y2 = box.astype(int)\n            if debug:\n                # Draw text on image\n                cv2.putText(draw, str(i), (x1,y1), cv2.FONT_HERSHEY_SIMPLEX, 1, 255, 2)\n                cv2.rectangle(draw, (x1,y1), (x2,y2), (255,0,0), 2)\n            \n            bar_len = get_chart_y_value(y1, min_y_ticks[1], max_y_ticks[1], min_y_ticks[2], max_y_ticks[2])\n\n            # find the corresponding x-label\n            mean_x = (x1 + x2) / 2\n            chosen_tick_text = 'None'\n            for text_id, (tick_coord_x, tick_coord_y, text) in enumerate(xticks):\n                if np.abs(tick_coord_x - mean_x) < 10:\n                    chosen_tick_text = text\n                    assinged_tick_ids.append(text_id)\n                    break\n\n            data_xs.append(chosen_tick_text)\n            data_ys.append(bar_len)\n            \n        if len(xticks) > len(bounding_boxes):\n            for text_id, (tick_coord_x, tick_coord_y, text) in enumerate(xticks):\n                if text_id not in assinged_tick_ids :\n                    data_xs.append(text)\n                    data_ys.append(0)\n                    assinged_tick_ids.append(text_id)\n                    \n#             print(assinged_tick_ids)\n            sorted_inds = np.argsort(assinged_tick_ids)\n            data_xs = np.array(data_xs)[sorted_inds].tolist()\n            data_ys = np.array(data_ys)[sorted_inds].tolist()\n            \n        # histogram\n        if np.mean([x == 'None' for x in data_xs]) > 0.8 and np.abs(bounding_boxes[0][2] - bounding_boxes[1][0]) < 5:\n            data_xs = []\n            data_ys = []\n            for i, box in enumerate(bounding_boxes):\n                x1,y1,x2,y2 = box.astype(int)\n                \n                bar_len = get_chart_y_value(y1, min_y_ticks[1], max_y_ticks[1], min_y_ticks[2], max_y_ticks[2])\n            \n                chosen_tick_text = 'None'\n                for text_id, (tick_coord_x, tick_coord_y, text) in enumerate(xticks):\n                    if np.abs(tick_coord_x - x1) < 10:\n                        chosen_tick_text = text\n                        break\n\n                data_xs.append(chosen_tick_text)\n                data_ys.append(bar_len)\n            \n            chosen_tick_text = 'None'\n            for text_id, (tick_coord_x, tick_coord_y, text) in enumerate(xticks):\n                if np.abs(tick_coord_x - x2) < 10:\n                    chosen_tick_text = text\n                    break\n            data_xs.append(chosen_tick_text)\n            \n        if debug:\n            plt.figure()\n            plt.imshow(draw)\n            plt.show()\n\n        return data_xs, data_ys","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.248735Z","iopub.execute_input":"2023-06-20T01:26:24.249257Z","iopub.status.idle":"2023-06-20T01:26:24.280194Z","shell.execute_reply.started":"2023-06-20T01:26:24.249220Z","shell.execute_reply":"2023-06-20T01:26:24.278964Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class LineReader:\n    \n    def read(xticks, yticks, X, Y, img, mask, debug=False):\n        numeric_yticks = []\n        for tick in yticks:\n            clean_tick =  re.sub( r\"[^\\d\\-.]\", \"\", tick[2])\n            if is_numeric_string(clean_tick):\n                numeric_yticks.append((tick[0], tick[1], float(clean_tick)))\n                \n        numeric_yticks = filter_tick_outliers(numeric_yticks)\n\n        if len(numeric_yticks) > 0:\n            min_y_ticks = min(numeric_yticks, key=lambda x: x[-1])\n            max_y_ticks = max(numeric_yticks, key=lambda x: x[-1])\n        else:\n            max_y_ticks = (0,0,0)\n            min_y_ticks = (0,0,0)\n        \n        if debug:\n            mask3c = np.stack([mask]+[np.zeros_like(mask)]*2, axis=-1)\n            blend = cv2.addWeighted(mask3c, 0.2, img, 0.8, 0.0)\n            plt.figure(figsize=(10,8))\n            plt.imshow(blend)\n            \n        data_xs = []\n        data_ys = []\n            \n        for xtick_x, _, xtext in xticks:\n            intersection = int(np.clip(xtick_x, 0, mask.shape[1]-1))\n            y_pixel = np.where(mask[:, int(xtick_x-6):int(xtick_x+6)])[0].mean()\n            if np.isnan(y_pixel):\n                continue\n                \n            yvalue = get_chart_y_value(y_pixel, min_y_ticks[1], max_y_ticks[1], min_y_ticks[2], max_y_ticks[2])\n            data_xs.append(xtext)\n            data_ys.append(yvalue)\n\n        return data_xs, data_ys\n    \n","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.283852Z","iopub.execute_input":"2023-06-20T01:26:24.284293Z","iopub.status.idle":"2023-06-20T01:26:24.300037Z","shell.execute_reply.started":"2023-06-20T01:26:24.284249Z","shell.execute_reply":"2023-06-20T01:26:24.298996Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_chart_x_value(pixel_xval, pixel_min, pixel_max, tick_min, tick_max ):\n    denominator = 1 if (pixel_min - pixel_max) == 0 else (pixel_min - pixel_max)\n    if pixel_min >= pixel_xval:\n        return (pixel_xval - pixel_min) / denominator * (tick_max - tick_min) + tick_min\n    else:\n        return (pixel_min - pixel_xval) / denominator * (tick_max - tick_min) + tick_min\n\n\nclass ScatterReader:\n    \n    def read(xticks, yticks, X, Y, img, instances, debug=False):\n        numeric_yticks = []\n        for tick in yticks:\n            clean_tick =  numeric_clean(tick[2])\n            if is_numeric_string(clean_tick):\n                numeric_yticks.append((tick[0], tick[1], float(clean_tick)))\n\n        numeric_xticks = []\n        for tick in xticks:\n            clean_tick =  numeric_clean(tick[2])\n            if is_numeric_string(clean_tick):\n                numeric_xticks.append((tick[0], tick[1], float(clean_tick)))\n\n        numeric_xticks = filter_tick_outliers(numeric_xticks)\n        \n        if len(numeric_yticks) > 0:\n            min_y_ticks = min(numeric_yticks, key=lambda x: x[-1])\n            max_y_ticks = max(numeric_yticks, key=lambda x: x[-1])\n        else:\n            max_y_ticks = (0,0,0)\n            min_y_ticks = (0,0,0)\n\n        if len(numeric_xticks) > 0:\n            min_x_ticks = min(numeric_xticks, key=lambda x: x[-1])\n            max_x_ticks = max(numeric_xticks, key=lambda x: x[-1])\n        else:\n            max_x_ticks = (0,0,0)\n            min_x_ticks = (0,0,0)\n            \n        bounding_boxes = instances['instances'].pred_boxes.tensor.cpu().numpy()\n\n        bounding_boxes = np.array( sorted(bounding_boxes, key=lambda x: x[0]) )\n\n        data_xs = []\n        data_ys = []\n            \n        draw = img.copy()\n        for i, box in enumerate(bounding_boxes):\n            x1,y1,x2,y2 = box\n            x_center = (x1+x2)/2\n            y_center = (y1+y2)/2\n            if debug:\n                # Draw text on image\n                x1,y1,x2,y2 = int(x1), int(y1), int(x2), int(y2)\n                cv2.putText(draw, str(i), (x1,y1), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,0,0), 1)\n                cv2.rectangle(draw, (x1,y1), (x2,y2), (0,255,0), 1)\n        \n            yvalue = get_chart_y_value(y_center, min_y_ticks[1], max_y_ticks[1], min_y_ticks[2], max_y_ticks[2])\n            xvalue = get_chart_x_value(x_center, min_x_ticks[0], max_x_ticks[0], min_x_ticks[2], max_x_ticks[2])\n            data_xs.append(xvalue)\n            data_ys.append(yvalue)\n            \n        if debug:\n#             print(2)\n            plt.figure(figsize=(20,10))\n            plt.imshow(draw)\n            plt.show()\n\n        return data_xs, data_ys\n","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.301962Z","iopub.execute_input":"2023-06-20T01:26:24.302789Z","iopub.status.idle":"2023-06-20T01:26:24.327973Z","shell.execute_reply.started":"2023-06-20T01:26:24.302746Z","shell.execute_reply":"2023-06-20T01:26:24.326540Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HBarReader:\n    \n    def read(xticks, yticks, X, Y, img, instances, debug=False):\n\n        bounding_boxes = instances['instances'].pred_boxes.tensor.cpu().numpy()\n\n        bounding_boxes = np.array( sorted(bounding_boxes, key=lambda x: x[1]) )\n\n        numeric_xticks = []\n        for tick in xticks:\n            clean_tick =  re.sub( r\"[^\\d\\-.]\", \"\", tick[2])\n            if is_numeric_string(clean_tick):\n                numeric_xticks.append((tick[0], tick[1], float(clean_tick)))\n                \n        numeric_xticks = filter_tick_outliers(numeric_xticks)\n\n        # print(numeric_xticks)\n\n        if len(numeric_xticks) > 0:\n            min_x_ticks = min(numeric_xticks, key=lambda x: x[-1])\n            max_x_ticks = max(numeric_xticks, key=lambda x: x[-1])\n        else:\n            max_x_ticks = (0,0,0)\n            min_x_ticks = (0,0,0)\n\n\n        draw = img.copy()\n\n        data_xs = []\n        data_ys = []\n        for i, box in enumerate(bounding_boxes):\n            x1,y1,x2,y2 = box\n            if debug:\n                # Draw text on image\n                x2, y2 = int(x2), int(y2)\n                cv2.putText(draw, str(i), (x2,y2), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,0,0), 2)\n        \n\n            bar_len = get_chart_x_value(x2, min_x_ticks[0], max_x_ticks[0], min_x_ticks[2], max_x_ticks[2])\n\n            # find the corresponding y-label\n            mean_y = (y1 + y2) / 2\n            chosen_tick_text = 'None'\n            for tick_coord_x, tick_coord_y, text in yticks:\n                if np.abs(tick_coord_y - mean_y) < 10:\n                    chosen_tick_text = text\n                    break\n\n            # print(tick_coord_x, tick_coord_y, text)\n            \n            data_ys.append(chosen_tick_text)\n            data_xs.append(bar_len)\n\n        if debug:\n            plt.figure()\n            plt.imshow(draw)\n            plt.show()\n\n        return data_xs, data_ys   \n\n","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.329773Z","iopub.execute_input":"2023-06-20T01:26:24.330489Z","iopub.status.idle":"2023-06-20T01:26:24.350115Z","shell.execute_reply.started":"2023-06-20T01:26:24.330445Z","shell.execute_reply":"2023-06-20T01:26:24.348689Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def calc_score(pred_xs, pred_ys, pred_chart_type, gt_ann):\n    gt_xs = [it['x'] for it in gt_ann['data-series']]\n    gt_ys = [it['y'] for it in gt_ann['data-series']]\n    gt_chart_type = gt_ann['chart-type']\n    if gt_chart_type in ['dot']:\n        gt_xs = cast_to_numeric_if_possible(gt_xs)\n        \n    pred_ys_copy = [y for y in pred_ys]\n    if len(pred_xs) > len(pred_ys):\n        if np.isnan(gt_ys[-1]) :\n            gt_ys[-1] = 0\n            pred_ys_copy.append(0)\n\n    ground_truth = pd.DataFrame.from_dict({\n            f'_x': (gt_xs, gt_chart_type),\n            f'_y': (gt_ys, gt_chart_type),\n        }, orient='index', columns=['data_series', 'chart_type']).rename_axis('id')\n\n    predictions = pd.DataFrame.from_dict({\n        f'_x': (pred_xs, pred_chart_type),\n        f'_y': (pred_ys_copy, pred_chart_type),\n    }, orient='index', columns=['data_series', 'chart_type']).rename_axis('id')\n\n    return benetech_score(ground_truth, predictions)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.352300Z","iopub.execute_input":"2023-06-20T01:26:24.354101Z","iopub.status.idle":"2023-06-20T01:26:24.367569Z","shell.execute_reply.started":"2023-06-20T01:26:24.353423Z","shell.execute_reply":"2023-06-20T01:26:24.366309Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plotbb_model = Unet(\n    encoder_name = PLOTBB_CFG.backbone,\n    encoder_depth = 5,\n    encoder_weights = None,\n    classes = 1,\n    activation = 'sigmoid',\n    aux_params = {\n        'classes': 5,\n        'activation':'softmax',\n    }\n)\n\nline_seg_model = Unet(\n    encoder_name = LINE_SEG_CFG.backbone,\n    encoder_depth = 5,\n    encoder_weights = None,\n    classes = 1,\n    activation = 'sigmoid',\n   \n)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:24.377213Z","iopub.execute_input":"2023-06-20T01:26:24.377539Z","iopub.status.idle":"2023-06-20T01:26:25.947223Z","shell.execute_reply.started":"2023-06-20T01:26:24.377509Z","shell.execute_reply":"2023-06-20T01:26:25.946041Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def modify_key_name_2(state_dict):\n    new_state_dict = {}\n    for key, value in state_dict.items():\n        new_key = key.replace('model.encoder.model.', 'encoder.model.')\n        new_key = new_key.replace('model.decoder.', 'decoder.')\n        new_key = new_key.replace('model.segmentation', 'segmentation')\n        new_key = new_key.replace('model.classification', 'classification')\n        new_state_dict[new_key] = value\n    return new_state_dict","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:25.949165Z","iopub.execute_input":"2023-06-20T01:26:25.949575Z","iopub.status.idle":"2023-06-20T01:26:25.957331Z","shell.execute_reply.started":"2023-06-20T01:26:25.949526Z","shell.execute_reply":"2023-06-20T01:26:25.955882Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plotbb_models = []\nfor path in glob.glob(PLOTBB_WEIGHT_PATHS+'/*.pth'):\n    plotbb_statedict = torch.load(path)['model']\n    plotbb_statedict = modify_key_name_2(plotbb_statedict)\n    plotbb_model.load_state_dict(plotbb_statedict)\n    plotbb_model.eval()\n    plotbb_model.to(DEVICE)\n    plotbb_models.append(plotbb_model)\n\nline_seg_statedict = torch.load(LINE_SEG_WEIGHT_PATH)['model']\nline_seg_statedict = modify_key_name(line_seg_statedict)\nline_seg_model.load_state_dict(line_seg_statedict)\n\nline_seg_model.eval()\nline_seg_model.to(DEVICE)\nprint()","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:25.959474Z","iopub.execute_input":"2023-06-20T01:26:25.960388Z","iopub.status.idle":"2023-06-20T01:26:45.021197Z","shell.execute_reply.started":"2023-06-20T01:26:25.960341Z","shell.execute_reply":"2023-06-20T01:26:45.019019Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_plotbb_XY_and_chart_type(img, models, transforms, device=DEVICE, draw=False):\n    X = transforms(image=img)['image']\n    X = X.to(device)\n    \n    mean_mask_pred = 0\n    mean_y_pred = 0\n    for model in models:\n        with torch.no_grad():\n            mask_pred, y_pred = model(X.unsqueeze(0))\n            mask_pred = mask_pred.cpu().numpy()[0][0]\n            mean_mask_pred += mask_pred\n            mean_y_pred += y_pred\n            \n    mean_mask_pred /= len(models)\n    mean_y_pred /= len(models)\n    \n    mean_mask_pred = cv2.resize(mean_mask_pred, (img.shape[1], img.shape[0]))\n    mean_mask_pred = np.where(mean_mask_pred > 0.5, 255, 0).astype('uint8')\n    \n    contours, _ = cv2.findContours(mean_mask_pred, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n    x, y, w, h = cv2.boundingRect(contours[0])\n\n    img_draw = None\n    if draw:\n        img_draw = cv2.rectangle(img.copy(), (x,y), (x+w, y+h), (0,0,255), 1)\n    chart_type = mean_y_pred.argmax().item()\n\n    return x, y+h, chart_type, img_draw","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:45.022631Z","iopub.execute_input":"2023-06-20T01:26:45.023045Z","iopub.status.idle":"2023-06-20T01:26:45.036164Z","shell.execute_reply.started":"2023-06-20T01:26:45.023012Z","shell.execute_reply":"2023-06-20T01:26:45.034928Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ocr = OCR()","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:45.038059Z","iopub.execute_input":"2023-06-20T01:26:45.038480Z","iopub.status.idle":"2023-06-20T01:26:48.334633Z","shell.execute_reply.started":"2023-06-20T01:26:45.038441Z","shell.execute_reply":"2023-06-20T01:26:48.332879Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"READER_MAP = {\n    'vertical_bar': VBarReader,\n    'dot': DotReader,\n    'line': LineReader,\n    'scatter': ScatterReader,\n    'horizontal_bar': HBarReader\n}","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:48.337063Z","iopub.execute_input":"2023-06-20T01:26:48.337829Z","iopub.status.idle":"2023-06-20T01:26:48.345837Z","shell.execute_reply.started":"2023-06-20T01:26:48.337778Z","shell.execute_reply":"2023-06-20T01:26:48.344641Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize a prediction","metadata":{}},{"cell_type":"code","source":"%%time\nimg_name= f'0000ae6cbdb1.jpg'\n\nDEBUG = True\n\nimg_path = os.path.join(TRAIN_DIR, img_name)\nimg_name_no_ext = img_name.split('.')[0]\n\nwith open(os.path.join(TRAIN_ANN_DIR, img_name_no_ext + '.json'), 'r') as f:\n    ann = json.load(f)\n\n\nimg = cv2.imread(img_path)\nimg = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\nX,Y,chart_type1, img_draw = get_plotbb_XY_and_chart_type(img, plotbb_models, PLOTBB_CFG.val_transforms, draw=DEBUG)\npred_chart_type = CHART_MAP_REV[chart_type1]\n\nif pred_chart_type == 'line':\n    line_seg_mask = get_chart_data_mask(img, line_seg_model, LINE_SEG_CFG.val_transforms)\nelse:\n    instances = predict_inst_seg(img[:,:,::-1], pred_chart_type)\n    instances = inst_seg_post_process(instances, INST_SEG_CONF[pred_chart_type], INST_SEG_NMS[pred_chart_type])\n\ntext_boxes, ocr_show = ocr(img, return_im_show=DEBUG)\n\nplot_axes = PlotAxes(X, Y, text_boxes, pred_chart_type)\n\nxticks = plot_axes.get_x_ticks()\nyticks = plot_axes.get_y_ticks()\n\nif pred_chart_type == 'line':\n    pred_xs, pred_ys = READER_MAP[pred_chart_type].read(xticks, yticks, X, Y, img, line_seg_mask, debug=DEBUG)\n    \nelse:\n    pred_xs, pred_ys = READER_MAP[pred_chart_type].read(xticks, yticks, X, Y, img, instances, debug=DEBUG)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:48.347480Z","iopub.execute_input":"2023-06-20T01:26:48.348635Z","iopub.status.idle":"2023-06-20T01:26:53.448538Z","shell.execute_reply.started":"2023-06-20T01:26:48.348595Z","shell.execute_reply":"2023-06-20T01:26:53.447299Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"score = calc_score(pred_xs, pred_ys, pred_chart_type, ann)\nscore","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:53.449745Z","iopub.execute_input":"2023-06-20T01:26:53.450106Z","iopub.status.idle":"2023-06-20T01:26:53.477307Z","shell.execute_reply.started":"2023-06-20T01:26:53.450068Z","shell.execute_reply":"2023-06-20T01:26:53.476492Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(ann['data-series']), len(pred_xs)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:53.478894Z","iopub.execute_input":"2023-06-20T01:26:53.480081Z","iopub.status.idle":"2023-06-20T01:26:53.488203Z","shell.execute_reply.started":"2023-06-20T01:26:53.480040Z","shell.execute_reply":"2023-06-20T01:26:53.487092Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(30,10))\nplt.imshow(ocr_show)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:53.489945Z","iopub.execute_input":"2023-06-20T01:26:53.490685Z","iopub.status.idle":"2023-06-20T01:26:54.205334Z","shell.execute_reply.started":"2023-06-20T01:26:53.490642Z","shell.execute_reply":"2023-06-20T01:26:54.204373Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img_draw2 = PlotAxes.draw_ticks(img_draw, xticks)\nimg_draw2 = PlotAxes.draw_ticks(img_draw2, yticks)\nplt.imshow(img_draw2)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:54.207103Z","iopub.execute_input":"2023-06-20T01:26:54.207862Z","iopub.status.idle":"2023-06-20T01:26:54.507468Z","shell.execute_reply.started":"2023-06-20T01:26:54.207822Z","shell.execute_reply":"2023-06-20T01:26:54.503248Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"xticks, yticks","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:54.509121Z","iopub.execute_input":"2023-06-20T01:26:54.509618Z","iopub.status.idle":"2023-06-20T01:26:54.519565Z","shell.execute_reply.started":"2023-06-20T01:26:54.509574Z","shell.execute_reply":"2023-06-20T01:26:54.518453Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_xs, pred_ys","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:54.521041Z","iopub.execute_input":"2023-06-20T01:26:54.522097Z","iopub.status.idle":"2023-06-20T01:26:54.530281Z","shell.execute_reply.started":"2023-06-20T01:26:54.522056Z","shell.execute_reply":"2023-06-20T01:26:54.529139Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ann['data-series']","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:54.531739Z","iopub.execute_input":"2023-06-20T01:26:54.532731Z","iopub.status.idle":"2023-06-20T01:26:54.541664Z","shell.execute_reply.started":"2023-06-20T01:26:54.532670Z","shell.execute_reply":"2023-06-20T01:26:54.540585Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Evaluate","metadata":{}},{"cell_type":"code","source":"meta_df = pd.read_csv('/kaggle/input/benetech-extract-meta/metadata.csv')\nval_df = meta_df[meta_df.fold==0]","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:54.543400Z","iopub.execute_input":"2023-06-20T01:26:54.543871Z","iopub.status.idle":"2023-06-20T01:26:54.685544Z","shell.execute_reply.started":"2023-06-20T01:26:54.543826Z","shell.execute_reply":"2023-06-20T01:26:54.684430Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"val_extracted = val_df[(val_df.source=='extracted') \n#                       & (val_df.chart_type=='horizontal_bar') \n                      ]","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:54.687118Z","iopub.execute_input":"2023-06-20T01:26:54.687534Z","iopub.status.idle":"2023-06-20T01:26:54.696875Z","shell.execute_reply.started":"2023-06-20T01:26:54.687483Z","shell.execute_reply":"2023-06-20T01:26:54.695740Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEBUG=False\n\nall_names = []\nall_chart_types = []\nall_scores = []\n\nfor img_name_no_ext in tqdm(val_extracted['name'].tolist()):\n    img_name = img_name_no_ext + '.jpg'\n    img_path = os.path.join(TRAIN_DIR, img_name)\n    \n    with open(os.path.join(TRAIN_ANN_DIR, img_name_no_ext + '.json'), 'r') as f:\n        ann = json.load(f)\n    \n    img = cv2.imread(img_path)\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n    X,Y,chart_type1, img_draw = get_plotbb_XY_and_chart_type(img, plotbb_models, PLOTBB_CFG.val_transforms, draw=DEBUG)\n    pred_chart_type = CHART_MAP_REV[chart_type1]\n\n    if pred_chart_type == 'line':\n        line_seg_mask = get_chart_data_mask(img, line_seg_model, LINE_SEG_CFG.val_transforms)\n    else:\n        instances = predict_inst_seg(img[:,:,::-1], pred_chart_type)\n        instances = inst_seg_post_process(instances, INST_SEG_CONF[pred_chart_type], INST_SEG_NMS[pred_chart_type])\n\n    text_boxes, ocr_show = ocr(img, return_im_show=DEBUG)\n\n    plot_axes = PlotAxes(X, Y, text_boxes, pred_chart_type)\n\n    xticks = plot_axes.get_x_ticks()\n    yticks = plot_axes.get_y_ticks()\n\n    if pred_chart_type == 'line':\n        pred_xs, pred_ys = READER_MAP[pred_chart_type].read(xticks, yticks, X, Y, img, line_seg_mask, debug=DEBUG)\n        \n    else:\n        pred_xs, pred_ys = READER_MAP[pred_chart_type].read(xticks, yticks, X, Y, img, instances, debug=DEBUG)\n    \n#     print(img_name)\n    try:\n        score = calc_score(pred_xs, pred_ys, pred_chart_type, ann)\n    except Exception as ex:\n        print(img_name)\n        score = 0\n        print(ex)\n        raise ex\n        \n        \n    all_names.append(img_name_no_ext)\n    all_scores.append(score)\n    all_chart_types.append(ann['chart-type'])\n    \n#     print('Cummu mean score:', np.mean(all_scores))\n\n","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:26:54.698510Z","iopub.execute_input":"2023-06-20T01:26:54.699395Z","iopub.status.idle":"2023-06-20T01:28:34.629730Z","shell.execute_reply.started":"2023-06-20T01:26:54.699355Z","shell.execute_reply":"2023-06-20T01:28:34.628544Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"score_df = pd.DataFrame({'name':all_names, 'chart_type':all_chart_types, 'score': all_scores})\nscore_df.groupby('chart_type').score.mean()","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:28:34.631576Z","iopub.execute_input":"2023-06-20T01:28:34.632360Z","iopub.status.idle":"2023-06-20T01:28:34.651156Z","shell.execute_reply.started":"2023-06-20T01:28:34.632315Z","shell.execute_reply":"2023-06-20T01:28:34.649768Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"score_df.score.mean()","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:28:34.653653Z","iopub.execute_input":"2023-06-20T01:28:34.654631Z","iopub.status.idle":"2023-06-20T01:28:34.664356Z","shell.execute_reply.started":"2023-06-20T01:28:34.654575Z","shell.execute_reply":"2023-06-20T01:28:34.662973Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"score_df.groupby('chart_type').score.sum() / len(score_df)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:28:34.666484Z","iopub.execute_input":"2023-06-20T01:28:34.667077Z","iopub.status.idle":"2023-06-20T01:28:34.680990Z","shell.execute_reply.started":"2023-06-20T01:28:34.666881Z","shell.execute_reply":"2023-06-20T01:28:34.679876Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_chart_type","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:28:34.682498Z","iopub.execute_input":"2023-06-20T01:28:34.683597Z","iopub.status.idle":"2023-06-20T01:28:34.693072Z","shell.execute_reply.started":"2023-06-20T01:28:34.683549Z","shell.execute_reply":"2023-06-20T01:28:34.691761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"score_df","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:28:34.694623Z","iopub.execute_input":"2023-06-20T01:28:34.695873Z","iopub.status.idle":"2023-06-20T01:28:34.719464Z","shell.execute_reply.started":"2023-06-20T01:28:34.695832Z","shell.execute_reply":"2023-06-20T01:28:34.718260Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Infer on test","metadata":{}},{"cell_type":"code","source":"DEBUG=False","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:28:34.720923Z","iopub.execute_input":"2023-06-20T01:28:34.721850Z","iopub.status.idle":"2023-06-20T01:28:34.727310Z","shell.execute_reply.started":"2023-06-20T01:28:34.721809Z","shell.execute_reply":"2023-06-20T01:28:34.725980Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"all_names = []\nall_preds = []\nall_chart_types = []\n\nfor img_name in tqdm(os.listdir(TEST_DIR)):\n    img_path = os.path.join(TEST_DIR, img_name)\n    img_name_no_ext = img_name.split('.')[0]\n    try:\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        X,Y,chart_type1, img_draw = get_plotbb_XY_and_chart_type(img, plotbb_models, PLOTBB_CFG.val_transforms, draw=DEBUG)\n        pred_chart_type = CHART_MAP_REV[chart_type1]\n        \n        if pred_chart_type == 'line':\n            line_seg_mask = get_chart_data_mask(img, line_seg_model, LINE_SEG_CFG.val_transforms)\n        else:\n            instances = predict_inst_seg(img[:,:,::-1], pred_chart_type)\n            instances = inst_seg_post_process(instances, INST_SEG_CONF[pred_chart_type], INST_SEG_NMS[pred_chart_type])\n\n        text_boxes, ocr_show = ocr(img, return_im_show=DEBUG)\n\n        plot_axes = PlotAxes(X, Y, text_boxes, pred_chart_type)\n\n        xticks = plot_axes.get_x_ticks()\n        yticks = plot_axes.get_y_ticks()\n\n        if pred_chart_type == 'line':\n            pred_xs, pred_ys = READER_MAP[pred_chart_type].read(xticks, yticks, X, Y, img, line_seg_mask, debug=DEBUG)\n            if DEBUG:\n                plt.figure()\n                plt.imshow(line_seg_mask)\n                plt.show()\n        else:\n            pred_xs, pred_ys = READER_MAP[pred_chart_type].read(xticks, yticks, X, Y, img, instances, debug=DEBUG)\n            \n        all_names.append(f'{img_name_no_ext}_x')\n        all_names.append(f'{img_name_no_ext}_y')\n        all_preds.append(';'.join([str(x) for x in pred_xs]))\n        all_preds.append(';'.join([str(y) for y in pred_ys]))\n        all_chart_types.append(pred_chart_type)\n        all_chart_types.append(pred_chart_type)\n\n    except Exception as ex:\n        print(img_name)\n        all_names.append(f'{img_name_no_ext}_x')\n        all_names.append(f'{img_name_no_ext}_y')\n        all_preds.append('abc;def')\n        all_preds.append('0.0;1.0')\n        all_chart_types.append('vertical_bar')\n        all_chart_types.append('vertical_bar')\n\n#     break","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:28:34.729405Z","iopub.execute_input":"2023-06-20T01:28:34.730581Z","iopub.status.idle":"2023-06-20T01:28:36.362501Z","shell.execute_reply.started":"2023-06-20T01:28:34.730538Z","shell.execute_reply":"2023-06-20T01:28:36.361145Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"try:\n    sub_df = pd.DataFrame({\n        'id': all_names,\n        'data_series': all_preds,\n        'chart_type': all_chart_types\n    })\n\n    sub_df.to_csv('submission.csv', index=False)\nexcept:\n    sub_df = pd.read_csv('/kaggle/input/benetech-making-graphs-accessible/sample_submission.csv')\n    sub_df.to_csv('submission.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:28:36.364589Z","iopub.execute_input":"2023-06-20T01:28:36.365469Z","iopub.status.idle":"2023-06-20T01:28:36.379171Z","shell.execute_reply.started":"2023-06-20T01:28:36.365416Z","shell.execute_reply":"2023-06-20T01:28:36.377785Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_df","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:28:36.381899Z","iopub.execute_input":"2023-06-20T01:28:36.382258Z","iopub.status.idle":"2023-06-20T01:28:36.398350Z","shell.execute_reply.started":"2023-06-20T01:28:36.382223Z","shell.execute_reply":"2023-06-20T01:28:36.397040Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# img = plt.imread(os.path.join(TEST_DIR, '007a18eb4e09.jpg'))\n# plt.imshow(img)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T01:28:36.400266Z","iopub.execute_input":"2023-06-20T01:28:36.400809Z","iopub.status.idle":"2023-06-20T01:28:36.408429Z","shell.execute_reply.started":"2023-06-20T01:28:36.400764Z","shell.execute_reply":"2023-06-20T01:28:36.406384Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}