{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# !pip install git+https://github.com/huggingface/transformers\n!pip uninstall -y transformers\n# !pip install --no-deps /kaggle/input/certifi/}certifi-2022.12.7-py3-none-any.whl\n!mkdir temp && cp -r /kaggle/input/transformers-pix2struct-fix temp/transformers && cd temp/transformers && python setup.py develop --no-deps","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:41:04.254200Z","iopub.execute_input":"2023-06-19T18:41:04.255137Z","iopub.status.idle":"2023-06-19T18:41:25.807295Z","shell.execute_reply.started":"2023-06-19T18:41:04.255097Z","shell.execute_reply":"2023-06-19T18:41:25.805993Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path.insert(0, \"/kaggle/working/temp/transformers/src/\")\n","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:41:25.810030Z","iopub.execute_input":"2023-06-19T18:41:25.810741Z","iopub.status.idle":"2023-06-19T18:41:25.817878Z","shell.execute_reply.started":"2023-06-19T18:41:25.810696Z","shell.execute_reply":"2023-06-19T18:41:25.816506Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# installation\n!cp -r /kaggle/input/bmgacode/bmgacode ./bmga","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-06-19T18:41:25.820793Z","iopub.execute_input":"2023-06-19T18:41:25.822092Z","iopub.status.idle":"2023-06-19T18:41:28.381450Z","shell.execute_reply.started":"2023-06-19T18:41:25.822051Z","shell.execute_reply":"2023-06-19T18:41:28.380070Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !tar -czvf ./pycocotools-2.0.6.tar.gz /kaggle/input/bmgadeps/pycocotools-2.0.6/pycocotools-2.0.6/* && pip install pycocotools-2.0.6.tar.gz\n# !cd /kaggle/input/bmgadeps/pycocotools-2.0.6/pycocotools-2.0.6 && pip install .\n# !cp -r /kaggle/input/bmgadeps/pycocotools-2.0.6/pycocotools-2.0.6 ./pycocotools-2.0.6 && cd pycocotools-2.0.6 && python setup.py develop\n!pip install --disable-pip-version-check --no-deps --root-user-action=ignore /kaggle/input/pycocotools202/pycocotools-2.0.2-cp37-cp37m-linux_x86_64.whl","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:41:28.385027Z","iopub.execute_input":"2023-06-19T18:41:28.385714Z","iopub.status.idle":"2023-06-19T18:41:30.574648Z","shell.execute_reply.started":"2023-06-19T18:41:28.385671Z","shell.execute_reply":"2023-06-19T18:41:30.573299Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pycocotools","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:41:30.576607Z","iopub.execute_input":"2023-06-19T18:41:30.577292Z","iopub.status.idle":"2023-06-19T18:41:30.588024Z","shell.execute_reply.started":"2023-06-19T18:41:30.577243Z","shell.execute_reply":"2023-06-19T18:41:30.586975Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install --disable-pip-version-check --no-deps --root-user-action=ignore /kaggle/input/bmgadeps2/segmentation_models_pytorch-0.3.2-py3-none-any.whl\n!pip install --disable-pip-version-check --no-deps --root-user-action=ignore /kaggle/input/bmgadeps2/pretrainedmodels-0.7.4-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:41:30.589908Z","iopub.execute_input":"2023-06-19T18:41:30.590404Z","iopub.status.idle":"2023-06-19T18:41:35.045184Z","shell.execute_reply.started":"2023-06-19T18:41:30.590363Z","shell.execute_reply":"2023-06-19T18:41:35.043870Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cp --recursive /kaggle/input/bmgadeps2/efficientnet_pytorch-0.7.1/ efficientnet_pytorch-0.7.1 && cd efficientnet_pytorch-0.7.1/efficientnet_pytorch-0.7.1 && pip install --disable-pip-version-check --no-deps --root-user-action=ignore .","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:41:35.047518Z","iopub.execute_input":"2023-06-19T18:41:35.047977Z","iopub.status.idle":"2023-06-19T18:41:38.837816Z","shell.execute_reply.started":"2023-06-19T18:41:35.047931Z","shell.execute_reply":"2023-06-19T18:41:38.836536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install --disable-pip-version-check --no-deps --root-user-action=ignore /kaggle/input/bmgadeps/loguru-0.6.0-py3-none-any.whl\n# !pip install /kaggle/input/bmgadeps/pycocotools-2.0.6.tar.gz\n!pip install --disable-pip-version-check --no-deps --root-user-action=ignore /kaggle/input/bmgadeps/lmdb-1.4.0-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl\n!pip install --disable-pip-version-check --no-deps --root-user-action=ignore /kaggle/input/bmgadeps/anyconfig-0.13.0-py2.py3-none-any.whl \n!pip install --disable-pip-version-check --no-deps --root-user-action=ignore /kaggle/input/bmgadeps/gevent-22.10.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n!pip install --disable-pip-version-check --no-deps --root-user-action=ignore /kaggle/input/bmgadeps/zope.event-4.6-py2.py3-none-any.whl\n!pip install --disable-pip-version-check --no-deps --root-user-action=ignore /kaggle/input/bmgadeps/zope.interface-6.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n!pip install --disable-pip-version-check --no-deps --root-user-action=ignore /kaggle/input/bmgadeps/gevent_websocket-0.10.1-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:41:38.840406Z","iopub.execute_input":"2023-06-19T18:41:38.841156Z","iopub.status.idle":"2023-06-19T18:41:55.614138Z","shell.execute_reply.started":"2023-06-19T18:41:38.841112Z","shell.execute_reply":"2023-06-19T18:41:55.612297Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nos.chdir(\"/kaggle/working/bmga\")","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:41:55.616275Z","iopub.execute_input":"2023-06-19T18:41:55.616667Z","iopub.status.idle":"2023-06-19T18:41:55.622527Z","shell.execute_reply.started":"2023-06-19T18:41:55.616626Z","shell.execute_reply":"2023-06-19T18:41:55.621136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cd text_detection/src/assets/ops/dcn/ && python setup.py build_ext --inplace","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:41:55.626156Z","iopub.execute_input":"2023-06-19T18:41:55.626834Z","iopub.status.idle":"2023-06-19T18:43:38.921065Z","shell.execute_reply.started":"2023-06-19T18:41:55.626778Z","shell.execute_reply":"2023-06-19T18:43:38.919787Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# %%\nTEST_MODE = True\n\nif TEST_MODE:\n DATA_FOLDER = \"/kaggle/input/benetech-making-graphs-accessible/test/\"\n WEIGHTS_FOLDER = \"/kaggle/input/bmgaweights/\"\n WEIGHTS_FULL_FOLDER = \"/kaggle/input/bmgaweightsfull/\"\nelse:\n DATA_FOLDER = \"./data/validation/\"\n WEIGHTS_FOLDER = \"./weights/\"\n\n# %%\n%load_ext autoreload\n%autoreload 2\n\n# export environment variables\nimport os\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n\nimport sys\n\nsys.path.insert(0, \"./text_detection/src\")\n# sys.path.append(\"./classification/src\")\nsys.path.append(\"./detection/src\")\nsys.path.append(\"./segmentation/src\")\nsys.path.append(\"./text_recognition/src\")\n\n# %% [markdown]\n# # Compute metrics code\n\n# %%\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nfrom rapidfuzz.distance.Levenshtein import distance as levenshtein\nfrom sklearn.metrics import r2_score\n\n# %%\ndef sigmoid(x):\n return 2 - 2 / (1 + np.exp(-x))\n\n\ndef 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 return sigmoid((1 - r2_score(y_true, y_pred)) ** 0.5)\n\n\ndef 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\ndef 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\ndef 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=False), predictions.itertuples(index=False))\n scores = []\n for (gt_series, gt_type), (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 scores.append(score_series(gt_series, pred_series))\n return np.mean(scores)\n\n# %% [markdown]\n# ### Inference Code\n\n# %%\nimport numpy as np\nimport pandas as pd\nimport os\n\nfrom classification.core import ClassificationModel\nfrom detection.core import ObjectDetectionModel\nfrom segmentation.core import SegmentationModel\nfrom text_recognition.core import TextRecognitionModel\nfrom text_detection.core import TextDetectionModel\nfrom postprocessing.core import Postprocessing\n\n# %%\nfrom matplotlib import pyplot as plt\nimport cv2\nimport numpy as np\n\n\n# %%\n# image_folder = \"./data/train/images\"\n# origin_image_paths = [os.path.join(image_folder, x) for x in os.listdir(image_folder) if \".jpg\" in x][:500]\n# image_paths = [os.path.join(image_folder, x) for x in os.listdir(image_folder) if \".jpg\" in x][:500]\nimage_folder = os.path.join(DATA_FOLDER, \"images\")\n\norigin_image_paths = [os.path.join(image_folder, x) for x in os.listdir(image_folder) if \".jpg\" in x]\nimage_paths = [os.path.join(image_folder, x) for x in os.listdir(image_folder) if \".jpg\" in x]\n\n# %%\n# graph_classfication_config = {\n# \"model_name\": \"resnet50\",\n# \"n_classes\": 5,\n# \"weights_path\": os.path.join(WEIGHTS_FOLDER, \"graph_classification.pth\"),\n# }\n\ngraph_classfication_config = {\n \"model_name\": \"tf_efficientnetv2_s_in21ft1k\",\n# \"model_name\": \"tf_efficientnetv2_m_in21ft1k\",\n \"n_classes\": 5,\n \"weights_path\": os.path.join(WEIGHTS_FULL_FOLDER, \"graph_classification_full2.pth\"),\n# \"weights_path\": os.path.join(\"/kaggle/input/bmga-graph-classification-weights/graph_classification_v2m_all.bin\"),\n# \"weights_path\": os.path.join(WEIGHTS_FOLDER, \"graph_classification2.pth\"),\n}\n\nx_type_classification_config = {\n \"model_name\": \"resnet50\",\n \"n_classes\": 2,\n \"weights_path\": os.path.join(WEIGHTS_FOLDER, \"x_type_classification.pth\"),\n}\n\ny_type_classification_config = {\n \"model_name\": \"resnet50\",\n \"n_classes\": 2,\n \"weights_path\": os.path.join(WEIGHTS_FOLDER, \"y_type_classification.pth\"),\n}\n\nold_line_segmentation_config = {\n \"weights_path\": os.path.join(WEIGHTS_FULL_FOLDER, \"cxl_512_ftune_full.pth\"),\n \"arch\": \"UnetPlusPlusFix\",\n \"encoder_name\": \"convnext_xlarge_in22ft1k\",\n \"drop_path\": 0,\n \"size\": 512,\n \"version\": \"old\",\n}\n\nline_segmentation_config = {\n \"weights_path\": os.path.join(WEIGHTS_FULL_FOLDER, \"outputs_cb_ftune.pth\"),\n \"arch\": \"UnetPlusPlusFix\",\n \"encoder_name\": \"convnext_base_384_in22ft1k\",\n \"drop_path\": 0,\n \"size\": 512,\n \"version\": \"new\",\n}\n\nvertical_bar_segmentation_config = {\n \"weights_path\": os.path.join(WEIGHTS_FOLDER, \"outputs_v2s_bar_512.pth\"),\n \"arch\": \"Unet\",\n \"encoder_name\": \"tf_efficientnetv2_s_in21ft1k\",\n \"drop_path\": 0,\n \"size\": 512,\n \"version\": \"old\",\n}\n\nkeypoint_detection_config = {\n \"name\": \"keypoint_detection\",\n \"experiment_path\": \"./detection/src/exps/example/custom/bmga.py\",\n # \"weights_path\": os.path.join(WEIGHTS_FOLDER, \"keypoint_detection.pth\"),\n \"weights_path\": os.path.join(WEIGHTS_FULL_FOLDER, \"keypoint_detection_yoloxx_xy_full.pth\"),\n \"classes\": [\"value\", \"x\", \"y\"], #, \"x_label\", \"y_label\"],\n \"conf_thre\": 0.05,\n \"nms_thre\": 0.25,\n \"test_size\": (640, 640),\n}\n\ntext_detection_config = {\n \"weights_path\": os.path.join(WEIGHTS_FOLDER, \"synthtext_finetune_ic19_res50_dcn_fpn_dbv2\"),\n \"config_path\": \"text_detection/src/experiments/ASF/td500_resnet50_deform_thre_asf_inference.yaml\",\n \"image_short_side\": 768,\n \"thresh\": 0.1,\n \"box_thresh\": 0.2,\n \"resize\": False,\n \"polygon\": True,\n}\n\nx_labels_text_detection_config = {\n # \"weights_path\": os.path.join(WEIGHTS_FOLDER, \"db_x_labels\"),\n \"weights_path\": os.path.join(WEIGHTS_FULL_FOLDER, \"db_x_labels_full2\"),\n \"config_path\": \"./text_detection/src/experiments/ASF/td500_resnet50_deform_thre_asf_inference.yaml\",\n \"image_short_side\": 896,\n \"thresh\": 0.2,\n \"box_thresh\": 0.4,\n \"resize\": False,\n \"polygon\": True,\n}\n\ny_labels_text_detection_config = {\n \"weights_path\": os.path.join(WEIGHTS_FULL_FOLDER, \"db_y_labels_full3\"),\n \"config_path\": \"./text_detection/src/experiments/ASF/td500_resnet50_deform_thre_asf_inference.yaml\",\n \"image_short_side\": 768,\n \"thresh\": 0.05,\n \"box_thresh\": 0.25,\n \"resize\": False,\n \"polygon\": True,\n}\n\ntext_recognition_config = {\n \"weights_path\": os.path.join(WEIGHTS_FOLDER, \"parseq-bb5792a6.pt\"),\n # \"weights_path\": \"baudm/parseq\",\n \"model_name\": \"parseq\",\n \"config_path\": os.path.join(WEIGHTS_FOLDER, \"parseq_hparams.json\"),\n}\n\ngraph_classification_model = ClassificationModel(**graph_classfication_config)\nx_type_classification_model = ClassificationModel(**x_type_classification_config)\ny_type_classification_model = ClassificationModel(**y_type_classification_config)\n# old_line_segmentation_model = SegmentationModel(**old_line_segmentation_config)\n# line_segmentation_model = SegmentationModel(**old_line_segmentation_config)\nline_segmentation_model = SegmentationModel(**line_segmentation_config)\nvertical_bar_segmentation_model = SegmentationModel(**vertical_bar_segmentation_config)\n\nkeypoint_detection_model = ObjectDetectionModel(**keypoint_detection_config)\n\n# keypoint_detection_config[\"weights_path\"] = os.path.join(WEIGHTS_FOLDER, \"keypoint_detection_yoloxx_xy.pth\")\n# xy_keypoint_detection_model = ObjectDetectionModel(**keypoint_detection_config)\n\nkeypoint_detection_config[\"conf_thre\"] = 0.15\nkeypoint_detection_config[\"nms_thre\"] = 0.25\nkeypoint_detection_config[\"weights_path\"] = os.path.join(WEIGHTS_FOLDER, \"keypoint_detection_yoloxx_line_value.pth\")\nline_value_keypoint_detection_model = ObjectDetectionModel(**keypoint_detection_config)\n\n# keypoint_detection_config[\"conf_thre\"] = 0.3\n# keypoint_detection_config[\"nms_thre\"] = 0.3\nkeypoint_detection_config[\"weights_path\"] = os.path.join(WEIGHTS_FULL_FOLDER, \"keypoint_detection_yoloxx_bar_value_full.pth\")\nbar_value_keypoint_detection_model = ObjectDetectionModel(**keypoint_detection_config)\n\nkeypoint_detection_config[\"conf_thre\"] = 0.3\nkeypoint_detection_config[\"nms_thre\"] = 0.45\n# keypoint_detection_config[\"test_size\"] = (768, 768)\nkeypoint_detection_config[\"weights_path\"] = os.path.join(WEIGHTS_FULL_FOLDER, \"keypoint_detection_yoloxx_scatter_value_full2.pth\")\nscatter_value_keypoint_detection_model = ObjectDetectionModel(**keypoint_detection_config)\n\ntext_detection_model = TextDetectionModel(**text_detection_config)\nx_labels_text_detection_model = TextDetectionModel(**x_labels_text_detection_config)\nx_labels_text_detection_config[\"image_short_side\"] = 1024\nx_labels_text_detection_model_2 = TextDetectionModel(**x_labels_text_detection_config)\ny_labels_text_detection_model = TextDetectionModel(**y_labels_text_detection_config)\ntext_recognition_model = TextRecognitionModel(**text_recognition_config)\ntext_recognition_model.parseq.eval()\nprint()\n\n# %%\n# read ground truth from /home/thanh/bmga/data/validation/metadata.jsonl\nimport json\n\nif not TEST_MODE:\n # with open(\"/home/thanh/bmga/data/train/metadata.jsonl\", \"r\") as f:\n with open(\"/home/thanh/bmga/data/validation/metadata.jsonl\", \"r\") as f:\n metadata = [json.loads(x) for x in f.readlines()]\n\n metadata_dict = {}\n for x in metadata:\n metadata_dict[x[\"file_name\"]] = x\n\n filtered_image_paths = []\n filtered_original_image_paths = []\n\n for image_path in image_paths:\n if \"images/\" + image_path.split(\"/\")[-1] not in metadata_dict.keys():\n continue\n filtered_image_paths.append(image_path)\n filtered_original_image_paths.append(image_path)\n\n image_paths = filtered_image_paths\n origin_image_paths = filtered_original_image_paths\n\n# %% [markdown]\n# ### Utility Functions\n\n# %%\n# function to convert polygon points to smallest 4 points polygon\ndef convert_polygon_to_min_rect(polygon):\n polygon = np.array(polygon)\n polygon = polygon.reshape(-1, 2)\n polygon = polygon.astype(np.float32)\n rect = cv2.minAreaRect(polygon)\n box = cv2.boxPoints(rect)\n box = np.int0(box)\n\n return box\n\ndef crop_polygon_from_image(image, polygon):\n polygon = convert_polygon_to_min_rect(polygon)\n mask = np.zeros(image.shape[:2], np.uint8)\n cv2.drawContours(mask, [polygon], 0, 255, -1, cv2.LINE_AA)\n out = 255 - np.zeros_like(image)\n out[mask == 255] = image[mask == 255]\n\n # return crop from image\n try:\n crop = out[np.min(polygon[:, 1]):np.max(polygon[:, 1]), np.min(polygon[:, 0]):np.max(polygon[:, 0])]\n except:\n crop = np.ones((32, 100, 3), dtype=np.uint8) * 255\n return crop\n\n\n# sample_image_path = image_paths[0]\n# sample_image = cv2.imread(sample_image_path)\n# sample_image = cv2.cvtColor(sample_image, cv2.COLOR_BGR2RGB)\n\n# sample_polygon = [[20, 20], [10, 100], [100, 200], [300, 40]]\n\n# # draw polygon\n# sample_image = cv2.polylines(sample_image, [np.array(sample_polygon)], True, (0, 255, 0), 2)\n# plt.imshow(sample_image)\n\n# %%\n# crop = crop_polygon_from_image(sample_image, sample_polygon)\n# plt.imshow(crop)\n\n# %%\ndef filter_x_polygons(polygons, img_height, img_path):\n# # calculate polygon area\n# from shapely.geometry import Polygon\n\n# polygon_areas = []\n# for polygon in polygons:\n# if polygon:\n# polygon_areas.append(Polygon(polygon).area)\n# else:\n# polygon_areas.append(0)\n\n# # remove \n# mean_area = np.mean(polygon_areas)\n# filtered_polygons = []\n# for polygon, area in zip(polygons, polygon_areas):\n# if area > mean_area * 0.2:\n# filtered_polygons.append(polygon)\n \n# polygons = filtered_polygons\n\n # first, draw a line along y axis then count the number of x_label_boxes that intersect with the line\n max_count = 0\n max_count_line_y = 0\n\n for line_y in range(img_height):\n count = 0\n for polygon in polygons:\n if polygon:\n min_y = min([x[1] for x in polygon])\n max_y = max([x[1] for x in polygon])\n else:\n min_y = 0\n max_y = 0\n\n if min_y <= line_y <= max_y:\n count += 1\n if count > max_count:\n max_count = count\n max_count_line_y = line_y\n\n # filter out y_label_boxes that intersect with the line\n filtered_x_label_polygons = []\n for polygon in polygons:\n if polygon:\n min_y = min([x[1] for x in polygon])\n max_y = max([x[1] for x in polygon])\n else:\n min_y = 0\n max_y = 0\n\n if min_y <= max_count_line_y <= max_y:\n filtered_x_label_polygons.append(polygon)\n\n return filtered_x_label_polygons\n\n\ndef filter_y_polygons(polygons, img_width, image):\n # first, draw a line along x axis then count the number of y_label_boxes that intersect with the line\n max_count = 0\n max_count_line_x = 0\n\n for line_x in range(img_width):\n count = 0\n for polygon in polygons:\n if polygon:\n min_x = min([x[0] for x in polygon])\n max_x = max([x[0] for x in polygon])\n else:\n min_x = 0\n max_x = 0\n w = max_x - min_x\n if min_x + w // 4 <= line_x <= max_x - w // 4:\n count += 1\n if count > max_count:\n max_count = count\n max_count_line_x = line_x\n\n # filter out y_label_boxes that intersect with the line\n filtered_y_label_polygons = []\n for polygon in polygons:\n if polygon:\n min_x = min([x[0] for x in polygon])\n max_x = max([x[0] for x in polygon])\n else:\n min_x = 0\n max_x = 0\n if min_x <= max_count_line_x <= max_x:\n filtered_y_label_polygons.append(polygon)\n\n return filtered_y_label_polygons\n # # second, do text recognition on y_label_boxes\n # crops = []\n # for polygon in filtered_y_label_polygons:\n # crop = crop_polygon_from_image(image, polygon)\n # crops.append(crop)\n\n # text_recognition_results = text_recognition_model.predict(crops)\n\n # # filter out those boxes that the values can't be converted to float: only case that y labels are numbers, have to update\n # filtered_y_label_boxes_2 = []\n # for i, box in enumerate(filtered_y_label_polygons):\n # try:\n # text = \"\".join([c for c in text_recognition_results[0][i][0] if c in \"0123456789.\"])\n # if not text:\n # float(text)\n # filtered_y_label_boxes_2.append(box)\n # except:\n # pass\n\n # return filtered_y_label_boxes_2\n\ndef calculate_iou(polygon1, polygon2, image):\n # calculate iou between two polygons\n polygon1 = np.array(polygon1)\n polygon2 = np.array(polygon2)\n polygon1 = polygon1.reshape(-1, 2)\n polygon2 = polygon2.reshape(-1, 2)\n polygon1 = polygon1.astype(np.float32)\n polygon2 = polygon2.astype(np.float32)\n\n rect1 = cv2.minAreaRect(polygon1)\n box1 = cv2.boxPoints(rect1)\n box1 = np.int0(box1)\n\n rect2 = cv2.minAreaRect(polygon2)\n box2 = cv2.boxPoints(rect2)\n box2 = np.int0(box2)\n\n mask1 = np.zeros(image.shape[:2], np.uint8)\n cv2.drawContours(mask1, [box1], 0, 255, -1, cv2.LINE_AA)\n mask2 = np.zeros(image.shape[:2], np.uint8)\n cv2.drawContours(mask2, [box2], 0, 255, -1, cv2.LINE_AA)\n\n intersection = np.logical_and(mask1, mask2)\n union = np.logical_or(mask1, mask2)\n iou_score = np.sum(intersection) / np.sum(union)\n\n return iou_score\n\ndef calculate_label_polygons_accuracy(pred_polygons, gt_polygons, image, is_x_label=True, iou_thre=0.5):\n if len(pred_polygons) != len(gt_polygons):\n return 0\n \n if is_x_label:\n gt_polygons = sorted(gt_polygons, key=lambda x: min([y[0] for y in x]) if x else 0)\n gt_polygons = sorted(pred_polygons, key=lambda x: min([y[0] for y in x]) if x else 0)\n else:\n gt_polygons = sorted(gt_polygons, key=lambda x: min([y[1] for y in x]) if x else 0)\n gt_polygons = sorted(pred_polygons, key=lambda x: min([y[1] for y in x]) if x else 0)\n\n iou_score = 0\n for i in range(len(gt_polygons)):\n iou = calculate_iou(gt_polygons[i], gt_polygons[i], image)\n if iou > iou_thre:\n iou_score += 1\n\n if iou_score == len(gt_polygons):\n return 1\n\n return 0\n\ndef visualize(image_path, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons):\n image = cv2.imread(image_path)\n\n if value_boxes is not None:\n for box in value_boxes:\n cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 0, 255), 2)\n\n if x_boxes is not None:\n for box in x_boxes:\n cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)\n\n if y_boxes is not None:\n for box in y_boxes:\n cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (255, 0, 0), 2) \n \n if x_labels_polygons is not None:\n # visualize x_label_boxes\n for polygon in x_labels_polygons:\n polygon = np.array(polygon)\n polygon = polygon.reshape(-1, 2)\n polygon = polygon.astype(np.int32)\n cv2.drawContours(image, [polygon], 0, (255, 255, 0), 2)\n\n if y_labels_polygons is not None:\n # visualize y_label_boxes\n for polygon in y_labels_polygons:\n polygon = np.array(polygon)\n polygon = polygon.reshape(-1, 2)\n polygon = polygon.astype(np.int32)\n cv2.drawContours(image, [polygon], 0, (0, 255, 255), 2)\n\n plt.figure(figsize=(8, 8))\n plt.imshow(image)\n\n# %% [markdown]\n# ### Graph classification model, x/y labels classification model\n\n# %%\ngraph_classes = ['dot', 'line', 'scatter', 'vertical_bar', \"horizontal_bar\"]\n\ngraph_type_predictions = graph_classification_model.predict(image_paths=image_paths, size=(640, 640))\n\n# convert predictions to graph type\ngraph_type_predictions = np.argmax(graph_type_predictions, axis=1)\ngraph_type_predictions = [graph_classes[i] for i in graph_type_predictions]\n\n# %%\nif not TEST_MODE:\n # convert predictions to graph type\n gt_classes = []\n\n for image_path in image_paths:\n gt_classes.append(metadata_dict[\"images/\" + image_path.split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"][\"class\"])\n\n # calculate accuracy\n acc = 0\n for idx in range(len(image_paths)):\n if graph_type_predictions[idx] == gt_classes[idx]:\n acc += 1\n\n print(\"acc: \", acc / len(image_paths))\n print(np.unique(gt_classes, return_counts=True))\n\n# %%\ntype_classes = [\"numerical\", \"categorical\"]\n\nx_type_predictions = x_type_classification_model.predict(image_paths=image_paths)\nx_type_predictions = np.argmax(x_type_predictions, axis=1)\nx_type_predictions = [type_classes[i] for i in x_type_predictions]\n\ny_type_predictions = y_type_classification_model.predict(image_paths=image_paths)\ny_type_predictions = np.argmax(y_type_predictions, axis=1)\ny_type_predictions = [type_classes[i] for i in y_type_predictions]\n\nif not TEST_MODE:\n x_type_gt_classes = []\n y_type_gt_classes = []\n\n for image_path in image_paths:\n x_type_gt_classes.append(metadata_dict[\"images/\" + image_path.split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"][\"x_type\"])\n y_type_gt_classes.append(metadata_dict[\"images/\" + image_path.split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"][\"y_type\"])\n\n # calculate accuracy\n x_type_acc = 0\n y_type_acc = 0\n for i in range(len(image_paths)):\n if x_type_predictions[i] == x_type_gt_classes[i]:\n x_type_acc += 1\n if y_type_predictions[i] == y_type_gt_classes[i]:\n y_type_acc += 1\n\n print(\"x_type_acc: \", x_type_acc / len(image_paths))\n print(\"y_type_acc: \", y_type_acc / len(image_paths))\n\n# %% [markdown]\n# ### X/Y labels detection using DB model\n\n# %%\n# x_labels_predictions = x_labels_text_detection_model.predict(image_paths=image_paths)\n# y_labels_predictions = y_labels_text_detection_model.predict(image_paths=image_paths)\n\n# %%\n# from tqdm import tqdm\n# # calucate accuracy\n# x_acc = 0\n# y_acc = 0\n\n# for idx in tqdm(range(len(image_paths))):\n# image = cv2.imread(image_paths[idx])\n\n# x_labels_polygons = x_labels_predictions[idx][0][0]\n# y_labels_polygons = y_labels_predictions[idx][0][0]\n\n# x_labels_polygons = filter_x_polygons(\n# x_labels_polygons,\n# image.shape[0],\n# image_paths[idx],\n# )\n\n# y_labels_polygons = filter_y_polygons(\n# y_labels_polygons,\n# image.shape[1],\n# image\n# )\n \n# x_acc += calculate_label_polygons_accuracy(\n# x_labels_polygons,\n# metadata_dict[\"images/\" + image_paths[idx].split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"][\"x_labels_polygons\"],\n# image=image,\n# is_x_label=True,\n# )\n \n# y_acc += calculate_label_polygons_accuracy(\n# y_labels_polygons,\n# metadata_dict[\"images/\" + image_paths[idx].split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"][\"y_labels_polygons\"],\n# image=image,\n# is_x_label=False, \n# )\n\n# print(\"x_acc: \", x_acc / len(image_paths))\n# print(\"y_acc: \", y_acc / len(image_paths))\n\n# # x_acc: 0.8872987477638641\n# # y_acc: 0.8461538461538461\n\n\n# %%\n# # visualize keypoint detection results, data is boxes\n# idx = 162\n# # idx = (idx + 1) % len(image_paths)\n# image = cv2.imread(image_paths[idx])\n# x_labels_polygons = x_labels_predictions[idx][0][0]\n# y_labels_polygons = y_labels_predictions[idx][0][0]\n\n# x_labels_polygons = filter_x_polygons(\n# x_labels_polygons,\n# image.shape[0],\n# image_paths[idx],\n# )\n\n# y_labels_polygons = filter_y_polygons(\n# y_labels_polygons,\n# image.shape[1],\n# image\n# )\n\n# # visualize x_label_boxes\n# image = cv2.imread(image_paths[idx])\n# for polygon in x_labels_polygons:\n# polygon = np.array(polygon)\n# polygon = polygon.reshape(-1, 2)\n# polygon = polygon.astype(np.int32)\n# cv2.drawContours(image, [polygon], 0, (0, 255, 0), 2)\n\n# # visualize y_label_boxes\n# for polygon in y_labels_polygons:\n# polygon = np.array(polygon)\n# polygon = polygon.reshape(-1, 2)\n# polygon = polygon.astype(np.int32)\n# cv2.drawContours(image, [polygon], 0, (0, 0, 255), 2)\n\n# plt.figure(figsize=(10, 10))\n# plt.imshow(image)\n# print(idx, image_paths[idx])\n\n# %%\n# # visualize ground truth\n# image = cv2.imread(image_paths[idx])\n# for polygon in metadata_dict[\"images/\" + image_paths[idx].split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"][\"y_labels_polygons\"]:\n# x0, y0, x1, y1, x2, y2, x3, y3 = polygon[\"x0\"], polygon[\"y0\"], polygon[\"x1\"], polygon[\"y1\"], polygon[\"x2\"], polygon[\"y2\"], polygon[\"x3\"], polygon[\"y3\"]\n# polygon = np.array([[x0, y0], [x1, y1], [x2, y2], [x3, y3]])\n# polygon = polygon.reshape(-1, 2)\n# polygon = polygon.astype(np.int32)\n# cv2.drawContours(image, [polygon], 0, (0, 0, 255), 2)\n\n# for polygon in metadata_dict[\"images/\" + image_paths[idx].split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"][\"x_labels_polygons\"]:\n# x0, y0, x1, y1, x2, y2, x3, y3 = polygon[\"x0\"], polygon[\"y0\"], polygon[\"x1\"], polygon[\"y1\"], polygon[\"x2\"], polygon[\"y2\"], polygon[\"x3\"], polygon[\"y3\"]\n# polygon = np.array([[x0, y0], [x1, y1], [x2, y2], [x3, y3]])\n# polygon = polygon.reshape(-1, 2)\n# polygon = polygon.astype(np.int32)\n# cv2.drawContours(image, [polygon], 0, (0, 255, 0), 2)\n# plt.imshow(image)\n# print(metadata_dict[\"images/\" + image_paths[idx].split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"][\"value\"])\n\n# %% [markdown]\n# ### Object detection model to detect point on graphs\n# 1. Detect x_labels and y_labels points\n# 2. Map these points with x_labels and y_labels texts\n# 3. Post processing depends on the graph type\n\n# %%\n# keypoint_predictions = keypoint_detection_model.predict(image_paths=image_paths)\n\n# %%\n# # sample image of horizontal bar chart\n# sample_image = cv2.imread(\"./data/validation/images/\" + image_paths[idx].split(\"/\")[-1])\n# sample_image = cv2.cvtColor(sample_image, cv2.COLOR_BGR2RGB)\n\n# # rotate image using cv2\n# sample_image = cv2.rotate(sample_image, cv2.ROTATE_90_COUNTERCLOCKWISE)\n\n# # horizontal flip\n# sample_image = cv2.flip(sample_image, 1)\n\n# # save temporary image\n# temp_path = \"./data/temp.jpg\"\n# cv2.imwrite(temp_path, sample_image)\n\n# plt.imshow(sample_image)\n\n# %%\n# x_labels_polygons = x_labels_text_detection_model.predict(image_paths=[temp_path])[0][0][0]\n# y_labels_polygons = y_labels_text_detection_model.predict(image_paths=[temp_path])[0][0][0]\n\n# x_labels_polygons = filter_x_polygons(\n# x_labels_polygons,\n# image.shape[0],\n# image_paths[idx],\n# )\n\n# y_labels_polygons = filter_y_polygons(\n# y_labels_polygons,\n# image.shape[1],\n# image\n# )\n\n# # visualize x_label_boxes\n# for polygon in x_labels_polygons:\n# polygon = np.array(polygon)\n# polygon = polygon.reshape(-1, 2)\n# polygon = polygon.astype(np.int32)\n# cv2.drawContours(sample_image, [polygon], 0, (0, 255, 0), 2)\n\n# # visualize y_label_boxes\n# for polygon in y_labels_polygons:\n# polygon = np.array(polygon)\n# polygon = polygon.reshape(-1, 2)\n# polygon = polygon.astype(np.int32)\n# cv2.drawContours(sample_image, [polygon], 0, (0, 0, 255), 2)\n\n# single_keypoint_predictions = keypoint_detection_model.predict(image_paths=[temp_path])\n# data = single_keypoint_predictions[0][0][0].cpu().numpy()\n\n# value_boxes = (data[data[:, 6] == 0][:, :4] / single_keypoint_predictions[1][0][\"ratio\"]).astype(int)\n# x_boxes = (data[data[:, 6] == 1][:, :4] / single_keypoint_predictions[1][0][\"ratio\"]).astype(int)\n# y_boxes = (data[data[:, 6] == 2][:, :4] / single_keypoint_predictions[1][0][\"ratio\"]).astype(int)\n\n\n# visualize(temp_path, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n\n# %%\ndef process_labels_polygons(idx, graph_type):\n TEMP_IMAGE_FOLDER = \"./data/temporary/\"\n if not os.path.exists(TEMP_IMAGE_FOLDER):\n os.makedirs(TEMP_IMAGE_FOLDER)\n\n if graph_type == \"horizontal_bar\":\n chosen_x_labels_text_detection_model = x_labels_text_detection_model_2\n chosen_y_labels_text_detection_model = y_labels_text_detection_model\n else:\n chosen_x_labels_text_detection_model = x_labels_text_detection_model\n chosen_y_labels_text_detection_model = y_labels_text_detection_model\n\n if graph_type == \"horizontal_bar\":\n try:\n sample_image = cv2.imread(os.path.join(image_folder, origin_image_paths[idx].split(\"/\")[-1]))\n\n # rotate image using cv2\n sample_image = cv2.rotate(sample_image, cv2.ROTATE_90_COUNTERCLOCKWISE)\n\n # horizontal flip\n sample_image = cv2.flip(sample_image, 1)\n\n # save temporary image\n temp_path = f\"{os.path.join(TEMP_IMAGE_FOLDER, os.path.basename(image_paths[idx]))}\"\n cv2.imwrite(temp_path, sample_image)\n\n image_paths[idx] = temp_path\n except:\n x_labels_polygons, y_labels_polygons = [], []\n\n try:\n x_labels_polygons = chosen_x_labels_text_detection_model.predict(image_paths=[image_paths[idx]])[0][0][0]\n y_labels_polygons = chosen_y_labels_text_detection_model.predict(image_paths=[image_paths[idx]])[0][0][0]\n except:\n print(\"can't predict original image, using resized image\")\n sample_image = cv2.imread(image_paths[idx])\n if graph_type == \"horizontal_bar\":\n # resize image to size = 512, 768\n sample_image = cv2.resize(sample_image, (512, 768))\n else:\n sample_image = cv2.resize(sample_image, (768, 512))\n\n temp_path = f\"{os.path.join(TEMP_IMAGE_FOLDER, os.path.basename(image_paths[idx]))}\"\n cv2.imwrite(temp_path, sample_image)\n image_paths[idx] = temp_path\n\n try:\n x_labels_polygons = chosen_x_labels_text_detection_model.predict(image_paths=[image_paths[idx]])[0][0][0]\n y_labels_polygons = chosen_y_labels_text_detection_model.predict(image_paths=[image_paths[idx]])[0][0][0]\n except:\n x_labels_polygons, y_labels_polygons = [], []\n\n image = cv2.imread(image_paths[idx])\n x_labels_polygons = filter_x_polygons(\n x_labels_polygons,\n image.shape[0],\n image_paths[idx],\n )\n\n y_labels_polygons = filter_y_polygons(\n y_labels_polygons,\n image.shape[1] // 2,\n image\n )\n\n try:\n # use general model x/y/value\n single_keypoint_predictions = keypoint_detection_model.predict(image_paths=[image_paths[idx]])\n data = single_keypoint_predictions[0][0][0].cpu().numpy()\n x_boxes = (data[data[:, 6] == 1][:, :4] / single_keypoint_predictions[1][0][\"ratio\"]).astype(int)\n y_boxes = (data[data[:, 6] == 2][:, :4] / single_keypoint_predictions[1][0][\"ratio\"]).astype(int)\n\n # use separate model x/y\n # xy_keypoint_predictions = xy_keypoint_detection_model.predict(image_paths=[image_paths[idx]])\n # xy_data = xy_keypoint_predictions[0][0][0].cpu().numpy()\n # x_boxes = (xy_data[xy_data[:, 6] == 1][:, :4] / xy_keypoint_predictions[1][0][\"ratio\"]).astype(int)\n # y_boxes = (xy_data[xy_data[:, 6] == 2][:, :4] / xy_keypoint_predictions[1][0][\"ratio\"]).astype(int)\n\n # use general model for all chart types\n # value_boxes = (data[data[:, 6] == 0][:, :4] / single_keypoint_predictions[1][0][\"ratio\"]).astype(int)\n\n # use separate model for each chart type\n if graph_type == \"line\":\n keypoint_predictions = line_value_keypoint_detection_model.predict(image_paths=[image_paths[idx]])\n value_data = keypoint_predictions[0][0][0].cpu().numpy()\n elif graph_type == \"scatter\":\n keypoint_predictions = scatter_value_keypoint_detection_model.predict(image_paths=[image_paths[idx]])\n value_data = keypoint_predictions[0][0][0].cpu().numpy()\n else:\n keypoint_predictions = bar_value_keypoint_detection_model.predict(image_paths=[image_paths[idx]])\n value_data = keypoint_predictions[0][0][0].cpu().numpy()\n\n value_boxes = (value_data[value_data[:, 6] == 0][:, :4] / keypoint_predictions[1][0][\"ratio\"]).astype(int)\n except:\n value_boxes, x_boxes, y_boxes = [], [], []\n\n return value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons\n\n\nif not TEST_MODE:\n # ------------ SELECT ONE SAMPLE ------------\n idx = 6\n # idx = (idx + 1) % len(image_paths)\n # while graph_type_predictions[idx] != \"line\":\n # idx = (idx + 1) % len(image_paths)\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons = process_labels_polygons(idx, graph_type_predictions[idx])\n print(\"-------- PREDICTION ---------\")\n print(\"len(x_labels_polygons): \", len(x_labels_polygons), \"len(x_boxes): \", len(x_boxes), \"len(value_boxes): \", len(value_boxes))\n\n visualize(image_paths[idx], value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n\n # ground truth\n print(\"-------- GROUND TRUTH ---------\")\n print(\"graph type: \", metadata_dict[\"images/\" + image_paths[idx].split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"][\"class\"])\n gt_values = metadata_dict[\"images/\" + image_paths[idx].split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"][\"value\"]\n for v in gt_values:\n print(v)\n\n print(\"len(gt_values): \", len(gt_values))\n\n print(idx, image_paths[idx].split(\"/\")[-1])\n\n word_polygons = text_detection_model.predict([image_paths[idx]])[0][0][0]\n # visualize on images\n new_image = cv2.imread(image_paths[idx])\n for word_polygon in word_polygons:\n word_polygon = np.array(word_polygon, dtype=np.int32)\n new_image = cv2.polylines(new_image, [word_polygon], True, (0, 255, 0), 2)\n\n plt.figure(figsize=(10, 10))\n plt.imshow(new_image)\n\n# %%\n# GENERAL RULE:\n# 1. Filter x_points, y_points by draw a line_y, line_x\ndef process_filter_xy_value_boxes(idx, x_boxes, y_boxes, value_boxes):\n def convert_4_points_box_to_polygon(box):\n x1, y1, x2, y2 = box\n return [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]\n\n image = cv2.imread(image_paths[idx])\n x_boxes_polygons = [convert_4_points_box_to_polygon(box) for box in x_boxes]\n x_boxes_polygons = filter_x_polygons(x_boxes_polygons, image.shape[0], image_paths[idx])\n x_boxes = [\n [\n min([p[0] for p in polygon]) if polygon else 0,\n min([p[1] for p in polygon]) if polygon else 0,\n max([p[0] for p in polygon]) if polygon else 0,\n max([p[1] for p in polygon]) if polygon else 0,\n ] \n for polygon in x_boxes_polygons\n ]\n\n y_boxes_polygons = [convert_4_points_box_to_polygon(box) for box in y_boxes]\n y_boxes_polygons = filter_y_polygons(y_boxes_polygons, image.shape[1] // 2, image)\n y_boxes = [\n [\n min([p[0] for p in polygon]) if polygon else 0,\n min([p[1] for p in polygon]) if polygon else 0,\n max([p[0] for p in polygon]) if polygon else 0,\n max([p[1] for p in polygon]) if polygon else 0,\n ]\n for polygon in y_boxes_polygons\n ]\n\n\n # draw Ox, Oy of the graph based on centers of x_boxes and y_boxes\n Oy = np.mean([(box[0] + box[2]) / 2 for box in y_boxes])\n Ox = np.mean([(box[1] + box[3]) / 2 for box in x_boxes])\n origin = (Ox, Oy)\n\n # filter out those value_boxes that are not in the graph\n filter_value_boxes = []\n for box in value_boxes:\n x1, y1, x2, y2 = box\n if x2 > origin[1] and y1 < origin[0] + 10:\n filter_value_boxes.append(box)\n\n return x_boxes, y_boxes, filter_value_boxes\n\n\nif not TEST_MODE:\n # visualize x_boxes and y_boxes\n x_boxes, y_boxes, value_boxes = process_filter_xy_value_boxes(idx, x_boxes, y_boxes, value_boxes)\n\n visualize(image_paths[idx], value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n\n# %%\ndef compute_iou_and_distance_one_direction(polygon, box, direction=\"x\"):\n x1, y1, x2, y2 = box\n if len(polygon) == 0:\n polygon_min_x, polygon_max_x, polygon_min_y, polygon_max_y = 0, 0, 0, 0\n else:\n polygon_min_x = min([p[0] for p in polygon])\n polygon_max_x = max([p[0] for p in polygon])\n polygon_min_y = min([p[1] for p in polygon])\n polygon_max_y = max([p[1] for p in polygon])\n\n\n if direction == \"x\":\n intersection = set(range(int(x1), int(x2))).intersection(set(range(int(polygon_min_x), int(polygon_max_x))))\n partly_union = set(range(int(x1), int(x2)))\n iou = len(intersection) / len(partly_union)\n distance = abs((y2 + y1) / 2 - (polygon_max_y + polygon_min_y) / 2)\n return iou, distance\n elif direction == \"y\":\n intersection = set(range(int(y1), int(y2))).intersection(set(range(int(polygon_min_y), int(polygon_max_y))))\n partly_union = set(range(int(y1), int(y2)))\n iou = len(intersection) / len(partly_union)\n distance = abs((x2 + x1) / 2 - (polygon_max_x + polygon_min_x) / 2)\n return iou, distance\n else:\n raise ValueError(\"direction must be x or y\")\n\ndef compute_iou_and_all_distances(polygon, box, direction=\"x\"):\n x1, y1, x2, y2 = box\n if len(polygon) == 0:\n polygon_min_x, polygon_max_x, polygon_min_y, polygon_max_y = 0, 0, 0, 0\n else:\n polygon_min_x = min([p[0] for p in polygon])\n polygon_max_x = max([p[0] for p in polygon])\n polygon_min_y = min([p[1] for p in polygon])\n polygon_max_y = max([p[1] for p in polygon])\n\n distance_x = abs((x2 + x1) / 2 - (polygon_max_x + polygon_min_x) / 2)\n distance_y = abs((y2 + y1) / 2 - (polygon_max_y + polygon_min_y) / 2)\n\n if direction == \"x\":\n intersection = set(range(int(x1), int(x2))).intersection(set(range(int(polygon_min_x), int(polygon_max_x))))\n partly_union = set(range(int(x1), int(x2)))\n if len(partly_union) == 0:\n return 0, 0, 0\n iou = len(intersection) / len(partly_union)\n return iou, distance_x, distance_y\n elif direction == \"y\":\n intersection = set(range(int(y1), int(y2))).intersection(set(range(int(polygon_min_y), int(polygon_max_y))))\n partly_union = set(range(int(y1), int(y2)))\n if len(partly_union) == 0:\n return 0, 0, 0\n iou = len(intersection) / len(partly_union)\n return iou, distance_x, distance_y\n else:\n raise ValueError(\"direction must be x or y\")\n\n\ndef mapping_labels_and_value(graph_type, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons):\n # if graph_type == \"horizontal_bar\":\n # # map x_labels with x_boxes because we don't need to be exact\n # # map based on best iou in Ox direction\n # # find all iou > o pairs between x_labels polygons and x_boxes\n # pairs = [] # (x_labels_index, x_boxes_index, similarity)\n # for i, x_label_polygon in enumerate(x_labels_polygons):\n # for j, x_box in enumerate(x_boxes):\n # iou, distance = compute_iou_and_distance_one_direction(x_label_polygon, x_box, direction=\"x\")\n # if iou > 0.2:\n # pairs.append((i, j, iou, distance))\n \n # # select only one pair for each y_label with the highest iou\n # pairs = sorted(pairs, key=lambda x: (x[2], x[3]), reverse=True)\n \n # filtered_pairs = []\n # existed_indices = set()\n # for pair in pairs:\n # if pair[0] not in existed_indices:\n # filtered_pairs.append(pair)\n # existed_indices.add(pair[0])\n\n # # get remaining y_labels_polygons and y_boxes\n # remaining_x_labels_polygons = [x_labels_polygons[p[0]] for p in filtered_pairs]\n # remaining_x_boxes = [x_boxes[p[1]] for p in filtered_pairs]\n\n # return remaining_x_boxes, y_boxes, remaining_x_labels_polygons, y_labels_polygons\n # else:\n # map y_labels with y_boxes because we don't need to be exact\n # map based on best iou in Oy direction\n # find all iou > o pairs between y_labels polygons and y_boxes\n pairs = [] # (y_labels_index, y_boxes_index, similarity)\n for i, y_label_polygon in enumerate(y_labels_polygons):\n for j, y_box in enumerate(y_boxes):\n iou, distance = compute_iou_and_distance_one_direction(y_label_polygon, y_box, direction=\"y\")\n if iou > 0.2:\n pairs.append((i, j, iou, distance))\n \n # select only one pair for each y_label with the highest iou\n pairs = sorted(pairs, key=lambda x: (x[2], x[3]), reverse=True)\n \n filtered_pairs = []\n existed_indices = set()\n for pair in pairs:\n if pair[0] not in existed_indices:\n filtered_pairs.append(pair)\n existed_indices.add(pair[0])\n\n # get remaining y_labels_polygons and y_boxes\n remaining_y_labels_polygons = [y_labels_polygons[p[0]] for p in filtered_pairs]\n remaining_y_boxes = [y_boxes[p[1]] for p in filtered_pairs]\n\n return x_boxes, remaining_y_boxes, x_labels_polygons, remaining_y_labels_polygons\n\nif not TEST_MODE:\n # visualize x_boxes and y_boxes\n graph_type = graph_type_predictions[idx]\n x_boxes, y_boxes, x_labels_polygons, y_labels_polygons = mapping_labels_and_value(graph_type, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n\n visualize(image_paths[idx], value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n\n# %%\n# ------------ VERTICAL BAR GRAPH -------------\n# Data: x will always be categorical, y will always be numerical\n# 1. we should prioritize value prediction, map 1-1 with x_points then if there is outlier x_points/values, ignore it, map with closest x2-x1 first, then y2-y1\n# if number of x_labels is equal to number of x_boxes, then map 1-1\n# if number of x_labels is different from number of x_boxes, then:\n# - get the rectangle of x_label\n# - draw a rhombus with the points is center of the rectangle edges\n# - draw a rectangle with the center be the highest point of the rhombus\n# - then map 1-1 with x_boxes\n\ndef map_x_labels_polygons_and_x_boxes(x_labels_polygons, x_boxes):\n if False: # len(x_labels_polygons) == len(x_boxes):\n # map 1-1\n indices_mapping = [(i, i) for i in range(len(x_labels_polygons))]\n else:\n # - get the min rectangle of x_label\n x_labels_boxes = []\n for x_label_polygon in x_labels_polygons:\n rect = cv2.minAreaRect(np.array(x_label_polygon))\n box = cv2.boxPoints(rect)\n box = np.intp(box)\n x_labels_boxes.append(box)\n \n # - draw a rhombus with the points is center of the rectangle edges\n x_labels_rhombuses = []\n for x_label_box in x_labels_boxes:\n x1, y1, x2, y2, x3, y3, x4, y4 = x_label_box.flatten()\n new_x1 = (x1 + x2) / 2\n new_y1 = (y1 + y2) / 2\n new_x2 = (x2 + x3) / 2\n new_y2 = (y2 + y3) / 2\n new_x3 = (x3 + x4) / 2\n new_y3 = (y3 + y4) / 2\n new_x4 = (x4 + x1) / 2\n new_y4 = (y4 + y1) / 2\n x_labels_rhombuses.append(np.array([[new_x1, new_y1], [new_x2, new_y2], [new_x3, new_y3], [new_x4, new_y4]]))\n\n # - draw a rectangle with the center be the highest point of the rhombus\n x_labels_rectangles = []\n for x_label_rhombus in x_labels_rhombuses:\n # highest point is the point that has minimum y\n highest_point = None\n for point in x_label_rhombus:\n if highest_point is None:\n highest_point = point\n else:\n if point[1] < highest_point[1]:\n highest_point = point\n x, y = highest_point\n w = 10\n h = 10 # x_boxes[0].shape[1]\n\n x1 = x - w\n y1 = y - h\n x2 = x1 + 2 * w\n y2 = y1 + 2 * h\n\n x_labels_rectangles.append(np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]]))\n\n # then map 1-1 with x_boxes\n pairs = []\n for i, x_label_rect in enumerate(x_labels_rectangles):\n for j, x_box in enumerate(x_boxes):\n iou, distance = compute_iou_and_distance_one_direction(x_label_rect, x_box, direction=\"x\")\n if iou > 0.2:\n pairs.append((i, j, iou, distance))\n\n # select only one pair for each y_label with the highest iou\n pairs = sorted(pairs, key=lambda x: (x[2], x[3]), reverse=True)\n \n filtered_pairs = []\n existed_labels_indices = set()\n existed_boxes_indices = set()\n for pair in pairs:\n if pair[0] not in existed_labels_indices and pair[1] not in existed_boxes_indices:\n filtered_pairs.append(pair)\n existed_labels_indices.add(pair[0])\n existed_boxes_indices.add(pair[1])\n\n # check if there is missing pairs because of segmentation issues\n remaining_x_labels_indices = set(range(len(x_labels_rectangles))).difference(set([p[0] for p in pairs]))\n remaining_x_boxes_indices = set(range(len(x_boxes))).difference(set([p[1] for p in pairs]))\n\n remaining_pairs = []\n for i in remaining_x_labels_indices:\n for j in remaining_x_boxes_indices:\n x_label_polygon = np.array(x_labels_polygons[i])\n x_label_rect = [\n [min(x_label_polygon[:, 0]), min(x_label_polygon[:, 1])],\n [max(x_label_polygon[:, 0]), min(x_label_polygon[:, 1])],\n [max(x_label_polygon[:, 0]), max(x_label_polygon[:, 1])],\n [min(x_label_polygon[:, 0]), max(x_label_polygon[:, 1])]\n ]\n x_box = x_boxes[j]\n iou, distance = compute_iou_and_distance_one_direction(x_label_rect, x_box, direction=\"x\")\n if iou > 0.9:\n remaining_pairs.append((i, j, iou, distance))\n\n for pair in remaining_pairs:\n if pair[0] not in existed_labels_indices and pair[1] not in existed_boxes_indices:\n filtered_pairs.append(pair)\n existed_labels_indices.add(pair[0])\n existed_boxes_indices.add(pair[1])\n\n\n indices_mapping = [(p[0], p[1]) for p in filtered_pairs]\n\n \n # get remaining y_labels_polygons and y_boxes\n remaining_x_labels_polygons = [x_labels_polygons[p[0]] for p in indices_mapping]\n remaining_x_boxes = [x_boxes[p[1]] for p in indices_mapping]\n\n return remaining_x_labels_polygons, remaining_x_boxes\n\n\ndef map_x_boxes_and_value_boxes(x_boxes, value_boxes, graph_type=\"\", mask_mapping=None):\n value_indices_mapping = []\n if len(x_boxes) == len(value_boxes) and graph_type in [\"vertical_bar\", \"horizontal_bar\", \"dot\"]:\n value_indices_mapping = [(i, i) for i in range(len(value_boxes))]\n else:\n # rely on the number of x_boxes\n # map 1-1 with x_boxes, if there is any missing, set the value to minimum value of y_boxes value\n value_x_box_pairs = []\n for i, x_box in enumerate(x_boxes):\n for j, value_box in enumerate(value_boxes):\n # convert x_box to polygon\n x1, y1, x2, y2 = x_box\n x_box_polygon = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]])\n # iou, distance = compute_iou_and_distance_one_direction(x_box_polygon, value_box, direction=\"x\")\n iou, distance_x, distance_y = compute_iou_and_all_distances(x_box_polygon, value_box, direction=\"x\")\n \n if mask_mapping is not None:\n value_box_y = (value_box[1] + value_box[3]) // 2\n x_box_x = (x1 + x2) // 2\n distance_y = abs(mask_mapping[x_box_x] - value_box_y) \n \n if iou > 0.2:\n value_x_box_pairs.append((i, j, iou, distance_y))\n \n # value_x_box_pairs = sorted(value_x_box_pairs, key=lambda x: (x[2], x[3]), reverse=True)\n value_x_box_pairs = sorted(value_x_box_pairs, key=lambda x: (-x[3], x[2]), reverse=True)\n \n filtered_pairs = []\n existed_boxes_indices = set()\n existed_values_indices = set()\n for pair in value_x_box_pairs:\n if pair[0] not in existed_boxes_indices and pair[1] not in existed_values_indices:\n filtered_pairs.append(pair)\n existed_boxes_indices.add(pair[0])\n existed_values_indices.add(pair[1])\n\n value_indices_mapping = [(p[0], p[1]) for p in filtered_pairs]\n\n return value_indices_mapping\n\n\ndef filter_non_numerical_boxes_and_polygons(image_path, boxes, polygons):\n origin_image = cv2.imread(image_path)\n origin_image = cv2.cvtColor(origin_image, cv2.COLOR_BGR2RGB)\n\n crops = []\n for polygon, box in zip(polygons, boxes):\n try:\n crop = crop_polygon_from_image(origin_image, polygon)\n except:\n crop = np.ones((32, 100, 3), dtype=np.uint8) * 255\n if min(crop.shape) == 0:\n # random white blank crop\n crop = np.ones((32, 100, 3), dtype=np.uint8) * 255\n\n if graph_type_predictions[idx] == \"horizontal_bar\":\n # horizontal flip\n crop = cv2.flip(crop, 1)\n # rotate 90 degree\n crop = cv2.rotate(crop, cv2.ROTATE_90_CLOCKWISE)\n crops.append(crop)\n \n filtered_texts, filtered_polygons, filtered_boxes = [], [], []\n\n crops_texts = [p[0] for p in text_recognition_model.predict(crops)[0]]\n\n processed_texts = []\n for t in crops_texts:\n for i in range(10):\n t = t.replace(f\",{str(i)}00\", f\"{str(i)}00\")\n t = t.replace(\",\", \".\")\n processed_texts.append(t)\n crops_texts = processed_texts\n\n for text, polygon, box in zip(crops_texts, polygons, boxes):\n try:\n # filter out those character that are not numerical and alphabets\n text = \"\".join([c for c in text if c in \"-0123456789.abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ\"])\n\n t = text.replace(\",\", \"\").replace(\"%\", \"\")\n if not t.startswith(\"-\"):\n t = t.replace(\"-\", \"\")\n if t.endswith(\"K\") or t.endswith(\"k\"):\n t = t[:-1] + \"000\"\n if t.lower() == \"o\":\n t = \"0\"\n float(t)\n\n filtered_texts.append(t)\n filtered_polygons.append(polygon)\n filtered_boxes.append(box)\n except:\n print(\"Not numerical\", text)\n\n return filtered_texts, filtered_polygons, filtered_boxes\n\n\ndef longest_increasing_subsequence(arr):\n n = len(arr)\n\n # Declare the list (array) for LIS and initialize LIS\n # values for all indexes\n lis = [1]*n\n\n prev = [0]*n\n for i in range(0, n):\n prev[i] = i\n\n # Compute optimized LIS values in bottom up manner\n for i in range (1 , n):\n for j in range(0 , i):\n if arr[i] > arr[j] and lis[i]< lis[j] + 1:\n lis[i] = lis[j]+1\n prev[i] = j\n\n # Initialize maximum to 0 to get the maximum of all\n # LIS\n maximum = 0\n idx = 0\n\n # Pick maximum of all LIS values\n for i in range(n):\n if maximum < lis[i]:\n maximum = lis[i]\n idx = i\n\n seq = [arr[idx]]\n while idx != prev[idx]:\n idx = prev[idx]\n seq.append(arr[idx])\n\n return (maximum, reversed(seq))\n\ndef get_pixel_to_value_pair(boxes, texts, direction=\"y\"):\n pixel_to_value_pairs = []\n\n for i, box in enumerate(boxes):\n value_text = texts[i].replace(\",\", \"\").replace(\"%\", \"\")\n if not value_text.startswith(\"-\"):\n value_text = value_text.replace(\"-\", \"\")\n if value_text.endswith(\"K\") or value_text.endswith(\"k\"):\n value_text = value_text.replace(\"K\", \"000\")\n\n # only keep numbers\n value_text = \"\".join([c for c in value_text if c in \"-0123456789.abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ\"])\n\n if direction == \"y\":\n pixel_to_value_pairs.append(((box[1] + box[3]) / 2, float(value_text)))\n else:\n pixel_to_value_pairs.append(((box[0] + box[2]) / 2, float(value_text)))\n\n if direction == \"y\":\n # find longest increasing subsequence in pixel_to_value_pairs according to value\n # first sort by pixel\n pixel_to_value_pairs = sorted(pixel_to_value_pairs, key=lambda x: x[0], reverse=True)\n values = [p[1] for p in pixel_to_value_pairs]\n\n new_values = list(longest_increasing_subsequence(values)[1])\n new_pixel_to_value_pairs = []\n for pixel, value in pixel_to_value_pairs:\n if new_values and value == new_values[0]:\n new_pixel_to_value_pairs.append((pixel, value))\n new_values = new_values[1:]\n\n pixel_to_value_pairs = new_pixel_to_value_pairs\n\n # set value above zero pixel to positive, and below zero pixel to negative\n zero_pixel = None\n for pixel, value in pixel_to_value_pairs:\n if value == 0:\n zero_pixel = pixel\n break\n\n if zero_pixel is not None:\n new_pixel_to_value_pairs = []\n for pixel, value in pixel_to_value_pairs:\n if pixel < zero_pixel:\n new_pixel_to_value_pairs.append((pixel, abs(value)))\n else:\n new_pixel_to_value_pairs.append((pixel, -abs(value)))\n pixel_to_value_pairs = new_pixel_to_value_pairs\n else:\n # find longest increasing subsequence in pixel_to_value_pairs according to value\n # first sort by pixel\n pixel_to_value_pairs = sorted(pixel_to_value_pairs, key=lambda x: x[0], reverse=False)\n values = [p[1] for p in pixel_to_value_pairs]\n\n new_values = list(longest_increasing_subsequence(values)[1])\n new_pixel_to_value_pairs = []\n for pixel, value in pixel_to_value_pairs:\n if new_values and value == new_values[0]:\n new_pixel_to_value_pairs.append((pixel, value))\n new_values = new_values[1:]\n\n pixel_to_value_pairs = new_pixel_to_value_pairs\n\n # set value above zero pixel to positive, and below zero pixel to negative\n zero_pixel = None\n for pixel, value in pixel_to_value_pairs:\n if value == 0:\n zero_pixel = pixel\n break\n\n if zero_pixel is not None:\n new_pixel_to_value_pairs = []\n for pixel, value in pixel_to_value_pairs:\n if pixel > zero_pixel:\n new_pixel_to_value_pairs.append((pixel, abs(value)))\n else:\n new_pixel_to_value_pairs.append((pixel, -abs(value)))\n pixel_to_value_pairs = new_pixel_to_value_pairs\n\n return pixel_to_value_pairs\n\n\ndef get_x_values_from_value_boxes(value_boxes, pixel_to_value_pairs, default_len):\n \"\"\"\n default_len: in case there is no pixel_to_value_pairs, we return `[0] * default_len` for all x labels\n \"\"\"\n if len(pixel_to_value_pairs) == 0:\n print(\"No y_boxes found!!!\")\n all_values = [0] * default_len\n else:\n min_value = min([p[1] for p in pixel_to_value_pairs]) if len(pixel_to_value_pairs) > 0 else 0\n\n if min_value < 0:\n new_pixel_to_value_pairs = []\n for pair in pixel_to_value_pairs:\n new_pixel_to_value_pairs.append((pair[0], pair[1] + abs(min_value)))\n pixel_to_value_pairs = new_pixel_to_value_pairs\n\n # print(pixel_value_pairs)\n # calculate the real value of value boxes based on Oy axis\n # sort value_boxes\n value_boxes = sorted(value_boxes, key=lambda x: x[0])\n all_values = []\n for value_box in value_boxes:\n value_x_pixel = (value_box[0] + value_box[2]) / 2\n\n # find 2 nearest pixel_value_pairs to value_y_pixel\n nearest_pixel_value_pairs = sorted(pixel_to_value_pairs, key=lambda x: abs(x[0] - value_x_pixel))[:2]\n if len(nearest_pixel_value_pairs) >= 2:\n x1_pixel, x1_value = nearest_pixel_value_pairs[0]\n x2_pixel, x2_value = nearest_pixel_value_pairs[1]\n else:\n # handle the case there is only one nearest y value -> use origin as 0 or use highest y_labels then map value\n x1_pixel, x1_value = nearest_pixel_value_pairs[0]\n x2_pixel, x2_value = nearest_pixel_value_pairs[0]\n\n # calculate the real value of value_box\n # handle the case value_y_pixel > y2_pixel\n if value_x_pixel > x1_pixel: # on the right of x1_pixel\n value_box_value = x1_value + abs((x2_value - x1_value) / (x2_pixel - x1_pixel) * (value_x_pixel - x1_pixel))\n # y1_value - abs((y2_value - y1_value) / (y2_pixel - y1_pixel) * (value_y_pixel - y1_pixel))\n else:\n value_box_value = x1_value - abs((x2_value - x1_value) / (x2_pixel - x1_pixel) * (value_x_pixel - x1_pixel))\n # y1_value + abs((y2_value - y1_value) / (y2_pixel - y1_pixel) * (value_y_pixel - y1_pixel))\n\n # print(\"----------------\")\n # print(\"y1_pixel =\", y1_pixel, \"y1_value =\", y1_value)\n # print(\"y2_pixel =\", y2_pixel, \"y2_value =\", y2_value)\n # print(\"value_y_pixel =\", value_y_pixel, \"value_box_value =\", value_box_value)\n if math.isnan(value_box_value):\n if min_value < 0:\n value_box_value = 0\n else:\n value_box_value = min_value\n # elif value_box_value < 0:\n # if min_value < 0:\n # value_box_value = 0\n # else:\n # value_box_value = min_value\n all_values.append(value_box_value)\n\n # set all infinite values to 0\n all_values = [0 if abs(v) == float(\"inf\") else v for v in all_values]\n\n if min_value < 0:\n all_values = [v + min_value for v in all_values]\n\n return all_values\n\n\ndef get_y_values_from_value_boxes(value_boxes, pixel_to_value_pairs, default_len):\n \"\"\"\n default_len: in case there is no pixel_to_value_pairs, we return `[0] * default_len` for all x labels\n \"\"\"\n if len(pixel_to_value_pairs) == 0:\n print(\"No y_boxes found!!!\")\n all_values = [0] * default_len\n else:\n min_value = min([p[1] for p in pixel_to_value_pairs]) if len(pixel_to_value_pairs) > 0 else 0\n\n if min_value < 0:\n new_pixel_to_value_pairs = []\n for pair in pixel_to_value_pairs:\n new_pixel_to_value_pairs.append((pair[0], pair[1] + abs(min_value)))\n pixel_to_value_pairs = new_pixel_to_value_pairs\n\n # print(pixel_value_pairs)\n # calculate the real value of value boxes based on Oy axis\n # sort value_boxes\n value_boxes = sorted(value_boxes, key=lambda x: x[0])\n all_values = []\n for value_box in value_boxes:\n value_y_pixel = (value_box[1] + value_box[3]) / 2\n\n # find 2 nearest pixel_value_pairs to value_y_pixel\n nearest_pixel_value_pairs = sorted(pixel_to_value_pairs, key=lambda x: abs(x[0] - value_y_pixel))[:2]\n if len(nearest_pixel_value_pairs) >= 2:\n y1_pixel, y1_value = nearest_pixel_value_pairs[0]\n y2_pixel, y2_value = nearest_pixel_value_pairs[1]\n else:\n # handle the case there is only one nearest y value -> use origin as 0 or use highest y_labels then map value\n y1_pixel, y1_value = nearest_pixel_value_pairs[0]\n y2_pixel, y2_value = nearest_pixel_value_pairs[0]\n\n # calculate the real value of value_box\n # handle the case value_y_pixel > y2_pixel\n if value_y_pixel > y1_pixel: # below the y1_pixel\n value_box_value = y1_value - abs((y2_value - y1_value) / (y2_pixel - y1_pixel) * (value_y_pixel - y1_pixel))\n else:\n value_box_value = y1_value + abs((y2_value - y1_value) / (y2_pixel - y1_pixel) * (value_y_pixel - y1_pixel))\n\n # print(\"----------------\")\n # print(\"y1_pixel =\", y1_pixel, \"y1_value =\", y1_value)\n # print(\"y2_pixel =\", y2_pixel, \"y2_value =\", y2_value)\n # print(\"value_y_pixel =\", value_y_pixel, \"value_box_value =\", value_box_value)\n if math.isnan(value_box_value):\n if min_value < 0:\n value_box_value = 0\n else:\n value_box_value = min_value\n # elif value_box_value < 0:\n # if min_value < 0:\n # value_box_value = 0\n # else:\n # value_box_value = min_value\n all_values.append(value_box_value)\n\n # set all infinite values to 0\n all_values = [0 if abs(v) == float(\"inf\") else v for v in all_values]\n\n if min_value < 0:\n all_values = [v + min_value for v in all_values]\n\n return all_values\n\n\ndef calculate_overlap_score(big_polygon, small_polygon, image):\n polygon1 = big_polygon\n polygon2 = small_polygon\n\n # calculate iou between two polygons\n polygon1 = np.array(polygon1)\n polygon2 = np.array(polygon2)\n polygon1 = polygon1.reshape(-1, 2)\n polygon2 = polygon2.reshape(-1, 2)\n polygon1 = polygon1.astype(np.float32)\n polygon2 = polygon2.astype(np.float32)\n\n rect1 = cv2.minAreaRect(polygon1)\n box1 = cv2.boxPoints(rect1)\n box1 = np.int0(box1)\n\n rect2 = cv2.minAreaRect(polygon2)\n box2 = cv2.boxPoints(rect2)\n box2 = np.int0(box2)\n\n mask1 = np.zeros(image.shape[:2], np.uint8)\n cv2.drawContours(mask1, [box1], 0, 255, -1, cv2.LINE_AA)\n mask2 = np.zeros(image.shape[:2], np.uint8)\n cv2.drawContours(mask2, [box2], 0, 255, -1, cv2.LINE_AA)\n\n intersection = np.logical_and(mask1, mask2)\n union = np.logical_or(mask2, mask2)\n overlap_score = np.sum(intersection) / np.sum(union)\n\n return overlap_score\n\n\ndef check_if_overlap(polygon1, polygon2, direction=\"x\"):\n box2 = [\n min([p[0] for p in polygon2]),\n min([p[1] for p in polygon2]),\n max([p[0] for p in polygon2]),\n max([p[1] for p in polygon2])\n ]\n if compute_iou_and_distance_one_direction(polygon1, box2, direction)[0] > 0.5:\n return True\n \n return False\n\n\ndef is_multiple_lines(polygons, is_horizontal_bar=False):\n check = False\n for i in range(len(polygons)):\n for j in range(i + 1, len(polygons)):\n if is_horizontal_bar:\n if check_if_overlap(polygons[i], polygons[j], \"y\") or check_if_overlap(polygons[j], polygons[i], \"y\"):\n check = True\n break\n else:\n # convert polygon j to box\n if check_if_overlap(polygons[i], polygons[j], \"x\") or check_if_overlap(polygons[j], polygons[i], \"x\"):\n check = True\n break\n return check\n\n\ndef read_text_from_polygons(image_path, polygons, graph_type, split_polygons=False):\n origin_image = cv2.imread(image_path)\n if split_polygons:\n texts = []\n\n try:\n word_polygons = text_detection_model.predict([image_path])[0][0][0]\n except:\n print(\"Can't predict on original image, try to pad image to multiples of 128\")\n temp_image = cv2.imread(image_paths[idx])\n h, w = temp_image.shape[:2]\n\n # find nearest padding size to multiples of 128\n h = int(np.ceil(h / 128) * 128)\n w = int(np.ceil(w / 128) * 128)\n\n # pad both side to 768\n temp_image = cv2.copyMakeBorder(temp_image, 0, h - temp_image.shape[0], 0, w - temp_image.shape[1], cv2.BORDER_CONSTANT, value=[255, 255, 255])\n cv2.imwrite(\"./data/temp.png\", temp_image)\n\n try:\n word_polygons = text_detection_model.predict([\"./data/temp.png\"])[0][0][0]\n except:\n word_polygons = []\n\n # find all word polygons that are inside the main polygon\n threshold = 0.7\n index_to_word_polygons = {i: [] for i in range(len(polygons))}\n for i, polygon in enumerate(polygons):\n for word_polygon in word_polygons:\n overlap_score = calculate_overlap_score(polygon, word_polygon, origin_image)\n if overlap_score > threshold:\n index_to_word_polygons[i].append(word_polygon)\n\n print(index_to_word_polygons.keys(), [len(index_to_word_polygons[i]) for i in index_to_word_polygons])\n # crop word polygons and read text\n count_single = 0\n for i, polygon in enumerate(polygons):\n word_polygons = index_to_word_polygons[i]\n if len(word_polygons) == 1:\n count_single += 1\n \n if count_single >= len(polygons) - 1:\n # set the remaining word polygon to the main polygon\n for i, polygon in enumerate(polygons):\n word_polygons = index_to_word_polygons[i]\n if len(word_polygons) != 1:\n index_to_word_polygons[i] = [polygon]\n\n for i, polygon in enumerate(polygons):\n word_polygons = index_to_word_polygons[i] if len(index_to_word_polygons[i]) > 1 else [polygon]\n\n # implement the case that the label has multiple lines\n if is_multiple_lines(word_polygons, is_horizontal_bar=(graph_type == \"horizontal_bar\")):\n if graph_type == \"horizontal_bar\":\n word_polygons = sorted(\n word_polygons, \n key=lambda x: - 4 * np.mean([p[0] for p in x]) - np.mean([p[1] for p in x])\n )\n else:\n is_rotated = False\n for word_polygon in word_polygons:\n min_x = min([p[0] for p in polygon])\n max_x = max([p[0] for p in polygon])\n min_y = min([p[1] for p in polygon])\n max_y = max([p[1] for p in polygon])\n\n if (max_y - min_y) / (max_x - min_x) > 3:\n is_rotated = True\n break\n \n # print(\"is_rotated\", is_rotated)\n if not is_rotated:\n word_polygons = sorted(\n word_polygons, \n key=lambda x: np.mean([p[0] for p in x]) + 4 * np.mean([p[1] for p in x])\n )\n else:\n word_polygons = sorted(word_polygons, key=lambda x: -min([p[1] for p in x]))\n else:\n if graph_type == \"horizontal_bar\":\n # sort word_polygons by max y\n word_polygons = sorted(word_polygons, key=lambda x: -max([p[1] for p in x]))\n else:\n # sort word_polygons by min x\n word_polygons = sorted(word_polygons, key=lambda x: min([p[0] for p in x]))\n\n crops = []\n for word_polygon in word_polygons:\n try:\n crop = crop_polygon_from_image(origin_image, word_polygon)\n except:\n crop = np.ones((32, 100, 3), dtype=np.uint8) * 255\n if min(crop.shape) == 0:\n # random white blank crop\n crop = np.ones((32, 100, 3), dtype=np.uint8) * 255\n if graph_type == \"horizontal_bar\":\n # horizontal flip\n crop = cv2.flip(crop, 1)\n # rotate 90 degree\n crop = cv2.rotate(crop, cv2.ROTATE_90_CLOCKWISE)\n crops.append(crop)\n\n texts.append([p[0] for p in text_recognition_model.predict(crops)[0]])\n\n texts = [\" \".join(t) for t in texts]\n else:\n crops = []\n for polygon in polygons:\n try:\n crop = crop_polygon_from_image(origin_image, polygon)\n except:\n crop = np.ones((32, 100, 3), dtype=np.uint8) * 255\n if min(crop.shape) == 0:\n # random white blank crop\n crop = np.ones((32, 100, 3), dtype=np.uint8) * 255\n if graph_type == \"horizontal_bar\":\n # horizontal flip\n crop = cv2.flip(crop, 1)\n # rotate 90 degree\n crop = cv2.rotate(crop, cv2.ROTATE_90_CLOCKWISE)\n crops.append(crop)\n\n texts = [p[0] for p in text_recognition_model.predict(crops)[0]]\n\n return texts\n\nif not TEST_MODE:\n new_x_labels_polygons, new_x_boxes = map_x_labels_polygons_and_x_boxes(x_labels_polygons, x_boxes)\n\n visualize(image_paths[idx], value_boxes, new_x_boxes, y_boxes, new_x_labels_polygons, y_labels_polygons)\n\n# %%\nimport math\n\ndef postprocess_bar_graph(idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons):\n x_labels_polygons, x_boxes = map_x_labels_polygons_and_x_boxes(x_labels_polygons, x_boxes)\n # new_x_labels_polygons, new_x_boxes = map_x_labels_polygons_and_x_boxes(x_labels_polygons, x_boxes)\n\n # if len(new_x_labels_polygons) != 0 and (len(x_labels_polygons) - len(new_x_labels_polygons) <= 2 or len(new_x_labels_polygons) / len(x_labels_polygons) >= 0.8):\n # x_labels_polygons, x_boxes = new_x_labels_polygons, new_x_boxes\n # else:\n # # calculate x_boxes based on x_labels_polygons\n # x_boxes = []\n # for x_labels_polygon in x_labels_polygons:\n # min_x = min(np.array(x_labels_polygon)[:, 0])\n # max_x = max(np.array(x_labels_polygon)[:, 0])\n # min_y = min(np.array(x_labels_polygon)[:, 1])\n # max_y = max(np.array(x_labels_polygon)[:, 1])\n # mean_x = (min_x + max_x) // 2\n # mean_y = (min_y + max_y) // 2\n\n # x_boxes.append([mean_x - 5, mean_y - 5, mean_x + 5, mean_y + 5])\n \n mask = vertical_bar_segmentation_model.predict([image_paths[idx]])[0][0]\n image = cv2.imread(image_paths[idx])\n\n mask = np.array(mask > 0.9, dtype=np.uint8)\n if graph_type_predictions[idx] == \"vertical_bar\":\n mask = cv2.resize(mask, image.shape[:2][::-1])\n else:\n mask = cv2.resize(mask, image.shape[:2][::-1])\n\n # get the highest point of the bar\n if x_boxes:\n Ox = (x_boxes[0][1] + x_boxes[0][3]) // 2\n else:\n Ox = image.shape[0] - 1\n highest_points = []\n for j in range(mask.shape[1]):\n found = False\n for k in range(mask.shape[0]):\n if mask[k, j] == 1:\n highest_points.append(k)\n found = True\n break\n\n if not found:\n if x_boxes:\n highest_points.append(Ox)\n else:\n highest_points.append(0)\n \n # TODO: check if vertical bar is a histogram\n # conditions: 99% continous value from first x_box to last x_box\n start = 0\n end = len(highest_points) - 1\n for i, p in enumerate(highest_points):\n if p < Ox:\n start = i\n break\n\n for i, p in enumerate(highest_points[::-1]):\n if p < Ox:\n end = len(highest_points) - i\n break\n \n is_histogram = False\n if graph_type_predictions[idx] == \"vertical_bar\" and x_boxes:\n first_x = min([(box[0] + box[2]) // 2 for box in x_boxes])\n last_x = max([(box[0] + box[2]) // 2 for box in x_boxes])\n \n total_positives = 0\n for i in range(first_x, last_x + 1):\n if highest_points[i] < Ox:\n total_positives += 1\n print(\"ratio =\", total_positives / (end - start + 1))\n if total_positives / (end - start + 1) > 0.95:\n is_histogram = True\n print(\"is_histogram =\", is_histogram, \"first_x - last_x:\", (first_x, last_x), \"start - end:\", (start, end))\n\n if is_histogram:\n # find value boxes in middle point of 2 neighbor x_boxes\n sorted_x_boxes = sorted(x_boxes, key=lambda box: box[0])\n \n value_boxes = []\n for i in range(len(sorted_x_boxes) - 1):\n first_box = sorted_x_boxes[i]\n second_box = sorted_x_boxes[i + 1]\n \n value_x = (first_box[0] + second_box[2]) // 2\n value_y = highest_points[value_x]\n \n value_boxes.append([value_x - 5, value_y - 5, value_x + 5, value_y + 5])\n else:\n value_indices_mapping = map_x_boxes_and_value_boxes(x_boxes, value_boxes, mask_mapping=highest_points) \n\n # add missing pairs with value_index -1 if missing value\n missing_x_indices = set(range(len(x_boxes))) - set([p[0] for p in value_indices_mapping])\n for missing_x_index in missing_x_indices:\n value_indices_mapping.append((missing_x_index, -1))\n\n # then add new boxes to value_boxes\n # TODO: update this using segmentation model and also handle histogram\n inserted_value_boxes = []\n value_indices_mapping.sort(key=lambda x: x[0])\n\n # print(len(highest_points), mask.shape[1])\n\n# plt.imshow(mask)\n# plt.plot(list(range(len(highest_points))), highest_points)\n# plt.show()\n\n for i, value_index in value_indices_mapping:\n if value_index != -1: \n inserted_value_boxes.append(value_boxes[value_index])\n else:\n center_x = (x_boxes[i][0] + x_boxes[i][2]) // 2\n # print(\"adding\", center_x, highest_points[center_x])\n inserted_value_boxes.append([x_boxes[i][0], highest_points[center_x] - 5, x_boxes[i][2], highest_points[center_x] + 5])\n\n # inserted_value_boxes.append(x_boxes[i])\n\n value_boxes = inserted_value_boxes\n x_boxes = [x_boxes[p[0]] for p in value_indices_mapping]\n\n # filter out those y boxes that are not numerical\n image_path = image_paths[idx]\n filtered_texts, filtered_y_labels_polygons, filtered_y_boxes = filter_non_numerical_boxes_and_polygons(image_path, y_boxes, y_labels_polygons)\n \n # get pixel to value pair\n pixel_to_value_pairs = get_pixel_to_value_pair(filtered_y_boxes, filtered_texts, direction=\"y\")\n print(\"pixel_to_value_pairs\", pixel_to_value_pairs)\n if len(pixel_to_value_pairs) == 1 and x_boxes:\n Ox_pixel = int(np.mean([(p[1] + p[3]) / 2 for p in x_boxes]))\n pixel_to_value_pairs.append((Ox_pixel, 0))\n # print(\"new pixel_to_value_pairs\", pixel_to_value_pairs)\n\n # get y values from value boxes\n all_values = get_y_values_from_value_boxes(value_boxes, pixel_to_value_pairs, default_len=len(x_labels_polygons))\n\n # predict text for x_labels_polygons and sort x_labels_polygons based on min x of x_labels_polygons\n x_labels_polygons = sorted(x_labels_polygons, key=lambda x: max([p[0] for p in x]))\n\n x_labels_texts = read_text_from_polygons(image_path, x_labels_polygons, graph_type_predictions[idx], split_polygons=True)\n if x_labels_texts:\n if x_labels_texts[0].lower() == 'o':\n x_labels_texts[0] = '0'\n\n if graph_type_predictions[idx] == \"horizontal_bar\":\n x_labels_texts = x_labels_texts[::-1]\n all_values = all_values[::-1]\n\n return value_boxes, x_boxes, filtered_y_boxes, x_labels_polygons, filtered_y_labels_polygons, x_labels_texts, all_values\n\nif not TEST_MODE:\n new_value_boxes, new_x_boxes, new_filtered_y_boxes, new_x_labels_polygons, new_filtered_y_labels_polygons, new_x_labels_texts, new_all_values = postprocess_bar_graph(idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n visualize(image_paths[idx], new_value_boxes, new_x_boxes, y_boxes, new_x_labels_polygons, y_labels_polygons)\n\n# %%\n# 1. get line, then map x pixel t y_pixel\n# 2. use missing x boxes to get remaining value -> value box with corresponding y pixel\n\n# find the line in line chart using opencv\ndef find_line(image_path, text_detection_prediction, erode_size=2):\n image = cv2.imread(image_path)\n\n # FIND THE LINE USING SEGMENTATION MODEL\n \n# old_mask = old_line_segmentation_model.predict([image_path])[0][0]\n# old_mask = old_mask > 0.995\n \n# mask = line_segmentation_model.predict([image_path])[0][0]\n# # apply threshold\n# mask = mask > 0.65\n# # mask = np.logical_and(mask, old_mask)\n\n# mask = mask.astype(np.uint8)\n# # resize mask to original image size\n# mask = cv2.resize(mask, (image.shape[1], image.shape[0]))\n# # round mask to integer\n# mask = np.round(mask).astype(np.uint8)\n \n mask = line_segmentation_model.predict([image_path])[0][0]\n mask = mask > 0.65\n mask = mask.astype(np.uint8)\n mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)\n mask = np.round(mask).astype(np.uint8)\n\n# mask_2 = line_segmentation_model_2.predict([image_path])[0][0]\n# mask_2 = mask_2 > 0.65\n# mask_2 = mask_2.astype(np.uint8)\n# mask_2 = cv2.resize(mask_2, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)\n# mask_2 = np.round(mask_2).astype(np.uint8)\n\n# mask_3 = line_segmentation_model_3.predict([image_path])[0][0]\n# mask_3 = mask_3 > 0.65\n# mask_3 = mask_3.astype(np.uint8)\n# mask_3 = cv2.resize(mask_3, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)\n# mask_3 = np.round(mask_3).astype(np.uint8)\n \n# mask = mask + mask_2 + mask_3\n# mask[mask < 2] = 0\n# mask[mask >= 2] = 1\n \n mask = mask * 255\n\n # mask0 = mask.copy()\n\n # # FIND THE LINE USING OPENCV\n # # convert to grayscale\n # image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n\n # # remove all pixels that are in text_detection_prediction polygons\n # for polygon in text_detection_prediction:\n # # if polygon is bigger than 0.2 * image size, then it is not a text\n # if not cv2.contourArea(np.array(polygon)) > 0.2 * image.shape[0] * image.shape[1]:\n # cv2.fillPoly(image, [np.array(polygon)], 255)\n\n # lower = np.array([0])\n # upper = np.array([180])\n # mask = cv2.inRange(image, lower, upper)\n\n # mask = cv2.erode(mask, np.ones((erode_size, erode_size), np.uint8), iterations=1)\n # mask = cv2.dilate(mask, np.ones((3, 3), np.uint8), iterations=1)\n\n # # remove all horizontal lines\n # horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (40, 1))\n # detected_lines = cv2.morphologyEx(mask, cv2.MORPH_OPEN, horizontal_kernel, iterations=2)\n # cnts = cv2.findContours(detected_lines, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n # cnts = cnts[0] if len(cnts) == 2 else cnts[1]\n\n # for c in cnts:\n # cv2.drawContours(mask, [c], -1, 0, -1)\n\n # # remove all vertical lines\n # vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40))\n # detected_lines = cv2.morphologyEx(mask, cv2.MORPH_OPEN, vertical_kernel, iterations=2)\n # cnts = cv2.findContours(detected_lines, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n # cnts = cnts[0] if len(cnts) == 2 else cnts[1]\n\n # for c in cnts:\n # cv2.drawContours(mask, [c], -1, 0, -1)\n\n # # # remove all connected components smaller than 16 pixels\n # # output = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S)\n # # num_labels = output[0]\n # # labels = output[1]\n # # stats = output[2]\n # # centroids = output[3]\n\n # # sizes = stats[1:, -1]\n # # num_labels = num_labels - 1\n\n # # min_size = 16\n\n # # for i in range(0, num_labels):\n # # if sizes[i] <= min_size:\n # # mask[labels == i + 1] = 0\n\n # mask = np.minimum(mask, mask0)\n return image, mask\n\ndef add_missing_value_box_for_line_graph(image_path, x_boxes, y_boxes, value_boxes, x_labels_polygons, value_indices_mapping, missing_x_indices, is_visualize=False):\n # remove all the text before find line\n try:\n text_detection_prediction = text_detection_model.predict([image_path])[0][0][0]\n except:\n try:\n print(\"Can't predict on original image, try to pad image to multiples of 128\")\n temp_image = cv2.imread(image_paths[idx])\n h, w = temp_image.shape[:2]\n\n # find nearest padding size to multiples of 128\n h = int(np.ceil(h / 128) * 128)\n w = int(np.ceil(w / 128) * 128)\n\n # pad both side to 768\n temp_image = cv2.copyMakeBorder(temp_image, 0, h - temp_image.shape[0], 0, w - temp_image.shape[1], cv2.BORDER_CONSTANT, value=[255, 255, 255])\n cv2.imwrite(\"./data/temp.png\", temp_image)\n\n text_detection_prediction = text_detection_model.predict([\"./data/temp.png\"])[0][0][0]\n except:\n text_detection_prediction = []\n\n image, mask = find_line(image_path, text_detection_prediction)\n\n if y_boxes and x_boxes:\n Oy = np.mean([(box[0] + box[2]) / 2 for box in y_boxes])\n Ox = np.mean([(box[1] + box[3]) / 2 for box in x_boxes])\n else:\n Oy = 0\n Ox = 0\n\n if len(y_boxes) == 0:\n min_y = 0\n else:\n min_y = min([(box[1] + box[3]) // 2 for box in y_boxes])\n\n if len(x_boxes) == 0:\n max_x = mask.shape[1]\n else:\n max_x = max([(box[0] + box[2]) // 2 for box in x_boxes])\n\n # # set value on the left of Ox to 0 and below of Oy to 0\n # mask[:, :int(Oy) + 3] = 0\n # mask[int(Ox) - 2:, :] = 0\n\n # # only apply this if there is no value box bigger than min_y\n # if not any([(box[1] + box[3]) / 2 < min_y for box in value_boxes]):\n # mask[:int(min_y) - 5, :] = 0\n # mask[:, int(max_x):] = 0\n if is_visualize:\n new_image = cv2.imread(image_path)\n new_image = cv2.cvtColor(new_image, cv2.COLOR_BGR2RGB)\n three_channels_mask = np.zeros((mask.shape[0], mask.shape[1], 3))\n three_channels_mask[:, :, 0] = mask\n three_channels_mask[:, :, 1] = 0\n three_channels_mask[:, :, 2] = 0\n three_channels_mask = three_channels_mask.astype(np.uint8)\n new_image = cv2.addWeighted(new_image, 0.7, three_channels_mask, 0.5, 0)\n plt.imshow(new_image)\n plt.show()\n\n # find and keep the largest connected component in mask\n # TODO: in the future, when we have a good line segmentation model, we don't need to use this `try-except` anymore, only use the raw output of segmentation model\n try:\n num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8)\n largest_label = np.argmax(stats[1:, 4]) + 1\n new_mask = np.zeros_like(mask)\n new_mask[labels == largest_label] = 255\n\n # using predicted mask, map to current image, then find largest connected component inside that mask region of the image\n # print(\"image info: \", image.shape, image.max(), image.min())\n masked_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n masked_image[new_mask != 255] = 255\n masked_image = 255 - masked_image\n\n # lower = np.array([0])\n # upper = np.array([120])\n # masked_image = cv2.inRange(masked_image, lower, upper)\n masked_image[masked_image < 50] = 0\n masked_image[masked_image >= 50] = 255\n\n # # find first and last x pixel that has value > 0\n # first_x = 0\n # last_x = 0\n # for i in range(masked_image.shape[1]):\n # if masked_image[:, i].sum() > 0:\n # first_x = i\n # break\n # for i in range(masked_image.shape[1] - 1, -1, -1):\n # if masked_image[:, i].sum() > 0:\n # last_x = i\n # break\n\n # mask[:, :first_x] = 0\n # mask[:, last_x:] = 0\n\n # new_mask = mask\n\n # dilate masked_image\n kernel = np.ones((2, 2), np.uint8)\n new_masked_image = cv2.dilate(masked_image, kernel, iterations=1)\n\n # masked_image = 255 - masked_image\n # print(np.unique(masked_image, return_counts=True))\n\n num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(new_masked_image, connectivity=8)\n largest_label = np.argmax(stats[1:, 4]) + 1\n smallest_mask = np.zeros_like(mask)\n smallest_mask[labels == largest_label] = 255\n\n # if horizontal length of new_mask is smaller than 0.2 * horizontal length of masked_image, then we don't use it\n horizontal_sum_smallest_mask = smallest_mask.sum(axis=1) > 0\n horizontal_sum_mask = mask.sum(axis=1) > 0\n \n if horizontal_sum_smallest_mask.sum() < 0.6 * horizontal_sum_mask.sum():\n smallest_mask = mask\n \n new_mask = smallest_mask\n except:\n new_mask = mask\n\n\n if is_visualize:\n plt.imshow(new_mask)\n plt.show()\n\n mask = new_mask # np.logical_or(mask, new_mask)\n\n if is_visualize:\n # show mask on the image\n # display mask on image with alpha 0.5\n image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n three_channels_mask = np.zeros((mask.shape[0], mask.shape[1], 3))\n three_channels_mask[:, :, 0] = mask\n three_channels_mask[:, :, 1] = 0\n three_channels_mask[:, :, 2] = 0\n three_channels_mask = three_channels_mask.astype(np.uint8)\n image = cv2.addWeighted(image, 0.7, three_channels_mask, 0.5, 0)\n plt.imshow(image)\n # plt.show()\n # save image to file\n os.makedirs(\"./data/line_mask_prediction\", exist_ok=True)\n plt.savefig(f\"./data/line_mask_prediction/{image_path.split('/')[-1]}\")\n\n\n # get the mapping x_pixel -> (x_value, y_pixel)\n mapping = {}\n for x_pixel in range(mask.shape[1]):\n # mean of y_pixel\n y_pixel = np.mean(np.where(mask[:, x_pixel] > 0)[0])\n if not math.isnan(y_pixel):\n mapping[x_pixel] = y_pixel\n # print(\"mapping: \", mapping)\n # there is still case that x_pixel is discontinuous, so we need to interpolate\n x_pixels = list(mapping.keys())\n y_pixels = [mapping[x_pixel] for x_pixel in x_pixels]\n\n if len(x_pixels) == 0:\n print(\"Not found x_pixels, try with erode = 1\")\n _, mask = find_line(image_path, text_detection_prediction, 1)\n\n # set value on the left of Ox to 0 and below of Oy to 0\n mask[:, :int(Oy) + 3] = 0\n mask[int(Ox) - 5:, :] = 0\n\n # only apply this if there is no value box bigger than min_y\n if not any([(box[1] + box[3]) / 2 < min_y for box in value_boxes]):\n mask[:int(min_y) - 5, :] = 0\n mask[:, int(max_x):] = 0\n if is_visualize:\n plt.imshow(mask)\n\n # get the mapping x_pixel -> (x_value, y_pixel)\n mapping = {}\n for x_pixel in range(mask.shape[1]):\n # mean of y_pixel\n y_pixel = np.mean(np.where(mask[:, x_pixel] > 0)[0])\n if not math.isnan(y_pixel):\n mapping[x_pixel] = y_pixel\n\n # there is still case that x_pixel is discontinuous, so we need to interpolate\n x_pixels = list(mapping.keys())\n y_pixels = [mapping[x_pixel] for x_pixel in x_pixels]\n\n return_value_boxes, return_x_boxes, return_x_labels_polygons = [], [], []\n if len(x_pixels) <= 1:\n # adding default values\n filtered_x_boxes = [x_boxes[p[0]] for p in value_indices_mapping]\n filtered_x_labels_polygons = [x_labels_polygons[p[0]] for p in value_indices_mapping]\n filtered_value_boxes = [value_boxes[p[1]] for p in value_indices_mapping]\n\n inserted_value_boxes = []\n inserted_x_boxes = []\n inserted_x_labels_polygons = []\n\n for index in missing_x_indices:\n if y_pixel is not None:\n if len(value_boxes):\n inserted_value_boxes.append(value_boxes[0])\n else:\n mid_x = mask.shape[1] // 2\n mid_y = mask.shape[0] // 2\n inserted_value_boxes.append([mid_x - 5, mid_y - 5, mid_x + 5, mid_y + 5])\n inserted_x_boxes.append(x_boxes[index])\n inserted_x_labels_polygons.append(x_labels_polygons[index])\n\n filtered_x_boxes.extend(inserted_x_boxes)\n filtered_value_boxes.extend(inserted_value_boxes)\n filtered_x_labels_polygons.extend(inserted_x_labels_polygons)\n else:\n all_x_pixels = []\n all_y_pixels = []\n for i in range(len(x_pixels) - 1):\n # insert all the missing x_pixels and y_pixels using linspace y_pixels\n new_y_pixels = np.linspace(y_pixels[i], y_pixels[i + 1], x_pixels[i + 1] - x_pixels[i] + 1)\n new_x_pixels = np.linspace(x_pixels[i], x_pixels[i + 1], x_pixels[i + 1] - x_pixels[i] + 1)\n \n all_x_pixels.extend(new_x_pixels)\n all_y_pixels.extend(new_y_pixels)\n\n # sort all_x_pixels and all_y_pixels by x_pixel\n x_pixels, y_pixels = zip(*sorted(zip(all_x_pixels, all_y_pixels)))\n mapping = {x: y for x, y in zip(x_pixels, y_pixels)}\n # draw the mapping\n # print(min(mapping.keys()), max(mapping.keys()))\n if is_visualize:\n # set figure size\n plt.plot(list(mapping.keys()), list(mapping.values()))\n\n inserted_value_boxes = []\n inserted_x_boxes = []\n inserted_x_labels_polygons = []\n\n # for index in missing_x_indices:\n # keep the value_indices_mapping that are in the line\n keep_value_indices_mapping = []\n for index in range(len(value_indices_mapping)):\n value_box = value_boxes[value_indices_mapping[index][1]]\n # check if there is any x, y pair in mapping that is in the value_box\n for x_pixel in range(int(value_box[0]), int(value_box[2])):\n y_pixel = mapping.get(x_pixel, None)\n if y_pixel is not None and y_pixel >= value_box[1] and y_pixel <= value_box[3]:\n keep_value_indices_mapping.append(value_indices_mapping[index])\n break\n \n # remove all the value_indices_mapping that are in keep_value_indices_mapping\n value_indices_mapping = [item for item in value_indices_mapping if item not in keep_value_indices_mapping]\n\n for index in list(missing_x_indices) + [p[0] for p in value_indices_mapping]:\n x_pixel = (x_boxes[index][0] + x_boxes[index][2]) // 2\n # print(\"x_pixel =\", x_pixel, \"Oy =\", Oy)\n x_pixel = int(Oy + 4) if x_pixel <= Oy + 3 else x_pixel\n x_pixel = int(max_x - 4) if x_pixel >= max_x - 3 else x_pixel\n # print(\"new x_pixel =\", x_pixel, \"Oy =\", Oy)\n for x_pixel in range(x_pixel, x_pixel + 3):\n y_pixel = mapping.get(x_pixel, None)\n if y_pixel is not None:\n inserted_value_boxes.append([int(x_boxes[index][0]), int(y_pixel - 5), int(x_boxes[index][2]), int(y_pixel + 5)])\n inserted_x_boxes.append(x_boxes[index])\n inserted_x_labels_polygons.append(x_labels_polygons[index])\n break\n\n filtered_x_boxes = [x_boxes[p[0]] for p in keep_value_indices_mapping]\n filtered_x_labels_polygons = [x_labels_polygons[p[0]] for p in keep_value_indices_mapping]\n filtered_value_boxes = [value_boxes[p[1]] for p in keep_value_indices_mapping]\n\n # filtered_x_boxes = []\n # filtered_x_labels_polygons = []\n # filtered_value_boxes = []\n\n filtered_x_boxes.extend(inserted_x_boxes)\n filtered_value_boxes.extend(inserted_value_boxes)\n filtered_x_labels_polygons.extend(inserted_x_labels_polygons)\n\n # # print(value_indices_mapping, len(filtered_x_boxes), len(filtered_value_boxes), len(filtered_x_labels_polygons))\n # # filter out those box that are not in predicted mapping\n # # print(\"mapping: \", mapping)\n # # remove those first mapping that has value > Ox's y\n # mapping_list = list(mapping.items())\n # mapping_list.sort(key=lambda x: x[0])\n \n # do_check = False\n # index = 0\n # for i, item in enumerate(mapping_list[:10]):\n # if item[1] >= Ox:\n # do_check = True\n # index = i\n # break\n \n # if do_check:\n # count = 0\n # stop = False\n # while not stop:\n # if mapping_list[0][1] >= Ox - 1 or count <= index:\n # mapping_list.pop(0)\n # count += 1\n # else:\n # stop = True\n\n # # print(\"count:\", count)\n # if count <= 50:\n # mapping = dict(mapping_list)\n\n min_x = min(mapping.keys())\n max_x = max(mapping.keys())\n\n # sort filtered_value_boxes, filtered_x_boxes, filtered_x_labels_polygons by x_box x\n if not len(filtered_value_boxes) or not len(filtered_x_boxes) or not len(filtered_x_labels_polygons):\n return [], [], []\n\n filtered_value_boxes, filtered_x_boxes, filtered_x_labels_polygons = zip(*sorted(zip(filtered_value_boxes, filtered_x_boxes, filtered_x_labels_polygons), key=lambda x: x[1][0]))\n for i, (value_box, x_box, x_labels_polygon) in enumerate(zip(filtered_value_boxes, filtered_x_boxes, filtered_x_labels_polygons)):\n x_pixel = (x_box[0] + x_box[2]) // 2\n value_x_pixel = (value_box[0] + value_box[2]) // 2\n margin = 3\n if min_x - margin <= x_pixel <= max_x + 2 and min_x - margin <= value_x_pixel <= max_x + 2: # or (i <= 1 and i!= (len(filtered_value_boxes) - 1)):\n return_value_boxes.append(value_box)\n return_x_boxes.append(x_box)\n return_x_labels_polygons.append(x_labels_polygon)\n\n # # remove first box if it is on Oy and Oy doesn't overlap with line mask\n # first_box = return_value_boxes[0]\n # if first_box[0] <= Oy <= first_box[2]:\n # Oy_mask = image.copy()\n # Oy_mask = cv2.cvtColor(Oy_mask, cv2.COLOR_BGR2GRAY)\n # min_Oy_x = min([x[0] for x in y_boxes])\n # max_Oy_x = max([x[2] for x in y_boxes])\n\n # Oy_mask[:, :min_Oy_x] = 255\n # Oy_mask[:, max_Oy_x:] = 255\n\n # Oy_mask[Oy_mask > 80] = 255\n # Oy_mask[Oy_mask <= 80] = 0\n\n # if visualize:\n # plt.imshow(Oy_mask)\n # plt.show()\n \n # # find and keep largest connected component in Oy mask\n # num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(255 - Oy_mask, connectivity=8)\n # largest_label = np.argmax(stats[1:, 4]) + 1\n # new_Oy_mask = np.zeros_like(Oy_mask)\n # new_Oy_mask[labels == largest_label] = 255\n\n # # check if new_Oy_mask and mask overlap\n # if np.sum(new_Oy_mask & mask) == 0:\n # return_value_boxes, return_x_boxes, return_x_labels_polygons = return_value_boxes[1:], return_x_boxes[1:], return_x_labels_polygons[1:]\n\n return return_value_boxes, return_x_boxes, return_x_labels_polygons\n\n\n# if not TEST_MODE:\n# # ---------- EXAMPLE ----------\n# new_x_labels_polygons, new_x_boxes = map_x_labels_polygons_and_x_boxes(x_labels_polygons, x_boxes)\n# if len(new_x_labels_polygons) != 0 and (len(x_labels_polygons) - len(new_x_labels_polygons) <= 2 or len(new_x_labels_polygons) / len(x_labels_polygons) >= 0.8):\n# x_labels_polygons, x_boxes = new_x_labels_polygons, new_x_boxes\n# else:\n# # calculate x_boxes based on x_labels_polygons\n# x_boxes = []\n# for x_labels_polygon in x_labels_polygons:\n# min_x = min(np.array(x_labels_polygon)[:, 0])\n# max_x = max(np.array(x_labels_polygon)[:, 0])\n# min_y = min(np.array(x_labels_polygon)[:, 1])\n# max_y = max(np.array(x_labels_polygon)[:, 1])\n# mean_x = (min_x + max_x) // 2\n# mean_y = (min_y + max_y) // 2\n\n# x_boxes.append([mean_x - 5, mean_y - 5, mean_x + 5, mean_y + 5])\n \n\n# visualize(image_paths[idx], value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n\n# value_indices_mapping = map_x_boxes_and_value_boxes(x_boxes, value_boxes, graph_type=graph_type_predictions[idx])\n\n# missing_x_indices = set(range(len(x_boxes))) - set([p[0] for p in value_indices_mapping])\n# print(value_indices_mapping, missing_x_indices)\n\n# new_value_boxes, new_x_boxes, new_x_labels_polygons = add_missing_value_box_for_line_graph(image_paths[idx], x_boxes, y_boxes, value_boxes, x_labels_polygons, value_indices_mapping, missing_x_indices, is_visualize=True)\n# visualize(image_paths[idx], new_value_boxes, new_x_boxes, y_boxes, new_x_labels_polygons, y_labels_polygons)\n\n\n# %%\ndef postprocess_line_graph(idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, is_visualize_line=False):\n new_x_labels_polygons, new_x_boxes = map_x_labels_polygons_and_x_boxes(x_labels_polygons, x_boxes)\n\n if len(new_x_labels_polygons) != 0 and (len(x_labels_polygons) - len(new_x_labels_polygons) <= 2 or len(new_x_labels_polygons) / len(x_labels_polygons) >= 0.8):\n x_labels_polygons, x_boxes = new_x_labels_polygons, new_x_boxes\n else:\n # calculate x_boxes based on x_labels_polygons\n x_boxes = []\n for x_labels_polygon in x_labels_polygons:\n min_x = min(np.array(x_labels_polygon)[:, 0])\n max_x = max(np.array(x_labels_polygon)[:, 0])\n min_y = min(np.array(x_labels_polygon)[:, 1])\n max_y = max(np.array(x_labels_polygon)[:, 1])\n mean_x = (min_x + max_x) // 2\n mean_y = (min_y + max_y) // 2\n\n x_boxes.append([mean_x - 5, mean_y - 5, mean_x + 5, mean_y + 5])\n \n value_indices_mapping = map_x_boxes_and_value_boxes(x_boxes, value_boxes)\n\n missing_x_indices = set(range(len(x_boxes))) - set([p[0] for p in value_indices_mapping])\n # if len(missing_x_indices) > 0:\n # if there are missing values, we use CV algorithm to project x_boxes to line to get value_boxes\n # value_boxes, x_boxes, x_labels_polygons = add_missing_value_box_for_line_graph(image_paths[idx], x_boxes, y_boxes, value_boxes, x_labels_polygons, {}, set(), is_visualize=is_visualize_line)\n value_boxes, x_boxes, x_labels_polygons = add_missing_value_box_for_line_graph(image_paths[idx], x_boxes, y_boxes, value_boxes, x_labels_polygons, value_indices_mapping, missing_x_indices, is_visualize=is_visualize_line)\n \n # filter out those y boxes that are not numerical\n image_path = image_paths[idx]\n filtered_texts, filtered_y_labels_polygons, filtered_y_boxes = filter_non_numerical_boxes_and_polygons(image_path, y_boxes, y_labels_polygons)\n\n # get pixel to value pair\n pixel_to_value_pairs = get_pixel_to_value_pair(filtered_y_boxes, filtered_texts, direction=\"y\")\n if len(pixel_to_value_pairs) == 1 and x_boxes:\n Ox_pixel = int(np.mean([(p[0] + p[2]) / 2 for p in x_boxes]))\n pixel_to_value_pairs.append((Ox_pixel, 0))\n\n # get y values from value boxes\n all_values = get_y_values_from_value_boxes(value_boxes, pixel_to_value_pairs, default_len=len(x_labels_polygons))\n\n # predict text for x_labels_polygons and sort x_labels_polygons based on min x of x_labels_polygons\n x_labels_polygons = sorted(x_labels_polygons, key=lambda x: max([p[0] for p in x]) if len(x) > 0 else 0)\n\n x_labels_texts = read_text_from_polygons(image_path, x_labels_polygons, graph_type_predictions[idx], split_polygons=True)\n if x_labels_texts:\n if x_labels_texts[0] == 'o':\n x_labels_texts[0] = '0'\n\n return value_boxes, x_boxes, filtered_y_boxes, x_labels_polygons, filtered_y_labels_polygons, x_labels_texts, all_values\n\n# %%\ndef postprocess_scatter_graph(idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons):\n x_labels_polygons, x_boxes = map_x_labels_polygons_and_x_boxes(x_labels_polygons, x_boxes)\n\n filtered_x_texts, filtered_x_labels_polygons, filtered_x_boxes = filter_non_numerical_boxes_and_polygons(image_paths[idx], x_boxes, x_labels_polygons)\n filtered_y_texts, filtered_y_labels_polygons, filtered_y_boxes = filter_non_numerical_boxes_and_polygons(image_paths[idx], y_boxes, y_labels_polygons)\n\n x_pixel_to_value_pairs = get_pixel_to_value_pair(filtered_x_boxes, filtered_x_texts, direction=\"x\")\n y_pixel_to_value_pairs = get_pixel_to_value_pair(filtered_y_boxes, filtered_y_texts, direction=\"y\")\n if len(y_pixel_to_value_pairs) == 1 and x_boxes:\n Ox_pixel = int(np.mean([(p[0] + p[2]) / 2 for p in x_boxes]))\n y_pixel_to_value_pairs.append((Ox_pixel, 0))\n\n\n all_x_values = get_x_values_from_value_boxes(value_boxes, x_pixel_to_value_pairs, default_len=0)\n all_y_values = get_y_values_from_value_boxes(value_boxes, y_pixel_to_value_pairs, default_len=0)\n\n # convert all values to float\n all_x_values = [float(x) for x in all_x_values]\n all_y_values = [float(x) for x in all_y_values]\n\n return value_boxes, filtered_x_boxes, filtered_y_boxes, filtered_x_labels_polygons, filtered_y_labels_polygons, all_x_values, all_y_values\n\n\n# %%\ndef postprocess_dot_graph(idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons):\n if x_type_predictions[idx] == \"categorical\":\n # if Ox is categorical, we use vertical bar postporcessing method\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values = postprocess_bar_graph(idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n else:\n # if Ox is numerical, we use scatter plot postprocessing method, then round y (and maybe x also) value to integer.\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values = postprocess_scatter_graph(idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n all_values = [round(y) for y in all_values]\n x_labels_texts = [round(x) for x in x_labels_texts]\n\n return value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values\n\n# %%\nif not TEST_MODE:\n if graph_type_predictions[idx] in [\"vertical_bar\", \"horizontal_bar\"]:\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values = postprocess_bar_graph(\n idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons\n )\n elif graph_type_predictions[idx] == \"line\":\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values = postprocess_line_graph(\n idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons\n )\n elif graph_type_predictions[idx] == \"scatter\":\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values = postprocess_scatter_graph(\n idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons\n )\n elif graph_type_predictions[idx] == \"dot\":\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values = postprocess_dot_graph(\n idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons\n )\n\n visualize(image_paths[idx], value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n\n# %%\nimport math\nimport uuid\n\n# compute metrics\ndef compute_metrics(idx, all_values, x_labels_texts, graph_type):\n gt_xs = []\n gt_ys = []\n\n gt_data = metadata_dict[\"images/\" + image_paths[idx].split(\"/\")[-1]][\"ground_truth\"][\"gt_parse\"]\n x_type = gt_data[\"x_type\"]\n y_type = gt_data[\"y_type\"]\n\n if gt_data[\"class\"] != graph_type:\n gt_xs, gt_ys = [], []\n else:\n for v in gt_data[\"value\"]:\n # check if v[\"y\"] is not a number\n if gt_data[\"class\"] == \"horizontal_bar\":\n if not math.isnan(float(v[\"x\"])):\n gt_ys.append(v[\"x\"])\n gt_xs.append(v[\"y\"])\n \n else:\n if not math.isnan(float(v[\"y\"])):\n gt_ys.append(v[\"y\"])\n gt_xs.append(v[\"x\"])\n \n if graph_type == \"horizontal_bar\":\n x_type, y_type = y_type, x_type\n if x_type == \"categorical\":\n gt_xs = [str(x) for x in gt_xs]\n else:\n gt_xs = [float(x) for x in gt_xs]\n\n if y_type == \"categorical\":\n gt_ys = [str(y) for y in gt_ys]\n else:\n gt_ys = [float(y) for y in gt_ys]\n\n # random_id = str(uuid.uuid4())[:10] # image_paths[\"idx\"].split(\"/\")[-1] # \n random_id = os.path.basename(image_paths[idx]).split(\".\")[0]\n\n ground_truth = pd.DataFrame.from_dict({\n f'{random_id}_x': (gt_xs, gt_data[\"class\"]),\n f'{random_id}_y': (gt_ys, gt_data[\"class\"]),\n }, orient='index', columns=['data_series', 'chart_type']).rename_axis('id')\n\n\n # --------- PREDICTION ------------\n pred_xs = []\n pred_ys = []\n\n if gt_data[\"class\"] != graph_type:\n pred_xs, pred_ys = [1], [1]\n else:\n # for x_label_text, value in zip(x_labels_texts, all_values):\n # pred_xs.append(x_label_text)\n # pred_ys.append(value)\n \n pred_xs = x_labels_texts\n pred_ys = all_values\n\n if x_type == \"categorical\":\n pred_xs = [str(x) for x in pred_xs]\n else:\n pred_xs = [float(x) for x in pred_xs]\n\n if y_type == \"categorical\":\n pred_ys = [str(y) for y in pred_ys]\n else:\n pred_ys = [float(y) for y in pred_ys]\n\n predictions = pd.DataFrame.from_dict({\n f'{random_id}_x': (pred_xs, graph_type),\n f'{random_id}_y': (pred_ys, graph_type),\n }, orient='index', columns=['data_series', 'chart_type']).rename_axis('id')\n\n print(\"prediction graph type: \", graph_type)\n print(\"ground truth graph type: \", gt_data[\"class\"])\n\n return benetech_score(ground_truth, predictions), ground_truth, predictions\n\nif not TEST_MODE:\n print(graph_type_predictions[idx])\n score, gt, pred = compute_metrics(idx, all_values, x_labels_texts, graph_type_predictions[idx])\n print(\"score =\", score)\n print(\"----------\")\n print(\"pred =\\n\", pred)\n print(\"----------\")\n print(\"gt =\\n\", gt)\n print(\"-------------------------------------\")\n print(pred.data_series.values[0])\n print(gt.data_series.values[0])\n print(pred.data_series.values[1])\n print(gt.data_series.values[1])\n\n\n\n# %%\nif not TEST_MODE:\n font_path = None\n deplot_weights_path = \"google/deplot\"\nelse:\n font_path = \"/kaggle/input/deplot-fonts/arial.ttf\"\n deplot_weights_path = \"/kaggle/input/deplot\"\n\nimport pandas as pd\nimport io\nimport torch\nfrom transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor\nfrom PIL import Image\nimport re\n\ndef display_deplot_output(deplot_output, visualize=True):\n '''\n The raw output of deplot\n TITLE | Rural population (%) long-run with 2050 projections<0x0A>(OWID) in Greece<0x0A>Years | Rural population<0x0A>1940 | 47.38<0x0A>1960 | 43.68<0x0A>1980 | 30.28<0x0A>...\n '''\n \n deplot_output = deplot_output.replace(\"<0x0A>\", \"\\n\").replace(\" | \", \"\\t\")\n\n second_a_index = [m.start() for m in re.finditer('\\t', deplot_output)][1]\n last_newline_index = deplot_output.rfind('\\n', 0, second_a_index) \n\n title = deplot_output[:last_newline_index]\n table = deplot_output[last_newline_index+1:]\n\n data = io.StringIO(table)\n df = pd.read_csv(data, sep='\\t')\n if visualize:\n display(df)\n \n return df\n\n\ndef deplot(path, model, processor, device, visualize=True):\n\n image = Image.open(path)\n if visualize:\n display(image)\n inputs = processor(images=image, text=\"Generate underlying data table of the figure below:\", return_tensors=\"pt\", font_path=font_path)\n\n # Move inputs to GPU\n inputs = {key: value.to(device) for key, value in inputs.items()}\n\n predictions = model.generate(**inputs, max_new_tokens=512)\n return processor.decode(predictions[0], skip_special_tokens=True)\n\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndeplot_model = Pix2StructForConditionalGeneration.from_pretrained(deplot_weights_path).to(device)\nprocessor = Pix2StructProcessor.from_pretrained(deplot_weights_path)\n\ndef deplot_inference(image_path, visualize):\n deplot_output = deplot(image_path, deplot_model, processor, device, visualize)\n return display_deplot_output(deplot_output, visualize)\n\ndef deplot_predict(image_path, idx):\n inference_df = deplot_inference(image_path, visualize=False)\n\n x_values = inference_df[inference_df.columns[0]].values\n y_values = inference_df[inference_df.columns[1]].values\n \n graph_type = graph_type_predictions[idx]\n # check value type for each graph type\n if graph_type in [\"horizontal_bar\", \"vertical_bar\", \"line\", \"dot\", \"scatter\"]:\n # x is categorical, y is numerical\n formatted_y_values = []\n for v in y_values:\n try:\n float(v)\n if math.isnan(float(v)):\n formatted_y_values.append(0)\n else:\n formatted_y_values.append(v)\n except:\n formatted_y_values.append(0)\n \n y_values = formatted_y_values\n \n if graph_type in [\"dot\", \"scatter\"]:\n # for dot graph, x could be categorical as well as numerical, how to handle this?\n # just leave all the value be numerical for now\n\n formatted_x_values = []\n for v in x_values:\n try:\n float(v)\n if math.isnan(float(v)):\n formatted_x_values.append(0)\n else:\n formatted_x_values.append(v)\n except:\n formatted_x_values.append(0)\n \n x_values = formatted_x_values\n\n if not TEST_MODE:\n score, gt, pred = compute_metrics(idx, y_values, x_values, graph_type)\n print(\"Score =\", score)\n return score, gt, pred\n else:\n length = min(len(x_values), len(y_values))\n image_id = os.path.basename(image_paths[idx]).split(\".\")[0]\n predictions = pd.DataFrame.from_dict({\n f'{image_id}_x': (\";\".join([str(x) for x in x_values]), graph_type),\n f'{image_id}_y': (\";\".join([str(x) for x in y_values]), graph_type),\n }, orient='index', columns=['data_series', 'chart_type']).rename_axis('id')\n\n return predictions\n\n\n# %%\ndef predict(idx, is_visualize=False, is_visualize_line=False):\n graph_type = graph_type_predictions[idx]\n\n try:\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons = process_labels_polygons(idx, graph_type)\n x_boxes, y_boxes, value_boxes = process_filter_xy_value_boxes(idx, x_boxes, y_boxes, value_boxes)\n x_boxes, y_boxes, x_labels_polygons, y_labels_polygons = mapping_labels_and_value(graph_type, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n\n if graph_type in [\"vertical_bar\", \"horizontal_bar\"]:\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values = postprocess_bar_graph(\n idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons\n )\n elif graph_type == \"line\":\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values = postprocess_line_graph(\n idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, is_visualize_line\n )\n elif graph_type == \"scatter\":\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values = postprocess_scatter_graph(\n idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons\n )\n elif graph_type == \"dot\":\n value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons, x_labels_texts, all_values = postprocess_dot_graph(\n idx, value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons\n )\n \n # post process x_labels_texts\n if graph_type != \"scatter\" and x_labels_texts and type(x_labels_texts[0]) == str:\n new_x_labels_texts = []\n for value in x_labels_texts:\n if len(value) <= 4:\n new_x_labels_texts.append(value)\n else:\n new_value = \"\"\n for i in range(len(value) - 1):\n if (value[i].isalpha() and (value[i + 1].isdigit() or value[i+1] in [\">\", \"<\", ])) or \\\n ((value[i].isdigit() or value[i] in [\">\", \"<\", \")\"]) and value[i + 1].isalpha()) or \\\n (value[i].islower() and value[i + 1].isupper()):\n new_value += value[i] + \" \"\n else:\n new_value += value[i]\n\n new_value += value[-1]\n new_x_labels_texts.append(new_value)\n \n x_labels_texts = new_x_labels_texts\n\n if not TEST_MODE:\n if is_visualize:\n visualize(image_paths[idx], value_boxes, x_boxes, y_boxes, x_labels_polygons, y_labels_polygons)\n score, gt, pred = compute_metrics(idx, all_values, x_labels_texts, graph_type)\n\n if len(pred.data_series[0]) == 0 or len(pred.data_series[1]) == 0:\n print(\"Predicting nothing, try using Deplot...\")\n return deplot_predict(image_paths[idx], idx)\n\n print(\"Score =\", score)\n return score, gt, pred\n else:\n# pred_xs = []\n# pred_ys = []\n\n# for x_label_text, value in zip(x_labels_texts, all_values):\n# pred_xs.append(x_label_text)\n# pred_ys.append(value)\n\n pred_xs = x_labels_texts\n pred_ys = all_values\n\n pred_xs = [str(x) for x in pred_xs]\n pred_ys = [str(y) for y in pred_ys]\n\n if graph_type == \"horizontal_bar\":\n pred_xs, pred_ys = pred_ys, pred_xs\n\n if len(pred_xs) == 0 or len(pred_ys) == 0:\n return deplot_predict(image_paths[idx], idx)\n\n image_id = os.path.basename(image_paths[idx]).split(\".\")[0]\n predictions = pd.DataFrame.from_dict({\n f'{image_id}_x': (\";\".join([str(x) for x in pred_xs]), graph_type),\n f'{image_id}_y': (\";\".join([str(x) for x in pred_ys]), graph_type),\n }, orient='index', columns=['data_series', 'chart_type']).rename_axis('id')\n return predictions\n except:\n image_id = os.path.basename(image_paths[idx]).split(\".\")[0]\n\n try: # using pretrained deplot model\n return deplot_predict(image_path, idx)\n except:\n predictions = pd.DataFrame.from_dict({\n f'{image_id}_x': (\"0;0\", graph_type),\n f'{image_id}_y': (\"0;0\", graph_type),\n }, orient='index', columns=['data_series', 'chart_type']).rename_axis('id')\n return predictions\n\n# %%\nUSE_MULTIPLE_SIZE_SCATTER = False\n\nif USE_MULTIPLE_SIZE_SCATTER:\n keypoint_detection_config[\"test_size\"] = (640, 640)\n scatter_value_keypoint_detection_model_1 = ObjectDetectionModel(**keypoint_detection_config)\n\n keypoint_detection_config[\"test_size\"] = (768, 768)\n scatter_value_keypoint_detection_model_2 = ObjectDetectionModel(**keypoint_detection_config)\n\n scatter_value_keypoint_detection_model = scatter_value_keypoint_detection_model_1\n\npreds = []\n\nif not TEST_MODE:\n gts = []\n\n idx = 0\n # idx += 1\n while idx < len(image_paths):\n # if graph_type_predictions[idx] == \"vertical_bar\":\n print(f\"------------ {idx} -------------\")\n score, gt, pred = predict(idx)\n\n if len(pred.data_series.values[0]) > 20 and graph_type_predictions[idx] == \"scatter\" and USE_MULTIPLE_SIZE_SCATTER:\n scatter_value_keypoint_detection_model = scatter_value_keypoint_detection_model_2\n score, gt, pred = predict(idx)\n scatter_value_keypoint_detection_model = scatter_value_keypoint_detection_model_1\n\n preds.append(pred)\n gts.append(gt)\n print(\"Len of prediction: \", len(pred.data_series.values[0]))\n print(\"Len of ground truth: \", len(gt.data_series.values[0]))\n\n# if score <= 0.95:\n# print(pred)\n# print(gt)\n# predict(idx, is_visualize=True, is_visualize_line=True)\n# # predict(idx, is_visualize=True, is_visualize_line=False)\n# idx += 1\n# break\n idx += 1 \n\n concat_preds = pd.concat(preds)\n concat_gts = pd.concat(gts)\n\n score = benetech_score(concat_gts, concat_preds)\n\n print(\"----------------------------------------------------------------------------\")\n print(\"AVERAGE Score =\", score)\nelse:\n idx = 0\n while idx < len(image_paths):\n # if graph_type_predictions[idx] == \"dot\":\n print(f\"------------ {idx} -------------\")\n pred = predict(idx)\n if len(pred.data_series.values[0]) > 20 and graph_type_predictions[idx] == \"scatter\" and USE_MULTIPLE_SIZE_SCATTER:\n scatter_value_keypoint_detection_model = scatter_value_keypoint_detection_model_2\n pred = predict(idx)\n scatter_value_keypoint_detection_model = scatter_value_keypoint_detection_model_1\n\n preds.append(pred)\n\n idx += 1\n\n concat_preds = pd.concat(preds)\n\n# %%\n# for idx in [75, 210, 217, 295, 466]:\n# score, gt, pred = predict(idx)\n# print(\"---------------\")\n\n# %%\nif not TEST_MODE:\n print(pred.data_series.values[0])\n print(gt.data_series.values[0])\n print(pred.data_series.values[1])\n print(gt.data_series.values[1])\n\n# %%\nif not TEST_MODE:\n # with open(\"chart_type_results.txt\", \"w\") as f:\n for chart_type in concat_preds[\"chart_type\"].unique():\n chart_gts = concat_gts[concat_gts[\"chart_type\"] == chart_type]\n chart_preds = concat_preds[concat_preds.index.isin(chart_gts.index)]\n\n print(\"chart_type =\", chart_type, \"score =\", benetech_score(chart_gts, chart_preds))\n # f.write(\"chart_type = \" + chart_type + \" score = \" + str(benetech_score(chart_gts, chart_preds)) + \"\\n\")\n\n# %%\nif not TEST_MODE:\n # join concat_preds with concat_gts using id column\n id_to_scores = {}\n for id in concat_preds.index:\n concat_preds.loc[id, 'chart_type'] = concat_gts.loc[id, 'chart_type']\n score = benetech_score(concat_gts[concat_gts.index == id], concat_preds[concat_preds.index == id])\n id_to_scores[id] = score\n\n result = concat_preds.join(concat_gts, on='id', how='left', lsuffix='_pred', rsuffix='_gt')\n result['score'] = result.index.map(id_to_scores)\n result.to_csv('result.csv')\n\n origin_result = pd.read_csv(\"result2.csv\")\n\n for i, (id, pred, origin_pred, gt, score, origin_score) in enumerate(zip(result.index, result.data_series_pred, origin_result.data_series_pred, origin_result.data_series_gt, result.score, origin_result.score)):\n if origin_score - score >= 0.02 and id.endswith(\"_x\"):\n print(\"-----------------------------\", i // 2, score, origin_score, \"-----------------------------\")\n print(\"pred: \", pred)\n print(\"origin_pred:\", origin_pred)\n print(\"gt: \", gt)\n print(\"\\n\")\n","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:43:38.925971Z","iopub.execute_input":"2023-06-19T18:43:38.928157Z","iopub.status.idle":"2023-06-19T18:45:25.000861Z","shell.execute_reply.started":"2023-06-19T18:43:38.928121Z","shell.execute_reply":"2023-06-19T18:45:24.999748Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!rm -rf /kaggle/working/*","metadata":{"execution":{"iopub.status.busy":"2023-06-18T17:38:27.928036Z","iopub.execute_input":"2023-06-18T17:38:27.928435Z","iopub.status.idle":"2023-06-18T17:38:29.241496Z","shell.execute_reply.started":"2023-06-18T17:38:27.928391Z","shell.execute_reply":"2023-06-18T17:38:29.239571Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"concat_preds = concat_preds.reset_index()\nconcat_preds.to_csv(\"/kaggle/working/submission.csv\", index=False)\nconcat_preds","metadata":{"execution":{"iopub.status.busy":"2023-06-19T18:45:29.241929Z","iopub.execute_input":"2023-06-19T18:45:29.243061Z","iopub.status.idle":"2023-06-19T18:45:29.376350Z","shell.execute_reply.started":"2023-06-19T18:45:29.243008Z","shell.execute_reply":"2023-06-19T18:45:29.375362Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import pandas as pd\n# pd.read_csv(\"/kaggle/working/submission.csv\")","metadata":{"execution":{"iopub.status.busy":"2023-06-18T17:38:29.400397Z","iopub.execute_input":"2023-06-18T17:38:29.400841Z","iopub.status.idle":"2023-06-18T17:38:29.508432Z","shell.execute_reply.started":"2023-06-18T17:38:29.400798Z","shell.execute_reply":"2023-06-18T17:38:29.507288Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# deplot_predict(image_paths[0], 0)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T17:38:29.510246Z","iopub.execute_input":"2023-06-18T17:38:29.510670Z","iopub.status.idle":"2023-06-18T17:38:29.663893Z","shell.execute_reply.started":"2023-06-18T17:38:29.510627Z","shell.execute_reply":"2023-06-18T17:38:29.662680Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}