{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip -q install transformers==4.25.1","metadata":{"execution":{"iopub.status.busy":"2023-06-09T19:33:46.108012Z","iopub.execute_input":"2023-06-09T19:33:46.109206Z","iopub.status.idle":"2023-06-09T19:34:12.669011Z","shell.execute_reply.started":"2023-06-09T19:33:46.109165Z","shell.execute_reply":"2023-06-09T19:34:12.667374Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip -q install pytorch-lightning==1.7.1","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!git clone https://github.com/clovaai/donut.git\n!cd donut && pip -q install .","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import re\nfrom donut import DonutModel\nimport torch\nimport os\nimport json\nfrom pathlib import Path\nimport shutil\nfrom PIL import Image","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# mounting drive\n\nfrom google.colab import drive\ndrive.mount('/content/drive')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cp drive/MyDrive/benetech-making-graphs-accessible.zip .\n!unzip -q benetech-making-graphs-accessible.zip","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = DonutModel.from_pretrained(\"naver-clova-ix/donut-base-finetuned-cord-v2\")\nif torch.cuda.is_available():\n    model.half()\n    device = torch.device(\"cuda\")\n    model.to(device)\nelse:\n    model.encoder.to(torch.bfloat16)\nmodel.eval()\nimage = Image.open(\"train/images/0000ae6cbdb1.jpg\").convert(\"RGB\")\noutput = model.inference(image=image, prompt=\"\")\noutput","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!rm -r dataset\n!mkdir dataset","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import random, os, shutil\nfrom tqdm import tqdm\n\nimg = 'images'\nant = 'annotations'\n\n# Set the path to the directory containing the data\nimg_dir = f\"train/{img}\"\nant_dir = f\"train/{ant}\"\n\n# Set the paths to the train, validation, and test directories\ntrain_img_dir = f\"dataset/train/{img}\"\nval_img_dir = f\"dataset/validation/{img}\"\ntest_img_dir = f\"dataset/test/{img}\"\ntrain_ant_dir = f\"dataset/train/{ant}\"\nval_ant_dir = f\"dataset/validation/{ant}\"\ntest_ant_dir = f\"dataset/test/{ant}\"\n\n\n\n# Create the train, validation, and test directories if they don't already exist\nos.makedirs(train_img_dir, exist_ok=True)\nos.makedirs(val_img_dir, exist_ok=True)\nos.makedirs(test_img_dir, exist_ok=True)\nos.makedirs(train_ant_dir, exist_ok=True)\nos.makedirs(val_ant_dir, exist_ok=True)\nos.makedirs(test_ant_dir, exist_ok=True)\n\n# Set the ratio of data to use for train, validation, and test sets\ntrain_ratio = 0.8\nval_ratio = 0.1\ntest_ratio = 0.1\n\n# Get a list of all the files in the data directory\nall_files = os.listdir(ant_dir)\n\n# Shuffle the list of files randomly\nrandom.shuffle(all_files)\n\n# Calculate the number of files to use for train, validation, and test sets\nnum_train = int(len(all_files) * train_ratio)\nnum_val = int(len(all_files) * val_ratio)\nnum_test = len(all_files) - num_train - num_val\n\n# Iterate over the shuffled list of files and move them to the appropriate directory\nfor i, filename in tqdm(enumerate(all_files)):\n    if filename.endswith('.json'):\n        if i < num_train:\n            shutil.copy(os.path.join(img_dir, filename[:-5]+'.jpg'), os.path.join(train_img_dir, filename[:-5]+'.jpg'))\n            shutil.copy(os.path.join(ant_dir, filename), os.path.join(train_ant_dir, filename))\n        elif i < num_train + num_val:\n            shutil.copy(os.path.join(img_dir, filename[:-5]+'.jpg'), os.path.join(val_img_dir, filename[:-5]+'.jpg'))\n            shutil.copy(os.path.join(ant_dir, filename), os.path.join(val_ant_dir, filename))\n        else:\n            shutil.copy(os.path.join(img_dir, filename[:-5]+'.jpg'), os.path.join(test_img_dir, filename[:-5]+'.jpg'))\n            shutil.copy(os.path.join(ant_dir, filename), os.path.join(test_ant_dir, filename))\n    else:\n        print('different file found')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for dir in ['train','test','validation']:\n    # define paths\n    base_path = Path(f\"dataset/{dir}\")\n    metadata_path = base_path.joinpath(\"annotations\")\n    image_path = base_path.joinpath(\"images\")\n\n    # define metadata list\n    metadata_list = []\n\n    # parse metadata\n    for file_name in metadata_path.glob(\"*.json\"):\n        with open(file_name, \"r\") as json_file:\n\n            # load json file\n            data = json.load(json_file)\n            # create \"text\" column with json string\n            proc = {}\n            for key,value in data.items():\n                if key in ['chart-type']:\n                    proc[key] = value\n                elif key in ['data-series']: \n                    i = 0 \n                    for coord in data['data-series']:             \n                        for key1,value1 in coord.items():\n                            if isinstance(value1, float):\n                                # Round the float value to 4 decimal places\n                                value1 = round(value1, 4)\n                            proc[f\"{key1}_{i}\"] = str(value1)\n                        i+=1\n                            \n\n            text = json.dumps(proc)\n            # add to metadata list if image exists\n            if image_path.joinpath(f\"{file_name.stem}.jpg\").is_file():\n              metadata_list.append({\"ground_truth\":\"{\\\"gt_parse\\\": \"+f\"{text}\"+\"}\",\"file_name\":f\"{file_name.stem}.jpg\"})\n              # delete json file\n\n\n    # write jsonline file\n    with open(image_path.joinpath('metadata.jsonl'), 'w') as outfile:\n        for entry in metadata_list:\n            json.dump(entry, outfile)\n            outfile.write('\\n')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for dir in ['train','test','validation']:\n    # Set the path to the directory to move\n    dir_to_move = f\"dataset/{dir}/images\"\n\n    # Set the path to the parent directory\n    parent_dir = f\"dataset/{dir}\"\n\n\n    # Move the contents of the directory to the parent directory\n    for filename in os.listdir(dir_to_move):\n        src = os.path.join(dir_to_move, filename)\n        dst = os.path.join(parent_dir, filename)\n        shutil.move(src, dst)\n\n    # Remove the original directory\n    shutil.rmtree(dir_to_move)\n    shutil.rmtree(f\"dataset/{dir}/annotations\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"('metadata.jsonl' in os.listdir('dataset/validation')) and ('metadata.jsonl' in os.listdir('dataset/train')) and ('metadata.jsonl' in os.listdir('dataset/test'))","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cd donut && python train.py --config config/train_cord.yaml","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from donut import DonutModel\nfrom PIL import Image\nimport torch\nmodel = DonutModel.from_pretrained(\"./donut/result/train_cord/20230331_204628\")\nif torch.cuda.is_available():\n    model.half()\n    device = torch.device(\"cuda\")\n    model.to(device)\nelse:\n    model.encoder.to(torch.bfloat16)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"Image.open(\"train/images/0000ae6cbdb1.jpg\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image = Image.open(\"train/images/0000ae6cbdb1.jpg\").convert(\"RGB\")\noutput = model.inference(image=image, prompt=\"\")\noutput","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import csv\n\ndata = {'predictions': [\n    {'x_1': 'Malaysia', 'y_1': '96314', 'x_2': 'Maldives', 'y_2': '77777.7777', 'x_3': 'Mali', 'y_3': '66493',\n     'x_4': 'Malta', 'y_4': '547295.4785', 'x_5': 'Martinique', 'y_5': '495998.6', 'x_6': 'Mauritania', 'y_6': '66493',\n     'x_7': 'Mauritius', 'y_7': '595998.6', 'x_8': 'Mayotte', 'y_8': '395998.6', 'x_9': 'Melanesia', 'y_9': '595998.6'}\n]}\n\n# Extract the predictions list from the dictionary\npredictions = data['predictions']\n\n# Define the field names for the CSV file\nfield_names = ['x_1', 'y_1', 'x_2', 'y_2', 'x_3', 'y_3', 'x_4', 'y_4', 'x_5', 'y_5', 'x_6', 'y_6', 'x_7', 'y_7', 'x_8', 'y_8', 'x_9', 'y_9']\n\n# Specify the output CSV file path\noutput_file = 'sumission.csv'\n\n# Write the dictionary to a CSV file\nwith open(output_file, 'w', newline='') as csvfile:\n    writer = csv.DictWriter(csvfile, fieldnames=field_names)\n    writer.writeheader()\n    writer.writerows(predictions)\n\nprint('CSV file created successfully.')","metadata":{"execution":{"iopub.status.busy":"2023-06-09T20:04:28.250562Z","iopub.execute_input":"2023-06-09T20:04:28.251012Z","iopub.status.idle":"2023-06-09T20:04:28.262227Z","shell.execute_reply.started":"2023-06-09T20:04:28.25097Z","shell.execute_reply":"2023-06-09T20:04:28.261038Z"},"trusted":true},"execution_count":null,"outputs":[]}]}