{"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":"For more context:https://www.kaggle.com/competitions/benetech-making-graphs-accessible/discussion/418389","metadata":{}},{"cell_type":"code","source":"import random\nimport cv2\nimport glob\nimport json\nimport os\nfrom tqdm import tqdm\nimport time\nfrom PIL import Image\nimport sys\nimport numpy as np\nimport pandas as pd","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":0.179974,"end_time":"2023-06-19T15:44:36.756250","exception":false,"start_time":"2023-06-19T15:44:36.576276","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:59.197447Z","iopub.execute_input":"2023-06-21T07:31:59.197985Z","iopub.status.idle":"2023-06-21T07:31:59.207991Z","shell.execute_reply.started":"2023-06-21T07:31:59.197946Z","shell.execute_reply":"2023-06-21T07:31:59.206750Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Yolov7 model","metadata":{}},{"cell_type":"code","source":"class Configs:\n    conf_thres=0.3\n    iou_thres=0.65\n    classes=[0,1,2,3,4,5,6,7]\n    agnostic_nms=False\n    augment=False\noptyolo = Configs","metadata":{"execution":{"iopub.status.busy":"2023-06-21T07:31:08.082258Z","iopub.execute_input":"2023-06-21T07:31:08.083116Z","iopub.status.idle":"2023-06-21T07:31:08.088294Z","shell.execute_reply.started":"2023-06-21T07:31:08.083071Z","shell.execute_reply":"2023-06-21T07:31:08.087441Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sys.path.append(\"/kaggle/input/yolov7\")\nfrom models.experimental import attempt_load\nfrom utils.datasets import LoadStreams, LoadImages\nfrom utils.general import check_file, check_img_size, check_requirements, \\\n    box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr\n\ndevice=\"cuda:0\"\n# half = device != 'cpu'  # half precision only supported on CUDA\nyolo_model = attempt_load(\"/kaggle/input/yolov7-weights/yolov7_best.pt\", map_location=device)  # load FP32 model\nstride = int(yolo_model.stride.max())  # model stride\nimgsz = check_img_size(640, s=stride)  # check img_size\nyolo_model.eval()\n1","metadata":{"papermill":{"duration":15.12419,"end_time":"2023-06-19T15:44:51.890204","exception":false,"start_time":"2023-06-19T15:44:36.766014","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:08.393481Z","iopub.execute_input":"2023-06-21T07:31:08.393849Z","iopub.status.idle":"2023-06-21T07:31:24.619743Z","shell.execute_reply.started":"2023-06-21T07:31:08.393809Z","shell.execute_reply":"2023-06-21T07:31:24.618769Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sys.path.remove(\"/kaggle/input/yolov7\")","metadata":{"papermill":{"duration":0.017879,"end_time":"2023-06-19T15:44:51.918090","exception":false,"start_time":"2023-06-19T15:44:51.900211","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:24.621787Z","iopub.execute_input":"2023-06-21T07:31:24.622700Z","iopub.status.idle":"2023-06-21T07:31:24.627833Z","shell.execute_reply.started":"2023-06-21T07:31:24.622664Z","shell.execute_reply":"2023-06-21T07:31:24.626784Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\ndef detect_yolo(imgpth):\n    dataset = LoadImages(imgpth, img_size=imgsz, stride=stride)\n    old_img_w = old_img_h = imgsz\n    old_img_b = 1\n    for path, img, im0, vid_cap in dataset:\n        img = torch.from_numpy(img).to(device)\n        img = img.float()  # uint8 to fp16/32\n        img /= 255.0  # 0 - 255 to 0.0 - 1.0\n        if img.ndimension() == 3:\n            img = img.unsqueeze(0)\n        with torch.no_grad():   # Calculating gradients would cause a GPU memory leak\n            pred = yolo_model(img.to(\"cuda:0\"), augment=optyolo.augment)[0]\n        pred = non_max_suppression(pred, optyolo.conf_thres, optyolo.iou_thres, classes=optyolo.classes, agnostic=optyolo.agnostic_nms)\n        for i, det in enumerate(pred):  # detections per image\n            if len(det):\n            # Rescale boxes from img_size to im0 size\n                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()\n                det[:,2] = det[:,2] - det[:,0]\n                det[:,3] = det[:,3] - det[:,1]\n                det[:,0] = det[:,0] + det[:,2]/2\n                det[:,1] = det[:,1] + det[:,3]/2\n\n        return pred[0].cpu()","metadata":{"execution":{"iopub.status.busy":"2023-06-21T07:31:24.629502Z","iopub.execute_input":"2023-06-21T07:31:24.629867Z","iopub.status.idle":"2023-06-21T07:31:24.643043Z","shell.execute_reply.started":"2023-06-21T07:31:24.629835Z","shell.execute_reply":"2023-06-21T07:31:24.642002Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## load test images ","metadata":{}},{"cell_type":"code","source":"import glob\nall_imgs=glob.glob(\"/kaggle/input/benetech-making-graphs-accessible/test/images/*.*\")","metadata":{"papermill":{"duration":0.020137,"end_time":"2023-06-19T15:44:52.004594","exception":false,"start_time":"2023-06-19T15:44:51.984457","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:24.647242Z","iopub.execute_input":"2023-06-21T07:31:24.647530Z","iopub.status.idle":"2023-06-21T07:31:24.661826Z","shell.execute_reply.started":"2023-06-21T07:31:24.647505Z","shell.execute_reply":"2023-06-21T07:31:24.660827Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load EasyOCR model","metadata":{}},{"cell_type":"code","source":"import sys\nsys.path.append(\"/kaggle/input/easyocr-clone\")\nsys.path.append(\"/kaggle/input/easyocr-clone/trainer\")","metadata":{"papermill":{"duration":0.017099,"end_time":"2023-06-19T15:44:52.058067","exception":false,"start_time":"2023-06-19T15:44:52.040968","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:24.664848Z","iopub.execute_input":"2023-06-21T07:31:24.665110Z","iopub.status.idle":"2023-06-21T07:31:24.671121Z","shell.execute_reply.started":"2023-06-21T07:31:24.665087Z","shell.execute_reply":"2023-06-21T07:31:24.668904Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from model import Model\nfrom trainer.utils import AttrDict\nimport yaml\nimport torch\nfrom trainer.utils import CTCLabelConverter, AttnLabelConverter, Averager","metadata":{"papermill":{"duration":0.065921,"end_time":"2023-06-19T15:44:52.133726","exception":false,"start_time":"2023-06-19T15:44:52.067805","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:24.673071Z","iopub.execute_input":"2023-06-21T07:31:24.673725Z","iopub.status.idle":"2023-06-21T07:31:24.755009Z","shell.execute_reply.started":"2023-06-21T07:31:24.673674Z","shell.execute_reply":"2023-06-21T07:31:24.754111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_config(file_path):\n    with open(file_path, 'r', encoding=\"utf8\") as stream:\n        opt = yaml.safe_load(stream)\n    opt = AttrDict(opt)\n    if opt.lang_char == 'None':\n        characters = ''\n        for data in opt['select_data'].split('-'):\n            csv_path = os.path.join(opt['train_data'], data, 'labels.csv')\n            df = pd.read_csv(csv_path, sep='^([^,]+),', engine='python', usecols=['filename', 'words'], keep_default_na=False)\n            all_char = ''.join(df['words'])\n            characters += ''.join(set(all_char))\n        characters = sorted(set(characters))\n        opt.character= ''.join(characters)\n    else:\n        opt.character = opt.number + opt.symbol + opt.lang_char\n    os.makedirs(f'./saved_models/{opt.experiment_name}', exist_ok=True)\n    return opt","metadata":{"papermill":{"duration":0.020551,"end_time":"2023-06-19T15:44:52.164327","exception":false,"start_time":"2023-06-19T15:44:52.143776","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:24.756758Z","iopub.execute_input":"2023-06-21T07:31:24.757024Z","iopub.status.idle":"2023-06-21T07:31:24.768523Z","shell.execute_reply.started":"2023-06-21T07:31:24.757001Z","shell.execute_reply":"2023-06-21T07:31:24.767497Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"opt = get_config(\"/kaggle/input/easyocr-finetuned-dataset/en_filtered_config.yaml\")\ndevice = torch.device('cpu')\nif 'CTC' in opt.Prediction:\n    converter = CTCLabelConverter(opt.character)\nelse:\n    converter = AttnLabelConverter(opt.character)\nopt.num_class = len(converter.character)","metadata":{"papermill":{"duration":0.027551,"end_time":"2023-06-19T15:44:52.201677","exception":false,"start_time":"2023-06-19T15:44:52.174126","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:24.771910Z","iopub.execute_input":"2023-06-21T07:31:24.772171Z","iopub.status.idle":"2023-06-21T07:31:24.791944Z","shell.execute_reply.started":"2023-06-21T07:31:24.772148Z","shell.execute_reply":"2023-06-21T07:31:24.791105Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ocr_model = Model(opt)\nstatedict = torch.load(\"/kaggle/input/easyocr-finetuned-dataset/best_accuracy.pth\",map_location=device)\nnewstatedict = {}\nfor k in statedict:\n    newk = k[7:]\n    newstatedict[newk] =statedict[k]\nocr_model.load_state_dict(newstatedict)\nocr_model.eval()\n1","metadata":{"execution":{"iopub.status.busy":"2023-06-21T07:31:27.190262Z","iopub.execute_input":"2023-06-21T07:31:27.190965Z","iopub.status.idle":"2023-06-21T07:31:27.835248Z","shell.execute_reply.started":"2023-06-21T07:31:27.190929Z","shell.execute_reply":"2023-06-21T07:31:27.834332Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision import transforms\nclass NormalizePAD(object):\n\n    def __init__(self, max_size, PAD_type='right'):\n        self.toTensor = transforms.ToTensor()\n        self.max_size = max_size\n        self.max_width_half = math.floor(max_size[2] / 2)\n        self.PAD_type = PAD_type\n\n    def __call__(self, img):\n        img = self.toTensor(img)\n        img.sub_(0.5).div_(0.5)\n        c, h, w = img.size()\n        Pad_img = torch.FloatTensor(*self.max_size).fill_(0)\n        Pad_img[:, :, :w] = img  # right pad\n        if self.max_size[2] != w:  # add border Pad\n            Pad_img[:, :, w:] = img[:, :, w - 1].unsqueeze(2).expand(c, h, self.max_size[2] - w)\n\n        return Pad_img","metadata":{"execution":{"iopub.status.busy":"2023-06-21T07:31:30.126497Z","iopub.execute_input":"2023-06-21T07:31:30.127130Z","iopub.status.idle":"2023-06-21T07:31:30.135833Z","shell.execute_reply.started":"2023-06-21T07:31:30.127099Z","shell.execute_reply":"2023-06-21T07:31:30.134832Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import math\nclass AlignCollate(object):\n\n    def __init__(self, imgH=64, imgW=64, keep_ratio_with_pad=False, contrast_adjust = 0.):\n        self.imgH = imgH\n        self.imgW = imgW\n        self.keep_ratio_with_pad = keep_ratio_with_pad\n        self.contrast_adjust = contrast_adjust\n\n    def __call__(self, images):\n\n        if self.keep_ratio_with_pad:  # same concept with 'Rosetta' paper\n            resized_max_w = self.imgW\n            input_channel = 3 if images[0].mode == 'RGB' else 1\n            transform = NormalizePAD((input_channel, self.imgH, resized_max_w))\n\n            resized_images = []\n            for image in images:\n                w, h = image.size\n\n                #### augmentation here - change contrast\n                if self.contrast_adjust > 0:\n                    image = np.array(image.convert(\"L\"))\n                    image = adjust_contrast_grey(image, target = self.contrast_adjust)\n                    image = Image.fromarray(image, 'L')\n\n                ratio = w / float(h)\n                if math.ceil(self.imgH * ratio) > self.imgW:\n                    resized_w = self.imgW\n                else:\n                    resized_w = math.ceil(self.imgH * ratio)\n\n                resized_image = image.resize((resized_w, self.imgH), Image.BICUBIC)\n                resized_images.append(transform(resized_image))\n                # resized_image.save('./image_test/%d_test.jpg' % w)\n\n            image_tensors = torch.cat([t.unsqueeze(0) for t in resized_images], 0)\n\n        else:\n            transform = ResizeNormalize((self.imgW, self.imgH))\n            image_tensors = [transform(image) for image in images]\n            image_tensors = torch.cat([t.unsqueeze(0) for t in image_tensors], 0)\n\n        return image_tensors\n    \nAlignCollate_valid = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD, contrast_adjust=opt.contrast_adjust)","metadata":{"execution":{"iopub.status.busy":"2023-06-21T07:31:30.805322Z","iopub.execute_input":"2023-06-21T07:31:30.805980Z","iopub.status.idle":"2023-06-21T07:31:30.821006Z","shell.execute_reply.started":"2023-06-21T07:31:30.805948Z","shell.execute_reply":"2023-06-21T07:31:30.819871Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_prediction_ocr(image):\n    timg = AlignCollate_valid([image.convert(\"L\")])\n    preds = ocr_model(timg,\"\")\n    preds_size = torch.IntTensor([preds.size(1)] * 1)\n    _, preds_index = preds.max(2)\n    preds_index = preds_index.view(-1)\n    preds_str = converter.decode_greedy(preds_index.data, preds_size.data)\n    return preds_str[0]","metadata":{"execution":{"iopub.status.busy":"2023-06-21T07:31:31.472665Z","iopub.execute_input":"2023-06-21T07:31:31.473307Z","iopub.status.idle":"2023-06-21T07:31:31.479282Z","shell.execute_reply.started":"2023-06-21T07:31:31.473274Z","shell.execute_reply":"2023-06-21T07:31:31.478142Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Helper functions for inference","metadata":{}},{"cell_type":"code","source":"NAMES = [\"x-label\",\"y-label\",\"diagram-tick\",'line','dot','vertical_bar','scatter','horizontal_bar']\ndef divide_by_class(tensor,confid,classes):\n    bbox_by_class = {i:[] for i in range(len(NAMES))}\n    for i,c,confid in zip(range(tensor.shape[0]),classes,confid):\n        bbox_by_class[int(c.item())].append(tensor[i].tolist()+[confid.item()])\n    return bbox_by_class","metadata":{"execution":{"iopub.status.busy":"2023-06-21T07:31:32.925764Z","iopub.execute_input":"2023-06-21T07:31:32.926128Z","iopub.status.idle":"2023-06-21T07:31:32.934260Z","shell.execute_reply.started":"2023-06-21T07:31:32.926098Z","shell.execute_reply":"2023-06-21T07:31:32.932597Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Remove x/y-label boxes enclosed within other boxes. \ndef remove_contained_boxes(boxes,THRES=0.95):\n    to_remove = []\n    for i in range(len(boxes)):\n        for j in range(i + 1, len(boxes)):\n            box1 = boxes[i]\n            box2 = boxes[j]\n\n            # Calculate the top-left and bottom-right coordinates of each box\n            box1_coords = [(box1[0] - box1[2]/2, box1[1] - box1[3]/2), \n                           (box1[0] + box1[2]/2, box1[1] + box1[3]/2)]\n            box2_coords = [(box2[0] - box2[2]/2, box2[1] - box2[3]/2), \n                           (box2[0] + box2[2]/2, box2[1] + box2[3]/2)]\n\n            # Calculate the overlap of the two boxes\n            dx = min(box1_coords[1][0], box2_coords[1][0]) - max(box1_coords[0][0], box2_coords[0][0])\n            dy = min(box1_coords[1][1], box2_coords[1][1]) - max(box1_coords[0][1], box2_coords[0][1])\n            if (dx>=0) and (dy>=0):\n                overlap_area = dx*dy\n                box1_area = box1[2]*box1[3]\n                box2_area = box2[2]*box2[3]\n                \n                # If box1 is contained within box2\n                if overlap_area >= THRES * box1_area:\n                    to_remove.append(i)\n                \n                # If box2 is contained within box1\n                elif overlap_area >= THRES * box2_area:\n                    to_remove.append(j)\n\n    # Remove the boxes that are contained within others\n    return [box for i, box in enumerate(boxes) if i not in to_remove]","metadata":{"papermill":{"duration":0.023326,"end_time":"2023-06-19T15:44:52.266545","exception":false,"start_time":"2023-06-19T15:44:52.243219","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:33.454776Z","iopub.execute_input":"2023-06-21T07:31:33.455129Z","iopub.status.idle":"2023-06-21T07:31:33.467163Z","shell.execute_reply.started":"2023-06-21T07:31:33.455101Z","shell.execute_reply":"2023-06-21T07:31:33.466017Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"papermill":{"duration":0.01766,"end_time":"2023-06-19T15:44:53.146505","exception":false,"start_time":"2023-06-19T15:44:53.128845","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:34.214150Z","iopub.execute_input":"2023-06-21T07:31:34.214510Z","iopub.status.idle":"2023-06-21T07:31:34.220553Z","shell.execute_reply.started":"2023-06-21T07:31:34.214482Z","shell.execute_reply":"2023-06-21T07:31:34.219588Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CLASS_NAMES = [\"x-label\",\"y-label\",\"diagram-tick\",'line','dot','vertical_bar','scatter','horizontal_bar']","metadata":{"papermill":{"duration":0.017035,"end_time":"2023-06-19T15:44:53.203995","exception":false,"start_time":"2023-06-19T15:44:53.186960","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:35.174836Z","iopub.execute_input":"2023-06-21T07:31:35.175178Z","iopub.status.idle":"2023-06-21T07:31:35.180685Z","shell.execute_reply.started":"2023-06-21T07:31:35.175150Z","shell.execute_reply":"2023-06-21T07:31:35.179648Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import re\ndef is_number(s: str):\n    s = re.sub('[,$% ]', '', s)\n    try:\n        float(s)\n        return True\n    except ValueError:\n        return False\ndef get_number(s):\n    if type(s)==str:\n        return float(re.sub('[,$% ]', '', s))\n    else:\n        return float(s)","metadata":{"papermill":{"duration":0.018581,"end_time":"2023-06-19T15:44:53.259546","exception":false,"start_time":"2023-06-19T15:44:53.240965","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:35.745086Z","iopub.execute_input":"2023-06-21T07:31:35.745446Z","iopub.status.idle":"2023-06-21T07:31:35.751726Z","shell.execute_reply.started":"2023-06-21T07:31:35.745414Z","shell.execute_reply":"2023-06-21T07:31:35.750406Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import math\n\n#Get data series information from x-tick/x-label values,y-tick/y-label values, and data-point coordinates for scatter plot\ndef scatter_inverse(result, x_tickinfo, y_tickinfo,charttype):\n    x_labels_numerical = all(is_number(label) for label in x_tickinfo.keys())\n    y_labels_numerical = all(is_number(label) for label in y_tickinfo.keys())\n\n    if x_labels_numerical:\n        x_ticks = {get_number(k): v for k, v in x_tickinfo.items()}\n        x_values = list(x_ticks.keys())\n        x_coords = list(x_ticks.values())\n    else:\n        x_mapping = {info['x']: label for label, info in x_tickinfo.items()}\n\n    if y_labels_numerical:\n        y_ticks = {get_number(k): v for k, v in y_tickinfo.items()}\n        y_values = list(y_ticks.keys())\n        y_coords = list(y_ticks.values())\n    else:\n        y_mapping = {info['y']: label for label, info in y_tickinfo.items()}\n\n    dataseries = []\n    for point in result:\n        x_coord, y_coord = point\n\n        if x_labels_numerical:\n            for i in range(len(x_values) - 1):\n                if x_coords[i+1]['x'] <= x_coord <= x_coords[i]['x']:\n                    ratio = (x_coord - x_coords[i]['x']) / (x_coords[i + 1]['x'] - x_coords[i]['x'])\n                    x_val = x_values[i] + ratio * (x_values[i + 1] - x_values[i])\n                    break\n            else:\n                ratio = (x_coord - x_coords[0]['x']) / (x_coords[-1]['x'] - x_coords[0]['x'])\n                x_val = x_values[0] + ratio * (x_values[-1] - x_values[0])\n        else:\n            x_val = x_mapping[min(x_mapping.keys(), key=lambda x: abs(x - x_coord))]\n\n\n        if y_labels_numerical:\n            for i in range(len(y_values) - 1):\n                if y_coords[i+1]['y'] <= y_coord <= y_coords[i]['y']:\n                    ratio = (y_coord - y_coords[i]['y']) / (y_coords[i + 1]['y'] - y_coords[i]['y'])\n                    y_val = y_values[i] + ratio * (y_values[i + 1] - y_values[i])\n                    break\n            else:\n                ratio = (y_coord - y_coords[0]['y']) / (y_coords[-1]['y'] - y_coords[0]['y'])\n                y_val = y_values[0] + ratio * (y_values[-1] - y_values[0])\n        else:\n            y_val = y_mapping[min(y_mapping.keys(), key=lambda x: abs(x - y_coord))]\n            \n\n        dataseries.append({'x': x_val, 'y': y_val})\n\n    return dataseries","metadata":{"papermill":{"duration":0.027761,"end_time":"2023-06-19T15:44:53.297188","exception":false,"start_time":"2023-06-19T15:44:53.269427","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:36.485961Z","iopub.execute_input":"2023-06-21T07:31:36.486307Z","iopub.status.idle":"2023-06-21T07:31:36.503372Z","shell.execute_reply.started":"2023-06-21T07:31:36.486279Z","shell.execute_reply":"2023-06-21T07:31:36.502444Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import math\n#Get data series information from x-tick/x-label values,y-tick/y-label values, and data-point coordinates for line, dot, and vertical bar plot\ndef other_inverse(result, x_tickinfo, y_tickinfo,charttype):\n    y_labels_numerical = all(is_number(label) for label in y_tickinfo.keys())\n    \n    x_mapping = {info['x']: label for label, info in x_tickinfo.items()}\n\n    if y_labels_numerical:\n        y_ticks = {get_number(k): v for k, v in y_tickinfo.items()}\n        y_values = list(y_ticks.keys())\n        y_coords = list(y_ticks.values())\n    else:\n        y_mapping = {info['y']: label for label, info in y_tickinfo.items()}\n\n    dataseries = []\n    for x_val,tick in x_tickinfo.items():\n        x_coord, y_coord = min(result,key=lambda x: abs(x[0]-tick[\"x\"]))\n        if y_labels_numerical:\n            for i in range(len(y_values) - 1):\n                if y_coords[i+1]['y'] <= y_coord <= y_coords[i]['y']:\n                    ratio = (y_coord - y_coords[i]['y']) / (y_coords[i + 1]['y'] - y_coords[i]['y'])\n                    y_val = y_values[i] + ratio * (y_values[i + 1] - y_values[i])\n                    break\n            else:\n                ratio = (y_coord - y_coords[0]['y']) / (y_coords[-1]['y'] - y_coords[0]['y'])\n                y_val = y_values[0] + ratio * (y_values[-1] - y_values[0])\n        else:\n            y_val = y_mapping[min(y_mapping.keys(), key=lambda x: abs(x - y_coord))]\n            \n        dataseries.append({'x': x_val, 'y': y_val})\n    return dataseries\n#Get data series information from x-tick/x-label values,y-tick/y-label values, and data-point coordinates for horizontal bar plot\n\ndef horizontalbar_inverse(result, x_tickinfo, y_tickinfo,charttype):\n    y_labels_numerical = all(is_number(label) for label in y_tickinfo.keys())\n    \n    x_mapping = {info['x']: label for label, info in x_tickinfo.items()}\n\n    if y_labels_numerical:\n        y_ticks = {get_number(k): v for k, v in y_tickinfo.items()}\n        y_values = list(y_ticks.keys())\n        y_coords = list(y_ticks.values())\n    else:\n        y_mapping = {info['y']: label for label, info in y_tickinfo.items()}\n\n    dataseries = []\n    for x_val,tick in x_tickinfo.items():\n        y_coord, x_coord = min(result,key=lambda x: abs(x[1]-tick[\"y\"]))\n        if y_labels_numerical:\n            for i in range(len(y_values) - 1):\n                if y_coords[i]['x'] <= y_coord <= y_coords[i+1]['x']:\n                    ratio = (y_coord - y_coords[i]['x']) / (y_coords[i + 1]['x'] - y_coords[i]['x'])\n                    y_val = y_values[i] + ratio * (y_values[i + 1] - y_values[i])\n                    break\n            else:\n                ratio = (y_coord - y_coords[0]['x']) / (y_coords[-1]['x'] - y_coords[0]['x'])\n                y_val = y_values[0] + ratio * (y_values[-1] - y_values[0])\n        else:\n            y_val = y_mapping[min(y_mapping.keys(), key=lambda x: abs(x - y_coord))]\n            \n        dataseries.append({'x': x_val, 'y': y_val})\n    return dataseries","metadata":{"papermill":{"duration":0.031468,"end_time":"2023-06-19T15:44:53.338356","exception":false,"start_time":"2023-06-19T15:44:53.306888","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:37.055293Z","iopub.execute_input":"2023-06-21T07:31:37.055758Z","iopub.status.idle":"2023-06-21T07:31:37.075419Z","shell.execute_reply.started":"2023-06-21T07:31:37.055727Z","shell.execute_reply":"2023-06-21T07:31:37.074341Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#Calculate the ticks from bounding box\ndef get_paired_bbox_ticks(bboxes,bbox,idx):\n    x,y,w,h = bbox\n    tickbox = {}\n    tcount = {}\n    for x1,y1,text in bboxes:\n        if text in tickbox:\n            tcount[text] +=1\n            text = text+str(tcount[text])\n        if idx==0:\n            tickbox[text] = {\"x\":x1,\"y\":y+h/2}\n        else:\n            tickbox[text] = {\"x\":x-w/2,\"y\":y1}\n        if text not in tcount:\n            tcount[text] = 0\n    return tickbox\n    \ndef process_predictions(classwise_preds,img,charttype,bbox):\n    x_bboxes = []\n    y_bboxes = []\n    for x,y,w,h,_ in classwise_preds[0]:\n        x1 = x - w / 2\n        y1 = y - h / 2\n        cropbox = img[int(y1):int(y1+h),int(x1):int(x1+w)]\n        result = get_prediction_ocr(Image.fromarray(cropbox))\n        x_bboxes.append([x,y,result])\n    for x,y,w,h,_ in classwise_preds[1]:\n        x1 = x - w / 2\n        y1 = y - h / 2\n        cropbox = img[int(y1):int(y1+h),int(x1):int(x1+w)]\n        result = get_prediction_ocr(Image.fromarray(cropbox))\n        y_bboxes.append([x,y,result])\n    if charttype==\"horizontal_bar\":\n        xidx,yidx=1,0\n    else:\n        xidx,yidx=0,1\n    x_tickinfo = get_paired_bbox_ticks(x_bboxes,bbox,xidx)\n    y_tickinfo = get_paired_bbox_ticks(y_bboxes,bbox,yidx)\n    \n    results = [tmp[:2] for tmp in classwise_preds[2]]\n    \n    if charttype==\"horizontal_bar\":\n        results = sorted(results,key=lambda ttmp: ttmp[1])\n        y_tickinfo = {k: v for k, v in sorted(y_tickinfo.items(), key=lambda item: item[1][\"x\"])}\n        x_tickinfo = {k: v for k, v in sorted(x_tickinfo.items(), key=lambda item: item[1][\"y\"])}        \n    else:\n        results = sorted(results,key=lambda ttmp: ttmp[0])\n        x_tickinfo = {k: v for k, v in sorted(x_tickinfo.items(), key=lambda item: item[1][\"x\"])}\n        y_tickinfo = {k: v for k, v in sorted(y_tickinfo.items(),reverse=True, key=lambda item: item[1][\"y\"])}\n    if charttype==\"scatter\":\n        answer = scatter_inverse(results,x_tickinfo,y_tickinfo,charttype)\n    elif charttype==\"horizontal_bar\":\n        answer = horizontalbar_inverse(results,x_tickinfo,y_tickinfo,charttype)\n    else:\n        answer = other_inverse(results,x_tickinfo,y_tickinfo,charttype)\n    return answer,charttype","metadata":{"papermill":{"duration":0.028497,"end_time":"2023-06-19T15:44:53.376564","exception":false,"start_time":"2023-06-19T15:44:53.348067","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:38.450830Z","iopub.execute_input":"2023-06-21T07:31:38.451191Z","iopub.status.idle":"2023-06-21T07:31:38.468493Z","shell.execute_reply.started":"2023-06-21T07:31:38.451164Z","shell.execute_reply":"2023-06-21T07:31:38.467447Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_predictions(classwise_preds,img):\n    scores = []\n    bbox = []\n    for i in range(3,8):\n        if len(classwise_preds[i])>0:\n            mostconf = max(classwise_preds[i],key = lambda x: x[-1])\n            scores.append(mostconf[-1])\n            bbox.append(mostconf[:4])\n        else:\n            scores.append(0)\n            bbox.append([])\n            \n    amx = np.argmax(scores)\n    if scores[amx]==0:\n        return None\n    charttype = CLASS_NAMES[3+amx]\n    bounding_box = bbox[amx]\n    x,y,w,h = bounding_box\n    xe = x+w/2\n    ye = y+h/2\n    x = x-w/2\n    y= y-h/2\n    yaxisidx,xaxisidx=1,0\n    if charttype==\"horizontal_bar\":\n        yaxisidx,xaxisidx=0,1\n    # Check1: Only include diagram ticks that are inside the bounding box\n    filtered = []\n    for idx, (x1, y1, w1, h1,conf) in enumerate(classwise_preds[2]):\n        if not ( x>(x1+w1/2) or y>(y1+h1/2) or xe<(x1-w1/2) or ye<(y1-h1/2)):\n            filtered.append([x1, y1, w1, h1,conf])\n    print(len(classwise_preds[2]),len(filtered))\n    classwise_preds[2] = remove_contained_boxes(filtered)\n    # CHeck2: Only keep yaxis labels on the left side of bounding box\n    filtered = []\n    for idx,(x1,y1,w1,h1,conf) in enumerate(classwise_preds[yaxisidx]):\n        if (x1-w1/2)<x and (y1-h1/2)<ye and (y1+h1/2)>y:\n            filtered.append([x1, y1, w1, h1,conf])\n    print(len(classwise_preds[yaxisidx]),len(filtered))\n    classwise_preds[yaxisidx] = remove_contained_boxes(filtered)\n    print(\"removed contained:\",len(classwise_preds[yaxisidx]),len(filtered))\n    \n    # Check 3: Only keep xaxis labels on the bottom side of bounding box\n    filtered = []\n    for idx,(x1,y1,w1,h1,conf) in enumerate(classwise_preds[xaxisidx]):\n        if (y1+h1/2)>ye and (x1-w1/2)<xe and (x1+w1/2)>x:\n            filtered.append([x1, y1, w1, h1,conf])\n    print(len(classwise_preds[xaxisidx]),len(filtered))\n    classwise_preds[xaxisidx] = remove_contained_boxes(filtered)\n    print(\"removed contained:\",len(classwise_preds[xaxisidx]),len(filtered))\n#     print(classwise_preds)\n    return process_predictions(classwise_preds,img,charttype,bounding_box)","metadata":{"papermill":{"duration":0.02678,"end_time":"2023-06-19T15:44:53.413233","exception":false,"start_time":"2023-06-19T15:44:53.386453","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:31:39.490775Z","iopub.execute_input":"2023-06-21T07:31:39.491124Z","iopub.status.idle":"2023-06-21T07:31:39.507366Z","shell.execute_reply.started":"2023-06-21T07:31:39.491094Z","shell.execute_reply":"2023-06-21T07:31:39.506330Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## make predictions, and save submission","metadata":{}},{"cell_type":"code","source":"ST= time.time()\npredictions = {}\ngroundtruth = {}\nec=0\nfor pth in tqdm(all_imgs):\n    img = cv2.cvtColor(cv2.imread(pth), cv2.COLOR_BGR2RGB)\n    results = detect_yolo(pth)\n    tmppred = [results[:,:4].cpu(),results[:,4].cpu(),results[:,5].cpu()]    \n    byclass = divide_by_class(*tmppred)\n    try:\n        preds,ctype =get_predictions(byclass,img)\n    except:\n        print(\"Error\")\n        preds =[{\"x\":0,\"y\":0}]\n        ctype = \"scatter\"\n        ec+=1\n    jsonname = os.path.split(pth)[-1].replace(\".jpg\",\"\")\n    predictions[f\"{jsonname}_x\"] = (\";\".join([str(pp[\"x\"]).replace(\";\",\"\") for pp in preds]),ctype)\n    predictions[f\"{jsonname}_y\"] = (\";\".join([str(pp[\"y\"]).replace(\";\",\"\") for pp in preds]),ctype)\ntotal = time.time() - ST","metadata":{"papermill":{"duration":6.454641,"end_time":"2023-06-19T15:45:00.045189","exception":false,"start_time":"2023-06-19T15:44:53.590548","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-06-21T07:32:06.190301Z","iopub.execute_input":"2023-06-21T07:32:06.190909Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"prediction_df = pd.DataFrame.from_dict(predictions, orient='index', columns=['data_series', 'chart_type']).rename_axis('id')","metadata":{"execution":{"iopub.execute_input":"2023-06-19T15:45:00.068551Z","iopub.status.busy":"2023-06-19T15:45:00.068243Z","iopub.status.idle":"2023-06-19T15:45:00.076931Z","shell.execute_reply":"2023-06-19T15:45:00.075578Z"},"papermill":{"duration":0.022804,"end_time":"2023-06-19T15:45:00.079149","exception":false,"start_time":"2023-06-19T15:45:00.056345","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"prediction_df.to_csv(\"submission.csv\",index=True)","metadata":{"execution":{"iopub.execute_input":"2023-06-19T15:45:00.102282Z","iopub.status.busy":"2023-06-19T15:45:00.101762Z","iopub.status.idle":"2023-06-19T15:45:00.110530Z","shell.execute_reply":"2023-06-19T15:45:00.109675Z"},"papermill":{"duration":0.022395,"end_time":"2023-06-19T15:45:00.112490","exception":false,"start_time":"2023-06-19T15:45:00.090095","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pd.read_csv(\"submission.csv\")","metadata":{"execution":{"iopub.execute_input":"2023-06-19T15:45:00.136053Z","iopub.status.busy":"2023-06-19T15:45:00.135228Z","iopub.status.idle":"2023-06-19T15:45:00.152871Z","shell.execute_reply":"2023-06-19T15:45:00.151850Z"},"papermill":{"duration":0.031475,"end_time":"2023-06-19T15:45:00.155085","exception":false,"start_time":"2023-06-19T15:45:00.123610","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]}]}