{"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":"markdown","source":"
\n This notebook is to predict plot-bb (bounding box of plot area).
\n Opencv image processing liblary is used for prediction.
\n When plot-bb of test images are known, you can narrow down the area to search data-series.

\n Plz upvoke if useful!\n
","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"markdown","source":"\n#
Import
","metadata":{}},{"cell_type":"code","source":"from glob import glob\nfrom glob import glob\nimport json\nimport os\nimport sys\n\nimport argparse\nimport cv2\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport plotly.express as px\nimport plotly.graph_objects as go\nfrom plotly.subplots import make_subplots\nfrom sklearn.metrics import mean_squared_error, r2_score\nfrom tqdm.notebook import tqdm","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:38:35.420432Z","iopub.execute_input":"2023-04-12T02:38:35.421381Z","iopub.status.idle":"2023-04-12T02:38:38.681093Z","shell.execute_reply.started":"2023-04-12T02:38:35.421325Z","shell.execute_reply":"2023-04-12T02:38:38.679906Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n#
Load data
","metadata":{}},{"cell_type":"code","source":"INPUT_DIR = '/kaggle/input/benetech-making-graphs-accessible'\ntrain_image_paths = sorted(glob(os.path.join(INPUT_DIR, 'train/images/*.jpg')))\ntrain_annotation_paths = sorted(glob(os.path.join(INPUT_DIR, 'train/annotations/*.json')))\nprint('# of train images:', len(train_image_paths))\nprint('# of train annotations:', len(train_annotation_paths))\n\n# Parameters\nargs = argparse.Namespace(\n sampling_size = 10000,\n thresholds = [9, 128]\n)","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:38:38.683207Z","iopub.execute_input":"2023-04-12T02:38:38.683537Z","iopub.status.idle":"2023-04-12T02:38:40.169442Z","shell.execute_reply.started":"2023-04-12T02:38:38.683506Z","shell.execute_reply":"2023-04-12T02:38:40.168202Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load json files to pd.DataFrame\nannotation_dict = {}\nfor annotation_path in tqdm(train_annotation_paths[:args.sampling_size]):\n with open(annotation_path, 'r') as f:\n json_str = f.read() \n id_ = annotation_path.split('/')[-1].split('.')[0]\n annotation_dict[id_] = json.loads(json_str)\n \nannotation_df = pd.DataFrame(annotation_dict).T\nprint('annotation_df:')\ndisplay(annotation_df.head(2)\\\n .style.set_properties(\n **{'background-color': '#212636','color':'white','border': '1.5px solid white'})\n .set_properties(\n **{\"background-color\": '#8a2be2'}, subset=['plot-bb']))\nprint(annotation_df.info())","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:38:40.170826Z","iopub.execute_input":"2023-04-12T02:38:40.171173Z","iopub.status.idle":"2023-04-12T02:38:46.733230Z","shell.execute_reply.started":"2023-04-12T02:38:40.171139Z","shell.execute_reply":"2023-04-12T02:38:46.731920Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# def plot_bb_on_image(image_id):\n\n# # load image\n# img = Image.open(os.path.join(INPUT_DIR, f'train/images/{image_id}.jpg'))\n# img_width = img.size[0]\n# img_height = img.size[1]\n# aspect_ratio = img_width/img_height\n# # load annotation\n# with open(os.path.join(INPUT_DIR, f'train/annotations/{image_id}.json'), 'r') as f:\n# annotation = json.load(f) \n \n# # create figure\n# fig = go.Figure()\n# # add invisible scatter trace\n# fig.add_trace(go.Scatter(x=[0, img_width],\n# y=[0, img_height],\n# mode='markers',\n# marker_opacity=0, \n# marker_color=('black'),\n# name=' '))\n# fig.update_xaxes(visible=False,\n# range=[0, img_width])\n# fig.update_yaxes(visible=False,\n# range=[0, img_height],\n# scaleanchor='x') # ensures that the aspect ratio stays constant\n \n# # add image\n# fig.add_layout_image(dict(\n# source=img,\n# x=0,\n# y=img_height,\n# sizex=img_width,\n# sizey=img_height,\n# xref='x',\n# yref='y',\n# opacity=1.0,\n# layer='below',\n# sizing='stretch'\n# ))\n \n# # add bounding box around plot ares\n# fig.add_shape(\n# type='rect',\n# x0 = annotation['plot-bb']['x0'], \n# y0 = -annotation['plot-bb']['y0'] + img_height, \n# x1 = annotation['plot-bb']['x0'] + annotation['plot-bb']['width'], \n# y1 = -(annotation['plot-bb']['y0'] + annotation['plot-bb']['height']) + img_height,\n# line = dict(color='#ff00ff'),\n# name = 'plot-bb'\n# )\n# # layout\n# fig.update_layout(\n# title_text = f'{image_id}.jpg ({annotation[\"chart-type\"]})',\n# template='plotly_dark',\n# width=800,\n# height=800 / aspect_ratio,\n# margin={'l': 0, 'r': 0, 't': 50, 'b': 50},\n# legend=dict(x=1, y=0.5, font=dict(size=9, color='white')),\n# showlegend=True\n# )\n \n# fig.show(config={'doubleClick': 'reset'})\n \n# # plot_bb on image\n# print('Examples:')\n# for chart_type in ['vertical_bar', 'horizontal_bar', 'line', 'scatter', 'dot']:\n# image_id = annotation_df[annotation_df['chart-type']==chart_type].index\n# if len(image_id)==0:\n# pass\n# else:\n# plot_bb_on_image(image_id[0])","metadata":{"_kg_hide-input":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2023-04-12T02:38:46.736000Z","iopub.execute_input":"2023-04-12T02:38:46.736333Z","iopub.status.idle":"2023-04-12T02:38:46.743383Z","shell.execute_reply.started":"2023-04-12T02:38:46.736301Z","shell.execute_reply":"2023-04-12T02:38:46.741922Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n#
Predict plot-bb
","metadata":{"execution":{"iopub.status.busy":"2023-04-04T02:23:37.886876Z","iopub.execute_input":"2023-04-04T02:23:37.887344Z","iopub.status.idle":"2023-04-04T02:23:37.893498Z","shell.execute_reply.started":"2023-04-04T02:23:37.887306Z","shell.execute_reply":"2023-04-04T02:23:37.891812Z"}}},{"cell_type":"markdown","source":" \n Predict","metadata":{}},{"cell_type":"code","source":"def line_detector(img, threshs=args.thresholds):\n \"\"\"\n Detects horizontal and vertical lines in a given image.\n\n This function takes an input image and detects horizontal and vertical lines\n within the image using the Sobel operator for edge detection and HoughLinesP \n function for line detection. It iterates through multiple threshold values \n to identify lines and returns the combined results.\n\n Args:\n img (numpy.ndarray): Input image, should be in BGR color format.\n\n Returns:\n tuple: A tuple containing two lists:\n x_lines (list): A list of detected horizontal lines, \n where each line is represented as [x1, y1, x2, y2].\n y_lines (list): A list of detected vertical lines, \n where each line is represented as [x1, y1, x2, y2].\n\n Examples:\n >>> img = cv2.imread('example_image.jpg')\n >>> x_lines, y_lines = line_detector(img)\n \"\"\"\n def get_xy_lines(img, x_or_y, thresh):\n img_height = img.shape[0]\n img_width = img.shape[1]\n gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n if x_or_y == 'x':\n sobel_edges = cv2.Sobel(gray_img, cv2.CV_64F, 0, 1, ksize=3)\n theta = np.pi/4\n elif x_or_y == 'y':\n sobel_edges = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=3)\n theta = np.pi/2\n sobel_abs = np.abs(sobel_edges)\n sobel_normalized = cv2.normalize(sobel_abs, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)\n _, edges = cv2.threshold(sobel_normalized, thresh, 255, cv2.THRESH_BINARY) \n lines = cv2.HoughLinesP(edges, \n rho=1,\n theta=theta, \n threshold=min(img_width, img_height)//10, \n minLineLength=min(img_width, img_height)//10, \n maxLineGap=1)\n return lines\n\n x_lines = None\n for thresh in threshs: \n x_lines_tmp = get_xy_lines(img, 'x', thresh)\n if x_lines_tmp is None: # not get case\n pass\n else:\n if x_lines is None:\n x_lines = x_lines_tmp\n else:\n x_lines = np.concatenate((np.array(x_lines),\n np.array(x_lines_tmp)))\n \n y_lines = None\n for thresh in threshs: \n y_lines_tmp = get_xy_lines(img, 'y', thresh)\n if y_lines_tmp is None: # not get case\n pass\n else:\n if y_lines is None:\n y_lines = y_lines_tmp\n else:\n y_lines = np.concatenate((np.array(y_lines),\n np.array(y_lines_tmp)))\n \n return x_lines, y_lines\n\n\ndef predict_plot_bb(image_id, threshs=args.thresholds):\n \"\"\"\n Predicts the bounding box of a plot in a given image.\n\n This function takes an image ID, reads the corresponding image, detects the\n horizontal and vertical lines using the line_detector function, and predicts\n the bounding box of the plot. The function trims the outer edges of the image\n to avoid detecting the image boundary as a plot boundary.\n\n Args:\n image_id (str): The image identifier used to locate the image file in the\n specified directory.\n\n Returns:\n dict: A dictionary containing the bounding box information of the plot, with\n keys 'height', 'width', 'x0', and 'y0'. 'height' and 'width' represent\n the height and width of the bounding box, while 'x0' and 'y0' represent\n the top-left corner coordinates.\n \"\"\"\n \n img_path = os.path.join(INPUT_DIR, f'train/images/{image_id}.jpg')\n img = cv2.imread(img_path)\n img_height = img.shape[0]\n img_width = img.shape[1]\n # trim edge frame so as not to judge the image outer line as plot_bb \n trim_edge = 5\n img = img[trim_edge:img_height-trim_edge, trim_edge:img_width-trim_edge, :]\n # get lines\n x_lines, y_lines = line_detector(img)\n x_lines = np.array(x_lines).squeeze().T\n y_lines = np.array(y_lines).squeeze().T \n plot_bb = {}\n # initial values\n plot_bb['height'] = img_height\n plot_bb['width'] = img_width\n plot_bb['x0'] = 0\n plot_bb['y0'] = 0\n \n if x_lines.ndim==1:\n x1, _, x2, _ = x_lines\n plot_bb['width'] = abs(x1-x2)\n plot_bb['x0'] = min(x1, x2)\n if x_lines.ndim==2:\n x1, _, x2, _ = x_lines\n width = max(abs(x1-x2))\n x0_1 = x1[np.argmax(abs(x1-x2))] + trim_edge \n x0_2 = x2[np.argmax(abs(x1-x2))] + trim_edge\n plot_bb['width'] = width\n plot_bb['x0'] = min(x0_1, x0_2)\n \n if y_lines.ndim==1:\n _, y1, _, y2 = y_lines\n plot_bb['height'] = abs(y1-y2)\n plot_bb['y0'] = min(y1, y2)\n if y_lines.ndim==2:\n _, y1, _, y2 = y_lines\n height = max(abs(y1-y2))\n y0_1 = y1[np.argmax(abs(y1-y2))] + trim_edge \n y0_2 = y2[np.argmax(abs(y1-y2))] + trim_edge\n plot_bb['height'] = height\n plot_bb['y0'] = min(y0_1, y0_2)\n \n return plot_bb","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:39:27.693691Z","iopub.execute_input":"2023-04-12T02:39:27.694121Z","iopub.status.idle":"2023-04-12T02:39:27.717190Z","shell.execute_reply.started":"2023-04-12T02:39:27.694085Z","shell.execute_reply":"2023-04-12T02:39:27.715940Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_ids = annotation_df.index\npred_plot_bb = {}\nfor image_id in tqdm(image_ids):\n pred_plot_bb[image_id] = predict_plot_bb(image_id)\n\nfor i, (k, v) in enumerate(pred_plot_bb.items()):\n print(f'{k}: {v}')\n if i>5:\n break","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:39:28.692799Z","iopub.execute_input":"2023-04-12T02:39:28.693357Z","iopub.status.idle":"2023-04-12T02:39:55.132076Z","shell.execute_reply.started":"2023-04-12T02:39:28.693307Z","shell.execute_reply":"2023-04-12T02:39:55.130284Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" \n Plot on image","metadata":{}},{"cell_type":"code","source":"def plot_bb_to_coordinates(plot_bb):\n x = [plot_bb['x0'], \n plot_bb['x0'],\n plot_bb['x0'] + plot_bb['width'],\n plot_bb['x0'] + plot_bb['width'],\n plot_bb['x0']]\n y = [plot_bb['y0'], \n plot_bb['y0'] + plot_bb['height'],\n plot_bb['y0'] + plot_bb['height'],\n plot_bb['y0'],\n plot_bb['y0']]\n return np.array(x), np.array(y)\n\n\ndef show_true_and_pred_plot_bb(image_id):\n\n # load image\n img = Image.open(os.path.join(INPUT_DIR, f'train/images/{image_id}.jpg'))\n img_width = img.size[0]\n img_height = img.size[1]\n aspect_ratio = img_width/img_height\n # load annotation\n with open(os.path.join(INPUT_DIR, f'train/annotations/{image_id}.json'), 'r') as f:\n annotation = json.load(f) \n \n # create figure\n fig = go.Figure()\n # add invisible scatter trace\n fig.add_trace(go.Scatter(x=[0, img_width],\n y=[0, img_height],\n mode='markers',\n marker_opacity=0, \n marker_color=('black'),\n name=' '))\n fig.update_xaxes(visible=False,\n range=[0, img_width])\n fig.update_yaxes(visible=False,\n range=[0, img_height],\n scaleanchor='x') # ensures that the aspect ratio stays constant\n \n # add image\n fig.add_layout_image(dict(\n source=img,\n x=0,\n y=img_height,\n sizex=img_width,\n sizey=img_height,\n xref='x',\n yref='y',\n opacity=1.0,\n layer='below',\n sizing='stretch'\n ))\n \n # add ture bounding box around plot ares\n x, y = plot_bb_to_coordinates(annotation['plot-bb'])\n fig.add_trace(\n go.Scatter(x=x, y=-y+img_height,\n line=dict(color='#ff00ff'),\n marker=dict(color='#ff00ff', size=1),\n name='true plot-bb')\n )\n \n # add pred bounding box around plot ares\n plot_bb_pred = predict_plot_bb(image_id)\n x, y = plot_bb_to_coordinates(plot_bb_pred)\n fig.add_trace(\n go.Scatter(x=x, y=-y+img_height,\n line=dict(color='skyblue', dash='dash'),\n marker=dict(color='skyblue', size=1),\n name='pred plot-bb')\n )\n # layout\n fig.update_layout(\n title_text = f'{image_id}.jpg ({annotation[\"chart-type\"]})',\n template='plotly_dark',\n width=800,\n height=800 / aspect_ratio,\n margin={'l': 0, 'r': 0, 't': 50, 'b': 50},\n legend=dict(x=1, y=0.5, font=dict(size=12, color='white')),\n showlegend=True\n )\n \n fig.show(config={'doubleClick': 'reset'})\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-04-12T02:39:55.134597Z","iopub.execute_input":"2023-04-12T02:39:55.134988Z","iopub.status.idle":"2023-04-12T02:39:55.152556Z","shell.execute_reply.started":"2023-04-12T02:39:55.134950Z","shell.execute_reply":"2023-04-12T02:39:55.151156Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# plot_bb on image\nfor chart_type in ['vertical_bar', 'horizontal_bar', 'line', 'scatter', 'dot']:\n try:\n idxs = annotation_df[annotation_df['chart-type']==chart_type].index\n show_true_and_pred_plot_bb(idxs[0])\n except:\n continue","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:39:55.154558Z","iopub.execute_input":"2023-04-12T02:39:55.155597Z","iopub.status.idle":"2023-04-12T02:39:56.988292Z","shell.execute_reply.started":"2023-04-12T02:39:55.155543Z","shell.execute_reply":"2023-04-12T02:39:56.987091Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" \n true vs pred","metadata":{}},{"cell_type":"markdown","source":"### R^2 score","metadata":{}},{"cell_type":"code","source":"pred_plot_bb_df = pd.DataFrame(pred_plot_bb).T\nprint('pred_plot_bb_df:')\ndisplay(pred_plot_bb_df.head(5).style.set_properties(\n **{'background-color': '#212636',\n 'color':'white',\n 'border': '1.5px solid white'}\n))","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:39:56.991784Z","iopub.execute_input":"2023-04-12T02:39:56.992342Z","iopub.status.idle":"2023-04-12T02:39:57.055017Z","shell.execute_reply.started":"2023-04-12T02:39:56.992290Z","shell.execute_reply":"2023-04-12T02:39:57.053643Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"true_plot_bb = [[d['height'], d['width'], d['x0'], d['y0']] \n for d in annotation_df['plot-bb'].values]\ntrue_plot_bb_df = pd.DataFrame(true_plot_bb,\n columns=['height', 'width', 'x0', 'y0'], \n index=annotation_df.index)\n\nprint('true_plot_bb_df:')\ndisplay(true_plot_bb_df.head(5).style.set_properties(\n **{'background-color': '#212636',\n 'color':'white',\n 'border': '1.5px solid white'}\n))","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:39:57.056414Z","iopub.execute_input":"2023-04-12T02:39:57.056887Z","iopub.status.idle":"2023-04-12T02:39:57.080240Z","shell.execute_reply.started":"2023-04-12T02:39:57.056832Z","shell.execute_reply":"2023-04-12T02:39:57.078916Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# R^2 score\nprint('R^2 score:')\nfor col in ['height', 'width', 'x0', 'y0']:\n score = r2_score(true_plot_bb_df[col], pred_plot_bb_df[col])\n print(f'{col}\\t: {score:.5f}')","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:39:57.082030Z","iopub.execute_input":"2023-04-12T02:39:57.082532Z","iopub.status.idle":"2023-04-12T02:39:57.096279Z","shell.execute_reply.started":"2023-04-12T02:39:57.082484Z","shell.execute_reply":"2023-04-12T02:39:57.094886Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Plot true vs pred","metadata":{}},{"cell_type":"code","source":"def plot_true_vs_pred(true_plot_bb_df, pred_plot_bb_df):\n for col in true_plot_bb_df.columns:\n fig = go.Figure()\n true = true_plot_bb_df[col]\n pred = pred_plot_bb_df[col]\n fig.add_trace(\n go.Scatter(x=true, \n y=pred,\n mode='markers',\n marker=dict(color='#00ff7f', size=5),\n name='x0')\n )\n\n fig.update_layout(\n title_text = f'{col}',\n template='plotly_dark',\n width=500,\n height=500,\n xaxis=dict(title='true'),\n yaxis=dict(title='predict'),\n showlegend=False\n )\n range_min = min(min(true), min(pred))\n range_max = max(max(true), max(pred))\n fig.update_xaxes(title_text='true', dtick=round(range_max//5,-1),\n range=(range_min, range_max))\n fig.update_yaxes(title_text='pred', dtick=round(range_max//5,-1),\n range=(range_min, range_max)) \n\n fig.show(config={'doubleClick': 'reset'})","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-04-12T02:39:57.098161Z","iopub.execute_input":"2023-04-12T02:39:57.098644Z","iopub.status.idle":"2023-04-12T02:39:57.110927Z","shell.execute_reply.started":"2023-04-12T02:39:57.098592Z","shell.execute_reply":"2023-04-12T02:39:57.109367Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_true_vs_pred(true_plot_bb_df, pred_plot_bb_df)","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:39:57.113363Z","iopub.execute_input":"2023-04-12T02:39:57.114014Z","iopub.status.idle":"2023-04-12T02:39:57.276519Z","shell.execute_reply.started":"2023-04-12T02:39:57.113960Z","shell.execute_reply":"2023-04-12T02:39:57.275646Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Search thresholds\n\n# import optuna\n \n# class Objective():\n \n# def __init__(self, annotation_df):\n# self.annotation_df = annotation_df\n \n# def __call__(self, trial):\n \n# thresh1 = trial.suggest_int('thresh1', 5, 15)\n# thresh2 = trial.suggest_int('thresh2', 100, 150)\n# threshs = [thresh1, thresh2]\n# image_ids = self.annotation_df.index\n# pred_plot_bb = {}\n# for image_id in tqdm(image_ids):\n# pred_plot_bb[image_id] = predict_plot_bb(image_id, threshs)\n\n# pred_plot_bb_df = pd.DataFrame(pred_plot_bb).T\n# true_plot_bb = [[d['height'], d['width'], d['x0'], d['y0']] \n# for d in self.annotation_df['plot-bb'].values]\n# true_plot_bb_df = pd.DataFrame(true_plot_bb,\n# columns=['height', 'width', 'x0', 'y0'], \n# index=image_ids)\n# score = 0\n# for col in ['height', 'width', 'x0', 'y0']:\n# score += r2_score(true_plot_bb_df[col], pred_plot_bb_df[col]) / 4\n# return score\n \n \n# objective = Objective(annotation_df) \n\n# study = optuna.create_study(\n# direction='maximize',\n# study_name='optimization line detector'\n# )\n# study.optimize(objective, n_trials=10, show_progress_bar=True)","metadata":{"execution":{"iopub.status.busy":"2023-04-12T02:39:57.277637Z","iopub.execute_input":"2023-04-12T02:39:57.278696Z","iopub.status.idle":"2023-04-12T02:39:57.285958Z","shell.execute_reply.started":"2023-04-12T02:39:57.278642Z","shell.execute_reply":"2023-04-12T02:39:57.284447Z"},"trusted":true},"execution_count":null,"outputs":[]}]}