{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport random\nimport numpy as np\nimport pandas as pd\nfrom glob import glob\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nimport pickle\nimport json\nimport cv2\nfrom PIL import Image\n\nfrom rapidfuzz.distance.Levenshtein import distance as levenshtein\nfrom sklearn.metrics import r2_score\n\nimport io\nimport torch\nfrom transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor\nimport re","metadata":{"execution":{"iopub.status.busy":"2023-06-27T06:30:42.866497Z","iopub.execute_input":"2023-06-27T06:30:42.866895Z","iopub.status.idle":"2023-06-27T06:30:42.873772Z","shell.execute_reply.started":"2023-06-27T06:30:42.866842Z","shell.execute_reply":"2023-06-27T06:30:42.872872Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def display_deplot_output(deplot_output, visualize=True):\n    '''\n    The raw output of deplot\n    TITLE | Rural population (%) long-run with 2050 projections<0x0A>(OWID) in Greece<0x0A>Years | Rural population<0x0A>1940 | 47.38<0x0A>1960 | 43.68<0x0A>1980 | 30.28<0x0A>...\n    '''\n    \n    deplot_output = deplot_output.replace(\"<0x0A>\", \"\\n\").replace(\" | \", \"\\t\")\n\n    second_a_index = [m.start() for m in re.finditer('\\t', deplot_output)][1]\n    last_newline_index = deplot_output.rfind('\\n', 0, second_a_index) \n\n    title = deplot_output[:last_newline_index]\n    table = deplot_output[last_newline_index+1:]\n\n    data = io.StringIO(table)\n    df = pd.read_csv(data, sep='\\t')\n    if visualize:\n        display(df)\n    \n    return df\n    \n\ndef deplot(path, model, processor, device, visualize=True):\n   \n    image = Image.open(path)\n    if visualize:\n        display(image)\n    inputs = processor(images=image, text=\"Generate underlying data table of the figure below:\", return_tensors=\"pt\", font_path=\"/kaggle/input/deplot-fonts/arial.ttf\")\n\n    # Move inputs to GPU\n    inputs = {key: value.to(device) for key, value in inputs.items()}\n\n    predictions = model.generate(**inputs, max_new_tokens=512)\n    print(predictions)\n    return processor.decode(predictions[0], skip_special_tokens=True)","metadata":{"execution":{"iopub.status.busy":"2023-06-27T07:13:33.694583Z","iopub.execute_input":"2023-06-27T07:13:33.695322Z","iopub.status.idle":"2023-06-27T07:13:33.705213Z","shell.execute_reply.started":"2023-06-27T07:13:33.695283Z","shell.execute_reply":"2023-06-27T07:13:33.704308Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = Pix2StructForConditionalGeneration.from_pretrained('/kaggle/input/load-google-deplot-model/deplot').to(device)\nprocessor = Pix2StructProcessor.from_pretrained('/kaggle/input/load-google-deplot-model/deplot')\n","metadata":{"execution":{"iopub.status.busy":"2023-06-27T06:37:17.447270Z","iopub.execute_input":"2023-06-27T06:37:17.447682Z","iopub.status.idle":"2023-06-27T06:37:33.622069Z","shell.execute_reply.started":"2023-06-27T06:37:17.447652Z","shell.execute_reply":"2023-06-27T06:37:33.620890Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def deplot_inference(image_path, visualize):\n    deplot_output=deplot(image_path, model, processor, device, visualize)\n    return display_deplot_output(deplot_output, visualize)","metadata":{"execution":{"iopub.status.busy":"2023-06-27T07:09:55.362823Z","iopub.execute_input":"2023-06-27T07:09:55.363209Z","iopub.status.idle":"2023-06-27T07:09:55.368173Z","shell.execute_reply.started":"2023-06-27T07:09:55.363178Z","shell.execute_reply":"2023-06-27T07:09:55.367049Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torchvision\nimport torch.utils.data\nfrom torchvision import transforms\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.optim import lr_scheduler\n\nimport cv2\nimport numpy as np","metadata":{"execution":{"iopub.status.busy":"2023-06-27T06:38:21.842169Z","iopub.execute_input":"2023-06-27T06:38:21.842536Z","iopub.status.idle":"2023-06-27T06:38:22.073295Z","shell.execute_reply.started":"2023-06-27T06:38:21.842503Z","shell.execute_reply":"2023-06-27T06:38:22.072289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"label_map = {'dot': 0, 'horizontal_bar' : 1, 'vertical_bar': 2, 'line': 3, 'scatter': 4}\nlabel_idx_to_classname = {v: k for k, v in label_map.items()}","metadata":{"execution":{"iopub.status.busy":"2023-06-27T06:38:30.051395Z","iopub.execute_input":"2023-06-27T06:38:30.051781Z","iopub.status.idle":"2023-06-27T06:38:30.057315Z","shell.execute_reply.started":"2023-06-27T06:38:30.051732Z","shell.execute_reply":"2023-06-27T06:38:30.056160Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"classification_model = torchvision.models.resnet50(pretrained=False)\n\nnum_features = classification_model.fc.in_features\nclassification_model.fc = nn.Linear(num_features, 5)\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nclassification_model = classification_model.to(device)\nclassification_model.eval()\n\nstate_dict = torch.load(\"/kaggle/input/model/Benetech _ResNet50_fold0.pth\")\nclassification_model.load_state_dict(state_dict)","metadata":{"execution":{"iopub.status.busy":"2023-06-27T06:46:35.419130Z","iopub.execute_input":"2023-06-27T06:46:35.419516Z","iopub.status.idle":"2023-06-27T06:46:36.913766Z","shell.execute_reply.started":"2023-06-27T06:46:35.419487Z","shell.execute_reply":"2023-06-27T06:46:36.912799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"val_transforms = transforms.Compose(\n    [transforms.ToTensor(),\n     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\ndef classification_inference(image_path):\n    img = cv2.imread(image_path)\n    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)\n    img = cv2.resize(img,(500,300))\n    img = img.astype(np.float32)/255.0\n    img = val_transforms(img)\n    \n    #print(img.shape)\n    \n#     inp = torch.stack([img.permute(1, 2, 0)]).cuda()\n    inp = torch.stack([img]).cuda()\n    out = classification_model(inp).detach().cpu().numpy()\n    #print(img.shape)\n    return label_idx_to_classname[np.argmax(out)]","metadata":{"execution":{"iopub.status.busy":"2023-06-27T07:02:59.954387Z","iopub.execute_input":"2023-06-27T07:02:59.954845Z","iopub.status.idle":"2023-06-27T07:02:59.962732Z","shell.execute_reply.started":"2023-06-27T07:02:59.954814Z","shell.execute_reply":"2023-06-27T07:02:59.961812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfrom tqdm import tqdm\nimport math\n\nIMAGE_FOLDER = \"/kaggle/input/benetech-making-graphs-accessible/test/images\"\n\nall_ids = []\nall_values = []\nall_chart_types = []\n\nfor image_name in tqdm(os.listdir(IMAGE_FOLDER)):\n    \n    if \".jpg\" not in image_name:\n        continue\n    try:\n        image_path = os.path.join(IMAGE_FOLDER, image_name)\n        print(image_path)\n        graph_type = classification_inference(image_path)\n        print(graph_type)\n        inference_df = deplot_inference(image_path, False)\n        \n        x_values = inference_df[inference_df.columns[0]].values\n        y_values = inference_df[inference_df.columns[1]].values\n        \n        # check value type for each graph type\n        if graph_type in [\"horizontal_bar\", \"vertical_bar\", \"line\", \"dot\", \"scatter\"]:\n            # x is categorical, y is numerical\n            formatted_y_values = []\n            for v in y_values:\n                try:\n                    float(v)\n                    if math.isnan(float(v)):\n                        formatted_y_values.append(0)\n                    else:\n                        formatted_y_values.append(v)\n                except:\n                    formatted_y_values.append(0)\n                    \n            y_values = formatted_y_values\n        \n        if graph_type in [\"dot\", \"scatter\"]:\n            # for dot graph, x could be categorical as well as numerical, how to handle this?\n            # just leave all the value be numerical for now\n\n            formatted_x_values = []\n            for v in x_values:\n                try:\n                    float(v)\n                    if math.isnan(float(v)):\n                        formatted_x_values.append(0)\n                    else:\n                        formatted_x_values.append(v)\n                except:\n                    formatted_x_values.append(0)\n                    \n            x_values = formatted_x_values\n            \n        \n        length = min(len(x_values), len(y_values))\n        x_values = \";\".join([str(v).strip() for v in x_values][:length])\n        y_values = \";\".join([str(v).strip() for v in y_values][:length])\n    except Exception as e:\n        print(\"Exception\", e)\n        graph_type = \"line\"\n        x_values = \"0;0\"\n        y_values = \"0;0\"\n    \n    image_id = image_name.split(\".\")[0]\n\n    all_ids.append(image_id + \"_x\")\n    all_values.append(x_values)\n    \n    all_ids.append(image_id + \"_y\")\n    all_values.append(y_values)\n    \n    all_chart_types.extend([graph_type, graph_type])","metadata":{"execution":{"iopub.status.busy":"2023-06-27T07:20:22.627362Z","iopub.execute_input":"2023-06-27T07:20:22.627800Z","iopub.status.idle":"2023-06-27T07:20:41.113566Z","shell.execute_reply.started":"2023-06-27T07:20:22.627761Z","shell.execute_reply":"2023-06-27T07:20:41.112583Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df = pd.DataFrame({\n    \"id\": all_ids,\n    \"data_series\": all_values,\n    \"chart_type\": all_chart_types\n})\n\nsubmission_df.to_csv(\"submission.csv\", index=False)\nsubmission_df","metadata":{"execution":{"iopub.status.busy":"2023-06-27T07:22:19.996825Z","iopub.execute_input":"2023-06-27T07:22:19.997383Z","iopub.status.idle":"2023-06-27T07:22:20.056286Z","shell.execute_reply.started":"2023-06-27T07:22:19.997335Z","shell.execute_reply":"2023-06-27T07:22:20.048827Z"},"trusted":true},"execution_count":null,"outputs":[]}]}