{"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":"## Setup","metadata":{}},{"cell_type":"code","source":"import os\nimport sys\nimport cv2\nimport glob\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nimport matplotlib.pyplot as plt","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:50:45.969351Z","iopub.execute_input":"2023-06-18T23:50:45.969772Z","iopub.status.idle":"2023-06-18T23:50:46.843616Z","shell.execute_reply.started":"2023-06-18T23:50:45.969737Z","shell.execute_reply":"2023-06-18T23:50:46.842554Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cp -r /kaggle/input/benetech-code-v2/src ./\n!cp -r /kaggle/input/benetech-code-v2/yolox/ ./\n!mv yolox YOLOX  #","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:50:46.849453Z","iopub.execute_input":"2023-06-18T23:50:46.851961Z","iopub.status.idle":"2023-06-18T23:50:50.810968Z","shell.execute_reply.started":"2023-06-18T23:50:46.851924Z","shell.execute_reply":"2023-06-18T23:50:50.809654Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nif len(os.listdir(\"/kaggle/input/benetech-making-graphs-accessible/test/images/\")) == 5:\n    !mkdir test\n    !cp /kaggle/input/benetech-making-graphs-accessible/test/images/* test\n    !cp /kaggle/input/benetech-weights/dots/1_svg.jpg test","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:50:50.813265Z","iopub.execute_input":"2023-06-18T23:50:50.814022Z","iopub.status.idle":"2023-06-18T23:50:53.796421Z","shell.execute_reply.started":"2023-06-18T23:50:50.813982Z","shell.execute_reply":"2023-06-18T23:50:53.795016Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### YoloX","metadata":{}},{"cell_type":"code","source":"!mkdir -p /tmp/pip/cache/\n!cp /kaggle/input/yolox-whl/* /tmp/pip/cache/","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:50:53.799490Z","iopub.execute_input":"2023-06-18T23:50:53.800210Z","iopub.status.idle":"2023-06-18T23:51:16.623270Z","shell.execute_reply.started":"2023-06-18T23:50:53.800171Z","shell.execute_reply":"2023-06-18T23:51:16.621962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install -qqq --no-index /kaggle/input/yolox-whl/loguru-0.7.0-py3-none-any.whl\n!pip install -qqq --no-index /kaggle/input/yolox-whl/thop-0.1.1.post2209072238-py3-none-any.whl\n!pip install -qqq --no-index /kaggle/input/yolox-whl/protobuf-4.23.2-cp37-abi3-manylinux2014_x86_64.whl\n!pip install -qqq --no-index /kaggle/input/yolox-whl/onnx-1.14.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n!pip install -qqq --no-index /kaggle/input/yolox-whl/onnx_simplifier-0.4.10-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n!pip install -qqq --no-index /kaggle/input/yolox-whl/Cython-0.29.35-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_24_x86_64.whl\n!pip install -qqq /kaggle/input/cached/pycocotools-2.0.6-cp310-cp310-linux_x86_64.whl","metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:51:16.627034Z","iopub.execute_input":"2023-06-18T23:51:16.627471Z","iopub.status.idle":"2023-06-18T23:53:07.103365Z","shell.execute_reply.started":"2023-06-18T23:51:16.627441Z","shell.execute_reply":"2023-06-18T23:53:07.102039Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cd /kaggle/working","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-06-18T23:53:07.105313Z","iopub.execute_input":"2023-06-18T23:53:07.105683Z","iopub.status.idle":"2023-06-18T23:53:07.114204Z","shell.execute_reply.started":"2023-06-18T23:53:07.105647Z","shell.execute_reply":"2023-06-18T23:53:07.113135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Cached","metadata":{}},{"cell_type":"code","source":"!pip install -q --no-index --find-links=file:///kaggle/input/cached addict terminaltables pycocotools mmcv_full mmdet","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-06-18T23:53:07.115870Z","iopub.execute_input":"2023-06-18T23:53:07.116231Z","iopub.status.idle":"2023-06-18T23:53:20.143023Z","shell.execute_reply.started":"2023-06-18T23:53:07.116198Z","shell.execute_reply":"2023-06-18T23:53:20.141850Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Classification","metadata":{}},{"cell_type":"code","source":"%%writefile inf_main.py\n\nimport glob\nimport json\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom torch.utils.data import DataLoader\n\nfrom data.dataset import ClsDataset\nfrom data.transforms import get_transfos\nfrom model_zoo.models import define_model\nfrom util.torch import load_model_weights\nfrom util.metrics import accuracy\n\n\nNUM_WORKERS = 2\n\n\ndef predict(model, dataset, loss_config, batch_size=64, device=\"cuda\", use_fp16=False):\n    \"\"\"\n    Perform model inference on a dataset.\n\n    Args:\n        model (nn.Module): Trained model for inference.\n        dataset (Dataset): Dataset to perform inference on.\n        loss_config (dict): Loss configuration.\n        batch_size (int, optional): Batch size for inference. Defaults to 64.\n        device (str, optional): Device to use for inference. Defaults to \"cuda\".\n        use_fp16 (bool, optional): Whether to use mixed precision inference. Defaults to False.\n\n    Returns:\n        preds (numpy.ndarray): Predicted probabilities of shape (num_samples, num_classes).\n        preds_aux (numpy.ndarray): Auxiliary predictions of shape (num_samples, num_aux_classes).\n    \"\"\"\n    model.eval()\n    preds = np.empty((0, model.num_classes))\n    preds_aux = []\n\n    loader = DataLoader(\n        dataset, batch_size=batch_size, shuffle=False, num_workers=NUM_WORKERS\n    )\n\n    with torch.no_grad():\n        for img, _, _ in tqdm(loader):\n            with torch.cuda.amp.autocast(enabled=use_fp16):\n                pred, pred_aux = model(img.cuda())\n\n            # Get probabilities\n            if loss_config[\"activation\"] == \"sigmoid\":\n                pred = pred.sigmoid()\n            elif loss_config[\"activation\"] == \"softmax\":\n                pred = pred.softmax(-1)\n\n            if loss_config.get(\"activation_aux\", \"softmax\") == \"sigmoid\":\n                pred_aux = pred_aux.sigmoid()\n            elif loss_config.get(\"activation_aux\", \"softmax\") == \"softmax\":\n                pred_aux = pred_aux.softmax(-1)\n\n            preds = np.concatenate([preds, pred.cpu().numpy()])\n            preds_aux.append(pred_aux.cpu().numpy())\n\n    return preds, np.concatenate(preds_aux)\n\n\nclass Config:\n    \"\"\"\n    Placeholder to load a config from a saved json\n    \"\"\"\n\n    def __init__(self, dic):\n        for k, v in dic.items():\n            setattr(self, k, v)\n\n\ndef kfold_inference(\n    df,\n    exp_folder,\n    debug=False,\n    use_tta=False,\n    use_fp16=False,\n):\n    \"\"\"\n    Perform k-fold cross-validation for model inference on the validation set.\n\n    Args:\n        df (pd.DataFrame): DataFrame containing the data.\n        exp_folder (str): Path to the experiment folder.\n        debug (bool, optional): Whether to run in debug mode. Defaults to False.\n        save (bool, optional): Whether to save the predictions. Defaults to True.\n        use_tta (bool, optional): Whether to use test time augmentation. Defaults to False.\n        use_fp16 (bool, optional): Whether to use mixed precision inference. Defaults to False.\n        train (bool, optional): Whether to perform inference on the training set. Defaults to False.\n        use_mt (bool, optional): Whether to use model teacher. Defaults to False.\n        distilled (bool, optional): Whether to use distilled model. Defaults to False.\n        n_soup (int, optional): Number of models to use for model soup. Defaults to 0.\n\n    Returns:\n        np.ndarray: Array containing the predicted probabilities for each class.\n    \"\"\"\n    assert not use_tta, \"TTA not implemented\"\n    predict_fct = predict  # predict_tta if use_tta else predict\n\n    config = Config(json.load(open(exp_folder + \"config.json\", \"r\")))\n\n    model = define_model(\n        config.name,\n        num_classes=config.num_classes,\n        num_classes_aux=config.num_classes_aux,\n        n_channels=config.n_channels,\n        drop_rate=config.drop_rate,\n        drop_path_rate=config.drop_path_rate,\n        use_gem=config.use_gem,\n        reduce_stride=config.reduce_stride,\n        verbose=(config.local_rank == 0),\n        pretrained=False,\n    )\n    model = model.cuda().eval()\n\n    preds = []\n    for weights in glob.glob(exp_folder + f\"{config.name}_fullfit_*\"):\n        print(weights)\n        model = load_model_weights(model, weights, verbose=1)\n\n        dataset = ClsDataset(\n            df,\n            transforms=get_transfos(augment=False, resize=config.resize),\n        )\n\n        pred, _ = predict_fct(\n            model,\n            dataset,\n            config.loss_config,\n            batch_size=8,\n            use_fp16=use_fp16,\n        )\n\n        if df['target'].max():\n            acc = accuracy(df[\"target\"].values, pred)\n            print(f\"\\n -> Accuracy : {acc:.4f}\")\n        \n        preds.append(pred)\n\n    return np.mean(preds, 0)\n","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:53:20.145173Z","iopub.execute_input":"2023-06-18T23:53:20.145548Z","iopub.status.idle":"2023-06-18T23:53:20.157428Z","shell.execute_reply.started":"2023-06-18T23:53:20.145512Z","shell.execute_reply":"2023-06-18T23:53:20.156234Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%python\n\nimport sys\nsys.path.append('/kaggle/working/src/')\nsys.path.append('/kaggle/input/timm-0-6-9/pytorch-image-models-master')\n\nimport os\nimport re\nimport glob\nimport timm\nimport numpy as np\nimport pandas as pd\n\nfrom params import CLASSES\nfrom inf_main import kfold_inference\n\n\nDATA_PATH = \"/kaggle/input/benetech-making-graphs-accessible/\"\n\nimages = glob.glob(DATA_PATH + 'test/images/*.jpg')\nif len(images) == 5:\n    images = glob.glob('/kaggle/working/test/*')\n\nif len(images) > 0:  # 6 to infer val\n    df = pd.DataFrame({\"path\": images})\n    df['target'] = 0\nelse:\n    df = pd.read_csv('/kaggle/input/benetech-weights/cls/33/df_val.csv')\n    df = df[df['chart-type'] != \"dot\"].reset_index(drop=True)\n    df['path'] = df['path'].apply(lambda x: re.sub('../input/../input/', DATA_PATH, x))\n\nEXP_FOLDERS = [\n    \"/kaggle/input/benetech-weights/cls/2023-06-08_6/\",\n    \"/kaggle/input/benetech-weights/cls/2023-06-08_7/\",\n]\n\npreds = []\nfor exp_folder in EXP_FOLDERS:\n    pred = kfold_inference(df, exp_folder)\n    preds.append(pred)\n\npred = np.mean(preds, 0)\npred = pred.argmax(-1)\ndf['pred'] = pred\ndf['chart-type'] = df['pred'].apply(lambda x: CLASSES[x])\ndf['id'] = df['path'].apply(lambda x: x[:-4].split('/')[-1])\ndf.to_csv('pred_cls.csv', index=False)\nprint('\\n-> Saved results to pred_cls.csv')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-06-18T23:53:20.159124Z","iopub.execute_input":"2023-06-18T23:53:20.159455Z","iopub.status.idle":"2023-06-18T23:53:49.757144Z","shell.execute_reply.started":"2023-06-18T23:53:20.159426Z","shell.execute_reply":"2023-06-18T23:53:49.756123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/working/src/')\n\nimport pandas as pd\nimport seaborn as sns\nimport matplotlib.pyplot as plt\n\nfrom params import CLASSES\nfrom util.plots import plot_confusion_matrix\n\ndf = pd.read_csv('pred_cls.csv')\n\nif df['target'].max():    \n    plot_confusion_matrix(df['pred'], df['target'], display_labels=CLASSES, normalize=None)\n    plt.xticks(rotation=45)\n    plt.title(f\"Acc = {(df['pred'] == df['target']).mean() :.4f}\")\n    plt.show()\nelse:\n    sns.countplot(x=df['chart-type'])\n    plt.show()\n    \ndf.head(10)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:53:49.760814Z","iopub.execute_input":"2023-06-18T23:53:49.761131Z","iopub.status.idle":"2023-06-18T23:53:50.118038Z","shell.execute_reply.started":"2023-06-18T23:53:49.761105Z","shell.execute_reply":"2023-06-18T23:53:50.117082Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dots","metadata":{}},{"cell_type":"code","source":"%%writefile dot_cfg.py\n\nclass ConfigDot:\n    selected_model = \"yolo\"\n    bbox_format = \"yolo\"\n    pred_format = \"pascal_voc\"\n\n    name = \"benetech_1_m_1\"\n    \n    cfg = \"/kaggle/input/benetech-yolox-weights/yolox_m/benetech_1_m_1.py\"\n    ckpt = \"/kaggle/input/benetech-yolox-weights/yolox_m/best_ckpt.pth\"\n\n    version = \"v13\"\n    labels = [\"point\"]\n\n    size = (1024, 1024)\n\n    # NMS\n    conf_thresh = 0.05\n    iou_thresh = 0.1\n    max_per_img = 500\n    min_per_img = 1\n    \n    val_bs = 1  # if size[0] > 1024 else 16\n    device = \"cuda\"\n    \nconfig_marker = ConfigDot","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:53:50.119524Z","iopub.execute_input":"2023-06-18T23:53:50.120015Z","iopub.status.idle":"2023-06-18T23:53:50.127770Z","shell.execute_reply.started":"2023-06-18T23:53:50.119977Z","shell.execute_reply":"2023-06-18T23:53:50.126551Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%python\n\nimport sys\nsys.path.append('/kaggle/working/src/')\nsys.path.append('/kaggle/working/YOLOX/')\n\nimport warnings\nwarnings.filterwarnings(\"ignore\", category=UserWarning)\n\nimport os\nimport re\nimport cv2\nimport sys\nimport glob\nimport torch\nimport numpy as np\nimport pandas as pd\nimport transformers\nimport matplotlib.pyplot as plt\n\nfrom tqdm import tqdm\nfrom pathlib import Path\nfrom transformers import TrOCRProcessor\nfrom transformers import VisionEncoderDecoderModel\nfrom mmdet.apis import init_detector, inference_detector\n\nfrom util.ocr import ocr\nfrom util.plots import plot_results\nfrom util.metrics import score_series\n\nfrom inference.yolox import retrieve_yolox_model, predict, YoloXWrapper\nfrom inference.utils import get_transfos, InferenceDataset\n\nfrom post_process.retrieve import retrieve_missing_boxes\nfrom post_process.in_graph import post_process_preds_dots\nfrom post_process.dots import constraint_size, restrict_labels_x, assign_dots, cluster_on_x\n\nfrom dot_cfg import config_marker\n\n\ntransformers.utils.logging.set_verbosity_error()\n\n\n# Data\nprint('- Prepare data')\n\ndf_target = pd.read_csv('/kaggle/input/benetech-weights/y_train.csv')\ndf = pd.read_csv('/kaggle/working/pred_cls.csv')\n\nif len(df) > 0:  # 6 for dot val\n    df = df[df['chart-type'] == \"dot\"].reset_index(drop=True)\n    df['gt_path'] = \"\"\nelse:\n    df_test = pd.DataFrame({\"path\": glob.glob('/kaggle/input/benetech-weights/dots/*')})\n    df_test['id'] = df_test['path'].apply(lambda x: Path(x).stem)\n    df_test['source'] = \"extracted\"\n    df_test['chart-type'] = \"dot\"\n    df_test['gt_path'] = \"\"\n    df_test.to_csv('df_test_dot.csv', index=False)\n    df = df_test\n    df_target = pd.read_csv(\"/kaggle/input/benetech-weights/dot_labels.csv\")\n\ntransforms = get_transfos(size=config_marker.size)\ndataset = InferenceDataset(df, transforms, pad=True)\n\n\n# Dot inf\nprint('- Dot inference')\n\nassert os.path.exists(config_marker.ckpt), \"Weights do not exist\"\nassert os.path.exists(config_marker.cfg), \"Config does not exist\"\n    \nif len(dataset):\n    model = retrieve_yolox_model(config_marker.cfg, config_marker.ckpt)\n    model = YoloXWrapper(model, config_marker)\n    meter, _ = predict(model, dataset, config_marker)\n    dataset = InferenceDataset(df, None)\n\n\n# Cached\nprint('- Retrieve CACHED model')\nwdir = '/kaggle/input/cached/work_dirs'\nconfig_file = '/kaggle/input/benetech-weights/custom.py'\ncheckpoint_file = wdir + '/cascade_rcnn_swin-t_fpn_LGF_VCE_PCE_coco_focalsmoothloss/checkpoint.pth'\nif len(dataset):\n    cached_model = init_detector(config_file, checkpoint_file, device='cuda')\n\n\n# OCR\nprint('- Prepare OCR model')\n\nif len(dataset):\n    processor = TrOCRProcessor.from_pretrained(\"/kaggle/input/benetech-weights/ocr/\")\n    config = torch.load(\"/kaggle/input/benetech-weights/ocr/config.pth\")\n    ocr_model = VisionEncoderDecoderModel(config)\n    ocr_model.load_state_dict(torch.load('/kaggle/input/benetech-weights/ocr/pytorch_model.bin'))\n    ocr_model = ocr_model.cuda().eval()\n\n    \n# Loop\nprint('\\n -> Done ! Starting loop...\\n')\nPLOT = len(dataset) < 5\n\nscores, preds_x, preds_y = [], {}, {}\nfor idx in range(len(dataset)):\n    try:\n        # Load image & update shape\n        img, gt, shape = dataset[idx]\n        if img.shape[1] > img.shape[0] * 1.4:\n    #         padding = img.shape[0] // 2\n            padding = int(img.shape[1] * 0.9) - img.shape[0]\n        else:\n            padding = 0\n        meter.preds[idx].update_shape((img.shape[0] + padding, img.shape[1]))\n        \n        # Cached\n        cached_result = inference_detector(cached_model, dataset.paths[idx])  # list[array]\n        score_th = min(0.1, cached_result[4][2, 4])\n\n        img, gt, shape = dataset[idx]\n\n        id_ = df.id[idx]\n\n        print(idx, id_[:10], end=\"\\t\")\n        title = f\"{id_}\"\n\n        # Preds\n        preds_ = [meter.preds[idx]['pascal_voc'][meter.labels[idx] == 0] for i in range(1)][0]\n        preds = [[], [], [], preds_]\n\n        # Override with cached\n        preds[1] = cached_result[4][cached_result[4][:, -1] > score_th][:, :4].astype(int)\n        preds[0] = cached_result[2][:1, :4].astype(int)\n\n        preds = post_process_preds_dots(preds, margin_pt=5, margin_text=5)\n#         preds[-1] = constraint_size(preds[-1], margin=2, coef=0.2)\n\n        margin = (img.shape[0] + img.shape[1]) / (2 * 20)\n        preds = restrict_labels_x(preds, margin=margin)\n\n        # Visual similarity\n        try:\n            retrieved_boxes = retrieve_missing_boxes(\n                preds, img, verbose=False, min_sim=0.85, seed=100, hw=None, max_retrieved=20, margin=-1\n            )\n            if len(retrieved_boxes):\n                print('RETRIEVED', len(retrieved_boxes), end=\"\\t\")\n                preds[-1] = np.concatenate([preds[-1], retrieved_boxes])\n        except:\n            pass\n\n        try:\n            centers, clusters = cluster_on_x(preds[-1], shape[1], plot=False)\n            centers = np.array([c for i, c in enumerate(centers) if clusters[i] > 0])\n        except:\n            centers, clusters = None, None\n\n        if len(preds[1]):\n            try:\n                xlabels = preds[1]\n                xlabels_loc = (xlabels[:, 0] + xlabels[:, 2]) / 2\n        \n                if centers is not None:\n                    mapping, retrieved_xlabels = assign_dots(preds[1], centers, retrieve_missing=True)\n                    if len(retrieved_xlabels):\n                        xlabels = np.concatenate([xlabels, retrieved_xlabels])\n                    xlabels_loc = (xlabels[:, 0] + xlabels[:, 2]) / 2\n                    preds[1] = xlabels\n\n                # OCR\n                x_texts = ocr(ocr_model, processor, img, preds[1], margin=1, plot=False)\n\n                xs, ys, locs = [], [], []\n                for i, txt in enumerate(x_texts):\n                    if clusters is not None:\n                        if i in mapping.keys():\n                            xs.append(txt)\n                            locs.append(xlabels_loc[i])\n                            ys.append(clusters.get(mapping[i], 0))\n                        else:\n                            if xlabels_loc[i] > preds[0][0][0]:\n                                xs.append(txt)\n                                locs.append(xlabels_loc[i])\n                                ys.append(0)\n                    else:\n                        xs.append(txt)\n                        locs.append(xlabels_loc[i])\n                        ys.append(0)\n\n            except:  # OCR or assignment fails, using dummy values\n                print('Error')\n                xs = [str(i) for i in range(len(centers))]\n                locs = centers\n                ys = list(clusters.values())\n        else:\n            xs = [str(i) for i in range(len(centers))]\n            locs = centers\n            ys = list(clusters.values())\n\n        if PLOT:\n            plot_results(img, preds, figsize=(12, 7), title=title)\n\n        pred = pd.DataFrame({\"x\": xs, \"y\": np.array(ys).astype(int), \"loc\": locs})\n        pred = pred.sort_values('loc').reset_index(drop=True)\n\n        preds_x[id_] = \";\".join(pred['x'].astype(str))\n        preds_y[id_] = \";\".join(pred['y'].astype(str))\n\n        if df_target is not None:\n            gt = df_target[df_target['id'] == id_].reset_index(drop=True)\n            if len(gt):\n                gt['y'] = gt[\"y\"].astype(int)\n                score_x = score_series(gt['x'].values, pred['x'].values)\n                score_y = score_series(gt['y'].values, pred['y'].values)\n                print(f\"Scores  -  x: {score_x:.3f}  - y: {score_y:.3f}\")\n                scores += [score_x, score_y]\n    except:\n        pass\n\nif len(scores):\n    print(f'\\n-> Dot CV : {np.mean(scores) :.3f}')\n    \ndf['pred_x'] = df['id'].apply(lambda x: preds_x.get(x, \"\"))\ndf['pred_y'] = df['id'].apply(lambda x: preds_y.get(x, \"\"))\n\ndf.to_csv('df_dot_preds.csv', index=False)\n\nprint('\\n\\n-> Saved results to df_dot_preds.csv')","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:53:50.129911Z","iopub.execute_input":"2023-06-18T23:53:50.130296Z","iopub.status.idle":"2023-06-18T23:54:55.951545Z","shell.execute_reply.started":"2023-06-18T23:53:50.130262Z","shell.execute_reply":"2023-06-18T23:54:55.950507Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### To submission format","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv('pred_cls.csv')[['path', 'chart-type', \"id\"]]\n\n# Merge dot\ndf_dot = pd.read_csv('df_dot_preds.csv')\ndf_dot = df_dot[['id', 'pred_x', 'pred_y']]\n\ndf = df.merge(df_dot, on=\"id\", how=\"left\")\n\ndf['pred_x'] = df['pred_x'].fillna(\"\")\ndf['pred_y'] = df['pred_y'].fillna(\"\")\n\n# Reformat & save\ndfx = df[[\"id\", \"pred_x\", \"chart-type\"]].copy()\ndfx['id'] += \"_x\"\ndfx.columns = [\"id\", \"data_series\", \"chart_type\"]\n\ndfy = df[[\"id\", \"pred_y\", \"chart-type\"]].copy()\ndfy['id'] += \"_y\"\ndfy.columns = [\"id\", \"data_series\", \"chart_type\"]\n\nsub = pd.concat([dfx, dfy]).sort_values('id').reset_index(drop=True)\n\nsub.to_csv('submission_dot.csv', index=False)\nprint('-> Saved results to submission_dot.csv')\n\ndisplay(sub[sub['chart_type'] == \"dot\"].head())\nsub.head(20)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:54:55.952792Z","iopub.execute_input":"2023-06-18T23:54:55.954313Z","iopub.status.idle":"2023-06-18T23:54:55.998624Z","shell.execute_reply.started":"2023-06-18T23:54:55.954280Z","shell.execute_reply":"2023-06-18T23:54:55.997600Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Scatter","metadata":{}},{"cell_type":"code","source":"%%writefile scatter_cfg.py\n\n\nclass ConfigMarker:\n    selected_model = \"yolo\"\n    bbox_format = \"yolo\"\n    pred_format = \"pascal_voc\"\n\n    name = \"benetech_1_m_1\"\n    cfg = \"/kaggle/input/benetech-yolox-weights/yolox_m/benetech_1_m_1.py\"\n    ckpt = \"/kaggle/input/benetech-yolox-weights/yolox_m/best_ckpt.pth\"\n    \n    version = \"v13\"\n#     version = \"v13_sim\"\n    labels = [\"point\"]\n\n    size = (1024, 1024)\n\n    # NMS\n    conf_thresh = 0.6\n    iou_thresh = 0.4\n    max_per_img = 500\n    min_per_img = 1\n    \n    val_bs = 1  # if size[0] > 1024 else 16\n    device = \"cuda\"\n    \nclass ConfigMarker2:\n    selected_model = \"yolo\"\n    bbox_format = \"yolo\"\n    pred_format = \"pascal_voc\"\n\n    name = \"benetech_1_l_1\"\n    cfg = \"/kaggle/input/benetech-yolox-weights/yolox_l/benetech_1_l_1.py\"\n    ckpt = \"/kaggle/input/benetech-yolox-weights/yolox_l/best_ckpt.pth\"\n    \n    version = \"v13\"\n#     version = \"v13_sim\"\n    labels = [\"point\"]\n\n    size = (1024, 1024)\n\n    # NMS\n    conf_thresh = 0.55\n    iou_thresh = 0.4\n    max_per_img = 500\n    min_per_img = 1\n    \n    val_bs = 1  # if size[0] > 1024 else 16\n    device = \"cuda\"","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:54:56.003204Z","iopub.execute_input":"2023-06-18T23:54:56.003481Z","iopub.status.idle":"2023-06-18T23:54:56.011129Z","shell.execute_reply.started":"2023-06-18T23:54:56.003459Z","shell.execute_reply":"2023-06-18T23:54:56.010124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# %%python\n\nimport sys\nsys.path.append('/kaggle/working/src/')\nsys.path.append('/kaggle/working/YOLOX/')\n\nimport os\nimport re\nimport cv2\nimport sys\nimport glob\nimport torch\nimport warnings\nimport transformers\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nfrom tqdm import tqdm\nfrom transformers import TrOCRProcessor\nfrom transformers import VisionEncoderDecoderModel\nfrom mmdet.apis import init_detector, inference_detector\n\nfrom inference.yolox import retrieve_yolox_model, predict, YoloXWrapper\nfrom inference.utils import get_transfos, InferenceDataset, nms\n\nfrom util.plots import plot_results\nfrom util.metrics import score_series\nfrom util.torch import seed_everything\nfrom util.ocr import ocr, post_process_texts\n\nfrom post_process.retrieve import retrieve_missing_boxes\nfrom post_process.remove import remove_outlier_boxes\nfrom post_process.reg import rounding, linear_regression\nfrom post_process.ticks import restrict_on_line, assign\nfrom post_process.in_graph import post_process_preds\nfrom post_process.tick_point import post_process_arrow, post_process_point_as_tick\n\nfrom scatter_cfg import ConfigMarker as config_marker\nfrom scatter_cfg import ConfigMarker2 as config_marker_2\n\ntransformers.utils.logging.set_verbosity_error()\nwarnings.filterwarnings(\"ignore\", category=UserWarning)\n\n# Data\ndf = pd.read_csv('/kaggle/working/pred_cls.csv')\ndf = df[df['chart-type'] == \"scatter\"].reset_index(drop=True)\ndf['gt_path'] = \"\"\n\nprint(f'- Found {len(df)} scatter plots')\n    \n# Marker inference\nprint('- Marker inference')\n\nassert os.path.exists(config_marker.ckpt), \"Weights do not exist\"\nassert os.path.exists(config_marker.cfg), \"Config does not exist\"\n\nmodel_marker = retrieve_yolox_model(config_marker.cfg, config_marker.ckpt)\nmodel_marker = YoloXWrapper(model_marker, config_marker)\n\nmodel_marker_2 = retrieve_yolox_model(config_marker_2.cfg, config_marker_2.ckpt)\nmodel_marker_2 = YoloXWrapper(model_marker_2, config_marker_2)\n    \ntransforms = get_transfos(size=config_marker.size)\ndataset = InferenceDataset(df, transforms, pad=True, pad_advanced=True)\n\nmeter, _ = predict(model_marker, dataset, config_marker, extract_fts=False)\nmeter_2, _ = predict(model_marker_2, dataset, config_marker, extract_fts=False)\n\ndataset = InferenceDataset(df, None)\n\n    \n# Cached\nprint('- Retrieve CACHED model')\nCACHED_CLASSES = [\n    'x_title', 'y_title', 'plot_area', 'other', 'xlabel', 'ylabel',\n    'chart_title', 'x_tick', 'y_tick', 'legend_patch', 'legend_label',\n    'legend_title', 'legend_area', 'mark_label', 'value_label',\n    'y_axis_area', 'x_axis_area', 'tick_grouping'\n]\n    \nwdir = '/kaggle/input/cached/work_dirs'\nconfig_file = '/kaggle/input/benetech-weights/custom.py'\ncheckpoint_file = wdir + '/cascade_rcnn_swin-t_fpn_LGF_VCE_PCE_coco_focalsmoothloss/checkpoint.pth'\nif len(dataset):\n    cached_model = init_detector(config_file, checkpoint_file, device='cuda')\n\n# OCR\nprint('- Prepare OCR model')\nprocessor = TrOCRProcessor.from_pretrained(\"/kaggle/input/benetech-weights/ocr/\")\n\nconfig = torch.load(\"/kaggle/input/benetech-weights/ocr/config.pth\")\nocr_model = VisionEncoderDecoderModel(config)\nocr_model.load_state_dict(torch.load('/kaggle/input/benetech-weights/ocr/pytorch_model.bin'))\nocr_model = ocr_model.cuda().eval()\n\n\n# Dataset & GT\nprint('- Prepare data')\ndataset = InferenceDataset(df, None)\nTO_REMOVE = [\"513147edc8a1\", \"a7a81c55df4c\", \"039d3e82ebaf\", \"82c3706f2698\", \"6d4d21bdc9a8\", \"ca30ad3528c4\", \"1ab7f626447d\"]\nFIXES = {\n    \"17000b60f53e\": [-2.36, -1.6, -1.3, -0.8, -0.5, 0.009, 0.416, 0.768, 1.296, 1.539, 2.027],\n    \"6d4d21bdc9a8\": [7400.0, 8100.0, 9300.0, 6300.0, 10000.0, 11800.0, 9800.0, 11700.0, 16128.0, 18823.0, 21519.0],\n    \"e93bed1228d6\": [ 5., 10., 15., 20., 22., 25., 26., 30., 32.],\n}\ndf_target = pd.read_csv('/kaggle/input/benetech-weights/y_train.csv')\n\nPLOT = len(df) < 5 \nDEBUG = False\n\nprint('\\n -> Done ! Starting loop...\\n')\n\nseed_everything(0)\nscores, preds_x, preds_y = [], {}, {}","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:54:56.012684Z","iopub.execute_input":"2023-06-18T23:54:56.013460Z","iopub.status.idle":"2023-06-18T23:55:26.594112Z","shell.execute_reply.started":"2023-06-18T23:54:56.013428Z","shell.execute_reply":"2023-06-18T23:55:26.592911Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"scores, preds_x, preds_y = [], {}, {}\nfor idx in range(len(df)):\n    try:\n        id_ = df.id[idx]\n        print(idx, id_, end=\"\\t\")\n        title = f\"{id_}\"\n        preds = [[], [], [], []]\n\n        # Update shapes\n        img, _, _ = dataset[idx]\n\n        if img.shape[1] > img.shape[0] * 1.4:\n            padding_bot = int(img.shape[1] * 0.9) - img.shape[0]\n        else:\n            padding_bot = 0\n        if img.shape[1] < img.shape[0] * 0.9:\n            padding_r = int(img.shape[0] * 1) - img.shape[1]\n        else:\n            padding_r = 0\n\n        meter.preds[idx].update_shape((img.shape[0] + padding_bot, img.shape[1] + padding_r))\n        meter_2.preds[idx].update_shape((img.shape[0] + padding_bot, img.shape[1] + padding_r))\n\n        # NMS\n        preds_1 = [meter.preds[idx]['pascal_voc'][meter.labels[idx] == i] for i in range(len(config_marker.labels))]\n        preds_2 = [meter_2.preds[idx]['pascal_voc'][meter_2.labels[idx] == i] for i in range(len(config_marker.labels))]\n\n        scores_1 = [meter.confidences[idx][meter.labels[idx] == i] for i in range(len(config_marker.labels))]\n        scores_2 = [meter_2.confidences[idx][meter_2.labels[idx] == i] for i in range(len(config_marker.labels))]\n\n        try:\n            boxes = np.concatenate([preds_1[0], preds_2[0]], 0)\n            confs = np.concatenate([scores_1[0], scores_2[0]], 0)\n\n            merged_boxes, merged_confs = nms(boxes, confs, threshold=0.4)\n            preds_marker = merged_boxes.copy()\n        except:\n            print('NMS error')\n            preds_marker = preds_2[0]\n\n        preds[3] = preds_marker\n\n        # Cached\n        if len(preds[-1]):\n            try:\n                cached_result = inference_detector(cached_model, dataset.paths[idx])  # list[array]\n\n                # Override with cached\n                score_th = min(0.1, cached_result[4][2, 4])\n                x_labels = cached_result[4][cached_result[4][:, -1] > score_th][:, :4].astype(int)\n\n                score_th = min(0.1, cached_result[5][2, 4])\n                y_labels = cached_result[5][cached_result[5][:, -1] > score_th][:, :4].astype(int)\n\n                score_th = min(0.1, cached_result[7][2, 4])\n                x_ticks = cached_result[7][cached_result[7][:, -1] > score_th][:, :4].astype(int)\n\n                score_th = min(0.1, cached_result[8][2, 4])\n                y_ticks = cached_result[8][cached_result[8][:, -1] > score_th][:, :4].astype(int)\n\n                preds[0] = cached_result[2][:1, :4].astype(int)\n                preds[1] = np.concatenate([x_labels, y_labels])\n                preds[2] = np.concatenate([x_ticks, y_ticks])\n            except:\n                print('CACHED Error')\n                pass\n\n        try:\n            preds = post_process_point_as_tick(\n                preds, marker_conf=merged_confs, th=0.5, max_dist=4, max_dist_o=10, verbose=DEBUG\n            )\n        except:\n            print('Error with post_process_point_as_tick')\n            pass\n\n        try:\n            preds = post_process_preds(preds)\n        except:\n            print('Error with post_process_preds')\n            pass\n\n        try:\n            preds = post_process_arrow(preds, verbose=1)\n        except:\n            print('Error with post_process_arrow')\n            pass\n\n        margin = (img.shape[0] + img.shape[1]) / (2 * 20)\n        preds = restrict_on_line(preds, margin=margin)\n\n        # Visual similarity\n        try:\n            retrieved_boxes = retrieve_missing_boxes(preds, img, verbose=False, min_sim=0.75, seed=0)\n            if len(retrieved_boxes):\n                preds[-1] = np.concatenate([preds[-1], retrieved_boxes])\n        except:\n            print('Error with retrieve_missing_boxes')\n            pass\n\n        if PLOT:\n            plot_results(img, preds, figsize=(12, 7), title=title)\n\n        # OCR\n        score_x, score_y = -1, -1\n        if len(preds[-1]):\n            try:\n                x_texts = ocr(ocr_model, processor, img, preds[1], margin=1, plot=DEBUG)\n                x_values, x_errors = post_process_texts(x_texts)\n                reg_x = linear_regression(preds[3], x_values, x_errors, preds[-1], mode=\"x\", verbose=DEBUG)\n                reg_x = np.round(reg_x, rounding(np.max(reg_x)))\n            except:  # OCR or reg fail for x, use dummy values\n                print(\"Error with x OCR\")\n                reg_x = np.arange(len(preds[-1]))\n\n            try:\n                y_texts = ocr(ocr_model, processor, img, preds[2], margin=3, plot=DEBUG)\n                y_values, y_errors = post_process_texts(y_texts)\n                reg_y = linear_regression(preds[4], y_values, y_errors, preds[-1], mode=\"y\", verbose=DEBUG)\n            except:\n                print(\"Error with y OCR\")\n                reg_y = np.arange(len(preds[-1]))\n\n            pred = pd.DataFrame({\"x\": reg_x, \"y\": reg_y})\n            pred = pred.sort_values(['x', 'y'], ignore_index=True)\n            preds_x[id_] = \";\".join(pred['x'].astype(str))\n            preds_y[id_] = \";\".join(pred['y'].astype(str))\n\n            gt = df_target[df_target['id'] == id_].reset_index(drop=True)\n            if len(gt):\n                gt[[\"x\", \"y\"]] = gt[[\"x\", \"y\"]].astype(float)\n                gt = gt.sort_values(['x', 'y'], ignore_index=True)\n\n                if id_ in FIXES:\n                    gt[\"y\"] = FIXES[id_]\n\n                score_x = score_series(gt['x'].values, pred['x'].values)\n                score_y = score_series(gt['y'].values, pred['y'].values)\n\n        if id_ not in TO_REMOVE:\n            scores += [max(0, score_x), max(0, score_y)]\n            if score_x >= 0:\n                print(f\"Scores  -  x: {score_x:.3f}  - y: {score_y:.3f}\")\n\n    except:\n        continue\n\n    if DEBUG:\n        break\n\nif len(scores):\n    if np.mean(scores):\n        print(f'\\n-> Scatter CV : {np.mean(scores) :.3f}')\n    \ndf['pred_x'] = df['id'].apply(lambda x: preds_x.get(x, \"\"))\ndf['pred_y'] = df['id'].apply(lambda x: preds_y.get(x, \"\"))\n\ndf.to_csv('df_scatter_preds.csv', index=False)\n\nprint('\\n-> Saved results to df_scatter_preds.csv')","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:55:26.596050Z","iopub.execute_input":"2023-06-18T23:55:26.596791Z","iopub.status.idle":"2023-06-18T23:55:30.987745Z","shell.execute_reply.started":"2023-06-18T23:55:26.596754Z","shell.execute_reply":"2023-06-18T23:55:30.986593Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### To submission format","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv('pred_cls.csv')[['path', 'chart-type', \"id\"]]\n# df['id'] = df['path'].apply(lambda x: x[:-4].split('/')[-1])\n\n# Merge scatter\ndf_scatter = pd.read_csv('df_scatter_preds.csv')\ndf_scatter = df_scatter[['id', 'pred_x', 'pred_y']]\n\ndf = df.merge(df_scatter, on=\"id\", how=\"left\")\n\ndf['pred_x'] = df['pred_x'].fillna(\"\")\ndf['pred_y'] = df['pred_y'].fillna(\"\")\n\n# Reformat & save\ndfx = df[[\"id\", \"pred_x\", \"chart-type\"]].copy()\ndfx['id'] += \"_x\"\ndfx.columns = [\"id\", \"data_series\", \"chart_type\"]\n\ndfy = df[[\"id\", \"pred_y\", \"chart-type\"]].copy()\ndfy['id'] += \"_y\"\ndfy.columns = [\"id\", \"data_series\", \"chart_type\"]\n\nsub = pd.concat([dfx, dfy]).sort_values('id').reset_index(drop=True)\n\nsub.to_csv('submission_scatter.csv', index=False)\nprint('-> Saved results to submission_scatter.csv')\n\ndisplay(sub[sub['chart_type'] == \"scatter\"].head())\nsub.head(20)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-18T23:55:30.989357Z","iopub.execute_input":"2023-06-18T23:55:30.989820Z","iopub.status.idle":"2023-06-18T23:55:31.031406Z","shell.execute_reply.started":"2023-06-18T23:55:30.989784Z","shell.execute_reply":"2023-06-18T23:55:31.030306Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Matcha","metadata":{}},{"cell_type":"code","source":"!pip install -qqq --no-index --find-links='/kaggle/input/transformers-deps' /kaggle/input/transformers-deps/transformers-4.30.0.dev0-py3-none-any.whl\n\nimport re\nimport torch\nimport warnings\nimport numpy as np\nimport pandas as pd\n\nfrom tqdm import tqdm\nfrom pathlib import Path\nfrom torch.utils.data import DataLoader\nfrom transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor\nimport pytorch_lightning as pl\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\nimport yaml\n\nwarnings.filterwarnings(\"ignore\", category=UserWarning)\n\ndata_dir = Path(\"/kaggle/input/benetech-making-graphs-accessible/test\")\nall_images = sorted((data_dir / \"images\").glob(\"*.jpg\"))\n\nif len(all_images) == 5:\n    image_dir = Path(\"test/\")\n    all_images = list(image_dir.glob(\"*\"))\n\n# Remove already computed\nsub_scatter = pd.read_csv('submission_scatter.csv').dropna(axis=0)\nscatters = sub_scatter['id'].apply(lambda x: x[:-2]).unique()\n\nsub_dot = pd.read_csv('submission_dot.csv').dropna(axis=0)\ndots = sub_dot['id'].apply(lambda x: x[:-2]).unique()\n\nimage_paths = [\n    img for img in all_images if\n    ((img.stem not in scatters) and (img.stem not in dots))\n]\n\nimage_ids = [img.stem for img in image_paths]\n\npredict_df = pd.DataFrame(zip(image_ids, image_paths), columns=['id', 'image_path'])","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:55:31.033202Z","iopub.execute_input":"2023-06-18T23:55:31.033588Z","iopub.status.idle":"2023-06-18T23:55:52.236118Z","shell.execute_reply.started":"2023-06-18T23:55:31.033553Z","shell.execute_reply":"2023-06-18T23:55:52.235029Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"MAX_LENGTH = 512\n\nX_START = \"<x_start>\"\nX_END = \"<x_end>\"\nY_START = \"<y_start>\"\nY_END = \"<y_end>\"\n\ndef string2preds(pred_string: str):\n\n    if pred_string.startswith(\"<dot>\"):\n        chart_type = \"dot\"\n    elif pred_string.startswith(\"<horizontal_bar>\"):\n        chart_type = \"horizontal_bar\"\n    elif pred_string.startswith(\"<vertical_bar>\"):\n        chart_type = \"vertical_bar\"\n    elif pred_string.startswith(\"<histogram>\"):\n        chart_type = \"histogram\"\n    else:\n        chart_type = \"line\"\n\n    x = pred_string.split(X_START)[1].split(X_END)[0].strip().split(\";\")\n    y = pred_string.split(Y_START)[1].split(Y_END)[0].strip().split(\";\")\n\n    if len(x) == 0:\n        x = []\n    if len(y) == 0:\n        y = []\n\n    return chart_type, x, y","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:55:52.237574Z","iopub.execute_input":"2023-06-18T23:55:52.237936Z","iopub.status.idle":"2023-06-18T23:55:52.250195Z","shell.execute_reply.started":"2023-06-18T23:55:52.237899Z","shell.execute_reply":"2023-06-18T23:55:52.246793Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MatchaModel(pl.LightningModule):\n    def __init__(self, processor, model, max_length):\n\n        super().__init__()\n        self.processor = processor\n        self.model = model\n        self.max_length = max_length\n\n    def predict_step(self, batch, batch_idx):\n        ids, images = batch\n\n        outputs = self.model.generate(**images, max_new_tokens=self.max_length)\n        preds = self.processor.batch_decode(outputs, skip_special_tokens=True)\n\n        return ids, preds\n\nclass PredictImageDataset(Dataset):\n    def __init__(self, df, processor):\n\n        self.processor = processor\n\n        self.ids = df['id'].values\n        self.image_paths = df['image_path'].values\n\n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, idx):\n\n        img = Image.open(self.image_paths[idx])\n        img = self.processor(img, return_tensors='pt')\n\n        return self.ids[idx], img\n\nclass PredictCustomCollator():\n\n    def __init__(self, pad_token_id):\n        self.pad_token_id = pad_token_id\n\n    def __call__(self, batch):\n        ids, images = [], []\n\n        for item in batch:\n            ids.append(item[0])\n            images.append(item[1])\n        \n        images = {\n            \"flattened_patches\": torch.tensor(np.vstack([img['flattened_patches'] for img in images])),\n            \"attention_mask\": torch.tensor(np.vstack([img['attention_mask'] for img in images])),\n        }\n\n        return ids, images\n    \npred_loader_params = {\n    'batch_size': 4,\n    'drop_last': False,\n    'num_workers': 2,\n    'pin_memory': False,\n    'shuffle': False,\n}\n\nclass ImageDataModule(pl.LightningDataModule):\n    def __init__(\n        self,\n        predict_df = None,\n        processor = None,\n    ):\n        super().__init__()\n        self.predict_df = predict_df\n        self.processor = processor\n\n    def setup(self, stage):\n        self.predict_dataset = PredictImageDataset(self.predict_df, self.processor)\n\n    def predict_dataloader(self):\n        custom_collator = PredictCustomCollator(self.processor.tokenizer.pad_token_id)\n        return DataLoader(self.predict_dataset, **pred_loader_params, collate_fn=custom_collator)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:55:52.252085Z","iopub.execute_input":"2023-06-18T23:55:52.252526Z","iopub.status.idle":"2023-06-18T23:55:52.269204Z","shell.execute_reply.started":"2023-06-18T23:55:52.252493Z","shell.execute_reply":"2023-06-18T23:55:52.268075Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"MODEL_DIR = '/kaggle/input/benetech-base-model'\nmodel = Pix2StructForConditionalGeneration.from_pretrained(MODEL_DIR)\nprocessor = Pix2StructProcessor.from_pretrained(MODEL_DIR)\n\nmatcha = MatchaModel(processor, model, MAX_LENGTH)\ndatamodule = ImageDataModule(predict_df=predict_df, processor=processor)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:55:52.270319Z","iopub.execute_input":"2023-06-18T23:55:52.272344Z","iopub.status.idle":"2023-06-18T23:56:05.728235Z","shell.execute_reply.started":"2023-06-18T23:55:52.272313Z","shell.execute_reply":"2023-06-18T23:56:05.727240Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dict_paths = [\n    \"/kaggle/input/benetech-base-ckpts/10ep_3e-5_aug_2_extracted_7.ckpt\",\n    \"/kaggle/input/benetech-base-ckpts/matcha_full_noisy.ckpt\",\n    \"/kaggle/input/benetech-base-ckpts/10ep_3e-5_aug_1_extracted_5.ckpt\",\n    \"/kaggle/input/benetech-base-ckpts/matcha_full_stable.ckpt\",\n    #\"/kaggle/input/benetech-base-ckpts/10ep_1e-5_aug_2_0_35_cosine_extracted_7.ckpt\",\n]\n\nd, ct, xs, ys = {}, {}, {}, {}\n\nfirst_pred = {}\n\nfor dict_path in dict_paths:\n    \n    preds, ids = [], []\n    \n    state_dict = torch.load(dict_path)['state_dict']\n    \n    for key in list(state_dict.keys()):\n        state_dict[key[6:]] = state_dict.pop(key)\n        \n    matcha.model.load_state_dict(state_dict, strict=True)\n\n    trainer = pl.Trainer(\n        logger = False,\n        accelerator = \"gpu\",\n        devices = 1,\n        precision = '16-mixed',\n    )\n\n    outputs = trainer.predict(matcha, datamodule=datamodule)\n    \n    for o in outputs:\n        ids.extend(o[0])\n        preds.extend(o[1])\n        \n        for curr_id, pred in zip(o[0],o[1]):\n            if curr_id not in first_pred:\n                first_pred[curr_id] = pred\n\n    for curr_id, pred in zip(ids, preds):\n        try:\n            chart_type, x, y = string2preds(pred)\n\n            if curr_id not in d:\n                d[curr_id] = []\n\n            \"\"\"\n            # use only categorical axis - more stable\n            if chart_type == 'horizontal_bar':\n                d[curr_id].append(len(y))\n            else:\n                d[curr_id].append(len(x))\n            \"\"\"\n            \n            d[curr_id].append(len(x))\n            if chart_type == 'histogram':\n                # ignore ys for hist - sometimes it mistakes it for vertical_bar\n                #d[curr_id].append(len(y)+1)\n                pass\n            else:\n                d[curr_id].append(len(y))\n\n            if curr_id not in ct:\n                ct[curr_id] = []\n                xs[curr_id] = []\n                ys[curr_id] = []\n\n            ct[curr_id].append(chart_type)\n            xs[curr_id].append(x)\n            ys[curr_id].append(y)\n                \n        except:\n            pass","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:56:05.729681Z","iopub.execute_input":"2023-06-18T23:56:05.730265Z","iopub.status.idle":"2023-06-18T23:57:43.437147Z","shell.execute_reply.started":"2023-06-18T23:56:05.730232Z","shell.execute_reply":"2023-06-18T23:57:43.435990Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def is_float(vals):\n    try:\n        vals = [float(x) for x in vals]\n    except:\n        return False\n    return True\n\ncnt = {}\nfor k, v in d.items():\n    cnt[k] = max(v,key=v.count)\n\nres_ids, chart_types, x_preds, y_preds = [], [], [], []\n\nfor curr_id in ct:\n    c, x, y = ct[curr_id], xs[curr_id], ys[curr_id]\n\n    if len(c) == 0:\n        continue\n\n    chart_type = max(c,key=c.count)\n    \n    hm = cnt[curr_id]\n    x = [a for a in x if len(a) == hm]\n    \n    # sometimes it mistakes histograms with vertical_bars\n    if chart_type == 'histogram' and any([a for a in y if len(a) == hm-1]):\n        hm -= 1\n        \n    y = [a for a in y if len(a) == hm]\n\n    if chart_type == 'horizontal_bar':\n        x, y = y, x\n\n    if len(x) == 0 or len(y) == 0:\n        continue\n\n    #best_x = max(x,key=x.count)\n    best_x = []\n\n    for i in range(len(x[0])):\n        curr_x = [a[i] for a in x]\n        best_x.append(max(curr_x,key=curr_x.count))\n    \"\"\"\n    new_ys = []\n\n    for xx, yy in zip(x, y):\n        if xx != best_x:\n            continue\n        new_ys.append(yy)\n    \n    y = new_ys\n    \"\"\"\n    y = [a for a in y if is_float(a)]\n\n    if len(x) == 0 or len(y) == 0:\n        continue\n\n    y = [[float(a) for a in yy] for yy in y]\n    \n    y = np.vstack(y).mean(axis=0)\n\n    x = best_x\n\n    if chart_type == 'horizontal_bar':\n        x, y = y, x\n\n    x_str = \";\".join(list(map(str, x)))\n    y_str = \";\".join(list(map(str, y)))\n    \n    res_ids.append(curr_id)\n    chart_types.append(chart_type)\n    x_preds.append(x_str)\n    y_preds.append(y_str)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:57:43.440483Z","iopub.execute_input":"2023-06-18T23:57:43.443114Z","iopub.status.idle":"2023-06-18T23:57:43.458185Z","shell.execute_reply.started":"2023-06-18T23:57:43.443072Z","shell.execute_reply":"2023-06-18T23:57:43.457101Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# fallback\nids_set = set(res_ids)\nmissing_ids = [curr_id for curr_id in predict_df['id'].values if curr_id not in ids_set]\n        \nfor curr_id in missing_ids:\n    try:\n        chart_type, xs, ys = string2preds(first_pred[curr_id])\n        x_str = \";\".join(list(map(str, xs)))\n        y_str = \";\".join(list(map(str, ys)))\n        \n        res_ids.append(curr_id)\n        chart_types.append(chart_type)\n        x_preds.append(x_str)\n        y_preds.append(y_str)       \n    except:   \n        res_ids.append(curr_id)\n        chart_types.append('line')\n        x_preds.append('0;0;0')\n        y_preds.append('0;0;0')","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:57:43.459889Z","iopub.execute_input":"2023-06-18T23:57:43.460912Z","iopub.status.idle":"2023-06-18T23:57:43.492185Z","shell.execute_reply.started":"2023-06-18T23:57:43.460874Z","shell.execute_reply":"2023-06-18T23:57:43.491299Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#replace histograms with vertical bars\ndef fix_hist(vals):\n    new_vals = []\n    for v in vals:\n        if v == 'histogram':\n            new_vals.append('vertical_bar')\n        else:\n            new_vals.append(v)\n    return new_vals\n\nchart_types = fix_hist(chart_types)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:57:43.493922Z","iopub.execute_input":"2023-06-18T23:57:43.494340Z","iopub.status.idle":"2023-06-18T23:57:43.508461Z","shell.execute_reply.started":"2023-06-18T23:57:43.494309Z","shell.execute_reply":"2023-06-18T23:57:43.507480Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sub_df = pd.DataFrame(\n    data={\n        \"id\": [f\"{id_}_x\" for id_ in res_ids] + [f\"{id_}_y\" for id_ in res_ids],\n        \"data_series\": x_preds + y_preds,\n        \"chart_type\": chart_types * 2,\n    }\n)\n\nsub_df.to_csv(\"submission_matcha.csv\", index=False)\nprint('-> Saved results to submission_matcha.csv')\nsub_df.head(20)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:57:43.510086Z","iopub.execute_input":"2023-06-18T23:57:43.510489Z","iopub.status.idle":"2023-06-18T23:57:43.538699Z","shell.execute_reply.started":"2023-06-18T23:57:43.510458Z","shell.execute_reply":"2023-06-18T23:57:43.537761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Sub","metadata":{}},{"cell_type":"code","source":"dfy = pd.read_csv('pred_cls.csv')[['chart-type', \"id\"]]\ndfx = dfy.copy()\ndfx['id'] += \"_x\"\ndfy['id'] += \"_y\"\nsub = pd.concat([dfx, dfy], ignore_index=True).rename(columns={'chart-type': \"chart_type\"})\n\nsub_scatter = pd.read_csv('submission_scatter.csv')[[\"id\", \"data_series\"]]\nsub_scatter = sub_scatter.dropna(axis=0)\n\nsub_dot = pd.read_csv('submission_dot.csv')[[\"id\", \"data_series\"]]\nsub_dot = sub_dot.dropna(axis=0)\n\nsub_matcha = pd.read_csv('submission_matcha.csv')[[\"id\", \"data_series\"]]\n\nresults = pd.concat([sub_scatter, sub_dot, sub_matcha])\nresults = results.drop_duplicates(keep=\"last\", subset=\"id\")\n\nsub = sub.merge(results, on=\"id\")\nsub = sub[[\"id\", \"data_series\", \"chart_type\"]]\nsub = sub.sort_values('id', ignore_index=True)\n\nsub.to_csv('submission.csv', index=False)\nprint('-> Saved results to submission.csv')\nsub.head(20)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T23:57:43.539956Z","iopub.execute_input":"2023-06-18T23:57:43.540888Z","iopub.status.idle":"2023-06-18T23:57:43.590908Z","shell.execute_reply.started":"2023-06-18T23:57:43.540857Z","shell.execute_reply":"2023-06-18T23:57:43.589723Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Done !","metadata":{}}]}