{"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":"import torch\nimport torch.nn as nn\nimport sys, os, shutil\nimport numpy as np\nfrom scipy.stats import mode\nimport cv2\nimport warnings\nimport gc\nimport torch.quantization\nfrom transformers import TrOCRProcessor, VisionEncoderDecoderModel\nfrom fastai.vision.all import *\nimport pandas as pd\nimport albumentations as A\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\nfrom torchvision.models.detection.faster_rcnn import FasterRCNN","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:55:45.779371Z","iopub.execute_input":"2023-06-18T11:55:45.780744Z","iopub.status.idle":"2023-06-18T11:55:45.796185Z","shell.execute_reply.started":"2023-06-18T11:55:45.780696Z","shell.execute_reply":"2023-06-18T11:55:45.794975Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install ../input/paddle-wheels/paddlepaddle-2.4.2-cp37-cp37m-manylinux1_x86_64.whl --no-index --find-links=file:///kaggle/input/paddle-wheels -q\n!pip install ../input/paddle-wheels/paddleocr-2.6.1.3-py3-none-any.whl --no-index --find-links=file:///kaggle/input/paddle-wheels -q","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:55:45.800494Z","iopub.execute_input":"2023-06-18T11:55:45.801927Z","iopub.status.idle":"2023-06-18T11:56:42.413856Z","shell.execute_reply.started":"2023-06-18T11:55:45.801884Z","shell.execute_reply":"2023-06-18T11:56:42.412343Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!python -m pip install --no-index --find-links=/kaggle/input/spell-check neuspell -q\n!mkdir /kaggle/working/bert-base-cased\n!cp /kaggle/input/bert-base-cased/* -r /kaggle/working/bert-base-cased","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:56:42.417940Z","iopub.execute_input":"2023-06-18T11:56:42.423105Z","iopub.status.idle":"2023-06-18T11:57:22.149868Z","shell.execute_reply.started":"2023-06-18T11:56:42.423054Z","shell.execute_reply":"2023-06-18T11:57:22.148288Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ObjectDetectorInference:\n    def __init__(self, checkpoint_file, transforms, tfms_list, h_after, w_after) -> None:\n        self.score_threshold = 0.5\n        self.transforms_compose = transforms\n        self.transforms_list = tfms_list\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.backbone = self.resnet50_fpn_fn(pretrained=False)\n        self.model = self.get_model(num_classes=2, backbone=self.backbone)\n        self.model.load_state_dict(torch.load(checkpoint_file, map_location=torch.device('cuda')))\n        self.model = self.model.to(self.device)\n        self.BBox = namedtuple(\"BBox\", \"xmin, ymin, xmax, ymax\")\n        self.ImgSize = namedtuple(\"ImgSize\", \"width, height\")\n        self.h_after = h_after\n        self.w_after = w_after\n    \n    def get_bboxes(self, img):\n        img_arr = np.array(img)\n        transformed = self.transforms_compose(image=img_arr)\n        batched = transformed['image'].unsqueeze(0)\n        batched = batched.to(self.device)\n        with torch.no_grad():\n            self.model.eval()\n            out = self.model(batched)\n        first_out = out[0]\n        first_out['boxes'] = first_out['boxes'][first_out['scores'] > self.score_threshold]\n        postprocessed_bboxes = []\n        postprocessed_scores = []  \n        for i, box in enumerate(first_out['boxes']):\n            bbox = self.BBox(box[0].item(), box[1].item(), box[2].item(), box[3].item())\n            new_bbox = self.postprocess_bbox(img_size=self.ImgSize(width=img.size[0], height=img.size[1]), \n                                             bbox=bbox, transforms=self.transforms_list, h_after=self.h_after, w_after=self.w_after)\n            postprocessed_bboxes.append(list(new_bbox))\n            postprocessed_scores.append(first_out['scores'][i].item()) \n        return postprocessed_bboxes, postprocessed_scores\n\n\n    def get_img_with_bbs(self, img, postprocessed_bboxes):\n        img_arr = np.array(img)\n        for i in range(len(postprocessed_bboxes)):\n            start_point = postprocessed_bboxes[i][0], postprocessed_bboxes[i][1]\n            end_point = postprocessed_bboxes[i][2], postprocessed_bboxes[i][3]\n            cv2.rectangle(img_arr, start_point, end_point, color=(0, 0, 255), thickness=1)\n        img_with_bbs = Image.fromarray(img_arr)\n        return img_with_bbs\n\n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/models/inference.py\n    def postprocess_bbox(self, img_size, bbox, transforms, h_after, w_after):\n        \"\"\"\n        Post-process predicted bbox to adjust coordinates to input image size.\n        \n        Parameters\n        ----------\n        img_size: original image size, before any model-pre-processing done.\n        bbox: predicted bbox.\n        transforms: list of model-pre-processing transforms.\n        h_after: height of image after model-pre-processing transforms.\n        w_after: width of image after model-pre-processing transforms.\n        \n        Returns\n        -------\n        Tuple with (xmin, ymin, xmax, ymax) rescaled and re-adjusted to match the original image size.\n        \"\"\"\n        h_after, w_after = self.get_size_without_padding(transforms, img_size, h_after, w_after)\n        pad = np.abs(h_after - w_after) // 2\n\n        h_scale, w_scale = h_after / img_size.height, w_after / img_size.width\n        if h_after < w_after:\n            xmin, xmax, ymin, ymax = (\n                int(bbox.xmin),\n                int(bbox.xmax),\n                int(bbox.ymin) - pad,\n                int(bbox.ymax) - pad,\n            )\n        else:\n            xmin, xmax, ymin, ymax = (\n                int(bbox.xmin) - pad,\n                int(bbox.xmax) - pad,\n                int(bbox.ymin),\n                int(bbox.ymax),\n            )\n\n        xmin, xmax, ymin, ymax = (\n            max(xmin, 0),\n            min(xmax, w_after),\n            max(ymin, 0),\n            min(ymax, h_after),\n        )\n        xmin, xmax, ymin, ymax = (\n            int(xmin / w_scale),\n            int(xmax / w_scale),\n            int(ymin / h_scale),\n            int(ymax / h_scale),\n        )\n\n        return xmin, ymin, xmax, ymax\n    \n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/tfms/albumentations/albumentations_helpers.py\n    def get_size_without_padding(self, tfms_list, before_tfm_img_size, height, width):\n        \"\"\"\n        Infer the height and width of the pre-processed image after removing padding.\n        \n        Parameters\n        ----------\n        tfms_list: list of albumentations transforms applied to the `before_tfm_img` image before passing it to the model for inference.\n        before_tfm_img_size: original image size before being pre-processed for inference.\n        height: height of output image from icevision `predict` function.\n        width: width of output image from icevision `predict` function.\n        \n        Returns\n        -------\n        height and width of the image coming out of the inference pipeline, after removing padding.\n        \"\"\"\n        if self.get_transform(tfms_list, \"Pad\") is not None:\n\n            t = self.get_transform(tfms_list, \"SmallestMaxSize\")\n            if t is not None:\n                presize = t.max_size\n                height, width = self.func_max_size(before_tfm_img_size.height, before_tfm_img_size.width, presize, min)\n\n            t = self.get_transform(tfms_list, \"LongestMaxSize\")\n            if t is not None:\n                size = t.max_size\n                height, width = self.func_max_size(before_tfm_img_size.height, before_tfm_img_size.width, size, max)\n\n        return height, width\n\n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/tfms/albumentations/albumentations_helpers.py\n    def func_max_size(self, height, width, max_size, func):\n        \"\"\"\n        Calculate rescaled height and width of the image in question w.r.t. to a specific size.\n        \n        Parameters\n        ----------\n        height: height of the image in question.\n        width: width of the image in question.\n        max_size: size w.r.t. the image needs to be rescaled (resized).\n        func: min/max. Whether to compare max_size to the smallest/longest of the image dims.\n        \n        Returns\n        -------\n        Rescaled height and width.\n        \"\"\"\n        scale = max_size / float(func(width, height))\n\n        if scale != 1.0:\n            height, width = tuple(self.py3round(dim * scale) for dim in (height, width))\n        return height, width\n\n\n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/tfms/albumentations/albumentations_helpers.py\n    def get_transform(self, tfms_list, t):\n        \"\"\"\n        Extract transform `t` from `tfms_list`.\n        \n        Parameters\n        ----------\n        tfms_list: list of albumentations transforms.\n        t: name (str) of the transform to look for and return from within `tfms_list`.\n        \n        Returns\n        -------\n        The `t` transform if found inside `tfms_list`, otherwise `None`.\n        \"\"\"\n        for el in tfms_list:\n            if t in str(type(el)):\n                return el\n        return None\n    \n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/tfms/albumentations/albumentations_helpers.py\n    def py3round(self, number):\n        \"\"\"\n        Unified rounding in all Python versions. Used by Albumentations.\n        \n        Parameters\n        ----------\n        number: float to round.\n        \n        Returns\n        -------\n        Rounded number.\n        \"\"\"\n        if abs(round(number) - number) == 0.5:\n            return int(2.0 * round(number / 2.0))\n\n        return int(round(number))\n    \n\n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/utils/torch_utils.py\n    def check_all_model_params_in_groups2(self, model, param_groups):\n        num_params = len([param for group in param_groups for param in group])\n        num_params_expected = len(list(model.parameters()))\n\n        if num_params != num_params_expected:\n            raise RuntimeError(f\"{num_params_expected} params in model, but only {num_params} in parameter group.\")\n        \n\n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/models/torchvision/faster_rcnn/backbones/resnet_fpn_utils.py\n    def param_groups(self, model):\n        body = model.body\n\n        layers = []\n        layers += [nn.Sequential(body.conv1, body.bn1)]\n        layers += [getattr(body, l) for l in list(body) if l.startswith(\"layer\")]\n        layers += [model.fpn]\n\n        _param_groups = [list(layer.parameters()) for layer in layers]\n        self.check_all_model_params_in_groups2(model, _param_groups)\n\n        return _param_groups\n    \n\n    def patch_param_groups(self, model):\n        model.param_groups = MethodType(self.param_groups, model)\n\n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/models/torchvision/faster_rcnn/backbones/resnet_fpn_configs.py\n    def _resnet_fpn(self, name, pretrained=True, **kwargs):\n        model = resnet_fpn_backbone(backbone_name=name, pretrained=pretrained, **kwargs)\n        self.patch_param_groups(model)\n\n        return model\n    \n\n    def resnet50_fpn_fn(self, pretrained=True, **kwargs):\n        return self._resnet_fpn(\"resnet50\", pretrained=pretrained, **kwargs)\n    \n\n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/models/torchvision/utils.py\n    def patch_param_groups2(self, model, head_layers, backbone_param_groups):\n        def param_groups2(model):\n            head_param_groups = [list(layer.parameters()) for layer in head_layers]\n\n            _param_groups = backbone_param_groups + head_param_groups\n            self.check_all_model_params_in_groups2(model, _param_groups)\n\n            return _param_groups\n\n        model.param_groups = MethodType(param_groups2, model)\n\n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/models/torchvision/utils.py\n    def patch_rcnn_param_groups(self, model):\n        return self.patch_param_groups2(model=model, head_layers=[model.rpn, model.roi_heads], backbone_param_groups=model.backbone.param_groups())\n    \n\n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/models/torchvision/utils.py\n    def _noop_normalize(self, image):\n        return image\n\n\n    def _noop_resize(self, image, target):\n        return image, target\n\n\n    def remove_internal_model_transforms(self, model):\n        model.transform.normalize = self._noop_normalize\n        model.transform.resize = self._noop_resize\n\n    # Original source: https://github.com/airctic/icevision/blob/master/icevision/models/torchvision/faster_rcnn/model.py\n    def get_model(self, num_classes, backbone, remove_internal_transforms=True, **faster_rcnn_kwargs):\n        model = FasterRCNN(backbone, num_classes=num_classes, **faster_rcnn_kwargs)\n\n        self.patch_rcnn_param_groups(model=model)\n\n        if remove_internal_transforms:\n            self.remove_internal_model_transforms(model)\n\n        return model","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:57:22.166241Z","iopub.execute_input":"2023-06-18T11:57:22.169014Z","iopub.status.idle":"2023-06-18T11:57:22.249878Z","shell.execute_reply.started":"2023-06-18T11:57:22.168963Z","shell.execute_reply":"2023-06-18T11:57:22.248786Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from albumentations.pytorch.transforms import ToTensorV2\n\ntick_mark_transforms = A.Compose([\n    A.LongestMaxSize(always_apply=False, p=1, max_size=640, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=640, min_width=640, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0), \n    ToTensorV2()\n])\n\ntick_mark_tfms_list = [\n    A.LongestMaxSize(always_apply=False, p=1, max_size=640, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=640, min_width=640, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0)\n]\n\ninfer_tick_mark = ObjectDetectorInference(checkpoint_file='/kaggle/input/b-m-g-a-extra-data/tick_mark_5.pth', \n                                          transforms=tick_mark_transforms, \n                                          tfms_list=tick_mark_tfms_list,\n                                          h_after=640,\n                                          w_after=640)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:57:22.252965Z","iopub.execute_input":"2023-06-18T11:57:22.254059Z","iopub.status.idle":"2023-06-18T11:57:25.132636Z","shell.execute_reply.started":"2023-06-18T11:57:22.254006Z","shell.execute_reply":"2023-06-18T11:57:25.131408Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"line_transforms = A.Compose([\n    A.LongestMaxSize(always_apply=False, p=1, max_size=640, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=640, min_width=640, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0), \n    ToTensorV2()\n])\n\nline_tfms_list = [\n    A.LongestMaxSize(always_apply=False, p=1, max_size=640, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=640, min_width=640, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0)\n]\n\ninfer_line = ObjectDetectorInference(checkpoint_file='/kaggle/input/b-m-g-a-extra-data/line__5.pth', transforms=line_transforms, tfms_list=line_tfms_list, h_after=640, w_after=640)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:57:25.134411Z","iopub.execute_input":"2023-06-18T11:57:25.135270Z","iopub.status.idle":"2023-06-18T11:57:27.749613Z","shell.execute_reply.started":"2023-06-18T11:57:25.135227Z","shell.execute_reply":"2023-06-18T11:57:27.748381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"horizontal_bar_tfms_list = [\n    A.LongestMaxSize(always_apply=False, p=1, max_size=896, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=896, min_width=896, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0)\n]\nhorizontal_bar_transforms = A.Compose([\n    A.LongestMaxSize(always_apply=False, p=1, max_size=896, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=896, min_width=896, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0), \n    ToTensorV2()\n])\n\ninfer_horizontal_bar = ObjectDetectorInference(checkpoint_file='/kaggle/input/b-m-g-a-extra-data/horizontal_bar_23.pth',transforms=horizontal_bar_transforms, tfms_list=horizontal_bar_tfms_list, h_after=896, w_after=896)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:57:27.754858Z","iopub.execute_input":"2023-06-18T11:57:27.755865Z","iopub.status.idle":"2023-06-18T11:57:30.714618Z","shell.execute_reply.started":"2023-06-18T11:57:27.755819Z","shell.execute_reply":"2023-06-18T11:57:30.709227Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"y_tick_label_transforms = A.Compose([\n    A.LongestMaxSize(always_apply=False, p=1, max_size=640, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=640, min_width=640, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0), \n    ToTensorV2()\n])\n\ny_tick_label_tfms_list = [\n    A.LongestMaxSize(always_apply=False, p=1, max_size=640, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=640, min_width=640, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0)\n]\n\ninfer_label = ObjectDetectorInference(checkpoint_file='/kaggle/input/b-m-g-a-extra-data/y_tick_label__6.pth', \n                                             transforms=y_tick_label_transforms, \n                                             tfms_list=y_tick_label_tfms_list,\n                                             h_after=640,\n                                             w_after=640)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:57:30.716607Z","iopub.execute_input":"2023-06-18T11:57:30.717387Z","iopub.status.idle":"2023-06-18T11:57:34.046450Z","shell.execute_reply.started":"2023-06-18T11:57:30.717345Z","shell.execute_reply":"2023-06-18T11:57:34.045203Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x_tick_label_transforms = A.Compose([\n    A.LongestMaxSize(always_apply=False, p=1, max_size=896, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=896, min_width=896, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0), \n    ToTensorV2()\n])\n\nx_tick_label_tfms_list = [\n    A.LongestMaxSize(always_apply=False, p=1, max_size=896, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=896, min_width=896, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0)\n]\n\ninfer_label_x = ObjectDetectorInference(checkpoint_file='/kaggle/input/b-m-g-a-extra-data/x_tick_label_6.pth', \n                                             transforms=x_tick_label_transforms, \n                                             tfms_list=x_tick_label_tfms_list,\n                                             h_after=896,\n                                             w_after=896)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:57:34.051712Z","iopub.execute_input":"2023-06-18T11:57:34.052780Z","iopub.status.idle":"2023-06-18T11:57:36.650107Z","shell.execute_reply.started":"2023-06-18T11:57:34.052726Z","shell.execute_reply":"2023-06-18T11:57:36.648768Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from albumentations.pytorch.transforms import ToTensorV2\n\ntick_mark_transforms = A.Compose([\n    A.LongestMaxSize(always_apply=False, p=1, max_size=640, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=640, min_width=640, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0), \n    ToTensorV2()\n])\n\ntick_mark_tfms_list = [\n    A.LongestMaxSize(always_apply=False, p=1, max_size=640, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=640, min_width=640, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0)\n]\n\ninfer_tick_mark = ObjectDetectorInference(checkpoint_file='/kaggle/input/b-m-g-a-extra-data/tick_mark_5.pth', \n                                          transforms=tick_mark_transforms, \n                                          tfms_list=tick_mark_tfms_list,\n                                          h_after=640,\n                                          w_after=640)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:57:36.655444Z","iopub.execute_input":"2023-06-18T11:57:36.655879Z","iopub.status.idle":"2023-06-18T11:57:39.344447Z","shell.execute_reply.started":"2023-06-18T11:57:36.655837Z","shell.execute_reply":"2023-06-18T11:57:39.343118Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"vertical_bar_transforms = A.Compose([\n    A.LongestMaxSize(always_apply=False, p=1, max_size=896, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=896, min_width=896, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0), \n    ToTensorV2()\n])\n\nvertical_bar_tfms_list = [\n    A.LongestMaxSize(always_apply=False, p=1, max_size=896, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=896, min_width=896, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0)\n]\n\ninfer_vertical_bar = ObjectDetectorInference(checkpoint_file='/kaggle/input/bmge-extra-data-2/vertical_bar_8.pth', \n                                             transforms=vertical_bar_transforms, \n                                             tfms_list=vertical_bar_tfms_list,\n                                             w_after=896,\n                                             h_after=896)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:57:39.353513Z","iopub.execute_input":"2023-06-18T11:57:39.353943Z","iopub.status.idle":"2023-06-18T11:57:42.068765Z","shell.execute_reply.started":"2023-06-18T11:57:39.353902Z","shell.execute_reply":"2023-06-18T11:57:42.067548Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dot_scatter_transforms = A.Compose([\n    A.LongestMaxSize(always_apply=False, p=1, max_size=640, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=640, min_width=640, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0), \n    ToTensorV2()\n])\n\ndot_scatter_tfms_list = [\n    A.LongestMaxSize(always_apply=False, p=1, max_size=640, interpolation=1),\n    A.PadIfNeeded(always_apply=False, p=1.0, min_height=640, min_width=640, pad_height_divisor=None, pad_width_divisor=None, border_mode=0, value=[124, 116, 104], mask_value=None),\n    A.Normalize(always_apply=False, p=1.0, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0)\n]\n\ninfer_dot_scatter = ObjectDetectorInference(checkpoint_file='/kaggle/input/b-m-g-a-extra-data/dot_scatter__8.pth', transforms=dot_scatter_transforms, tfms_list=dot_scatter_tfms_list, h_after=640, w_after=640)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:57:42.073872Z","iopub.execute_input":"2023-06-18T11:57:42.074847Z","iopub.status.idle":"2023-06-18T11:57:45.014545Z","shell.execute_reply.started":"2023-06-18T11:57:42.074801Z","shell.execute_reply":"2023-06-18T11:57:45.013167Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from neuspell import BertChecker\nchecker = BertChecker()\nchecker.from_pretrained(\"/kaggle/input/bert-spell-model/\")","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:57:45.020097Z","iopub.execute_input":"2023-06-18T11:57:45.020541Z","iopub.status.idle":"2023-06-18T11:58:10.786256Z","shell.execute_reply.started":"2023-06-18T11:57:45.020498Z","shell.execute_reply":"2023-06-18T11:58:10.785015Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport shutil\nimport statistics\nimport sys\nimport re\nfrom collections import namedtuple\nfrom types import MethodType\n\nimport albumentations as A\nimport cv2\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom PIL import Image, ImageEnhance\n\n\nclass VerticalBar:\n    def __init__(self,\n                 rotation_learn,\n                 infer_tick_mark,\n                 infer_label_y,\n                 infer_label_x,\n                 infer_vertical_bar,\n                 trocr_processor,\n                 trocr_model,\n                 paddle_ocr,\n                 donut_processor,\n                 donut_model):\n\n        # Classifier Models\n        self.rotation_learn = rotation_learn\n\n        # Object detection model object\n        self.infer_tick_mark = infer_tick_mark\n        self.infer_label_y = infer_label_y\n        self.infer_label_x = infer_label_x\n        self.infer_vertical_bar = infer_vertical_bar\n\n        # PaddleOcr and TrOCR\n        self.trocr_processor = trocr_processor\n        self.trocr_model = trocr_model\n        self.paddle_ocr = paddle_ocr\n\n        # Donut Models\n        self.donut_processor = donut_processor\n        self.donut_model = donut_model\n        self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    # TrOCR Inference\n    def _enhance_image(self, img):\n        try:\n            enhancer = ImageEnhance.Contrast(img)\n            img = enhancer.enhance(1.4)\n            enhancer1 = ImageEnhance.Sharpness(img)\n            img = enhancer1.enhance(1.3)\n            width, height = img.size\n            new_width, new_height = width * 2, height * 2\n            resized_img = img.resize((new_width, new_height))\n            return resized_img\n        except:\n            return img\n\n    def paddle_ocr_call(self, snippet):\n        result = self.paddle_ocr.ocr(snippet, cls=False)\n        if len(result[0]) < 1:\n            return [snippet]\n        else:\n            image_list = []\n            polygon_points = []\n            for i in range(len(result[0])):\n                x0, y0, x1, y1, x2, y2, x3, y3 = np.array(\n                    result[0][i][0]).flatten()\n                xmin = min(x0, x1, x2, x3)\n                xmax = max(x0, x1, x2, x3)\n                ymin = min(y0, y1, y2, y3)\n                ymax = max(y0, y1, y2, y3)\n                crop = snippet[int(ymin):int(ymax), int(xmin):int(xmax)]\n                image_list.append(crop)\n        return image_list\n\n    def snippets_to_single_image(self, images):\n        patch_width = 20\n        heights, widths, depth = zip(*(i.shape for i in images))\n        total_width = sum(widths) + patch_width * (len(images) - 1)\n        max_height = max(heights)\n        patch_height = max_height\n        new_im = Image.new('RGB', (total_width, max_height), color='white')\n        x_offset = 0\n        for im in images:\n            im = Image.fromarray(im)\n            box = (x_offset, 0, x_offset + im.size[0], im.size[1])\n            new_im.paste(im, box)\n            x_offset += im.size[0] + patch_width\n        return new_im\n    \n    def fix_rotation(self, image):\n        if 2 * image.width <= image.height:\n            print(\"rotating image\")\n            fixed_image = image.rotate(-90, expand=True)\n            return fixed_image\n        else:\n            # No rotation needed\n            return image\n\n    def trocr_inference(self, img, bboxes, trocr_processor, trocr_model):\n        gen_text = []\n        txt_box = []\n        for box in bboxes:\n            xmin, ymin, xmax, ymax = box\n            crop = img.crop((xmin, ymin, xmax, ymax))\n            crop = self._enhance_image(crop)\n\n            if abs(ymax - ymin) > 18:\n                crop = self.fix_rotation(crop)\n                crop = np.array(crop)\n                image_list = self.paddle_ocr_call(crop)\n                single_img = self.snippets_to_single_image(image_list)\n                pixel_values = self.trocr_processor(\n                    images=single_img, return_tensors=\"pt\").pixel_values\n            else:\n                crop = self.fix_rotation(crop)\n                pixel_values = trocr_processor(\n                    images=crop, return_tensors=\"pt\").pixel_values\n\n            generated_ids = trocr_model.generate(pixel_values.to(\n                'cuda'), max_length=64, num_beams=2, early_stopping=True)\n            generated_text = trocr_processor.batch_decode(\n                generated_ids, skip_special_tokens=True)[0]\n            gen_text.append(generated_text.capitalize())\n            txt_box.append([xmin, ymin, xmax, ymax])\n        return gen_text, txt_box\n    \n    def trocr_inference_y(self, img, bboxes, trocr_processor, trocr_model):\n        gen_text = []\n        txt_box = []\n        for box in bboxes:\n            xmin, ymin, xmax, ymax = box\n            crop = img.crop((xmin, ymin, xmax, ymax))\n            crop = self._enhance_image(crop)\n\n            if abs(ymax - ymin) > 18:\n                crop = np.array(crop)\n                pixel_values = self.trocr_processor(\n                    images=crop, return_tensors=\"pt\").pixel_values\n            else:\n                pixel_values = trocr_processor(\n                    images=crop, return_tensors=\"pt\").pixel_values\n\n            generated_ids = trocr_model.generate(pixel_values.to(\n                'cuda'), max_length=64, num_beams=2, early_stopping=True)\n            generated_text = trocr_processor.batch_decode(\n                generated_ids, skip_special_tokens=True)[0]\n            gen_text.append(generated_text.capitalize())\n            txt_box.append([xmin, ymin, xmax, ymax])\n        return gen_text, txt_box\n\n    # Donut Inference\n    def get_text(self, text):\n        match = re.search(\n            r'<s_X_ticks>(.+?)</s_X_ticks>|<s_X_ticks>(.+?)</s_X_label>', text)\n        if match:\n            result = match.group(1) or match.group(2)\n        else:\n            result = \"a<xt>b<xt>c<xt>d<xt>e\"\n        text_list = [i.strip() for i in result.split(\"<xt>\")]\n        string = \";\".join(text_list)\n        string = re.sub(r'<[^>]*s[^>]*>', '', string)\n\n        return string\n\n    def donut_predict(self, crop):\n        pixel_values = self.donut_processor(\n            crop.convert(\"RGB\"), return_tensors=\"pt\").pixel_values\n        pixel_values = pixel_values.to(device)\n        task_prompt = \"<s>\"\n        decoder_input_ids = self.donut_processor.tokenizer(\n            task_prompt, add_special_tokens=False, return_tensors=\"pt\").input_ids\n        decoder_input_ids = decoder_input_ids.to(device)\n        outputs = self.donut_model.generate(\n            pixel_values,\n            decoder_input_ids=decoder_input_ids,\n            max_length=512,\n            early_stopping=True,\n            pad_token_id=processor.tokenizer.pad_token_id,\n            eos_token_id=processor.tokenizer.eos_token_id,\n            use_cache=True,\n            num_beams=4,\n            bad_words_ids=[[processor.tokenizer.unk_token_id]],\n            return_dict_in_generate=True,\n            no_repeat_ngram_size=3,\n            do_sample=False,\n        )\n        seq = processor.batch_decode(outputs.sequences)[0]\n        seq = seq.replace(processor.tokenizer.eos_token, \"\").replace(\n            processor.tokenizer.pad_token, \"\")\n        seq = seq.replace(\"<one>\", \"1\")\n        seq = re.sub(r\"<.*?>\", \"\", seq, count=1).strip()\n        seq = self.get_text(seq)\n        return seq\n\n    def vertical_group_boxes(self, tick_marks, xtick_boxes, vertical_bar_boxes):\n        \"\"\"\n        The function groups ticks, labels, and bars in an image based on their x-coordinate and removes any\n        duplicates.\n\n        :param tick_marks: A list of tuples, where each tuple contains the x-coordinate of a tick mark and\n        its corresponding y-coordinate\n        :param xtick_boxes: A list of bounding boxes for the tick labels in the image\n        :param vertical_bar_boxes: A list of bounding boxes for the vertical bars in the image\n        \"\"\"\n        def identify_x(tick_boxes):\n            \"\"\"\n            The function identifies the y-coordinate of x-axis.\n            \"\"\"\n            center_points = [(box[1] + box[3]) / 2 for box in tick_boxes]\n            y = np.median(sorted(center_points)[-3:-1])\n            return y\n\n        def calculate_intersection_area(box1, box2):\n            \"\"\"\n            Calculate intersection area between two boxes\n            \"\"\"\n            x1_min, y1_min, x1_max, y1_max = box1\n            x2_min, y2_min, x2_max, y2_max = box2\n\n            intersection_x_min = max(x1_min, x2_min)\n            intersection_y_min = max(y1_min, y2_min)\n            intersection_x_max = min(x1_max, x2_max)\n            intersection_y_max = min(y1_max, y2_max)\n\n            width = max(0, intersection_x_max - intersection_x_min)\n            height = max(0, intersection_y_max - intersection_y_min)\n\n            intersection_area = width * height\n            return intersection_area\n\n        # Check Intersection\n        def check_intersection(box1, box2, tolerance=0.02):\n            x_min1, y_min1, x_max1, y_max1 = box1\n            x_min2, y_min2, x_max2, y_max2 = box2\n            if x_min2 < x_min1 < (1-tolerance)*x_max2:\n                return True\n            if x_min1 < x_min2 < (1-tolerance)*x_max1:\n                return True\n\n        # Filter by bar overlapping\n        def filter_groups_bar_overlaps(groups):\n            boxes = [box for group in groups for box in group.get('bar', [])]\n\n            non_overlapping_boxes = []\n            for box1 in boxes:\n                is_smaller = any(\n                    check_intersection(box1, box2)\n                    and (box1[2] - box1[0]) * (box1[3] - box1[1]) < (box2[2] - box2[0]) * (box2[3] - box2[1])\n                    for box2 in boxes if box1 != box2\n                )\n                if not is_smaller:\n                    non_overlapping_boxes.append(box1)\n\n            filtered_groups = [group for group in groups if group.get(\n                'bar', [])[0] in non_overlapping_boxes]\n\n            return filtered_groups\n\n        # Calculate distance between two points\n        def calculate_distance(point1, point2):\n            x1, y1 = point1\n            x2, y2 = point2\n            return math.sqrt((x2 - x1)**2 + (y2 - y1)**2)\n\n        # Remove overlapping boxes based on a threshold\n        def remove_overlapping_boxes(boxes):\n            # Sort the boxes by their y-coordinate (top-left corner)\n            sorted_boxes = sorted(boxes, key=lambda box: box[1])\n            non_overlapping_boxes = []\n\n            # Iterate over the sorted boxes\n            for box in sorted_boxes:\n                # Check if the current box overlaps with any of the non-overlapping boxes\n                overlaps = False\n                for non_overlapping_box in non_overlapping_boxes:\n                    intersection_area = calculate_intersection_area(\n                        box, non_overlapping_box)\n                    if intersection_area > threshold:\n                        overlaps = True\n                        break\n\n                # If the box does not overlap, add it to the non-overlapping boxes list\n                if not overlaps:\n                    non_overlapping_boxes.append(box)\n\n            return non_overlapping_boxes\n\n        # Remove groups with two or more None values\n        def remove_two_none(groups):\n            i = 0\n            while i < len(groups):\n                group = groups[i]\n                if (group['ticks'] is None and group['label_box'] is None) or \\\n                        (group['ticks'] is None and group['bar'] is None) or \\\n                        (group['label_box'] is None and group['bar'] is None):\n                    groups.pop(i)\n                else:\n                    i += 1\n            return groups\n\n        # Filter groups based on bar width\n        def filter_by_barwidth(groups):\n            try:\n                bar_widths = [abs(bar[2] - bar[0])\n                              for group in groups if group['bar'] is not None for bar in group['bar']]\n                median_bar_width = statistics.median(bar_widths)\n\n                filtered_groups = []\n                for group in groups:\n                    if group['bar'] is None:\n                        filtered_groups.append(group)\n                    else:\n                        group_bar_widths = [abs(bar[2] - bar[0])\n                                            for bar in group['bar']]\n                        if all(0.8 * median_bar_width <= abs(width) <= 1.25 * median_bar_width for width in group_bar_widths):\n                            filtered_groups.append(group)\n                return filtered_groups\n            except Exception as e:\n                print('Filter by Barwidth:: ', groups)\n                return groups\n\n        # Remove duplicate groups with same labels\n        def remove_duplicate_label_boxes(groups):\n            unique_label_boxes = []\n            groups_filtered = []\n\n            for group in groups:\n                if group['label_box'] is None:\n                    groups_filtered.append(group)\n                else:\n                    label_box = tuple(group['label_box'])\n                    # Check if the label_box is already seen\n                    if label_box not in unique_label_boxes:\n                        unique_label_boxes.append(label_box)\n                        groups_filtered.append(group)\n            return groups_filtered\n\n        # Remove duplicate groups with same bars\n        def remove_duplicate_bar(groups):\n            unique_bar_boxes = []\n            groups_filtered = []\n\n            for group in groups:\n                if group['bar'] is None:\n                    groups_filtered.append(group)\n                else:\n                    bar_box = tuple(group['bar'])\n                    # Check if the bar_box is already seen\n                    if bar_box not in unique_bar_boxes:\n                        unique_bar_boxes.append(bar_box)\n                        groups_filtered.append(group)\n            return groups_filtered\n\n        def final_filter_groups(groups):\n            groups_with_same_bars = []\n\n            for i, group in enumerate(groups):\n                current_bars = group['bar']\n                for j in range(i + 1, len(groups)):\n                    next_bars = groups[j]['bar']\n                    if (current_bars is not None) and (next_bars is not None):\n                        if current_bars == next_bars:\n                            groups_with_same_bars.append((i, j))\n\n            if len(groups_with_same_bars) > 0:\n                for group_indices in groups_with_same_bars:\n                    index1, index2 = group_indices\n                    group1, group2 = groups[index1], groups[index2]\n                    if group1['label_box'][0] is None:\n                        groups.pop(index1)\n                    elif group2['label_box'][0] is None:\n                        groups.pop(index2)\n                    else:\n                        bar1, label_box1 = group1['bar'][0], group1['label_box'][0]\n                        bar2, label_box2 = group2['bar'][0], group2['label_box'][0]\n                        distance1 = calculate_distance(\n                            bar1[:2], label_box1[:2])\n                        distance2 = calculate_distance(\n                            bar2[:2], label_box2[:2])\n                        groups.pop(\n                            index1) if distance1 < distance2 else groups.pop(index2)\n            return groups\n\n        # Initialize groups, tick_boxes, and x-axis position\n        groups = []\n        tick_boxes = []\n        for i in tick_marks:\n            x_axis = identify_x(tick_marks)\n            if i[3] > x_axis:\n                tick_boxes.append(i)\n\n        # Group ticks\n        for label_box in xtick_boxes:\n            group = {'ticks': None, 'label_box': None, 'bar': None}\n\n            label_x_min, label_y_min, label_x_max, label_y_max = label_box\n            label_center = (label_x_min + label_x_max) / 2\n\n            for tick_box in tick_boxes:\n                tick_x_min, tick_y_min, tick_x_max, tick_y_max = tick_box\n                tick_center = (tick_x_min + tick_x_max) / 2\n\n                if label_x_min <= tick_center <= label_x_max:\n                    if group['ticks'] is None:\n                        group['ticks'] = [tick_box]\n                        break\n\n            for idx, bar_box in enumerate(vertical_bar_boxes):\n                bar_x_min, bar_y_min, bar_x_max, bar_y_max = bar_box\n                bar_center = (bar_x_min + bar_x_max) / 2\n\n                if label_x_min <= bar_center <= label_x_max:\n                    if group['label_box'] is None:\n                        group['label_box'] = [label_box]\n                    if group['bar'] is None:\n                        group['bar'] = [bar_box]\n                    break\n\n            # Additional checks\n            if group['label_box'] is not None and group['bar'] is not None:\n                groups.append(group)\n\n        # Handle ungrouped bar boxes\n        ungrouped_bars = [bar_box for bar_box in vertical_bar_boxes if\n                          [bar_box] not in [group['bar'] for group in groups]]\n        ungrouped_labels = [label_box for label_box in xtick_boxes if\n                            [label_box] not in [group['label_box'] for group in groups]]\n        ungrouped_ticks = [tick_box for tick_box in tick_boxes if\n                           [tick_box] not in [group['ticks'] for group in groups]]\n\n        for bar_box in ungrouped_bars:\n            group = {'ticks': None, 'label_box': None, 'bar': None}\n            bar_x_min, bar_y_min, bar_x_max, bar_y_max = bar_box\n            for tick_box in ungrouped_ticks:\n                tick_x_min, tick_y_min, tick_x_max, tick_y_max = tick_box\n                tick_center = (tick_x_min + tick_x_max) / 2\n                if bar_x_min <= tick_center <= bar_x_max:\n                    group['bar'] = [bar_box]\n                    group['ticks'] = [tick_box]\n\n            for label_box in ungrouped_labels:\n                label_x_min, label_y_min, label_x_max, label_y_max = label_box\n                label_center = (label_x_min + label_x_max) / 2\n                if bar_x_min <= label_center <= bar_x_max:\n                    if group['bar'] is None:\n                        group['bar'] = [bar_box]\n                        group['label_box'] = [label_box]\n                    break\n\n            groups.append(group)\n\n        ungrouped_labels = [label_box for label_box in xtick_boxes if\n                            [label_box] not in [group['label_box'] for group in groups]]\n        ungrouped_ticks = [tick_box for tick_box in tick_boxes if\n                           [tick_box] not in [group['ticks'] for group in groups]]\n\n        for tick_box in ungrouped_ticks:\n            group = {'ticks': None, 'label_box': None, 'bar': None}\n            tick_x_min, tick_y_min, tick_x_max, tick_y_max = tick_box\n            tick_center = (tick_x_min + tick_x_max) / 2\n            for label_box in ungrouped_labels:\n                label_x_min, label_y_min, label_x_max, label_y_max = label_box\n                label_center = (label_x_min + label_x_max) / 2\n                if label_x_min <= tick_center <= label_x_max:\n                    group['ticks'] = [tick_box]\n                    group['label_box'] = [label_box]\n                    break\n\n            groups.append(group)\n\n        for i in groups:\n            if i['label_box'] != None:\n                if len(i['label_box']) > 1:\n                    i['label_box'] = [\n                        remove_overlapping_boxes(i['label_box'])[0]]\n            if i['ticks'] != None:\n                if len(i['ticks']) > 1:\n                    i['ticks'] = [remove_overlapping_boxes(i['ticks'])[0]]\n            if i['bar'] != None:\n                if len(i['bar']) > 1 or i['bar'] != None:\n                    i['bar'] = [remove_overlapping_boxes(i['bar'])[0]]\n\n        groups = remove_two_none(groups)\n        groups = filter_by_barwidth(groups)\n        groups = remove_duplicate_label_boxes(groups)\n        groups = remove_duplicate_bar(groups)\n\n        height_list = []\n        for group in groups:\n            label_box = group['label_box']\n            if group['label_box'] is not None:\n                for box in label_box:\n                    height_list.append(box[3] - box[1])\n        try:\n            height_avg = statistics.median(height_list)\n        except:\n            height_avg = 6\n\n        for item in groups:\n            if item['bar'] is None:\n                try:\n                    item['bar'] = [[item['label_box'][0][0], item['label_box'][0][1] - 1,\n                                    item['label_box'][0][2] - 1, item['ticks'][0][1] - 1]]\n                except Exception as e:\n                    item['bar'] = [[item['ticks'][0][0], item['ticks'][0][1] - 1,\n                                    item['ticks'][0][2] - 2, x_axis - 1]]\n                    print(e)\n\n            if item['label_box'] is None:\n                try:\n                    min_xmin = min(item['bar'][0][0], item['ticks'][0][0])\n                    min_ymin = item['ticks'][0][3]\n                    max_xmax = max(item['bar'][0][2], item['ticks'][0][2])\n\n                    item['label_box'] = [[min_xmin, item['bar'][0]\n                                          [3], max_xmax, item['bar'][0][3]+int(height_avg)]]\n                except Exception as e:\n                    raise (e)\n                    item['label_box'] = [[0, 0, 0, 0]]\n        old_len = len(groups)\n        groups = filter_groups_bar_overlaps(groups)\n        if old_len != len(groups):\n            print(\"overlapping filter applied\")\n        groups = final_filter_groups(groups)\n\n        for group in groups:\n            if group['ticks'] is None:\n                group['ticks'] = [[1, 1, 1, 1]]\n        try:\n            label_box_list = sum((group['label_box'] for group in groups), [])\n            bar_list = sum((group['bar'] for group in groups), [])\n            ticks_list = sum((group['ticks'] for group in groups), [])\n        except Exception as e:\n            print(e)\n            print(groups)\n        return label_box_list, bar_list, ticks_list, groups\n\n    def classify_image(self, vertical_bar_boxes, tick_mark_boxes, label_boxes):\n        def identify_y(boxes):\n            # Identify the tick position based on the center points of the boxes\n            center_points = [(box[0] + box[2]) / 2 for box in boxes]\n            x = np.median(sorted(center_points)[0:3])\n            return x\n\n        def identify_x(boxes):\n            # Identify the tick position based on the center points of the boxes\n            center_points = [(box[1] + box[3]) / 2 for box in boxes]\n            y = np.median(sorted(center_points)[-3:-1])\n            return y\n        \n        def are_bars_closely_arranged(bar_boxes, threshold):\n            bar_boxes = sorted(bar_boxes, key=lambda box: box[0])  # Sort by x-coordinate\n            for i in range(len(bar_boxes) - 1):\n                current_box = bar_boxes[i]\n                next_box = bar_boxes[i + 1]\n                if next_box[0] - current_box[2] > threshold:  # Check if the gap exceeds the threshold\n                    return False\n            return True\n\n        tick_boxes = []\n        aligned_ticks = 0\n        y_mark = identify_x(tick_mark_boxes)\n        x_mark = identify_y(tick_mark_boxes)\n        for tick in tick_mark_boxes:\n            if tick[3] > y_mark and tick[0] > x_mark:\n                tick_boxes.append(tick)\n\n        total_ticks = len(tick_boxes)\n\n        vertical_bar_widths = []\n        tick_box_widths = []\n\n        for box in vertical_bar_boxes:\n            width = box[2] - box[0]\n            vertical_bar_widths.append(width)\n\n        for box in tick_boxes:\n            width = box[2] - box[0]\n            tick_box_widths.append(width)\n        try:\n            median_vertical_bar_width = statistics.median(vertical_bar_widths)\n        except:\n            median_vertical_bar_width = 0\n        try:\n            median_tick_box_width = statistics.median(tick_box_widths)\n        except:\n            median_tick_box_width = 0\n            \n        if median_vertical_bar_width < median_tick_box_width:\n            return 'vertical_bar'\n\n        for bar in vertical_bar_boxes:\n            for tick in tick_boxes:\n                tick_x_min, tick_y_min, tick_x_max, tick_y_max = tick\n                if tick_x_min <= bar[2] <= tick_x_max:\n                    aligned_ticks += 1\n\n        alignment_percentage = (aligned_ticks / total_ticks) * 100\n\n        if alignment_percentage >= 60:\n            if are_bars_closely_arranged(vertical_bar_boxes, 5):\n                print(\"predicted histogram\")\n                return 'histogram'\n            else: \n                print('bars not closely arranged so its vertical')\n                return 'vertical_bar'\n        else:\n            return 'vertical_bar'\n    \n    def remove_abnormally_wide_bars(self, bar_boxes, threshold):\n            def calculate_width(box):\n                x_min, y_min, x_max, y_max = box\n                width = x_max - x_min\n                return width\n            \n            widths = [calculate_width(box) for box in bar_boxes]\n            try:\n                median_width = statistics.median(widths)\n            except:\n                median_width = 1\n            filtered_boxes = [box for box in bar_boxes if calculate_width(box) <= median_width * threshold]\n            if len(bar_boxes) != len(filtered_boxes):\n                print('wide_boxes_filter_applied')\n            return filtered_boxes\n\n    def histogram_group_boxes(self, vertical_bar_boxes, tick_boxes, xtick_boxes):\n        def remove_duplicates(lst):\n            unique_list = []\n            seen_items = []\n            for item in lst:\n                if item not in seen_items:\n                    unique_list.append(item)\n                    seen_items.append(item)\n            return unique_list\n\n        threshold = 1.8  # Adjust the threshold value as per your requirement\n        vertical_bar_boxes = self.remove_abnormally_wide_bars(vertical_bar_boxes, threshold)\n        \n        groups = []\n        for bar in vertical_bar_boxes:\n            group = {'ticks': [], 'label_box': [], 'bar': []}\n            for tick_box in tick_boxes:\n                tick_x_min, tick_y_min, tick_x_max, tick_y_max = tick_box\n                if tick_x_min <= bar[2] <= tick_x_max or tick_x_min <= bar[0] <= tick_x_max:\n                    group['ticks'].append(tick_box)\n                    group['bar'].append(bar)\n                    break\n            for label in xtick_boxes:\n                label_x_min, label_y_min, label_x_max, label_y_max = label\n                if label_x_min <= bar[2] <= label_x_max or label_x_min <= bar[0] <= label_x_max:\n                    group['label_box'].append(label)\n            groups.append(group)\n\n        label_box_list = remove_duplicates(sum((group['label_box'] for group in groups), []))\n        bar_list = remove_duplicates(sum((group['bar'] for group in groups), []))\n        ticks_list = remove_duplicates(sum((group['ticks'] for group in groups), []))\n\n        return label_box_list, bar_list, ticks_list\n\n    def fix_xseries_length(self, x_preds, y_preds):\n        \"\"\"\n        The function `fix_xseries_length` takes two lists as input and ensures that they have the same\n        length by either extending the shorter list with its last value or adding zeros to the end of the\n        longer list.\n        \"\"\"\n        if len(x_preds) < len(y_preds):\n            missing_count = len(y_preds) - len(x_preds)\n            x_preds.extend([x_preds[-1] for i in range(missing_count)])\n        if len(x_preds) > len(y_preds):\n            missing_count = len(y_preds) - len(x_preds)\n            y_preds.extend([str(0) for i in range(missing_count)])\n        return x_preds, y_preds\n\n    def is_convertible_to_float(self, string):\n        try:\n            float(string)\n            return True\n        except ValueError:\n            return False\n\n    def replace_k_with_thousands(self, text):\n        return re.sub(r'(\\d)k', r'\\g<1>000', text, flags=re.IGNORECASE)\n\n    def remove_special_char(self, text):\n        return re.sub(r'[^a-zA-Z0-9.\\s-]', '', text)\n\n    def postprocess_text(self, text):\n        text = self.replace_k_with_thousands(text)\n        text = self.remove_special_char(text)\n        return text\n\n    def get_yscale(self, img, y_label_boxes):\n        try:\n            try:\n                y_label_boxes = y_label_boxes[0:5]\n            except:\n                y_label_boxes = y_label_boxes[0]\n            y_preds, y_pred_boxes = self.trocr_inference_y(\n                img, y_label_boxes, self.trocr_processor, self.trocr_model)\n            y_preds = [self.postprocess_text(pred) for pred in y_preds]\n            selected_box = [(float(output), bbox) for output, bbox in zip(\n                y_preds, y_pred_boxes) if self.is_convertible_to_float(output)][0:2]\n            selected_box = sorted(selected_box, key=lambda x: x[1][1])\n            y_diff = selected_box[0][0] - selected_box[1][0]\n            pixel_diff = abs(selected_box[0][1][1] - selected_box[1][1][1])\n            y_scale = y_diff/pixel_diff\n            return y_scale\n        except Exception as e:\n            print(f\"error in Vertical_bar get_yscale {e}\")\n\n    def get_y_preds(self, vertical_bar_boxes, ymark_value, y_scale):\n        \"\"\"\n        This function takes in vertical bar boxes, a ymark value, and a y scale, and returns a list of\n        interpolated values based on the difference between the ymark value and the top of each vertical\n        bar.\n        \"\"\"\n        interpolated_values = []\n        try:\n            minimum_y_value = float(minimum_y_value[0])\n        except:\n            try:\n                minimum_y_value = float(minimum_y_value[1])\n            except:\n                minimum_y_value = 0\n        vertical_bars = sorted(vertical_bar_boxes, key=lambda x: x[0])\n        for bar in vertical_bars:\n            bar_top = bar[1]\n            interpolated_value = (ymark_value - bar_top) * y_scale\n            if self.is_convertible_to_float(interpolated_value):\n                interpolated_values.append(str(interpolated_value))\n        return interpolated_values\n\n    def identify_xaxis(self, tick_boxes):\n        \"\"\"\n        The function identifies the y-coordinate of x-axis.\n        \"\"\"\n        center_points = [(box[1] + box[3]) / 2 for box in tick_boxes]\n        y = np.median(sorted(center_points)[-3:-1])\n        return y\n    \n    def final_text_data(self, data):\n        cleaned_data = []\n        for item in data:\n            cleaned_item = re.sub(r'\\s*-\\s*', '-', item)\n            cleaned_item = re.sub(r'\\s*\\.\\s*', '.', cleaned_item)\n            cleaned_data.append(cleaned_item)\n        return cleaned_data\n    \n    def get_predictions(self, image_path):\n        img = Image.open(image_path).convert(\"RGB\")\n        width, height = img.size\n        tick_boxes, tick_scores = self.infer_tick_mark.get_bboxes(img)\n        ymark = self.identify_xaxis(tick_boxes)\n        vertical_bar_boxes, vertical_bar_scores = self.infer_vertical_bar.get_bboxes(\n            img)\n        y_label_boxes, y_label_scores = self.infer_label_y.get_bboxes(img)\n        y_scale = self.get_yscale(img, y_label_boxes)\n        with self.rotation_learn.no_bar():\n            rotation_status = self.rotation_learn.predict(image_path)\n        if rotation_status[0] == \"rotated\":\n            try:\n                crop = img.crop((0, ymark, width, height))\n            except Exception as e:\n                crop = img\n            \n            x_preds = self.donut_predict(crop)\n            x_preds = [self.postprocess_text(pred)\n                       for pred in x_preds.split(';')]\n            vertical_bar_boxes= self.remove_abnormally_wide_bars(vertical_bar_boxes, 2)\n#             display(self.infer_label_x.get_img_with_bbs(img, vertical_bar_boxes+tick_boxes))\n\n            y_preds = self.get_y_preds(vertical_bar_boxes, ymark, y_scale)\n            x_preds, y_preds = self.fix_xseries_length(x_preds, y_preds)\n\n        else:\n            x_label_boxes, x_label_scores = self.infer_label_x.get_bboxes(img)\n            bar_type = self.classify_image(\n                vertical_bar_boxes, tick_boxes, x_label_boxes)\n            if bar_type == 'histogram':\n                x_label_box_list, bar_list, ticks_list = self.histogram_group_boxes(tick_boxes=tick_boxes,\n                                                                                    xtick_boxes=x_label_boxes,\n                                                                                    vertical_bar_boxes=vertical_bar_boxes)\n                # display(self.infer_label_x.get_img_with_bbs(img, x_label_box_list+bar_list+ticks_list))\n            else:\n                x_label_box_list, bar_list, ticks_list, grps = self.vertical_group_boxes(tick_marks=tick_boxes,\n                                                                                         xtick_boxes=x_label_boxes,\n                                                                                          vertical_bar_boxes=vertical_bar_boxes)\n                \n                # display(self.infer_label_x.get_img_with_bbs(img, x_label_box_list+bar_list+ticks_list))\n\n\n            label_boxes = sorted(x_label_box_list, key=lambda box: box[0])\n            tick_mark_boxes = sorted(ticks_list, key=lambda box: box[0])\n            vertical_bar_boxes = sorted(bar_list, key=lambda box: box[0])\n\n            x_preds, x_pred_boxes = self.trocr_inference(\n                img, label_boxes, self.trocr_processor, self.trocr_model)\n            x_preds = [self.postprocess_text(pred) for pred in x_preds]\n            y_preds = self.get_y_preds(vertical_bar_boxes, ymark, y_scale)\n            x_preds, y_preds = self.fix_xseries_length(x_preds, y_preds)\n\n        x_preds = [value if value != '' else 'abc' for value in x_preds]\n        y_preds = ['0' if float(value) < 0 else value for value in y_preds]\n        x_preds = checker.correct_strings(x_preds)\n        x_preds = self.final_text_data(x_preds)\n\n        return ';'.join(x_preds), ';'.join(y_preds)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:10.792697Z","iopub.execute_input":"2023-06-18T11:58:10.793135Z","iopub.status.idle":"2023-06-18T11:58:11.136430Z","shell.execute_reply.started":"2023-06-18T11:58:10.793092Z","shell.execute_reply":"2023-06-18T11:58:11.135235Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from PIL import Image, ImageEnhance\nimport numpy as np\nimport torch\nimport re\n\nclass HorizontalBar:\n    def __init__(self, \n                 infer_tick_mark, \n                 infer_label_y, \n                 infer_label_x, \n                 infer_horizontal_bar,\n                 trocr_processor,\n                 troocr_model,\n                 paddle_ocr,\n                 donut_processor,\n                 donut_model\n                 ):\n\n        ## Object detection model object\n        self.infer_tick_mark = infer_tick_mark\n        self.infer_label = infer_label_y\n        self.infer_label_x = infer_label_x\n        self.infer_horizontal_bar = infer_horizontal_bar\n\n        ## paddle_ocr and TrOCR\n        self.trocr_processor = trocr_processor\n        self.trocr_model = troocr_model\n        self.paddle_ocr = paddle_ocr\n        ## Donut Models\n        self.donut_processor = donut_processor\n        self.donut_model = donut_model\n        self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    def get_predictions(self, img_path):\n            img = Image.open(img_path)\n            img_arr = np.array(img)\n\n            # Get bounding boxes for tick marks, labels, and x-axis tick labels\n            tick_mark_boxes, tick_mark_scores = self.infer_tick_mark.get_bboxes(img)\n            label_boxes, label_boxes_scores = self.infer_label.get_bboxes(img)\n            x_tick_label_boxes, x_tick_label_scores = self.infer_label_x.get_bboxes(img)\n            \n            # Apply non-maximum suppression to suppress overlapping bounding boxes\n            tick_mark_boxes_suppressed = self.non_max_suppression_fast(np.array(tick_mark_boxes), 0.1)\n            x_tick_label_boxes_suppressed = self.non_max_suppression_fast(np.array(x_tick_label_boxes), 0.1)\n\n            horizontal_bar_boxes, horizontal_bar_scores = self.infer_horizontal_bar.get_bboxes(img)\n            \n#             display(img, horizontal_bar_boxes)\n            x_axis = self.identify_x(tick_mark_boxes_suppressed)\n            y_axis = self.identify_y(tick_mark_boxes_suppressed)\n\n            tick_mark_boxes_suppressed = self.filter_tick_marks(tick_mark_boxes_suppressed, y_axis, x_axis)\n            label_boxes = self.filter_labels(label_boxes, y_axis, x_axis)\n            x_tick_label_boxes_suppressed = self.filter_labels(x_tick_label_boxes_suppressed, y_axis, x_axis)\n\n#             display(img, horizontal_bar_boxes + tick_mark_boxes_suppressed + x_tick_label_boxes_suppressed)\n\n            # Get the image height\n            image_width = img.size[0]\n            image_height = img.size[1]\n\n            ## For Y \n            if len(horizontal_bar_boxes) > 1:\n                # If there are more than one horizontal bar boxes detected\n\n                # Sort label boxes based on their y-coordinate using horizontal bar boxes as reference\n                sorted_bboxes_label, merged_boxes_label = self.get_y_tick_label_bboxes(img, label_boxes, sorted(horizontal_bar_boxes, key=lambda x: x[1]))\n\n                bbox_snippets = []\n\n                # Extract snippets from the merged label boxes for further processing\n                for each_bbox in merged_boxes_label:\n                    x1, y1, x2, y2 = each_bbox\n                    snippet = img_arr[y1:y2, x1:x2]\n\n                    if abs(y1 - y2) > 18:\n                        image_list = self.paddle_ocr_call(snippet)\n\n                        single_img = self.snippets_to_single_image(image_list)\n                        pixel_values = self.trocr_processor(images=single_img, return_tensors=\"pt\").pixel_values\n                        bbox_snippets.append(pixel_values)\n                    else:\n                        pixel_values = self.trocr_processor(images=snippet, return_tensors=\"pt\").pixel_values\n                        bbox_snippets.append(pixel_values)\n\n                # Stack the snippets into a batch tensor\n                batch = torch.stack(bbox_snippets).squeeze(1)\n\n                # Perform OCR to predict text from the snippets\n                predicted_text = self.TrOCR_predict(batch, self.trocr_processor, self.trocr_model)\n            else:\n                predicted_text = ['a', 'b', 'c']\n\n            clean_texts = []\n\n            if len(predicted_text) < 1:\n                # If no text is predicted, fallback to default labels\n                clean_texts = ['a', 'b', 'c']\n\n            for j in predicted_text:\n                if self.contains_only_special_characters(j):\n                    # Remove predicted text if it contains only special characters\n                    predicted_text.remove(j)\n                else:\n                    clean_texts.append(j.capitalize())\n\n            if len(clean_texts) < 1:\n                # If no clean labels are obtained, use default labels\n                text = \"a;b;c\"\n            else:\n                # Join the clean labels into a single string separated by ';'\n                text = ';'.join(clean_texts)\n            \n            x_tick_label_boxes_suppressed = sorted(x_tick_label_boxes_suppressed, key=lambda x: x[0])\n            if len(x_tick_label_boxes_suppressed) >= 2:                \n                ## For X\n                if len(x_tick_label_boxes_suppressed) == 2:\n                    label_boxes = x_tick_label_boxes_suppressed\n                else:\n                    label_boxes = x_tick_label_boxes_suppressed[1:3]\n\n                bbox_snippets = [] \n                for each_bbox in label_boxes:\n                    x1, y1, x2, y2 = each_bbox\n                    snippet = img_arr[y1:y2, x1:x2]\n                    pixel_values = self.trocr_processor(images=snippet, return_tensors=\"pt\").pixel_values\n                    bbox_snippets.append(pixel_values)\n    \n                # Stack the snippets into a batch tensor\n                batch = torch.stack(bbox_snippets).squeeze(1)\n\n                # Perform OCR to predict text from the snippets\n                predicted_text = self.TrOCR_predict(batch, self.trocr_processor, self.trocr_model)\n                predicted_text = [self.remove_special_characters(i) for i in predicted_text]\n                if not False in [self.is_convertible_to_float(i) for i in predicted_text]:    \n\n                    try:\n                        value_difference = self.remove_special_characters_and_calculate_difference(predicted_text)\n                        pixel_difference = (x_tick_label_boxes_suppressed[1][0] + x_tick_label_boxes_suppressed[1][2])/2 - (x_tick_label_boxes_suppressed[0][0] + x_tick_label_boxes_suppressed[0][2])/2\n\n                        per_pixel_values = value_difference/pixel_difference\n\n                        predicted_values = []\n                        for i in range(len(merged_boxes_label)):\n                            label_box = merged_boxes_label[i]\n                            horizontal_bar_boxes = sorted(horizontal_bar_boxes, key=lambda x: x[1])\n                            distance_difference = [self.get_distance_label_bar(label_box, horizontal_bar_box) for horizontal_bar_box in horizontal_bar_boxes]\n\n                            # Find the index of the horizontal bar with the closest vertical position to the label\n\n                            if distance_difference[np.argmin(distance_difference)] < 8:\n                                box_given = horizontal_bar_boxes[np.argmin(distance_difference)]\n                                predicted_values.append(str((box_given[2] - y_axis) * per_pixel_values))\n                            else:\n                                predicted_values.append(str(0))\n                    except:\n                        i = 0\n                        for j in range(len(merged_boxes_label)):\n                            predicted_values.append(i)\n                            i += 10\n                else:\n                    i = 0\n                    for j in range(len(merged_boxes_label)):\n                        predicted_values.append(i)\n                        i += 10\n        \n#             text = text.split(';')\n#             print(\"Before: \", text)\n\n#             text = checker.correct_strings(text)\n#             text = self.final_text_data(text)\n            \n#             print(\"After: \", text)\n#             print(text, ';'.join(predicted_values))\n            return text, ';'.join(predicted_values)\n    \n    def final_text_data(self, data):\n        cleaned_data = []\n        for item in data:\n            cleaned_item = re.sub(r'\\s*-\\s*', '-', item)\n            cleaned_item = re.sub(r'\\s*\\.\\s*', '.', cleaned_item)\n            cleaned_data.append(cleaned_item)\n        return cleaned_data\n    def non_max_suppression_fast(self, boxes, overlapThresh):\n        # if there are no boxes, return an empty list\n        if len(boxes) == 0:\n            return []\n\n        # if the bounding boxes integers, convert them to floats --\n        # this is important since we'll be doing a bunch of divisions\n        if boxes.dtype.kind == \"i\":\n            boxes = boxes.astype(\"float\")\n        #  \n        # initialize the list of picked indexes   \n        pick = []\n\n        # grab the coordinates of the bounding boxes\n        x1 = boxes[:,0]\n        y1 = boxes[:,1]\n        x2 = boxes[:,2]\n        y2 = boxes[:,3]\n\n        # compute the area of the bounding boxes and sort the bounding\n        # boxes by the bottom-right y-coordinate of the bounding box\n        area = (x2 - x1 + 1) * (y2 - y1 + 1)\n        idxs = np.argsort(y2)\n\n        # keep looping while some indexes still remain in the indexes\n        # list\n        while len(idxs) > 0:\n            # grab the last index in the indexes list and add the\n            # index value to the list of picked indexes\n            last = len(idxs) - 1\n            i = idxs[last]\n            pick.append(i)\n\n            # find the largest (x, y) coordinates for the start of\n            # the bounding box and the smallest (x, y) coordinates\n            # for the end of the bounding box\n            xx1 = np.maximum(x1[i], x1[idxs[:last]])\n            yy1 = np.maximum(y1[i], y1[idxs[:last]])\n            xx2 = np.minimum(x2[i], x2[idxs[:last]])\n            yy2 = np.minimum(y2[i], y2[idxs[:last]])\n\n            # compute the width and height of the bounding box\n            w = np.maximum(0, xx2 - xx1 + 1)\n            h = np.maximum(0, yy2 - yy1 + 1)\n\n            # compute the ratio of overlap\n            overlap = (w * h) / area[idxs[:last]]\n\n            # delete all indexes from the index list that have\n            idxs = np.delete(idxs, np.concatenate(([last],\n                np.where(overlap > overlapThresh)[0])))\n\n        # return only the bounding boxes that were picked using the\n        # integer data type\n        return boxes[pick].astype(\"int\")\n\n    def identify_y(self, boxes):\n        # Identify the tick position based on the center points of the boxes\n        center_points = [(box[0] + box[2]) / 2 for box in boxes]\n        x = np.median(sorted(center_points)[0:3])\n        return x\n\n    def identify_x(self, boxes):\n        # Identify the tick position based on the center points of the boxes\n        center_points = [(box[1] + box[3]) / 2 for box in boxes]\n        y = np.median(sorted(center_points)[-3:-1])\n        return y\n\n    def filter_tick_marks(self, boxes, given_x, given_y):\n        filtered_boxes = []\n        for box in boxes:\n            x_min, y_min, x_max, y_max = box\n            if x_min < given_x < x_max or y_min < given_y < y_max:\n                filtered_boxes.append(box)\n        return filtered_boxes\n\n    def filter_labels(self, boxes, given_x, given_y):\n        filtered_boxes = []\n        for box in boxes:\n            x_min, y_min, x_max, y_max = box\n            if given_x + 10 >= x_max or given_y  <= y_max:\n                filtered_boxes.append(box)\n        return filtered_boxes\n\n    def is_square(self, box, tolerance=0.3):\n        x1, y1, x2, y2 = box\n        width = x2 - x1\n        height = y2 - y1\n        return abs(width - height) <= tolerance * max(width, height)\n\n    def identify_ticksv2(self,boxes):\n        center_points = [(box[0] + box[2]) / 2 for box in boxes]\n        x = np.median(sorted(center_points)[-3:-1])\n        return x\n\n    def identify_y(self,boxes):\n        # Identify the tick position based on the center points of the boxes\n        center_points = [(box[0] + box[2]) / 2 for box in boxes]\n        x = np.median(sorted(center_points)[0:3])\n        return x\n\n    def identify_x(self,boxes):\n        # Identify the tick position based on the center points of the boxes\n        center_points = [(box[1] + box[3]) / 2 for box in boxes]\n        y = np.median(sorted(center_points)[-3:-1])\n        return y\n\n    def remove_boxes(self, boxes):\n        x = self.identify_ticksv2(boxes)\n        median = np.median([box[2] - x   for box in boxes])\n        filtered_boxes = [box for box in boxes if abs((box[2] - x)/median) < 1.5]\n        return filtered_boxes\n\n    def is_small(self, box, threshold=20):\n        x1, y1, x2, y2 = box\n        width = x2 - x1\n        height = y2 - y1\n        return width <= threshold and height <= threshold\n\n    def is_overlapping(self, box1, box2, tolerance=0.0):\n        x1, y1, x2, y2 = box1\n        a1, b1, a2, b2 = box2\n        return not (x2 < a1 + tolerance or a2 < x1 + tolerance or y2 < b1 + tolerance or b2 < y1 + tolerance)\n\n    def merge_boxes(self, boxes):\n        x_coordinates = [box[0][0] for box in boxes] + [box[0][2] for box in boxes]\n        y_coordinates = [box[0][1] for box in boxes] + [box[0][3] for box in boxes]\n        return (min(x_coordinates), min(y_coordinates), max(x_coordinates), max(y_coordinates))\n\n    def merge_bbox(self, boxes, mapped_labels):\n        len(boxes), len(mapped_labels)\n        non_overlapping_boxes = []\n        overlapping_boxes = []\n        for i in range(len(boxes)):\n            overlapping = False\n            if not (self.is_square(boxes[i]) or self.is_small(boxes[i])):\n                for j in range(len(boxes)):\n                    if i != j and self.is_overlapping(boxes[i], boxes[j], tolerance=-0.9):\n                        overlapping_boxes.append([\n                            boxes[i],\n                            mapped_labels[i]\n                            ])\n                        overlapping_boxes.append([\n                            boxes[j],\n                            mapped_labels[j]\n                        ])\n                        overlapping = True\n                        break\n            if not overlapping:\n                non_overlapping_boxes.append(boxes[i])\n\n\n        overlapping_groups = []\n        for box in overlapping_boxes:\n            found_group = False\n            for group in overlapping_groups:\n                if any([(self.is_overlapping(box1=box_in_group[0] , box2=box[0], tolerance = -0.9) and box_in_group[1]==box[1]) for box_in_group in group]):\n                    group.append(box)\n                    found_group = True\n                    break\n            if not found_group:\n                overlapping_groups.append([box])\n\n        enclosing_boxes = []\n        for group in overlapping_groups:\n            enclosing_box = self.merge_boxes(group)\n            enclosing_boxes.append(enclosing_box)\n\n        all_boxes = non_overlapping_boxes + enclosing_boxes\n        sorted_bboxes = sorted(all_boxes, key=lambda x: x[1])\n        return sorted_bboxes\n    \n    def map_labels_with_bars(self, img, label_boxes, horizontal_bar_boxes):\n        mapped_labels = []\n        for i in range(len(label_boxes)):\n            label_box = label_boxes[i]\n            label_y_min = label_box[1]\n            label_y_max = label_box[3]\n            \n            label_middle = (label_y_min + label_y_max)/2\n            \n            distance_differnce = []\n            for j in range(len(horizontal_bar_boxes)):\n                horizontal_bar_box = horizontal_bar_boxes[j]\n                horizontal_y_min = horizontal_bar_box[1]\n                horizontal_y_max = horizontal_bar_box[3]\n                horizontal_middle = (horizontal_y_min + horizontal_y_max)/2\n\n                distance_differnce.append(abs(label_middle - horizontal_middle))\n                \n            mapped_labels.append(np.argmin(distance_differnce))\n        return mapped_labels\n    \n    def get_y_tick_label_bboxes(self, img, postprocessed_bboxes, horizontal_bar_boxes):\n        filtered_boxes = self.remove_boxes(postprocessed_bboxes)\n        sorted_bboxes = sorted(filtered_boxes, key=lambda x: x[1])\n        mapped_labels = self.map_labels_with_bars(img, sorted_bboxes, horizontal_bar_boxes)\n        img_arr = np.array(img)\n        merged_bboxes = self.merge_bbox(sorted_bboxes, mapped_labels)\n        return sorted_bboxes, merged_bboxes\n    \n    def paddle_ocr_call(self, snippet):\n        result = self.paddle_ocr.ocr(snippet, cls=False)\n        if len(result[0]) < 1:\n            return [snippet]\n        else:\n            image_list = []\n            polygon_points = []\n            for i in range(len(result[0])):\n                x0, y0, x1, y1, x2, y2, x3, y3 = np.array(\n                    result[0][i][0]).flatten()\n                xmin = min(x0, x1, x2, x3)\n                xmax = max(x0, x1, x2, x3)\n                ymin = min(y0, y1, y2, y3)\n                ymax = max(y0, y1, y2, y3)\n                crop = snippet[int(ymin):int(ymax), int(xmin):int(xmax)]\n                image_list.append(crop)\n        return image_list\n    \n    def snippets_to_single_image(self, images):\n        patch_width = 20\n        heights, widths, depth = zip(*(i.shape for i in images))\n        total_width = sum(widths) + patch_width * (len(images) - 1)\n        max_height = max(heights)\n        patch_height = max_height\n        new_im = Image.new('RGB', (total_width, max_height), color='white')\n        x_offset = 0\n        for im in images:\n            im = Image.fromarray(im)\n            box = (x_offset, 0, x_offset + im.size[0], im.size[1])\n            new_im.paste(im, box)\n            x_offset += im.size[0] + patch_width\n    #     display(new_im)\n        return new_im\n    \n    def TrOCR_predict(self, pixel_values, processor, model):\n        pixel_values = pixel_values.to('cuda')\n        generated_ids = model.generate(pixel_values,output_scores=True,return_dict_in_generate=True, max_length = 64)\n        predicted_text = processor.batch_decode(generated_ids[0], skip_special_tokens=True)\n        return predicted_text\n    \n    def contains_only_special_characters(self, string):\n        pattern = r'^[^\\w\\s]*$'\n        return re.match(pattern, string) is not None\n\n\n    def remove_special_characters_and_calculate_difference(self, strings):\n        values = []\n        \n        # Remove special characters and convert strings to float\n        for string in strings:\n            cleaned_string = re.sub('[^0-9.]', '', string)\n            value = float(cleaned_string)\n            values.append(value)\n        \n        # Calculate the difference between the values\n        difference = values[1] - values[0]\n        \n        return difference\n\n    def is_convertible_to_float(self, string):\n        try:\n            float(string)\n            return True\n        except ValueError:\n            return False\n        \n    def remove_special_characters(self, string):\n        # Remove all non-numeric characters\n        modified_string = re.sub(r'[^0-9]', '', string)\n        return modified_string\n\n    def get_distance_label_bar(self, label_box, horizontal_bar_box):\n        label_y_min = label_box[1]\n        label_y_max = label_box[3]\n\n        # Calculate the middle position of the label box vertically\n        label_middle = (label_y_min + label_y_max) / 2\n        \n        horizontal_y_min = horizontal_bar_box[1]\n        horizontal_y_max = horizontal_bar_box[3]\n                \n        # Calculate the middle position of the horizontal bar box vertically\n        horizontal_middle = (horizontal_y_min + horizontal_y_max) / 2\n        \n        return abs(label_middle - horizontal_middle)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:11.142099Z","iopub.execute_input":"2023-06-18T11:58:11.144841Z","iopub.status.idle":"2023-06-18T11:58:11.276028Z","shell.execute_reply.started":"2023-06-18T11:58:11.144797Z","shell.execute_reply":"2023-06-18T11:58:11.274990Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from PIL import Image, ImageEnhance\nimport numpy as np\nimport torch\nimport re\n\nclass ScatterPlot:\n    def __init__(self, \n                 infer_tick_mark, \n                 infer_label_y, \n                 infer_label_x, \n                 infer_dot_scatter,\n                 trocr_processor,\n                 troocr_model,\n                 donut_processor,\n                 donut_model\n                 ):\n\n        ## Object detection model object\n        self.infer_tick_mark = infer_tick_mark\n        self.infer_label = infer_label_y\n        self.infer_label_x = infer_label_x\n        self.infer_dot_scatter = infer_dot_scatter\n\n        ## paddle_ocr and TrOCR\n        self.trocr_processor = trocr_processor\n        self.trocr_model = troocr_model\n        \n        ## Donut Models\n        self.donut_processor = donut_processor\n        self.donut_model = donut_model\n        self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    def get_predictions(self, img_path):\n        try :\n            img = Image.open(img_path).convert(\"RGB\")\n            img_arr = np.array(img)\n            \n            tick_mark_boxes, tick_mark_boxes_scores = self.infer_tick_mark.get_bboxes(img)\n            label_boxes, label_boxes_scores = self.infer_label.get_bboxes(img)\n\n            label_boxes_suppressed = self.non_max_suppression_fast(np.array(label_boxes), 0.2)\n            x_tick_label_boxes, x_tick_label_boxes_scores = self.infer_label_x.get_bboxes(img)\n\n            label_boxes_suppressed = self.non_max_suppression_fast(np.array(label_boxes), 0.2)\n\n            x_axis = self.identify_x(tick_mark_boxes)\n            y_axis = self.identify_y(tick_mark_boxes)\n\n            scatter_boxes, scatter_boxes_score = self.infer_dot_scatter.get_bboxes(img)\n            \n            x_scale, x_origin, x_origin_label_box = self.get_xscale(img, x_tick_label_boxes)\n            y_scale, y_origin, y_origin_label_box = self.get_yscale(img, label_boxes)\n        \n            x_values = []\n            y_values = []\n            scatter_boxes = sorted(scatter_boxes, key=lambda x: x[0])\n\n            for scatter_box in scatter_boxes:\n                x_middle = (scatter_box[0] + scatter_box[2])/2\n                y_middle = (scatter_box[1] + scatter_box[3])/2\n\n                y_difference = x_axis - y_middle\n                x_difference = x_middle - y_axis\n\n                if x_origin_label_box[2] >= scatter_box[2]:\n                    x_difference_label = x_origin_label_box[2] - y_axis                \n                    x_values.append(str((x_difference/x_difference_label)*x_origin))\n                else:\n                    x_values.append(str((x_difference*x_scale) + x_origin))\n                    \n                if y_origin_label_box[3] >= scatter_box[3]:\n                    y_values.append(str((y_difference*y_scale) + y_origin))\n                else:\n                    y_values.append(str((y_difference*y_scale) + 0))\n\n        except Exception as e:\n            x_values = ['1', '2', '3', '4', '5']\n            y_values = ['1', '2', '3', '4', '5']\n\n            print(f'Error Occured in LinePlot for this file {img_path}')\n            print('Error: ', e)\n        \n        return ';'.join(x_values), ';'.join(y_values)\n\n    def non_max_suppression_fast(self, boxes, overlapThresh):\n        # if there are no boxes, return an empty list\n        if len(boxes) == 0:\n            return []\n\n        # if the bounding boxes integers, convert them to floats --\n        # this is important since we'll be doing a bunch of divisions\n        if boxes.dtype.kind == \"i\":\n            boxes = boxes.astype(\"float\")\n\n        # initialize the list of picked indexes   \n        pick = []\n\n        # grab the coordinates of the bounding boxes\n        x1 = boxes[:,0]\n        y1 = boxes[:,1]\n        x2 = boxes[:,2]\n        y2 = boxes[:,3]\n\n        # compute the area of the bounding boxes and sort the bounding\n        # boxes by the bottom-right y-coordinate of the bounding box\n        area = (x2 - x1 + 1) * (y2 - y1 + 1)\n        idxs = np.argsort(y2)\n\n        # keep looping while some indexes still remain in the indexes\n        # list\n        while len(idxs) > 0:\n            # grab the last index in the indexes list and add the\n            # index value to the list of picked indexes\n            last = len(idxs) - 1\n            i = idxs[last]\n            pick.append(i)\n\n            # find the largest (x, y) coordinates for the start of\n            # the bounding box and the smallest (x, y) coordinates\n            # for the end of the bounding box\n            xx1 = np.maximum(x1[i], x1[idxs[:last]])\n            yy1 = np.maximum(y1[i], y1[idxs[:last]])\n            xx2 = np.minimum(x2[i], x2[idxs[:last]])\n            yy2 = np.minimum(y2[i], y2[idxs[:last]])\n\n            # compute the width and height of the bounding box\n            w = np.maximum(0, xx2 - xx1 + 1)\n            h = np.maximum(0, yy2 - yy1 + 1)\n\n            # compute the ratio of overlap\n            overlap = (w * h) / area[idxs[:last]]\n\n            # delete all indexes from the index list that have\n            idxs = np.delete(idxs, np.concatenate(([last],\n                np.where(overlap > overlapThresh)[0])))\n\n        # return only the bounding boxes that were picked using the\n        # integer data type\n        return boxes[pick].astype(\"int\")\n    \n    def identify_y(self, boxes):\n        # Identify the tick position based on the center points of the boxes\n        center_points = [(box[0] + box[2]) / 2 for box in boxes]\n        x = np.median(sorted(center_points)[0:3])\n        return x\n\n\n    def identify_x(self, boxes):\n        # Identify the tick position based on the center points of the boxes\n        center_points = [(box[1] + box[3]) / 2 for box in boxes]\n        y = np.median(sorted(center_points)[-3:-1])\n        return y\n\n    def get_yscale(self, img, y_label_boxes):\n        try:\n            minimum_y = sorted(y_label_boxes, key=lambda x: x[1], reverse=True)\n            y_preds, y_pred_boxes = self.trocr_inference(img, y_label_boxes, self.trocr_processor, self.trocr_model)\n            y_preds = [self.postprocess_text(pred) for pred in y_preds]\n            \n            selected_box = [(float(output),bbox) for output, bbox in zip(y_preds, y_pred_boxes) if self.is_convertible_to_float(output)]\n            selected_box = sorted(selected_box, key=lambda x: x[1][1], reverse=True)\n            if len(selected_box) > 2:\n                origin_value = selected_box[0][0]\n                origin_label_box = selected_box[0][1]\n                selected_box = selected_box[1:3]\n            elif len(selected_box) == 2:\n                origin_value = selected_box[0][0]\n                origin_label_box = selected_box[0][1]\n                selected_box = selected_box[0:2]\n            else:\n                return 0.5, 10, [1000, 1000, 1000, 1000]\n            \n            y_diff = abs(selected_box[0][0] - selected_box[1][0])\n            pixel_diff = abs(selected_box[0][1][1] - selected_box[1][1][1])\n                \n            y_scale = y_diff/pixel_diff\n            return y_scale, origin_value, origin_label_box\n        except Exception as e:\n            print(y_preds)\n            print(y_pred_boxes)\n            print(e)\n        \n\n    def get_xscale(self, img, x_label_boxes):\n        try:\n            minimum_x = sorted(x_label_boxes, key=lambda x: x[0], reverse=False)            \n            x_preds, x_pred_boxes = self.trocr_inference(img, x_label_boxes, self.trocr_processor, self.trocr_model)\n            x_preds = [self.postprocess_text(pred) for pred in x_preds]\n            selected_box = [(float(output),bbox) for output, bbox in zip(x_preds, x_pred_boxes) if self.is_convertible_to_float(output)]\n            selected_box = sorted(selected_box, key=lambda x: x[1][0])\n            if len(selected_box) > 2:\n                origin_value = selected_box[0][0]\n                origin_label_box = selected_box[0][1]\n                selected_box = selected_box[1:3]\n            elif len(selected_box) == 2:\n                origin_value = selected_box[0][0]\n                selected_box = selected_box[0:2]\n                origin_label_box = selected_box[0][1]\n            else:\n                return 0.5, 10, [1000, 1000, 1000, 1000]\n\n            y_diff = abs(selected_box[0][0] - selected_box[1][0]        )\n            pixel_diff = abs(selected_box[0][1][0] - selected_box[1][1][0])\n            y_scale = y_diff/pixel_diff\n            return y_scale, origin_value, origin_label_box\n        except Exception as e:\n            print(e)\n\n    def postprocess_text(self, text):\n        text = self.replace_k_with_thousands(text)\n        text = self.remove_special_char(text)\n        return text\n\n    def is_convertible_to_float(self, string):\n        try:\n            float(string)\n            return True\n        except ValueError:\n            return False\n\n    def replace_k_with_thousands(self, text):\n        return re.sub(r'(\\d)K', r'\\g<1>000', text)\n\n    def remove_special_char(self,text):\n        return re.sub(r'[^a-zA-Z0-9.\\s-]', '', text)\n    \n\n    def trocr_inference(self, img ,bboxes, trocr_processor, trocr_model):\n        gen_text=[]\n        txt_box=[]\n        for box in bboxes:\n            xmin, ymin, xmax, ymax = box\n            crop = img.crop((xmin, ymin, xmax, ymax))\n            crop = self._enhance_image(crop)\n            pixel_values = trocr_processor(images=crop, return_tensors=\"pt\").pixel_values\n            generated_ids = trocr_model.generate(pixel_values.to('cuda'), max_length=64, num_beams=2, early_stopping=True)\n            generated_text = trocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]\n            gen_text.append(generated_text.capitalize())\n            txt_box.append([xmin, ymin, xmax, ymax])\n        return gen_text, txt_box\n    \n    def _enhance_image(self, img):\n        enhancer = ImageEnhance.Contrast(img)\n        img = enhancer.enhance(1.4)\n        enhancer1 = ImageEnhance.Sharpness(img)\n        img = enhancer1.enhance(1.3)\n        width, height = img.size\n        new_width, new_height = width * 2, height * 2\n        resized_img = img.resize((new_width, new_height))\n        return resized_img","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:11.280943Z","iopub.execute_input":"2023-06-18T11:58:11.283500Z","iopub.status.idle":"2023-06-18T11:58:11.359992Z","shell.execute_reply.started":"2023-06-18T11:58:11.283459Z","shell.execute_reply":"2023-06-18T11:58:11.358972Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"score_threshold = 0.5","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:11.364647Z","iopub.execute_input":"2023-06-18T11:58:11.367182Z","iopub.status.idle":"2023-06-18T11:58:11.373540Z","shell.execute_reply.started":"2023-06-18T11:58:11.367123Z","shell.execute_reply":"2023-06-18T11:58:11.372514Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"rotation_learn = load_learner('/kaggle/input/b-m-g-a-extra-data/rotation_learner.pkl')","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:11.375041Z","iopub.execute_input":"2023-06-18T11:58:11.375878Z","iopub.status.idle":"2023-06-18T11:58:12.501737Z","shell.execute_reply.started":"2023-06-18T11:58:11.375831Z","shell.execute_reply":"2023-06-18T11:58:12.493314Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nlarge_processor = TrOCRProcessor.from_pretrained('/kaggle/input/ocr-model', local_files_only = True)\nl_model = VisionEncoderDecoderModel.from_pretrained('/kaggle/input/ocr-model', local_files_only=True).to(device)\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:12.508091Z","iopub.execute_input":"2023-06-18T11:58:12.508830Z","iopub.status.idle":"2023-06-18T11:58:38.338768Z","shell.execute_reply.started":"2023-06-18T11:58:12.508779Z","shell.execute_reply":"2023-06-18T11:58:38.337522Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"learn = load_learner('/kaggle/input/bmge-extra-data-2/classifier_learner.pkl')","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:38.343860Z","iopub.execute_input":"2023-06-18T11:58:38.346485Z","iopub.status.idle":"2023-06-18T11:58:39.190327Z","shell.execute_reply.started":"2023-06-18T11:58:38.346438Z","shell.execute_reply":"2023-06-18T11:58:39.188939Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.copy(\"/kaggle/input/paddle-wheels/en_PP-OCRv3_det_infer.tar\", \"/kaggle/working/en_PP-OCRv3_det_infer.tar\")\nshutil.copy(\"/kaggle/input/paddle-wheels/ch_ppocr_mobile_v2.0_cls_infer.tar\", \"/kaggle/working/ch_ppocr_mobile_v2.0_cls_infer.tar\")\nshutil.copy(\"/kaggle/input/paddle-wheels/en_PP-OCRv3_rec_infer.tar\", \"/kaggle/working/en_PP-OCRv3_rec_infer.tar\")\n\nimport tarfile\ndef extract(tar_file, path):\n    opened_tar = tarfile.open(tar_file)\n    if tarfile.is_tarfile(tar_file):\n        opened_tar.extractall(path)\n    else:\n        print(\"The tar file you entered is not a tar file\")\n\nextract(\"/kaggle/working/ch_ppocr_mobile_v2.0_cls_infer.tar\" , \"/kaggle/working/\")\nextract(\"/kaggle/working/en_PP-OCRv3_det_infer.tar\" , \"/kaggle/working/\")\nextract(\"/kaggle/working/en_PP-OCRv3_rec_infer.tar\" , \"/kaggle/working/\")","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:39.195586Z","iopub.execute_input":"2023-06-18T11:58:39.198515Z","iopub.status.idle":"2023-06-18T11:58:39.484821Z","shell.execute_reply.started":"2023-06-18T11:58:39.198467Z","shell.execute_reply":"2023-06-18T11:58:39.483515Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---","metadata":{}},{"cell_type":"code","source":"def non_max_suppression_fast(boxes, overlapThresh):\n   # if there are no boxes, return an empty list\n   if len(boxes) == 0:\n      return []\n\n   # if the bounding boxes integers, convert them to floats --\n   # this is important since we'll be doing a bunch of divisions\n   if boxes.dtype.kind == \"i\":\n      boxes = boxes.astype(\"float\")\n#  \n   # initialize the list of picked indexes   \n   pick = []\n\n   # grab the coordinates of the bounding boxes\n   x1 = boxes[:,0]\n   y1 = boxes[:,1]\n   x2 = boxes[:,2]\n   y2 = boxes[:,3]\n\n   # compute the area of the bounding boxes and sort the bounding\n   # boxes by the bottom-right y-coordinate of the bounding box\n   area = (x2 - x1 + 1) * (y2 - y1 + 1)\n   idxs = np.argsort(y2)\n\n   # keep looping while some indexes still remain in the indexes\n   # list\n   while len(idxs) > 0:\n      # grab the last index in the indexes list and add the\n      # index value to the list of picked indexes\n      last = len(idxs) - 1\n      i = idxs[last]\n      pick.append(i)\n\n      # find the largest (x, y) coordinates for the start of\n      # the bounding box and the smallest (x, y) coordinates\n      # for the end of the bounding box\n      xx1 = np.maximum(x1[i], x1[idxs[:last]])\n      yy1 = np.maximum(y1[i], y1[idxs[:last]])\n      xx2 = np.minimum(x2[i], x2[idxs[:last]])\n      yy2 = np.minimum(y2[i], y2[idxs[:last]])\n\n      # compute the width and height of the bounding box\n      w = np.maximum(0, xx2 - xx1 + 1)\n      h = np.maximum(0, yy2 - yy1 + 1)\n\n      # compute the ratio of overlap\n      overlap = (w * h) / area[idxs[:last]]\n\n      # delete all indexes from the index list that have\n      idxs = np.delete(idxs, np.concatenate(([last],\n         np.where(overlap > overlapThresh)[0])))\n\n   # return only the bounding boxes that were picked using the\n   # integer data type\n   return boxes[pick].astype(\"int\")\n","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:39.486601Z","iopub.execute_input":"2023-06-18T11:58:39.486971Z","iopub.status.idle":"2023-06-18T11:58:39.509438Z","shell.execute_reply.started":"2023-06-18T11:58:39.486931Z","shell.execute_reply":"2023-06-18T11:58:39.508261Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def is_square(box, tolerance=0.3):\n    # Check if the box is approximately square within the given tolerance\n    x1, y1, x2, y2 = box\n    width = x2 - x1\n    height = y2 - y1\n    return abs(width - height) <= tolerance * max(width, height)\n\n\ndef identify_y(boxes):\n    # Identify the tick position based on the center points of the boxes\n    center_points = [(box[0] + box[2]) / 2 for box in boxes]\n    x = np.median(sorted(center_points)[0:3])\n    return x\n\n\ndef identify_x(boxes):\n    # Identify the tick position based on the center points of the boxes\n    center_points = [(box[1] + box[3]) / 2 for box in boxes]\n    y = np.median(sorted(center_points)[-3:-1])\n    return y\n\n\ndef remove_boxes(boxes):\n    # Remove boxes based on their distance from the identified tick position\n    y = identify_x(boxes)\n    median = np.median([box[3] - y for box in boxes])\n    filtered_boxes = [box for box in boxes if box[3] > median]\n    return filtered_boxes\n\n\ndef is_small(box, threshold=20):\n    # Check if the box is smaller than the specified threshold\n    x1, y1, x2, y2 = box\n    width = x2 - x1\n    height = y2 - y1\n    return width <= threshold and height <= threshold\n\n\ndef is_overlapping(box1, box2, tolerance=0.0):\n    # Check if two boxes are overlapping within the given tolerance\n    x1, y1, x2, y2 = box1\n    a1, b1, a2, b2 = box2\n    return not (x2 < a1 + tolerance or a2 < x1 + tolerance or y2 < b1 + tolerance or b2 < y1 + tolerance)\n\n\ndef merge_boxes(boxes):\n    # Merge a list of boxes by finding the minimum and maximum coordinates\n    x_coordinates = [box[0] for box in boxes] + [box[2] for box in boxes]\n    y_coordinates = [box[1] for box in boxes] + [box[3] for box in boxes]\n    return (min(x_coordinates), min(y_coordinates), max(x_coordinates), max(y_coordinates))\n\n\ndef merge_bbox(boxes):\n    # Perform post-processing of object detection model results\n    non_overlapping_boxes = []\n    overlapping_boxes = []\n    \n    # Iterate over each box\n    for i in range(len(boxes)):\n        overlapping = False\n        \n        # Check if the box is square or small\n        if not (is_square(boxes[i]) or is_small(boxes[i])):\n            for j in range(len(boxes)):\n                # Check for overlapping with other boxes\n                if i != j and is_overlapping(boxes[i], boxes[j], tolerance=-0.9):\n                    overlapping_boxes.append(boxes[i])\n                    overlapping_boxes.append(boxes[j])\n                    overlapping = True\n                    break\n        \n        if not overlapping:\n            non_overlapping_boxes.append(boxes[i])\n    \n    overlapping_groups = []\n    \n    # Group overlapping boxes that belong to the same label\n    for box in overlapping_boxes:\n        found_group = False\n        for group in overlapping_groups:\n            if any([(is_overlapping(box1=box_in_group, box2=box, tolerance=-0.9)) for box_in_group in group]):\n                group.append(box)\n                found_group = True\n                break\n        if not found_group:\n            overlapping_groups.append([box])\n\n    enclosing_boxes = []\n    for group in overlapping_groups:\n        enclosing_box = merge_boxes(group)\n        enclosing_boxes.append(enclosing_box)\n\n    all_boxes = non_overlapping_boxes + enclosing_boxes\n    sorted_bboxes = sorted(all_boxes, key=lambda x: x[0])\n    return sorted_bboxes\n\n\ndef get_x_tick_label_bboxes(img, label_boxes):\n    filtered_boxes = remove_boxes(label_boxes)\n    sorted_bboxes = sorted(label_boxes, key=lambda x: x[0])\n    merged_bboxes = merge_bbox(sorted_bboxes)\n    return sorted_bboxes, merged_bboxes","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:39.513726Z","iopub.execute_input":"2023-06-18T11:58:39.514048Z","iopub.status.idle":"2023-06-18T11:58:39.554882Z","shell.execute_reply.started":"2023-06-18T11:58:39.514019Z","shell.execute_reply":"2023-06-18T11:58:39.553728Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def final_text_data(data):\n    cleaned_data = []\n    for item in data:\n        cleaned_item = re.sub(r'\\s*-\\s*', '-', item)\n        cleaned_item = re.sub(r'\\s*\\.\\s*', '.', cleaned_item)\n        cleaned_data.append(cleaned_item)\n    return cleaned_data","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:39.559809Z","iopub.execute_input":"2023-06-18T11:58:39.562713Z","iopub.status.idle":"2023-06-18T11:58:39.571015Z","shell.execute_reply.started":"2023-06-18T11:58:39.562672Z","shell.execute_reply":"2023-06-18T11:58:39.569710Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from PIL import Image, ImageEnhance\nimport numpy as np\nimport torch\nimport re\n\nclass DotPlot:\n    def __init__(self, \n                 infer_tick_mark, \n                 infer_label_x, \n                 infer_dot_scatter,\n                 trocr_processor,\n                 troocr_model,\n                 paddle_ocr,\n                 ):\n\n        ## Object detection model object\n        self.infer_tick_mark = infer_tick_mark\n        self.infer_label_x = infer_label_x\n        self.infer_dot_scatter = infer_dot_scatter\n\n        ## paddle_ocr and TrOCR\n        self.trocr_processor = trocr_processor\n        self.trocr_model = troocr_model\n        self.paddle_ocr = paddle_ocr\n\n        self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    def get_predictions(self, img_path):\n        img = Image.open(img_path).convert(\"RGB\")\n        width, height = img.size\n        x_label_boxes = self.infer_label_x.get_bboxes(img)\n        x_label_boxes = sorted(x_label_boxes[0], key=lambda box: box[0])\n\n        x_tick_marks = self.infer_tick_mark.get_bboxes(img)\n        x_tick_boxes = sorted(x_tick_marks[0], key=lambda box: box[0])\n        filter_x_tick_boxes = self.filter_y_ticks(x_tick_boxes)\n\n        dot_boxes = self.infer_dot_scatter.get_bboxes(img)\n        filtered_dots = self.remove_x_axis_dot_boxes(x_label_boxes, filter_x_tick_boxes,  dot_boxes[0] )\n        filtred_sorted_dot_boxes = sorted(filtered_dots, key=lambda box: box[1], reverse= True)\n\n        last_line_dots_sorted = self.get_last_line_dot_boxes(filtred_sorted_dot_boxes)\n\n        mean_diff_last_line_dot = self.get_mean_diff_dots(last_line_dots_sorted)\n\n        total_line_wise_dots = self.get_dot_count_each_last_line(last_line_dots_sorted, filtred_sorted_dot_boxes, mean_diff_last_line_dot)\n\n        y_dot_str = \";\".join(str(i) for i in total_line_wise_dots)\n\n        pred, pred_list = self.trocr_inference(img, x_label_boxes, self.trocr_processor, self.trocr_model)\n\n        x_dot_str = self.get_x_labels(x_label_boxes, total_line_wise_dots, pred_list, pred,last_line_dots_sorted)\n\n        return x_dot_str, y_dot_str\n    \n    def identify_x(self, tick_boxes):\n        center_points = [(box[1] + box[3]) / 2 for box in tick_boxes]\n        y = np.median(sorted(center_points)[-3:-1])\n        return y\n    \n    def filter_y_ticks(self, tick_marks):\n        tick_boxes = []\n        for i in tick_marks:\n            x_axis = self.identify_x(tick_marks)\n            if i[3] > x_axis:\n                tick_boxes.append(i)\n        return tick_boxes\n    \n    def remove_x_axis_dot_boxes(self, x_label_boxes, x_tick_boxes, dots_sorted):\n        filtered_dots = []\n        for dot in dots_sorted:\n            dot_x1, dot_y1, dot_x2, dot_y2 = dot\n            is_along_x_axis = False\n            # Check if the dot is along the x-axis (within the x_label_boxes or x_tick_boxes)\n            for box in x_label_boxes + x_tick_boxes:\n                box_x1, box_y1, box_x2, box_y2 = box\n                if dot_y1 >= box_y1 and dot_y2 <= box_y2:\n                    is_along_x_axis = True\n                    break\n            # Add the dot to the filtered dots list if it is not along the x-axis\n            if not is_along_x_axis:\n                filtered_dots.append(dot)\n        return filtered_dots\n\n    def _enhance_image(self, img):\n        enhancer = ImageEnhance.Contrast(img)\n        img = enhancer.enhance(1.3)\n        enhancer1 = ImageEnhance.Sharpness(img)\n        img = enhancer1.enhance(1.3)\n        width, height = img.size\n        new_width, new_height = width * 2, height * 2\n        resized_img = img.resize((new_width, new_height))\n        return resized_img\n\n    def snippets_to_single_image(self,image_list):\n        images = []\n        for element in image_list:\n            pil_image=Image.fromarray(element)\n            images.append(pil_image)\n\n        patch_width = 20\n        widths, heights = zip(*(i.size for i in images))\n        total_width = sum(widths) + patch_width * (len(images) - 1)\n        max_height = max(heights)\n        patch_height = max_height\n        new_im = Image.new('RGB', (total_width, max_height), color='white')\n        x_offset = 0\n        for im in images:\n            new_im.paste(im, (x_offset, 0))\n            x_offset += im.size[0] + patch_width\n        return new_im\n\n    def paddle_ocr_call(self, image):\n        snippet=np.array(image)\n        result = self.paddle_ocr.ocr(snippet, cls=True)\n        image_list = []\n        for i in result[0]:\n            array = i[0]\n            xmin_l = []\n            ymin_l = []\n            for id, pos in enumerate(array):\n                xmin_l.append(pos[0])\n                x_min = min(xmin_l)\n                x_max = max(xmin_l)\n                ymin_l.append(pos[1])\n                y_min = min(ymin_l)\n                y_max = max(ymin_l)\n            image_list.append(snippet[int(y_min):int(y_max), int(x_min):int(x_max)])\n        return image_list\n\n    def trocr_inference(self, img ,bboxes, trocr_processor, trocr_model):\n        gen_text=[]\n        for box in bboxes:\n            xmin, ymin, xmax, ymax = box\n            crop = img.crop((xmin, ymin, xmax, ymax))\n            if abs(ymax - ymin) > 40:\n                image_list = self.paddle_ocr_call(crop)\n                if len(image_list) > 1:\n                    single_img = self.snippets_to_single_image(image_list)\n                else:\n                    single_img = crop\n                crop = single_img\n            crop = self._enhance_image(crop)\n            pixel_values = trocr_processor(images=crop, return_tensors=\"pt\").pixel_values\n            generated_ids = trocr_model.generate(pixel_values.to('cuda'), max_length=64, num_beams=2, early_stopping=True)\n            generated_text = trocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]\n            gen_text.append(generated_text.capitalize())\n        predicted_text = \";\".join(gen_text)\n        return predicted_text, gen_text\n    \n    def get_last_line_dot_boxes(self, filtred_sorted_dot_boxes):\n        last_line_bbox = filtred_sorted_dot_boxes[0]\n        y_center_last = (last_line_bbox[1] + last_line_bbox[3]) // 2\n        last_line_dots = []\n        for each_dot in filtred_sorted_dot_boxes:\n            if each_dot[1]  <= y_center_last <= each_dot[3]:\n                last_line_dots.append(each_dot)\n        last_line_dots_sorted = sorted(last_line_dots, key=lambda box: box[0])\n        return last_line_dots_sorted\n    \n    def get_mean_diff_dots(self, last_line_dots_sorted):\n        diff_array = []\n        for i in range(len(last_line_dots_sorted)-1):\n            diff = last_line_dots_sorted[i][2] - last_line_dots_sorted[i+1][0]\n            diff_array.append(abs(diff))\n        mean_diff_last_line_dot = np.mean(diff_array)\n        return mean_diff_last_line_dot\n    \n    def get_dot_count_each_last_line(self, last_line_dots_sorted, filtred_sorted_dot_boxes, mean_diff_last_line_dot):\n        total_line_wise_dots = []\n        for each_last_line_dot in last_line_dots_sorted:\n            each_last_line_dot_center = (each_last_line_dot[0] + each_last_line_dot[2]) // 2\n            line_wise_dots = []\n            for each_dot in filtred_sorted_dot_boxes:\n                if each_dot[0]  <= each_last_line_dot_center <= each_dot[2]:\n                    line_wise_dots.append(each_dot)\n            total_line_wise_dots.append(str(len(line_wise_dots )))\n\n        for i in range(len(last_line_dots_sorted) -1):\n            each_last_line_dot_center =last_line_dots_sorted[i][2]\n            second_dot_center = last_line_dots_sorted[i+1][0]\n            number_of_zeros = round((abs(second_dot_center - each_last_line_dot_center))/mean_diff_last_line_dot)\n            if number_of_zeros>1:\n                zeros_list = [0] * (number_of_zeros - 1)\n                insert_position = i + 1\n                total_line_wise_dots[insert_position:insert_position] = zeros_list\n            \n        return total_line_wise_dots     \n    \n    def get_x_labels(self, x_label_boxes, total_line_wise_dots, pred_list, pred,last_line_dots_sorted):\n        if len(total_line_wise_dots) > len(pred_list):\n            start = int(pred_list[0])\n            end = int(pred_list[1])\n            start_box = x_label_boxes[0]\n            end_box = x_label_boxes[1]\n            num_boxes_between = 0\n\n            for i in range(len(last_line_dots_sorted)):\n                box_middle = (last_line_dots_sorted[i][0] + last_line_dots_sorted[i][2])/2\n\n                if start_box[0] < box_middle < end_box[0]:\n                    each_last_line_dot_center =last_line_dots_sorted[i][2]\n                    second_dot_center = last_line_dots_sorted[i+1][0]\n                    number_of_zeros = round((abs(second_dot_center - each_last_line_dot_center))/mean_diff_last_line_dot)\n                    num_boxes_between += number_of_zeros\n\n            range_value = (end - start)/num_boxes_between\n            label_value = []\n            single_value = int(pred_list[-1])\n            label_value.append(single_value)\n\n            for i in range(len(total_line_wise_dots) - 1):\n                single_value -= range_value\n                label_value.append(int(single_value))\n\n            label_value.reverse()  \n            final_x_pred = \";\".join(str(j) for j in label_value)\n        else:\n            final_x_pred = pred\n        return final_x_pred","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:39.576679Z","iopub.execute_input":"2023-06-18T11:58:39.579530Z","iopub.status.idle":"2023-06-18T11:58:39.650510Z","shell.execute_reply.started":"2023-06-18T11:58:39.579488Z","shell.execute_reply":"2023-06-18T11:58:39.649477Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_y_tick_mark_boxes(tick_mark_boxes):\n    x2_values = [box[1] for box in tick_mark_boxes]\n    x2_mode = mode(x2_values)[0]\n    filtered_boxes = [box for box in tick_mark_boxes if box[2] <= x2_mode+15]\n    return filtered_boxes\n","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:39.654982Z","iopub.execute_input":"2023-06-18T11:58:39.657899Z","iopub.status.idle":"2023-06-18T11:58:39.667049Z","shell.execute_reply.started":"2023-06-18T11:58:39.657857Z","shell.execute_reply":"2023-06-18T11:58:39.665855Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"learn = load_learner('/kaggle/input/bmge-extra-data-2/classifier_learner.pkl')","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:39.672317Z","iopub.execute_input":"2023-06-18T11:58:39.675290Z","iopub.status.idle":"2023-06-18T11:58:39.802448Z","shell.execute_reply.started":"2023-06-18T11:58:39.675248Z","shell.execute_reply":"2023-06-18T11:58:39.801241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.copy(\"/kaggle/input/paddle-wheels/en_PP-OCRv3_det_infer.tar\", \"/kaggle/working/en_PP-OCRv3_det_infer.tar\")\nshutil.copy(\"/kaggle/input/paddle-wheels/ch_ppocr_mobile_v2.0_cls_infer.tar\", \"/kaggle/working/ch_ppocr_mobile_v2.0_cls_infer.tar\")\nshutil.copy(\"/kaggle/input/paddle-wheels/en_PP-OCRv3_rec_infer.tar\", \"/kaggle/working/en_PP-OCRv3_rec_infer.tar\")\n\nimport tarfile\ndef extract(tar_file, path):\n    opened_tar = tarfile.open(tar_file)\n    if tarfile.is_tarfile(tar_file):\n        opened_tar.extractall(path)\n    else:\n        print(\"The tar file you entered is not a tar file\")\n\nextract(\"/kaggle/working/ch_ppocr_mobile_v2.0_cls_infer.tar\" , \"/kaggle/working/\")\nextract(\"/kaggle/working/en_PP-OCRv3_det_infer.tar\" , \"/kaggle/working/\")\nextract(\"/kaggle/working/en_PP-OCRv3_rec_infer.tar\" , \"/kaggle/working/\")","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:39.807616Z","iopub.execute_input":"2023-06-18T11:58:39.810595Z","iopub.status.idle":"2023-06-18T11:58:39.926096Z","shell.execute_reply.started":"2023-06-18T11:58:39.810551Z","shell.execute_reply":"2023-06-18T11:58:39.924913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from paddleocr import PaddleOCR\n\nocr = PaddleOCR(use_angle_cls=False, lang='en', show_log=False, det_model_dir =\"/kaggle/working/en_PP-OCRv3_det_infer\",rec_model_dir =\"/kaggle/working/en_PP-OCRv3_rec_infer\",cls_model_dir = \"/kaggle/working/ch_ppocr_mobile_v2.0_cls_infer\")","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:39.936040Z","iopub.execute_input":"2023-06-18T11:58:39.937233Z","iopub.status.idle":"2023-06-18T11:58:40.924246Z","shell.execute_reply.started":"2023-06-18T11:58:39.937188Z","shell.execute_reply":"2023-06-18T11:58:40.922987Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def TrOCR_predict(pixel_values, processor, model):\n    pixel_values = pixel_values.to('cuda')\n    generated_ids = model.generate(pixel_values,output_scores=True,return_dict_in_generate=True, max_length = 64)\n    predicted_text = processor.batch_decode(generated_ids[0], skip_special_tokens=True)\n    return predicted_text","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:40.925968Z","iopub.execute_input":"2023-06-18T11:58:40.926757Z","iopub.status.idle":"2023-06-18T11:58:40.934454Z","shell.execute_reply.started":"2023-06-18T11:58:40.926715Z","shell.execute_reply":"2023-06-18T11:58:40.933217Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def snippets_to_single_image(images):\n    patch_width = 20\n    heights, widths, depth = zip(*(i.shape for i in images))\n    total_width = sum(widths) + patch_width * (len(images) - 1)\n    max_height = max(heights)\n    patch_height = max_height\n    new_im = Image.new('RGB', (total_width, max_height), color='white')\n    x_offset = 0\n    for im in images:\n        im = Image.fromarray(im)\n        box = (x_offset, 0, x_offset + im.size[0], im.size[1])\n        new_im.paste(im, box)\n        x_offset += im.size[0] + patch_width\n#     display(new_im)\n    return new_im","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:40.939419Z","iopub.execute_input":"2023-06-18T11:58:40.941591Z","iopub.status.idle":"2023-06-18T11:58:40.954846Z","shell.execute_reply.started":"2023-06-18T11:58:40.941549Z","shell.execute_reply":"2023-06-18T11:58:40.953732Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import logging\nfrom ppocr.utils.logging import get_logger as ppocr_get_logger\nppocr_get_logger().setLevel(logging.ERROR)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:40.960327Z","iopub.execute_input":"2023-06-18T11:58:40.963889Z","iopub.status.idle":"2023-06-18T11:58:40.971762Z","shell.execute_reply.started":"2023-06-18T11:58:40.963847Z","shell.execute_reply":"2023-06-18T11:58:40.970408Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def paddle_ocr_call(self, image):\n    snippet=np.array(image)\n    result = self.paddle_ocr.ocr(snippet, cls=True)\n    image_list = []\n    for i in result[0]:\n        array = i[0]\n        xmin_l = []\n        ymin_l = []\n        for id, pos in enumerate(array):\n            xmin_l.append(pos[0])\n            x_min = min(xmin_l)\n            x_max = max(xmin_l)\n            ymin_l.append(pos[1])\n            y_min = min(ymin_l)\n            y_max = max(ymin_l)\n        image_list.append(snippet[int(y_min):int(y_max), int(x_min):int(x_max)])\n    return image_list\n","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:40.977294Z","iopub.execute_input":"2023-06-18T11:58:40.979419Z","iopub.status.idle":"2023-06-18T11:58:40.992682Z","shell.execute_reply.started":"2023-06-18T11:58:40.979379Z","shell.execute_reply":"2023-06-18T11:58:40.991544Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## Dot Scatter Utility Function\n\ndef check_dot_scatter(x_labels, dot_boxes, x_axis, height):\n    output_dict = {}\n    highest_label = (0, 0)\n    for i in range(len(x_labels)):\n        label = x_labels[i]\n        x_min, y_min, x_max, y_max = label\n        \n        x_middle = (x_min + x_max)/2\n        \n        mapped_dots_boxes = []\n        for single_dot_box in dot_boxes:\n            if single_dot_box[0] <= x_middle <= single_dot_box[2]:\n                mapped_dots_boxes.append(single_dot_box)\n        \n        if highest_label[1] < len(mapped_dots_boxes):\n            highest_label = (i, len(mapped_dots_boxes))\n            \n        output_dict[i] = mapped_dots_boxes\n    \n    if highest_label[1] == 0:\n        return False\n    \n    max_dots_selected = output_dict[highest_label[0]]\n    \n    distance_list = []\n    distance_from_x_list = []\n    max_dots_selected = sorted(max_dots_selected, key=lambda x: x[1], reverse=True)\n    \n    for i in range(len(max_dots_selected) - 1):\n        current_box_middle = (max_dots_selected[i][1] + max_dots_selected[i][3])/2 \n        next_box_middle = (max_dots_selected[i + 1][1] + max_dots_selected[i + 1][3])/2\n        \n        distance= current_box_middle - next_box_middle\n        distance_list.append(distance)\n        \n    for key in output_dict.keys():\n        boxes = output_dict[key]\n        if boxes:\n            boxes = sorted(boxes, key=lambda x: x[1], reverse=True)\n            current_box_middle = (boxes[0][1] + boxes[0][3])/2 \n            distance_from_x = abs(x_axis - current_box_middle)/height\n            distance_from_x_list.append(distance_from_x)\n    \n    std_dist = np.std(distance_list)\n    std_dist_from_x = np.std(distance_from_x_list)\n            \n    if len(distance_list) < 2:\n        return False\n\n    if len(x_labels) > 2:\n        if len(distance_from_x_list)/len(x_labels) < 0.6:\n            return False\n    else:\n        return False\n\n    if (std_dist < 1.5) and (std_dist_from_x < 0.06):\n        return True\n    else:\n        return False\n\n    return False ","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:40.998323Z","iopub.execute_input":"2023-06-18T11:58:41.000957Z","iopub.status.idle":"2023-06-18T11:58:41.025334Z","shell.execute_reply.started":"2023-06-18T11:58:41.000917Z","shell.execute_reply":"2023-06-18T11:58:41.024132Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def filter_x_tick_marks(boxes, given_x, given_y):\n    filtered_boxes = []\n    for box in boxes:\n        x_min, y_min, x_max, y_max = box\n        if y_min < given_y < y_max:\n            filtered_boxes.append(box)\n    return filtered_boxes\n\ndef filter_labels(boxes, given_x, given_y):\n    filtered_boxes = []\n    for box in boxes:\n        x_min, y_min, x_max, y_max = box\n        if given_x + 10 >= x_max or given_y  <= y_max:\n            filtered_boxes.append(box)\n    return filtered_boxes\n\ndef get_nearest_label(filtered_line_intersection_box, labels_list):\n    intersection_box_middle = (filtered_line_intersection_box[0] + filtered_line_intersection_box[2])/2\n    \n    distances = []\n    for label in labels_list:\n        label_middle = (label[0] + label[2])/2\n        \n        distances.append(abs(intersection_box_middle - label_middle))\n    \n    nearest_idx = np.argmin(distances)\n    \n    return labels_list[nearest_idx]","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:41.030597Z","iopub.execute_input":"2023-06-18T11:58:41.033870Z","iopub.status.idle":"2023-06-18T11:58:41.047629Z","shell.execute_reply.started":"2023-06-18T11:58:41.033830Z","shell.execute_reply":"2023-06-18T11:58:41.046468Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def filter_tick_label(tick_boxes, x_label_boxes):\n    filtered_labels = []\n    filtered_ticks = []\n    for tick_box in tick_boxes:\n        tick_x_min, tick_y_min, tick_x_max, tick_y_max = tick_box\n        tick_center = (tick_x_min + tick_x_max) / 2\n\n        for label_box in x_label_boxes:\n            label_x_min, label_y_min, label_x_max, label_y_max = label_box\n            label_center = (label_x_min + label_x_max) / 2\n\n            if (label_x_min <= tick_center <= label_x_max or tick_x_min <= label_center <= tick_x_max) and label_box not in filtered_labels:\n                filtered_labels.append(label_box)\n                filtered_ticks.append(tick_box)\n\n    return filtered_labels, filtered_ticks\n\ndef fix_rotation(image):\n    if 2 * image.width <= image.height:\n        fixed_image = image.rotate(-90, expand=True)\n        display(image)\n        return fixed_image\n    else:\n        # No rotation needed\n        return image\n","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:41.053291Z","iopub.execute_input":"2023-06-18T11:58:41.056612Z","iopub.status.idle":"2023-06-18T11:58:41.069103Z","shell.execute_reply.started":"2023-06-18T11:58:41.056571Z","shell.execute_reply":"2023-06-18T11:58:41.067962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def draw_vertical_lines(img, boxes):\n    draw = ImageDraw.Draw(img)\n    line_color = (255, 0, 0)  # Red color (R, G, B)\n    line_width = 3\n\n    for box in boxes:\n        x_min, y_min, x_max, y_max = box\n        x = (x_min + x_max) // 2  # x-coordinate for the vertical line\n        line_coordinates = [(x, y_min), (x, y_max)]\n        draw.line(line_coordinates, fill=line_color, width=line_width)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:41.074628Z","iopub.execute_input":"2023-06-18T11:58:41.078147Z","iopub.status.idle":"2023-06-18T11:58:41.087882Z","shell.execute_reply.started":"2023-06-18T11:58:41.078106Z","shell.execute_reply":"2023-06-18T11:58:41.086736Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import DonutProcessor, VisionEncoderDecoderModel\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nprocessor = DonutProcessor.from_pretrained(\"/kaggle/input/rotated-v4-5\")\nmodel = VisionEncoderDecoderModel.from_pretrained(\"/kaggle/input/rotated-v4-5\").to(device)\n\nimport re\n\ndef get_text(s):\n    match = re.search(r'<s_X_ticks>(.+?)</s_X_ticks>|<s_X_ticks>(.+?)</s_X_label>', s)\n    if match:\n        result = match.group(1) or match.group(2)\n    else:\n        result = \"a<xt>b<xt>c<xt>d<xt>e\"\n    text_list = [i.strip() for i in result.split(\"<xt>\")]\n    string = \";\".join(text_list)\n    string = re.sub(r'<[^>]*s[^>]*>', '', string)\n    \n    return string\n\ndef donut_predict(crop):\n    pixel_values = processor(crop.convert(\"RGB\"), return_tensors=\"pt\").pixel_values\n    pixel_values = pixel_values.to(device)\n    task_prompt = \"<s>\"\n    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors=\"pt\").input_ids\n    decoder_input_ids = decoder_input_ids.to(device)\n    outputs = model.generate(\n            pixel_values,\n            decoder_input_ids=decoder_input_ids,\n            max_length=512,\n            early_stopping=True,\n            pad_token_id=processor.tokenizer.pad_token_id,\n            eos_token_id=processor.tokenizer.eos_token_id,\n            use_cache=True,\n            num_beams=4,\n            bad_words_ids=[[processor.tokenizer.unk_token_id]],\n            return_dict_in_generate=True,\n            no_repeat_ngram_size=3,\n            do_sample=False,\n            output_scores=True\n        )\n    seq = processor.batch_decode(outputs.sequences)[0]\n    seq = seq.replace(processor.tokenizer.eos_token, \"\").replace(processor.tokenizer.pad_token, \"\")\n    seq = seq.replace(\"<one>\", \"1\")\n    seq = re.sub(r\"<.*?>\", \"\", seq, count=1).strip()\n    seq = get_text(seq)\n    return seq","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:41.095327Z","iopub.execute_input":"2023-06-18T11:58:41.095990Z","iopub.status.idle":"2023-06-18T11:58:56.643991Z","shell.execute_reply.started":"2023-06-18T11:58:41.095951Z","shell.execute_reply":"2023-06-18T11:58:56.642747Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import re\n\ndef is_convertible_to_float(string):\n    try:\n        float(string)\n        return True\n    except ValueError:\n        return False\n\ndef replace_k_with_thousands(text):\n    return re.sub(r'(\\d)K', r'\\g<1>000', text)\n\ndef remove_special_char(text):\n    return re.sub(r'[^a-zA-Z0-9.\\s-]', '', text)\n\ndef postprocess_text(text):\n    text = replace_k_with_thousands(text)\n    text = remove_special_char(text)\n    return text","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:56.645707Z","iopub.execute_input":"2023-06-18T11:58:56.646356Z","iopub.status.idle":"2023-06-18T11:58:56.654766Z","shell.execute_reply.started":"2023-06-18T11:58:56.646314Z","shell.execute_reply":"2023-06-18T11:58:56.653617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pathlib import Path\nfrom tqdm.auto import tqdm\nfrom PIL import ImageDraw\n\npath = Path('/kaggle/input/benetech-making-graphs-accessible')\ntest_path = path/'test'\n# test_path = Path('/kaggle/input/bmga-tv-data/valid')\nimages_path = test_path/'images'\n\nimages_data = os.listdir(images_path)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:56.658296Z","iopub.execute_input":"2023-06-18T11:58:56.659633Z","iopub.status.idle":"2023-06-18T11:58:56.672434Z","shell.execute_reply.started":"2023-06-18T11:58:56.659589Z","shell.execute_reply":"2023-06-18T11:58:56.670940Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"vertical_predictor = VerticalBar(\n                    rotation_learn = rotation_learn, \n                    infer_tick_mark = infer_tick_mark, \n                    infer_label_y = infer_label, \n                    infer_label_x = infer_label_x, \n                    infer_vertical_bar = infer_vertical_bar,\n                    trocr_processor = large_processor,\n                    trocr_model = l_model,\n                    paddle_ocr = ocr,\n                    donut_processor = processor,\n                    donut_model = model)\n\nscatter_predictor = ScatterPlot(\n                infer_tick_mark, \n                infer_label, \n                infer_label_x, \n                infer_dot_scatter,\n                large_processor,\n                l_model,\n                processor,\n                model)\n\nhorizontal_predictor = HorizontalBar(\n                         infer_tick_mark, \n                         infer_label, \n                         infer_label_x, \n                         infer_horizontal_bar,\n                         large_processor,\n                         l_model,\n                         ocr,\n                         processor,\n                         model\n                         )\n\ndot_predictor = DotPlot(\n                 infer_tick_mark, \n                 infer_label_x, \n                 infer_dot_scatter,\n                large_processor,\n                l_model,\n                 ocr)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:56.674135Z","iopub.execute_input":"2023-06-18T11:58:56.676836Z","iopub.status.idle":"2023-06-18T11:58:56.903262Z","shell.execute_reply.started":"2023-06-18T11:58:56.676794Z","shell.execute_reply":"2023-06-18T11:58:56.902008Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Disable warnings\nwarnings.filterwarnings(\"ignore\")\n\n# Create an empty DataFrame to store the submission data\nsubmission_data = pd.DataFrame(columns=['id', 'data_series', 'chart_type'])\n\n# Iterate over the files in the valid_images_path directory\nfor i, filename in tqdm(enumerate(images_data)):\n    test_file = os.path.join(images_path, filename)\n\n    # Perform prediction on the test file to determine the graph type\n    with learn.no_bar():\n        graph_type = learn.predict(test_file)  # ['dot', 'horizontal_bar', 'line', 'scatter', 'vertical_bar']\n\n    # Create a dictionary for x-axis data series with default labels\n    x_data_series = {\n        'id': filename[:-4] + \"_x\",\n        'data_series': 'a;b;c;d;e;f;g',\n        'chart_type': graph_type[0] \n    }\n\n    # Create a dictionary for y-axis data series with default values\n    y_data_series = {\n        'id': filename[:-4] + \"_y\",\n        'data_series': '0;1;2;3;4;5;6',\n        'chart_type': graph_type[0] \n    }\n\n    if graph_type[0] == 'scatter':\n        try:\n            img = Image.open(test_file).convert(\"RGB\")\n            width, height = img.size\n            x_label_boxes = infer_label_x.get_bboxes(img)\n            x_label_boxes = sorted(x_label_boxes[0], key=lambda box: box[0])\n\n            x_tick_marks = infer_tick_mark.get_bboxes(img)\n            x_tick_boxes = sorted(x_tick_marks[0], key=lambda box: box[0])\n            filter_x_tick_boxes = dot_predictor.filter_y_ticks(x_tick_boxes)\n\n            dot_boxes = infer_dot_scatter.get_bboxes(img)\n            filtered_dots = dot_predictor.remove_x_axis_dot_boxes(x_label_boxes, filter_x_tick_boxes,  dot_boxes[0] )\n            filtred_sorted_dot_boxes = sorted(filtered_dots, key=lambda box: box[1], reverse= True)\n\n            x_axis = dot_predictor.identify_x(x_tick_boxes)\n            is_dot  = check_dot_scatter(x_label_boxes, filtred_sorted_dot_boxes, x_axis, height=img.size[1]) \n\n            if is_dot:\n                x_preds, y_preds = dot_predictor.get_predictions(test_file)\n            else:                \n                x_preds, y_preds = scatter_predictor.get_predictions(test_file)\n\n            x_data_series = {\n                'id': filename[:-4] + \"_x\",\n                'data_series': x_preds,\n                'chart_type': graph_type[0] \n            }\n\n            y_data_series = {\n                'id': filename[:-4] + \"_y\",\n                'data_series': y_preds,\n                'chart_type': graph_type[0] \n            }\n            \n        except Exception as e:\n            print('Error while processing image:', filename)\n            print(e)\n            \n    if graph_type[0] == 'line':\n        try:\n            img = Image.open(test_file)\n            img_arr = np.array(img)\n            draw = ImageDraw.Draw(img)\n\n            # Get bounding boxes for tick marks, labels, and x-axis tick labels\n            tick_mark_boxes, tick_mark_boxes_scores = infer_tick_mark.get_bboxes(img)\n            label_boxes, label_boxes_scores = infer_label.get_bboxes(img)\n            label_boxes_suppressed = non_max_suppression_fast(np.array(label_boxes), 0.2)\n            x_axis = identify_x(tick_mark_boxes)\n            y_axis = identify_y(tick_mark_boxes)\n            \n            # Perform prediction on the test file to determine the graph type\n            with rotation_learn.no_bar():\n                rotation_status = rotation_learn.predict(test_file)  # ['dot', 'horizontal_bar', 'line', 'scatter', 'vertical_bar']\n\n            if rotation_status[0] != \"rotated\":  \n                # Get bounding boxes for tick marks, labels, and x-axis tick labels\n                x_tick_label_boxes, x_tick_label_boxes_scores = infer_label_x.get_bboxes(img)\n\n                x_ticks = filter_x_tick_marks(tick_mark_boxes, y_axis, x_axis)\n                filtered_labels, filtered_ticks = filter_tick_label(x_ticks, x_tick_label_boxes)\n\n                line_color = (255, 0, 0)  # Red color (R, G, B)\n                line_width = 3\n                filtered_labels = list(non_max_suppression_fast(np.array(filtered_labels), 0.8))\n\n                for box in filtered_labels:\n                    x_min, y_min, x_max, y_max = box\n                    x_middle = (x_min + x_max) / 2 \n                    y_middle = (y_min + y_max) / 2 \n\n                    if x_middle < y_axis:\n                        x_middle = y_axis\n\n                    line_coordinates = [(x_middle, 0), (x_middle, x_axis - 3)]\n                    draw.line(line_coordinates, fill=line_color, width=line_width)\n\n                line_intersection_boxes, line_intersection_boxes_scores = infer_line.get_bboxes(img)\n\n                filtered_line_intersection_boxes = []\n                filtered_ticks_without_line = []\n\n                for i, filtered_tick in enumerate(filtered_labels):\n                    inside_box = []\n                    for line_intersection_box, line_intersection_boxes_score in zip(line_intersection_boxes, line_intersection_boxes_scores):\n                        line_box_middle_x = (line_intersection_box[0] + line_intersection_box[2]) / 2\n                        if filtered_tick[0] < line_box_middle_x < filtered_tick[2]:\n                            inside_box.append((line_intersection_box, line_intersection_boxes_score))\n\n                    if len(inside_box) > 1:\n                        element_with_highest_score = max(inside_box, key=lambda x: x[1])\n                        filtered_line_intersection_boxes.append(element_with_highest_score[0]) \n                        filtered_ticks_without_line.append(filtered_tick)\n                    elif len(inside_box) == 1:\n                        filtered_line_intersection_boxes.append(inside_box[0][0]) \n                        filtered_ticks_without_line.append(filtered_tick)\n\n                filtered_line_intersection_boxes = sorted(filtered_line_intersection_boxes, key=lambda x: x[0])\n                filtered_line_intersection_boxes = list(non_max_suppression_fast(np.array(filtered_line_intersection_boxes), 0.1))\n\n                filtered_labels = sorted(filtered_labels, key=lambda x: x[0])\n\n                bbox_snippets = []\n\n                if len(filtered_line_intersection_boxes) == len(filtered_labels):\n                    for filtered_label in filtered_labels:\n                        x1, y1, x2, y2 = filtered_label\n                        snippet = img_arr[y1:y2, x1:x2]\n                        pixel_values = large_processor(images=snippet, return_tensors=\"pt\").pixel_values\n                        bbox_snippets.append(pixel_values)\n                else:\n                    nearest_labels = []\n                    for filtered_line_intersection_box in filtered_line_intersection_boxes:\n                        label_nearest = get_nearest_label(filtered_line_intersection_box, filtered_labels)\n                        nearest_labels.append(label_nearest)\n                    nearest_labels = sorted(nearest_labels, key=lambda x: x[0])\n                    for near_label in nearest_labels:\n                        x1, y1, x2, y2 = near_label\n                        snippet = img_arr[y1:y2, x1:x2]\n                        img_pil = Image.fromarray(snippet)\n                        snippet = fix_rotation(img_pil)\n                        snippet = np.array(snippet)\n                        pixel_values = large_processor(images=snippet, return_tensors=\"pt\").pixel_values\n                        bbox_snippets.append(pixel_values)\n\n                # Stack the snippets into a batch tensor\n                if len(bbox_snippets) != 0:\n                    batch = torch.stack(bbox_snippets).squeeze(1)\n\n                    # Perform OCR to predict text from the snippets\n                    predicted_text = TrOCR_predict(batch, large_processor, l_model)\n                    predicted_text = [i.capitalize() for i in predicted_text]\n\n            else:\n                width, height = img.size\n\n                x_ticks = filter_x_tick_marks(tick_mark_boxes, y_axis, x_axis)\n\n                line_color = (255, 0, 0)  # Red color (R, G, B)\n                line_width = 3\n\n                for box in x_ticks:\n                    x_min, y_min, x_max, y_max = box\n                    x_middle = (x_min + x_max) / 2 \n                    y_middle = (y_min + y_max) / 2 \n\n                    line_coordinates = [(x_middle, 0), (x_middle, x_axis - 3)]\n                    draw.line(line_coordinates, fill=line_color, width=line_width)\n\n                line_intersection_boxes, line_intersection_boxes_scores = infer_line.get_bboxes(img)\n\n                filtered_line_intersection_boxes = []\n                filtered_ticks_without_line = []\n\n                for i, filtered_tick in enumerate(x_ticks):\n                    inside_box = []\n                    for line_intersection_box, line_intersection_boxes_score in zip(line_intersection_boxes, line_intersection_boxes_scores):\n                        line_box_middle_x = (line_intersection_box[0] + line_intersection_box[2]) / 2\n                        if filtered_tick[0] < line_box_middle_x < filtered_tick[2]:\n                            inside_box.append((line_intersection_box, line_intersection_boxes_score))\n                    if len(inside_box) > 1:\n                        element_with_highest_score = max(inside_box, key=lambda x: x[1])\n                        filtered_line_intersection_boxes.append(element_with_highest_score[0]) \n                        filtered_ticks_without_line.append(filtered_tick)\n                    elif len(inside_box) == 1:\n                        filtered_line_intersection_boxes.append(inside_box[0][0]) \n                        filtered_ticks_without_line.append(filtered_tick)\n\n                filtered_line_intersection_boxes = sorted(filtered_line_intersection_boxes, key=lambda x: x[0])\n                filtered_line_intersection_boxes = list(non_max_suppression_fast(np.array(filtered_line_intersection_boxes), 0.1))\n\n                try:\n                    crop = img.crop((0, x_axis, width, height))\n                except Exception as e:\n                    print('Error while cropping image:', e)\n                    crop = img\n\n                pred = donut_predict(crop)\n                predicted_text = pred.split(';')\n            \n                groups = {}\n                filtered_ticks_without_line = sorted(filtered_ticks_without_line, key=lambda x: x[0])\n                if (len(filtered_ticks_without_line) >  len(filtered_line_intersection_boxes)) and (len(predicted_text) == len(filtered_ticks_without_line)): \n                    groups = []\n                    for i in range(len(filtered_ticks_without_line)):\n                        \n                        intersection_box_detected = None\n                        for filtered_line_intersection_box in filtered_line_intersection_boxes:\n                            intersection_box_middle = (filtered_line_intersection_box[0] + filtered_line_intersection_box[2])/2\n                            if filtered_ticks_without_line[i][0] < intersection_box_middle < filtered_ticks_without_line[i][2]:\n                                intersection_box_detected = filtered_line_intersection_box\n                            \n                        if len(predicted_text) < i:\n                            break\n                            \n                        group = {\n                            'tick': filtered_ticks_without_line[i],\n                            'label': predicted_text[i],\n                            'intersect': intersection_box_detected \n                        }\n                        \n                        groups.append(group)\n            \n                    predicted_text = []\n                    \n                    for group in groups:\n                        if group['intersect'] is not None:\n                            predicted_text.append(group['label'])\n                        \n            # Get y-axis label bounding boxes and snippets for OCR\n            label_boxes_suppressed = sorted(label_boxes_suppressed, key=lambda x: x[1], reverse=True)\n            bbox_snippets_y = []\n            for box in label_boxes_suppressed:\n                x1, y1, x2, y2 = box\n                snippet = img_arr[y1:y2, x1:x2]\n                pixel_values = large_processor(images=snippet, return_tensors=\"pt\").pixel_values\n                bbox_snippets_y.append(pixel_values)\n\n            batch = torch.stack(bbox_snippets_y).squeeze(1)\n\n            # Perform OCR to predict text from the snippets\n            try:\n                label_text = TrOCR_predict(batch, large_processor, l_model)\n                label_text = [postprocess_text(i) for i in label_text]\n            except Exception as e:\n                print('Error while performing OCR on y-axis labels:', e)\n                label_text = []\n\n            # After getting label text, process the y-axis data\n            selected_box = [(float(output), bbox) for output, bbox in zip(label_text, label_boxes_suppressed) if is_convertible_to_float(output)]        \n            origin_value = selected_box[0][0]\n            y_diff = abs(selected_box[1][0] - selected_box[0][0])\n            pixel_diff = abs((selected_box[1][1][1] + selected_box[1][1][3]) / 2 - (selected_box[0][1][1] + selected_box[0][1][3]) / 2)\n            y_scale = y_diff / pixel_diff\n\n            y_data = []\n            filtered_line_intersection_boxes = sorted(filtered_line_intersection_boxes, key=lambda x: x[0])\n            for filtered_line_intersection_box in filtered_line_intersection_boxes:\n                box_y_middle = (filtered_line_intersection_box[1] + filtered_line_intersection_box[3]) / 2\n                difference = abs(box_y_middle - x_axis)\n                total_difference_value = difference * y_scale\n                y_data.append(str(total_difference_value + origin_value))\n\n            if y_data:\n                y_data_series = {\n                    'id': filename[:-4] + \"_y\",\n                    'data_series': ';'.join(list(y_data)),\n                    'chart_type': graph_type[0] \n                }\n\n            if predicted_text:\n                if y_data:\n                    if len(y_data) > len(predicted_text):\n                        predicted_text.extend(['abc'] * (len(y_data) - len(predicted_text)))\n                        \n                predicted_text = checker.correct_strings(predicted_text)\n                predicted_text = final_text_data(predicted_text)\n                x_data_series = {\n                    'id': filename[:-4] + \"_x\",\n                    'data_series': ';'.join(predicted_text),\n                    'chart_type': graph_type[0] \n                }\n            \n        except Exception as e:\n            print('Error while processing image:', filename)\n            print(e)\n        \n    if graph_type[0] == 'vertical_bar':\n        try:\n\n            x_preds, y_preds = vertical_predictor.get_predictions(test_file) \n            x_data_series = {\n                'id': filename[:-4] + \"_x\",\n                'data_series': x_preds,\n                'chart_type': graph_type[0] \n            }\n\n            y_data_series = {\n                'id': filename[:-4] + \"_y\",\n                'data_series': y_preds,\n                'chart_type': graph_type[0] \n            }\n        except Exception as e:\n            print('Error while processing image:', filename)\n            print(e)\n            \n    if graph_type[0] == 'horizontal_bar':\n        try:\n            x_preds, y_preds = horizontal_predictor.get_predictions(test_file) \n            x_data_series = {\n                'id': filename[:-4] + \"_x\",\n                'data_series': x_preds,\n                'chart_type': graph_type[0] \n            }\n\n            y_data_series = {\n                'id': filename[:-4] + \"_y\",\n                'data_series': y_preds,\n                'chart_type': graph_type[0] \n            }\n        except Exception as e:\n            print('Error while processing image:', filename)\n            print(e)\n\n    if graph_type[0] == 'dot':\n        try:\n            x_preds, y_preds = dot_predictor.get_predictions(test_file)\n            x_data_series = {\n                'id': filename[:-4] + \"_x\",\n                'data_series': x_preds,\n                'chart_type': graph_type[0] \n            }\n\n            y_data_series = {\n                'id': filename[:-4] + \"_y\",\n                'data_series': y_preds,\n                'chart_type': graph_type[0] \n            }\n\n        except Exception as e:\n            print('Error while processing image:', filename)\n            print(e)\n\n    # Append the x-axis and y-axis data series dictionaries to the submission data DataFrame\n    submission_data = submission_data.append(x_data_series, ignore_index=True)\n    submission_data = submission_data.append(y_data_series, ignore_index=True)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:58:56.907302Z","iopub.execute_input":"2023-06-18T11:58:56.908018Z","iopub.status.idle":"2023-06-18T11:59:10.398491Z","shell.execute_reply.started":"2023-06-18T11:58:56.907968Z","shell.execute_reply":"2023-06-18T11:59:10.397285Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_data[submission_data['chart_type'] == 'vertical_bar']","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.403426Z","iopub.execute_input":"2023-06-18T11:59:10.406573Z","iopub.status.idle":"2023-06-18T11:59:10.430695Z","shell.execute_reply.started":"2023-06-18T11:59:10.406524Z","shell.execute_reply":"2023-06-18T11:59:10.429541Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_data[submission_data['chart_type'] == 'line']","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.432410Z","iopub.execute_input":"2023-06-18T11:59:10.433210Z","iopub.status.idle":"2023-06-18T11:59:10.450727Z","shell.execute_reply.started":"2023-06-18T11:59:10.433148Z","shell.execute_reply":"2023-06-18T11:59:10.448546Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_data[submission_data['chart_type'] == 'scatter']","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.452729Z","iopub.execute_input":"2023-06-18T11:59:10.453700Z","iopub.status.idle":"2023-06-18T11:59:10.474578Z","shell.execute_reply.started":"2023-06-18T11:59:10.453648Z","shell.execute_reply":"2023-06-18T11:59:10.473117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_data[submission_data['chart_type'] == 'horizontal_bar']","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.476490Z","iopub.execute_input":"2023-06-18T11:59:10.477268Z","iopub.status.idle":"2023-06-18T11:59:10.489432Z","shell.execute_reply.started":"2023-06-18T11:59:10.477227Z","shell.execute_reply":"2023-06-18T11:59:10.488086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_data[submission_data['chart_type'] == 'dot']","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.491142Z","iopub.execute_input":"2023-06-18T11:59:10.491928Z","iopub.status.idle":"2023-06-18T11:59:10.504061Z","shell.execute_reply.started":"2023-06-18T11:59:10.491890Z","shell.execute_reply":"2023-06-18T11:59:10.502980Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_data.head()","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.505866Z","iopub.execute_input":"2023-06-18T11:59:10.506611Z","iopub.status.idle":"2023-06-18T11:59:10.521599Z","shell.execute_reply.started":"2023-06-18T11:59:10.506573Z","shell.execute_reply":"2023-06-18T11:59:10.520120Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_data.shape","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.523287Z","iopub.execute_input":"2023-06-18T11:59:10.523992Z","iopub.status.idle":"2023-06-18T11:59:10.531403Z","shell.execute_reply.started":"2023-06-18T11:59:10.523952Z","shell.execute_reply":"2023-06-18T11:59:10.530183Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_data.to_csv('submission.csv', index=None)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.533572Z","iopub.execute_input":"2023-06-18T11:59:10.534481Z","iopub.status.idle":"2023-06-18T11:59:10.541631Z","shell.execute_reply.started":"2023-06-18T11:59:10.534440Z","shell.execute_reply":"2023-06-18T11:59:10.540384Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import json\n# gt_annotation = \"/kaggle/input/bmga-tv-data/valid/annotations\"\n# gt_images = \"/kaggle/input/bmga-tv-data/valid/images\"\n# xseries = []\n# yseries = []\n# ground_truth_data = pd.DataFrame(columns=['id', 'data_series', 'chart_type'])\n\n# for jsn in os.listdir(gt_images):\n#     jsn_file = os.path.join(gt_annotation,jsn[:-4] + \".json\")\n#     with open(jsn_file) as f:\n#         data=json.load(f)\n#     series = data['data-series']\n#     ctype = data['chart-type']\n#     xs=[]\n#     ys=[]\n#     for value in series:\n#         if ctype == 'horizontal_bar':\n#             xs.append(str(value['y']))\n#             ys.append(str(value['x']))\n#         else:\n#             xs.append(str(value['x']))\n#             ys.append(str(value['y']))\n#     xser = ';'.join(xs)\n#     yser = ';'.join(ys)\n    \n#     gt_x_data_series = {\n#         'id': jsn[:-4] + \"_x\",\n#         'data_series': xser,\n#         'chart_type': ctype \n#     }\n    \n#     gt_y_data_series = {\n#         'id': jsn[:-4] + \"_y\",\n#         'data_series': yser,\n#         'chart_type': ctype\n#     }\n#     ground_truth_data = ground_truth_data.append(gt_x_data_series, ignore_index=True)\n#     ground_truth_data = ground_truth_data.append(gt_y_data_series, ignore_index=True)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.543737Z","iopub.execute_input":"2023-06-18T11:59:10.544644Z","iopub.status.idle":"2023-06-18T11:59:10.551658Z","shell.execute_reply.started":"2023-06-18T11:59:10.544604Z","shell.execute_reply":"2023-06-18T11:59:10.550487Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import matplotlib.pyplot as plt\n# import numpy as np\n# import pandas as pd\n# from rapidfuzz.distance.Levenshtein import distance as levenshtein\n# from sklearn.metrics import r2_score\n\n# def sigmoid(x):\n#     return 2 - 2 / (1 + np.exp(-x))\n\n\n# def normalized_rmse(y_true, y_pred):\n#     # The argument to the sigmoid transform is equal to \n#     # rmse(y_true, y_pred) / rmse(y_true, np.mean(y_true))\n    \n# #     return sigmoid((1 - r2_score(y_true, y_pred)) ** 0.5)\n\n\n# def normalized_levenshtein_score(y_true, y_pred):\n#     total_distance = np.sum([levenshtein(yt, yp) for yt, yp in zip(y_true, y_pred)])\n#     length_sum = np.sum([len(yt) for yt in y_true])\n#     return sigmoid(total_distance / length_sum)\n\n\n# def score_series(y_true, y_pred):\n#     if len(y_true) != len(y_pred):\n#         return 0.0\n#     if isinstance(y_true[0], str):\n#         return normalized_levenshtein_score(y_true, y_pred)\n#     else:\n#         return normalized_rmse(y_true, y_pred)\n\n\n# def benetech_score(ground_truth: pd.DataFrame, predictions: pd.DataFrame) -> float:\n#     \"\"\"Evaluate predictions using the metric from the Benetech - Making Graphs Accessible.\n    \n#     Parameters\n#     ----------\n#     ground_truth: pd.DataFrame\n#         Has columns `[data_series, chart_type]` and an index `id`. Values in `data_series` \n#         should be either arrays of floats or arrays of strings.\n    \n#     predictions: pd.DataFrame\n#     \"\"\"\n#     if not ground_truth.index.equals(predictions.index):\n#         raise ValueError(\"Must have exactly one prediction for each ground-truth instance.\")\n#     if not ground_truth.columns.equals(predictions.columns):\n#         raise ValueError(f\"Predictions must have columns: {ground_truth.columns}.\")\n#     pairs = zip(ground_truth.itertuples(index=True), predictions.itertuples(index=True))\n#     scores = []\n#     not_matching = 0\n#     low_score_images_list = []\n#     for (index, gt_series, gt_type), (index, pred_series, pred_type) in pairs:\n#         if gt_type != pred_type:  # Check chart_type condition                \n#             scores.append(0.0)\n#         else:  # Score with RMSE or Levenshtein as appropriate\n#             if gt_type == 'line':\n#                 if score_series(gt_series, pred_series) == 0:\n#                     low_score_images_list.append(index.split('_')[0] + '.jpg')\n#             scores.append(score_series(gt_series, pred_series))\n#     return np.mean(scores), low_score_images_list\n","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.553859Z","iopub.execute_input":"2023-06-18T11:59:10.554760Z","iopub.status.idle":"2023-06-18T11:59:10.563434Z","shell.execute_reply.started":"2023-06-18T11:59:10.554643Z","shell.execute_reply":"2023-06-18T11:59:10.562230Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# submission_data['chart_type'].value_counts()","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.565389Z","iopub.execute_input":"2023-06-18T11:59:10.566341Z","iopub.status.idle":"2023-06-18T11:59:10.576240Z","shell.execute_reply.started":"2023-06-18T11:59:10.566301Z","shell.execute_reply":"2023-06-18T11:59:10.574872Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# submission_data_type = submission_data[submission_data['chart_type'] == 'horizontal_bar']\n# ground_truth_data_type = ground_truth_data[ground_truth_data['chart_type'] == 'horizontal_bar']","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.577668Z","iopub.execute_input":"2023-06-18T11:59:10.578913Z","iopub.status.idle":"2023-06-18T11:59:10.585868Z","shell.execute_reply.started":"2023-06-18T11:59:10.578873Z","shell.execute_reply":"2023-06-18T11:59:10.584611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# gt = ground_truth_data.set_index(ground_truth_data.columns.tolist()[0])\n# pt = submission_data.set_index(submission_data.columns.tolist()[0])\n\n# # Get the intersection of indices\n# common_indices = gt.index.intersection(pt.index)\n\n# # Remove records not in common indices from both dataframes\n# gt = gt.loc[common_indices]\n# pt = pt.loc[common_indices]\n\n# # gt = gt[gt.index.str.contains('y')]\n# # pt = pt[pt.index.str.contains('y')]\n\n\n# for index, row in gt.iterrows():\n#     if 'y' in index or row['chart_type'] == 'scatter':\n#         row['data_series'] = [float(i) for i in row['data_series'].split(';')]\n#     else:\n#         row['data_series'] = row['data_series'].split(';')\n\n# for index, row in pt.iterrows():\n#     if 'y' in index or row['chart_type'] == 'scatter':\n#         try:\n#             row['data_series'] = [float(i) for i in row['data_series'].split(';')]\n#         except Exception as e:\n#             print(row['data_series'])\n#             print(row)\n#             print(e)\n#             row['data_series'] = row['data_series'].split(';')\n#     else:\n#         row['data_series'] = row['data_series'].split(';')\n\n# print(len(gt), len(pt))\n# benetech_score(gt, pt)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-18T11:59:10.587938Z","iopub.execute_input":"2023-06-18T11:59:10.588910Z","iopub.status.idle":"2023-06-18T11:59:10.600902Z","shell.execute_reply.started":"2023-06-18T11:59:10.588868Z","shell.execute_reply":"2023-06-18T11:59:10.599772Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---","metadata":{}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}