{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.10","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":43873,"databundleVersionId":5585780,"sourceType":"competition"},{"sourceId":5214849,"sourceType":"datasetVersion","datasetId":3033230},{"sourceId":7620814,"sourceType":"datasetVersion","datasetId":4438991},{"sourceId":7622443,"sourceType":"datasetVersion","datasetId":4440082},{"sourceId":7622485,"sourceType":"datasetVersion","datasetId":4440118},{"sourceId":7622541,"sourceType":"datasetVersion","datasetId":4440164},{"sourceId":7626918,"sourceType":"datasetVersion","datasetId":4443425}],"dockerImageVersionId":30476,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Indroduction\n\nThe competition is focused on making graphs accessible to people with visual impairments or other disabilities that make it difficult to read and interpret the values displayed in graphs. In this competition, we are given a set of graphs, along with some data points such as x and y coordinates, x and y-axis values, and bounding boxes. However, the values in the graph images are not always clear, making it challenging to extract accurate data from them.\n\n\nTo address this challenge, I have fine-tuned an OCR model using EasyOCR. While pre-trained OCR models may not work accurately in all cases, fine-tuning an OCR model on the dataset provided by the competition can help us achieve better results.\n\n\nIn this notebook, I will be sharing my approach and implementation of the OCR model, along with the data pre-processing techniques steps used to extract accurate data from the graph images.\n\n\nThe goal of this notebook is to provide a detailed guide for anyone interested in using OCR to make graphs accessible to a wider audience. By sharing my approach, I hope to contribute to the development of more effective methods for making visual data accessible to all.","metadata":{}},{"cell_type":"markdown","source":"# ","metadata":{}},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport cv2","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:05:53.534446Z","iopub.execute_input":"2024-02-14T22:05:53.534709Z","iopub.status.idle":"2024-02-14T22:05:53.731407Z","shell.execute_reply.started":"2024-02-14T22:05:53.534685Z","shell.execute_reply":"2024-02-14T22:05:53.730400Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!git clone https://github.com/JaidedAI/EasyOCR.git","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:05:53.733019Z","iopub.execute_input":"2024-02-14T22:05:53.733340Z","iopub.status.idle":"2024-02-14T22:06:04.454126Z","shell.execute_reply.started":"2024-02-14T22:05:53.733312Z","shell.execute_reply":"2024-02-14T22:06:04.453077Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!cp -r /kaggle/working/EasyOCR/trainer /kaggle/working/","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:04.455766Z","iopub.execute_input":"2024-02-14T22:06:04.456141Z","iopub.status.idle":"2024-02-14T22:06:05.427615Z","shell.execute_reply.started":"2024-02-14T22:06:04.456102Z","shell.execute_reply":"2024-02-14T22:06:05.426206Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"files = os.listdir(\"/kaggle/working/trainer\")\nfor file in files:\n    !cp -r /kaggle/working/trainer/{file} /kaggle/working/","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:05.430616Z","iopub.execute_input":"2024-02-14T22:06:05.430963Z","iopub.status.idle":"2024-02-14T22:06:17.095024Z","shell.execute_reply.started":"2024-02-14T22:06:05.430931Z","shell.execute_reply":"2024-02-14T22:06:17.093783Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!rm -r /kaggle/working/EasyOCR\n!rm -r /kaggle/working/trainer","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:17.099400Z","iopub.execute_input":"2024-02-14T22:06:17.099676Z","iopub.status.idle":"2024-02-14T22:06:19.116831Z","shell.execute_reply.started":"2024-02-14T22:06:17.099647Z","shell.execute_reply":"2024-02-14T22:06:19.115522Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.mkdir(\"/kaggle/working/all_data/en_train_filtered\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.118320Z","iopub.execute_input":"2024-02-14T22:06:19.118642Z","iopub.status.idle":"2024-02-14T22:06:19.123994Z","shell.execute_reply.started":"2024-02-14T22:06:19.118611Z","shell.execute_reply":"2024-02-14T22:06:19.123016Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# %%time\n\n# !cp -r /kaggle/input/ocr-dataset/train/train /kaggle/working/all_data\n# !cp -r /kaggle/input/ocr-dataset/valid/valid /kaggle/working/all_data","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.125261Z","iopub.execute_input":"2024-02-14T22:06:19.125587Z","iopub.status.idle":"2024-02-14T22:06:19.133042Z","shell.execute_reply.started":"2024-02-14T22:06:19.125552Z","shell.execute_reply":"2024-02-14T22:06:19.132032Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# a = os.listdir(\"/kaggle/input/ocr-dataset/valid/valid\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.134526Z","iopub.execute_input":"2024-02-14T22:06:19.134819Z","iopub.status.idle":"2024-02-14T22:06:19.143231Z","shell.execute_reply.started":"2024-02-14T22:06:19.134784Z","shell.execute_reply":"2024-02-14T22:06:19.142300Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# os.listdir(\"all_data/train\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.146717Z","iopub.execute_input":"2024-02-14T22:06:19.147033Z","iopub.status.idle":"2024-02-14T22:06:19.153433Z","shell.execute_reply.started":"2024-02-14T22:06:19.147000Z","shell.execute_reply":"2024-02-14T22:06:19.152417Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# os.rename(\"/kaggle/working/all_data/train\",\"/kaggle/working/all_data/en_train_filtered\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.154603Z","iopub.execute_input":"2024-02-14T22:06:19.154960Z","iopub.status.idle":"2024-02-14T22:06:19.163371Z","shell.execute_reply.started":"2024-02-14T22:06:19.154929Z","shell.execute_reply":"2024-02-14T22:06:19.162410Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.mkdir(\"/kaggle/working/all_data/en_train_filtered/__results___files\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.164590Z","iopub.execute_input":"2024-02-14T22:06:19.165189Z","iopub.status.idle":"2024-02-14T22:06:19.172391Z","shell.execute_reply.started":"2024-02-14T22:06:19.165156Z","shell.execute_reply":"2024-02-14T22:06:19.171454Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# os.listdir(\"/kaggle/input/ocr-dataset/valid/valid\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.173583Z","iopub.execute_input":"2024-02-14T22:06:19.173946Z","iopub.status.idle":"2024-02-14T22:06:19.181088Z","shell.execute_reply.started":"2024-02-14T22:06:19.173913Z","shell.execute_reply":"2024-02-14T22:06:19.180289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# \"labels.csv\" in os.listdir(\"all_data/en_train_filtered\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.182187Z","iopub.execute_input":"2024-02-14T22:06:19.182909Z","iopub.status.idle":"2024-02-14T22:06:19.190378Z","shell.execute_reply.started":"2024-02-14T22:06:19.182877Z","shell.execute_reply":"2024-02-14T22:06:19.189513Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !mv /kaggle/working/all_data/en_train_filtered/labels.csv /kaggle/working/all_data/en_train_filtered/__results___files/labels.csv","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.191319Z","iopub.execute_input":"2024-02-14T22:06:19.191610Z","iopub.status.idle":"2024-02-14T22:06:19.200561Z","shell.execute_reply.started":"2024-02-14T22:06:19.191586Z","shell.execute_reply":"2024-02-14T22:06:19.199647Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# os.rename(\"/kaggle/working/all_data/valid\",\"/kaggle/working/all_data/en_val\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.201807Z","iopub.execute_input":"2024-02-14T22:06:19.202152Z","iopub.status.idle":"2024-02-14T22:06:19.209604Z","shell.execute_reply.started":"2024-02-14T22:06:19.202125Z","shell.execute_reply":"2024-02-14T22:06:19.208595Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.mkdir(\"/kaggle/working/all_data/en_val\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.210828Z","iopub.execute_input":"2024-02-14T22:06:19.211118Z","iopub.status.idle":"2024-02-14T22:06:19.218657Z","shell.execute_reply.started":"2024-02-14T22:06:19.211094Z","shell.execute_reply":"2024-02-14T22:06:19.217777Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.mkdir(\"/kaggle/working/all_data/en_val/__results___files\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.219789Z","iopub.execute_input":"2024-02-14T22:06:19.220071Z","iopub.status.idle":"2024-02-14T22:06:19.227810Z","shell.execute_reply.started":"2024-02-14T22:06:19.220048Z","shell.execute_reply":"2024-02-14T22:06:19.226809Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !mv /kaggle/working/all_data/en_val/labels.csv /kaggle/working/all_data/en_val/__results___files/labels.csv","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:19.229328Z","iopub.execute_input":"2024-02-14T22:06:19.229613Z","iopub.status.idle":"2024-02-14T22:06:19.235634Z","shell.execute_reply.started":"2024-02-14T22:06:19.229579Z","shell.execute_reply":"2024-02-14T22:06:19.234609Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n\nimport os\nimport pandas as pd\nimport cv2\nimport json\nimport matplotlib.pyplot as plt\nBASE_DIR = '/kaggle/input/benetech-making-graphs-accessible'\nfiles = os.listdir(f\"{BASE_DIR}/train/images\")\nf = open(f\"{BASE_DIR}/train/annotations/0000ae6cbdb1.json\")\nannotated_data = json.load(f)\n# annotated_data['text'][10]annotated_data['text'][10]\ndef hwlt2ltrb(coor):\n    left = coor[2]\n    top = coor[3]\n    right = left + coor[1]\n    bottom = top + coor[0]\n    \n    return (left,top,right,bottom)\n\n\ndef get_ltrb(source,idx):\n    l,t,r,b = float(\"inf\"),float(\"inf\"),float(\"-inf\"),float(\"-inf\")\n    data = source[idx][\"polygon\"]\n    x_values = (data['x0'],data['x1'],data['x2'],data['x3'])\n    y_values = (data['y0'],data['y1'],data['y2'],data['y3'])\n    l = min(l,min(x_values))\n    t = min(t,min(y_values))\n    r = max(r,max(x_values))\n    b = max(b,max(y_values))\n    text = source[idx]['text']\n    return (l,t,r,b),text\n# !mkdir ./output/train\n# !mkdir ./output/valid\n# mode = 'train'\n\nclasses = {\n    'line': 0,\n    'scatter':1,\n    'dot':2,\n    'vertical_bar':3,\n    'horizontal_bar':4,\n    'X-axis':5,\n    'y-axis':6\n}\n\ntrain_image_names = []\ntrain_texts = []\n\nvalid_image_names = []\nvalid_texts = []\n\nmode = \"train\"\n\nfor i,file in enumerate(files):\n    img = cv2.imread(f\"{BASE_DIR}/train/images/{file}\")\n    image = img.copy()\n    h,w,_ = img.shape\n    f = open(f\"{BASE_DIR}/train/annotations/{file.replace('jpg','json')}\")\n    annotated_data = json.load(f)\n    \n    indicies = list(map(lambda x: x[\"id\"], annotated_data[\"axes\"][\"x-axis\"][\"ticks\"]))\n    \n    indicies.extend(list(map(lambda x: x[\"id\"], annotated_data[\"axes\"][\"y-axis\"][\"ticks\"])))\n    \n#     plot_l,plot_t,plot_r,plot_b = hwlt2ltrb((plot_bb[\"height\"],plot_bb[\"width\"],plot_bb[\"x0\"],plot_bb[\"y0\"]))\n#     try:\n    for idx in indicies:\n        try:\n            coor,text = get_ltrb(annotated_data['text'],idx)\n            l,t,r,b = coor\n            canvas = cv2.rectangle(image,(l,t),(r,b),(0,255,0),1)\n            text_crop = img[t:b,l:r]\n            \n            text_img_name = f\"{file.replace('.jpg','')}_{idx}.jpg\"\n            if mode == \"train\":\n                cv2.imwrite(f\"/kaggle/working/all_data/en_train_filtered/__results___files/{text_img_name}\",text_crop)\n                train_image_names.append(text_img_name)\n                train_texts.append(text)\n            else:\n                cv2.imwrite(f\"/kaggle/working/all_data/en_val/__results___files/{text_img_name}\",text_crop) \n                valid_image_names.append(text_img_name)\n                valid_texts.append(text)\n\n        except:\n            break\n    \n    if i == 50000:\n        mode = \"valid\"\n    \n    if i % 1000 == 0:\n        print(f\"Current Iter: {i}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = pd.DataFrame(list(zip(train_image_names, train_texts)),\n               columns =['filename', 'words'])\n\nvalid_df = pd.DataFrame(list(zip(valid_image_names, valid_texts)),\n               columns =['filename', 'words'])\ntrain_df.to_csv('/kaggle/working/all_data/en_train_filtered/__results___files/labels.csv',index=False)\nvalid_df.to_csv('/kaggle/working/all_data/en_val/__results___files/labels.csv',index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = pd.read_csv(\"/kaggle/working/all_data/en_train_filtered/__results___files/labels.csv\", sep='^([^,]+),', engine='python', usecols=['filename', 'words'], keep_default_na=False)\ntrain_df = train_df.dropna().reset_index(drop=True)\ntrain_df.to_csv(\"/kaggle/working/all_data/en_train_filtered/__results___files/labels.csv\",index=False)\n\nvalid_df = pd.read_csv(\"/kaggle/working/all_data/en_val/__results___files/labels.csv\", sep='^([^,]+),', engine='python', usecols=['filename', 'words'], keep_default_na=False)\nvalid_df = valid_df.dropna().reset_index(drop=True)\nvalid_df.to_csv(\"/kaggle/working/all_data/en_val/__results___files/labels.csv\",index=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.listdir(\"all_data\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sys.path.append(\"/kaggle/working/\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install natsort\n\n\n\nimport os\nimport torch.backends.cudnn as cudnn\nimport yaml\n# from train import train\nfrom utils import AttrDict\nimport pandas as pd","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:33.202167Z","iopub.execute_input":"2024-02-14T22:06:33.203116Z","iopub.status.idle":"2024-02-14T22:06:49.692468Z","shell.execute_reply.started":"2024-02-14T22:06:33.203079Z","shell.execute_reply":"2024-02-14T22:06:49.691513Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cudnn.benchmark = True\ncudnn.deterministic = False","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:49.694122Z","iopub.execute_input":"2024-02-14T22:06:49.694681Z","iopub.status.idle":"2024-02-14T22:06:49.699184Z","shell.execute_reply.started":"2024-02-14T22:06:49.694638Z","shell.execute_reply":"2024-02-14T22:06:49.698110Z"},"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":{"execution":{"iopub.status.busy":"2024-02-14T22:06:49.700592Z","iopub.execute_input":"2024-02-14T22:06:49.700952Z","iopub.status.idle":"2024-02-14T22:06:49.711079Z","shell.execute_reply.started":"2024-02-14T22:06:49.700918Z","shell.execute_reply":"2024-02-14T22:06:49.710154Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%writefile config_files/en_filtered_config.yaml\nnumber: '0123456789'\nsymbol: \"!\\\"#$%&'()*+,-./:;<=>?@[\\\\]^_`{|}~ €\"\nlang_char: 'ؤ إئاجحخدذرزسشصضطظعغفقكلمنهويآأتىبثة'\nexperiment_name: 'en_filtered'\ntrain_data: '/kaggle/input/preprocessed-with-valid/processed images'\nvalid_data: '/kaggle/input/preprocessed-with-valid/processed images/valid'\nmanualSeed: 1111\nworkers: 2\nbatch_size: 16 # 32\nnum_iter: 25000\nvalInterval: 1000\nsaved_model: '' #'saved_models/en_filtered/iter_300000.pth'\nFT: False\noptim: False # default is Adadelta\nlr: 1\nbeta1: 0.9\nrho: 0.95\neps: 0.00000001\ngrad_clip: 5\n#Data processing\nselect_data: 'processed images' # this is dataset folder in train_data\nbatch_ratio: '1' \ntotal_data_usage_ratio: 1.0\nbatch_max_length: 34 \nimgH: 256\nimgW: 256\nrgb: False\ncontrast_adjust: False\nsensitive: False\nPAD: False\ncontrast_adjust: 0.0\ndata_filtering_off: False\n# Model Architecture\nTransformation: 'None'\nFeatureExtraction: 'VGG'\nSequenceModeling: 'BiLSTM'\nPrediction: 'CTC'\nnum_fiducial: 20\ninput_channel: 1\noutput_channel: 256\nhidden_size: 256\ndecode: 'greedy'\nnew_prediction: True\nfreeze_FeatureFxtraction: False\nfreeze_SequenceModeling: False","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:49.713110Z","iopub.execute_input":"2024-02-14T22:06:49.713434Z","iopub.status.idle":"2024-02-14T22:06:49.726384Z","shell.execute_reply.started":"2024-02-14T22:06:49.713390Z","shell.execute_reply":"2024-02-14T22:06:49.725499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"opt = get_config(\"config_files/en_filtered_config.yaml\")","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:49.727478Z","iopub.execute_input":"2024-02-14T22:06:49.727746Z","iopub.status.idle":"2024-02-14T22:06:49.745062Z","shell.execute_reply.started":"2024-02-14T22:06:49.727724Z","shell.execute_reply":"2024-02-14T22:06:49.744387Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# a = os.listdir('/kaggle/working/all_data/en_train_filtered')","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:49.746124Z","iopub.execute_input":"2024-02-14T22:06:49.746414Z","iopub.status.idle":"2024-02-14T22:06:49.753321Z","shell.execute_reply.started":"2024-02-14T22:06:49.746390Z","shell.execute_reply":"2024-02-14T22:06:49.752279Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# %%writefile /kaggle/working/dataset.py\n\n\nimport os\nimport sys\nimport re\nimport six\nimport math\nimport torch\nimport pandas  as pd\n\nfrom natsort import natsorted\nfrom PIL import Image\nimport numpy as np\nfrom torch.utils.data import Dataset, ConcatDataset, Subset\nfrom torch._utils import _accumulate\nimport torchvision.transforms as transforms\n\ndef contrast_grey(img):\n    high = np.percentile(img, 90)\n    low  = np.percentile(img, 10)\n    return (high-low)/(high+low), high, low\n\ndef adjust_contrast_grey(img, target = 0.4):\n    contrast, high, low = contrast_grey(img)\n    if contrast < target:\n        img = img.astype(int)\n        ratio = 200./(high-low)\n        img = (img - low + 25)*ratio\n        img = np.maximum(np.full(img.shape, 0) ,np.minimum(np.full(img.shape, 255), img)).astype(np.uint8)\n    return img\n\n\nclass Batch_Balanced_Dataset(object):\n\n    def __init__(self, opt):\n        \"\"\"\n        Modulate the data ratio in the batch.\n        For example, when select_data is \"MJ-ST\" and batch_ratio is \"0.5-0.5\",\n        the 50% of the batch is filled with MJ and the other 50% of the batch is filled with ST.\n        \"\"\"\n        log = open(f'./saved_models/{opt.experiment_name}/log_dataset.txt', 'a')\n        dashed_line = '-' * 80\n        print(dashed_line)\n        log.write(dashed_line + '\\n')\n        print(f'dataset_root: {opt.train_data}\\nopt.select_data: {opt.select_data}\\nopt.batch_ratio: {opt.batch_ratio}')\n        log.write(f'dataset_root: {opt.train_data}\\nopt.select_data: {opt.select_data}\\nopt.batch_ratio: {opt.batch_ratio}\\n')\n        assert len(opt.select_data) == len(opt.batch_ratio)\n\n        _AlignCollate = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD, contrast_adjust = opt.contrast_adjust)\n        self.data_loader_list = []\n        self.dataloader_iter_list = []\n        batch_size_list = []\n        Total_batch_size = 0\n        for selected_d, batch_ratio_d in zip(opt.select_data, opt.batch_ratio):\n            _batch_size = max(round(opt.batch_size * float(batch_ratio_d)), 1)\n            print(dashed_line)\n            log.write(dashed_line + '\\n')\n            _dataset, _dataset_log = hierarchical_dataset(root=opt.train_data, opt=opt, select_data=[selected_d])\n            total_number_dataset = len(_dataset)\n            log.write(_dataset_log)\n\n            \"\"\"\n            The total number of data can be modified with opt.total_data_usage_ratio.\n            ex) opt.total_data_usage_ratio = 1 indicates 100% usage, and 0.2 indicates 20% usage.\n            See 4.2 section in our paper.\n            \"\"\"\n            number_dataset = int(total_number_dataset * float(opt.total_data_usage_ratio))\n            dataset_split = [number_dataset, total_number_dataset - number_dataset]\n            indices = range(total_number_dataset)\n            _dataset, _ = [Subset(_dataset, indices[offset - length:offset])\n                           for offset, length in zip(_accumulate(dataset_split), dataset_split)]\n            selected_d_log = f'num total samples of {selected_d}: {total_number_dataset} x {opt.total_data_usage_ratio} (total_data_usage_ratio) = {len(_dataset)}\\n'\n            selected_d_log += f'num samples of {selected_d} per batch: {opt.batch_size} x {float(batch_ratio_d)} (batch_ratio) = {_batch_size}'\n            print(selected_d_log)\n            log.write(selected_d_log + '\\n')\n            batch_size_list.append(str(_batch_size))\n            Total_batch_size += _batch_size\n\n            _data_loader = torch.utils.data.DataLoader(\n                _dataset, batch_size=_batch_size,\n                shuffle=True,\n                num_workers=int(opt.workers), #prefetch_factor=2,persistent_workers=True,\n                collate_fn=_AlignCollate, pin_memory=True)\n            self.data_loader_list.append(_data_loader)\n            self.dataloader_iter_list.append(iter(_data_loader))\n\n        Total_batch_size_log = f'{dashed_line}\\n'\n        batch_size_sum = '+'.join(batch_size_list)\n        Total_batch_size_log += f'Total_batch_size: {batch_size_sum} = {Total_batch_size}\\n'\n        Total_batch_size_log += f'{dashed_line}'\n        opt.batch_size = Total_batch_size\n\n        print(Total_batch_size_log)\n        log.write(Total_batch_size_log + '\\n')\n        log.close()\n\n    def get_batch(self):\n        balanced_batch_images = []\n        balanced_batch_texts = []\n\n        for i, data_loader_iter in enumerate(self.dataloader_iter_list):\n            try:\n                image,text = next(iter(data_loader_iter))\n                balanced_batch_images.append(image)\n                balanced_batch_texts += text\n            except StopIteration:\n                self.dataloader_iter_list[i] = iter(self.data_loader_list[i])\n                image, text = next(iter(self.dataloader_iter_list[i]))\n                balanced_batch_images.append(image)\n                balanced_batch_texts += text\n            except ValueError:\n                pass\n\n        balanced_batch_images = torch.cat(balanced_batch_images, 0)\n\n        return balanced_batch_images, balanced_batch_texts\n\n\ndef hierarchical_dataset(root, opt, select_data='/'):\n    \"\"\" select_data='/' contains all sub-directory of root directory \"\"\"\n    dataset_list = []\n    dataset_log = f'dataset_root:    {root}\\t dataset: {select_data[0]}'\n    print(dataset_log)\n    dataset_log += '\\n'\n    for dirpath, dirnames, filenames in os.walk(root+'/'):\n        if not dirnames:\n            select_flag = False\n            for selected_d in select_data:\n                if selected_d in dirpath:\n                    select_flag = True\n                    break\n\n            if select_flag:\n                dataset = OCRDataset(dirpath, opt)\n                sub_dataset_log = f'sub-directory:\\t/{os.path.relpath(dirpath, root)}\\t num samples: {len(dataset)}'\n                print(sub_dataset_log)\n                dataset_log += f'{sub_dataset_log}\\n'\n                dataset_list.append(dataset)\n\n    concatenated_dataset = ConcatDataset(dataset_list)\n\n    return concatenated_dataset, dataset_log\n\nclass OCRDataset(Dataset):\n\n    def __init__(self, root, opt):\n\n        self.root = root\n        self.opt = opt\n        print(root)\n        self.df = pd.read_csv(os.path.join(root,'labels.csv'), sep='^([^,]+),', engine='python', usecols=['filename', 'words'], keep_default_na=False)\n        self.nSamples = len(self.df)\n\n        if self.opt.data_filtering_off:\n            self.filtered_index_list = [index + 1 for index in range(self.nSamples)]\n        else:\n            self.filtered_index_list = []\n            for index in range(self.nSamples):\n                label = self.df.at[index,'words']\n                try:\n                    if len(label) > self.opt.batch_max_length:\n                        continue\n                except:\n                    print(label)\n                out_of_char = f'[^{self.opt.character}]'\n                if re.search(out_of_char, label.lower()):\n                    continue\n                self.filtered_index_list.append(index)\n            self.nSamples = len(self.filtered_index_list)\n\n    def __len__(self):\n        return self.nSamples\n\n    def __getitem__(self, index):\n        index = self.filtered_index_list[index]\n        img_fname = self.df.at[index,'filename']\n        img_fpath = os.path.join(self.root, img_fname)\n        label = self.df.at[index,'words']\n\n        if self.opt.rgb:\n            img = Image.open(img_fpath).convert('RGB')  # for color image\n        else:\n            img = Image.open(img_fpath).convert('L')\n\n        if not self.opt.sensitive:\n            label = label.lower()\n\n        # We only train and evaluate on alphanumerics (or pre-defined character set in train.py)\n        out_of_char = f'[^{self.opt.character}]'\n        label = re.sub(out_of_char, '', label)\n\n        return (img, label)\n\nclass ResizeNormalize(object):\n\n    def __init__(self, size, interpolation=Image.BICUBIC):\n        self.size = size\n        self.interpolation = interpolation\n        self.toTensor = transforms.ToTensor()\n\n    def __call__(self, img):\n        img = img.resize(self.size, self.interpolation)\n        img = self.toTensor(img)\n        img.sub_(0.5).div_(0.5)\n        return img\n\n\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\n\n\nclass AlignCollate(object):\n\n    def __init__(self, imgH=32, imgW=100, 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, batch):\n        batch = filter(lambda x: x is not None, batch)\n        images, labels = zip(*batch)\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, labels\n\n\ndef tensor2im(image_tensor, imtype=np.uint8):\n    image_numpy = image_tensor.cpu().float().numpy()\n    if image_numpy.shape[0] == 1:\n        image_numpy = np.tile(image_numpy, (3, 1, 1))\n    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0\n    return image_numpy.astype(imtype)\n\n\ndef save_image(image_numpy, image_path):\n    image_pil = Image.fromarray(image_numpy)\n    image_pil.save(image_path)","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:49.754966Z","iopub.execute_input":"2024-02-14T22:06:49.755290Z","iopub.status.idle":"2024-02-14T22:06:50.106183Z","shell.execute_reply.started":"2024-02-14T22:06:49.755232Z","shell.execute_reply":"2024-02-14T22:06:50.105389Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# %%writefile /kaggle/working/train.py\n\nimport os\nimport sys\nimport time\nimport random\nimport torch\nimport torch.backends.cudnn as cudnn\nimport torch.nn as nn\nimport torch.nn.init as init\nimport torch.optim as optim\nimport torch.utils.data\nfrom torch.cuda.amp import autocast, GradScaler\nimport numpy as np\n\nfrom utils import CTCLabelConverter, AttnLabelConverter, Averager\n# from dataset import hierarchical_dataset, AlignCollate, Batch_Balanced_Dataset\nfrom model import Model\nfrom test import validation\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\ndef count_parameters(model):\n    print(\"Modules, Parameters\")\n    total_params = 0\n    for name, parameter in model.named_parameters():\n        if not parameter.requires_grad: continue\n        param = parameter.numel()\n        #table.add_row([name, param])\n        total_params+=param\n        print(name, param)\n    print(f\"Total Trainable Params: {total_params}\")\n    return total_params\n\ndef train(opt, show_number = 2, amp=False):\n    \"\"\" dataset preparation \"\"\"\n    if not opt.data_filtering_off:\n        print('Filtering the images containing characters which are not in opt.character')\n        print('Filtering the images whose label is longer than opt.batch_max_length')\n\n    opt.select_data = opt.select_data.split('-')\n    opt.batch_ratio = opt.batch_ratio.split('-')\n    train_dataset = Batch_Balanced_Dataset(opt)\n    \n    log = open(f'./saved_models/{opt.experiment_name}/log_dataset.txt', 'a', encoding=\"utf8\")\n    AlignCollate_valid = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD, contrast_adjust=opt.contrast_adjust)\n    valid_dataset, valid_dataset_log = hierarchical_dataset(root=opt.valid_data, opt=opt)\n    valid_loader = torch.utils.data.DataLoader(\n        valid_dataset, batch_size=min(32, opt.batch_size),\n        shuffle=True,  # 'True' to check training progress with validation function.\n        num_workers=int(opt.workers), prefetch_factor=512,\n        collate_fn=AlignCollate_valid, pin_memory=True)\n    log.write(valid_dataset_log)\n    print('-' * 80)\n    log.write('-' * 80 + '\\n')\n    log.close()\n    \n    \"\"\" model configuration \"\"\"\n    if 'CTC' in opt.Prediction:\n        converter = CTCLabelConverter(opt.character)\n    else:\n        converter = AttnLabelConverter(opt.character)\n    opt.num_class = len(converter.character)\n\n    if opt.rgb:\n        opt.input_channel = 3\n    model = Model(opt)\n    print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial, opt.input_channel, opt.output_channel,\n          opt.hidden_size, opt.num_class, opt.batch_max_length, opt.Transformation, opt.FeatureExtraction,\n          opt.SequenceModeling, opt.Prediction)\n\n    if opt.saved_model != '':\n        pretrained_dict = torch.load(opt.saved_model)\n        if opt.new_prediction:\n            model.Prediction = nn.Linear(model.SequenceModeling_output, len(pretrained_dict['module.Prediction.weight']))  \n        \n        model = torch.nn.DataParallel(model).to(device) \n        print(f'loading pretrained model from {opt.saved_model}')\n        if opt.FT:\n            model.load_state_dict(pretrained_dict, strict=False)\n        else:\n            model.load_state_dict(pretrained_dict)\n        if opt.new_prediction:\n            model.module.Prediction = nn.Linear(model.module.SequenceModeling_output, opt.num_class)  \n            for name, param in model.module.Prediction.named_parameters():\n                if 'bias' in name:\n                    init.constant_(param, 0.0)\n                elif 'weight' in name:\n                    init.kaiming_normal_(param)\n            model = model.to(device) \n    else:\n        # weight initialization\n        for name, param in model.named_parameters():\n            if 'localization_fc2' in name:\n                print(f'Skip {name} as it is already initialized')\n                continue\n            try:\n                if 'bias' in name:\n                    init.constant_(param, 0.0)\n                elif 'weight' in name:\n                    init.kaiming_normal_(param)\n            except Exception as e:  # for batchnorm.\n                if 'weight' in name:\n                    param.data.fill_(1)\n                continue\n        model = torch.nn.DataParallel(model).to(device)\n    \n    model.train() \n    print(\"Model:\")\n    print(model)\n    count_parameters(model)\n    \n    \"\"\" setup loss \"\"\"\n    if 'CTC' in opt.Prediction:\n        criterion = torch.nn.CTCLoss(zero_infinity=True).to(device)\n    else:\n        criterion = torch.nn.CrossEntropyLoss(ignore_index=0).to(device)  # ignore [GO] token = ignore index 0\n    # loss averager\n    loss_avg = Averager()\n\n    # freeze some layers\n    try:\n        if opt.freeze_FeatureFxtraction:\n            for param in model.module.FeatureExtraction.parameters():\n                param.requires_grad = False\n        if opt.freeze_SequenceModeling:\n            for param in model.module.SequenceModeling.parameters():\n                param.requires_grad = False\n    except:\n        pass\n    \n    # filter that only require gradient decent\n    filtered_parameters = []\n    params_num = []\n    for p in filter(lambda p: p.requires_grad, model.parameters()):\n        filtered_parameters.append(p)\n        params_num.append(np.prod(p.size()))\n    print('Trainable params num : ', sum(params_num))\n    # [print(name, p.numel()) for name, p in filter(lambda p: p[1].requires_grad, model.named_parameters())]\n\n    # setup optimizer\n    if opt.optim=='adam':\n        #optimizer = optim.Adam(filtered_parameters, lr=opt.lr, betas=(opt.beta1, 0.999))\n        optimizer = optim.Adam(filtered_parameters)\n    else:\n        optimizer = optim.Adadelta(filtered_parameters, lr=opt.lr, rho=opt.rho, eps=opt.eps)\n    print(\"Optimizer:\")\n    print(optimizer)\n\n    \"\"\" final options \"\"\"\n    # print(opt)\n    with open(f'./saved_models/{opt.experiment_name}/opt.txt', 'a', encoding=\"utf8\") as opt_file:\n        opt_log = '------------ Options -------------\\n'\n        args = vars(opt)\n        for k, v in args.items():\n            opt_log += f'{str(k)}: {str(v)}\\n'\n        opt_log += '---------------------------------------\\n'\n        print(opt_log)\n        opt_file.write(opt_log)\n\n    \"\"\" start training \"\"\"\n    start_iter = 0\n    if opt.saved_model != '':\n        try:\n            start_iter = int(opt.saved_model.split('_')[-1].split('.')[0])\n            print(f'continue to train, start_iter: {start_iter}')\n        except:\n            pass\n\n    start_time = time.time()\n    best_accuracy = -1\n    best_norm_ED = -1\n    i = start_iter\n\n    scaler = GradScaler()\n    t1= time.time()\n        \n    while(True):\n        # train part\n        optimizer.zero_grad(set_to_none=True)\n        \n        if amp:\n            with autocast():\n                image_tensors, labels = train_dataset.get_batch()\n                image = image_tensors.to(device)\n                text, length = converter.encode(labels, batch_max_length=opt.batch_max_length)\n                batch_size = image.size(0)\n\n                if 'CTC' in opt.Prediction:\n                    preds = model(image, text).log_softmax(2)\n                    preds_size = torch.IntTensor([preds.size(1)] * batch_size)\n                    preds = preds.permute(1, 0, 2)\n                    torch.backends.cudnn.enabled = False\n                    cost = criterion(preds, text.to(device), preds_size.to(device), length.to(device))\n                    torch.backends.cudnn.enabled = True\n                else:\n                    preds = model(image, text[:, :-1])  # align with Attention.forward\n                    target = text[:, 1:]  # without [GO] Symbol\n                    cost = criterion(preds.view(-1, preds.shape[-1]), target.contiguous().view(-1))\n            scaler.scale(cost).backward()\n            scaler.unscale_(optimizer)\n            torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip)\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            image_tensors, labels = train_dataset.get_batch()\n            image = image_tensors.to(device)\n            text, length = converter.encode(labels, batch_max_length=opt.batch_max_length)\n            batch_size = image.size(0)\n            if 'CTC' in opt.Prediction:\n                preds = model(image, text).log_softmax(2)\n                preds_size = torch.IntTensor([preds.size(1)] * batch_size)\n                preds = preds.permute(1, 0, 2)\n                torch.backends.cudnn.enabled = False\n                cost = criterion(preds, text.to(device), preds_size.to(device), length.to(device))\n                torch.backends.cudnn.enabled = True\n            else:\n                preds = model(image, text[:, :-1])  # align with Attention.forward\n                target = text[:, 1:]  # without [GO] Symbol\n                cost = criterion(preds.view(-1, preds.shape[-1]), target.contiguous().view(-1))\n            cost.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip) \n            optimizer.step()\n        loss_avg.add(cost)\n\n        # validation part\n        if (i % opt.valInterval == 0) and (i!=0):\n            print('training time: ', time.time()-t1)\n            t1=time.time()\n            elapsed_time = time.time() - start_time\n            # for log\n            with open(f'./saved_models/{opt.experiment_name}/log_train.txt', 'a', encoding=\"utf8\") as log:\n                model.eval()\n                with torch.no_grad():\n                    valid_loss, current_accuracy, current_norm_ED, preds, confidence_score, labels,\\\n                    infer_time, length_of_data = validation(model, criterion, valid_loader, converter, opt, device)\n                model.train()\n\n                # training loss and validation loss\n                loss_log = f'[{i}/{opt.num_iter}] Train loss: {loss_avg.val():0.5f}, Valid loss: {valid_loss:0.5f}, Elapsed_time: {elapsed_time:0.5f}'\n                loss_avg.reset()\n\n                current_model_log = f'{\"Current_accuracy\":17s}: {current_accuracy:0.3f}, {\"Current_norm_ED\":17s}: {current_norm_ED:0.4f}'\n\n                # keep best accuracy model (on valid dataset)\n                if current_accuracy > best_accuracy:\n                    best_accuracy = current_accuracy\n                    torch.save(model.state_dict(), f'/kaggle/working/best_accuracy.pth')\n                if current_norm_ED > best_norm_ED:\n                    best_norm_ED = current_norm_ED\n                    torch.save(model.state_dict(), f'/kaggle/working/best_norm_ED.pth')\n                best_model_log = f'{\"Best_accuracy\":17s}: {best_accuracy:0.3f}, {\"Best_norm_ED\":17s}: {best_norm_ED:0.4f}'\n\n                loss_model_log = f'{loss_log}\\n{current_model_log}\\n{best_model_log}'\n                print(loss_model_log)\n                log.write(loss_model_log + '\\n')\n\n                # show some predicted results\n                dashed_line = '-' * 80\n                head = f'{\"Ground Truth\":25s} | {\"Prediction\":25s} | Confidence Score & T/F'\n                predicted_result_log = f'{dashed_line}\\n{head}\\n{dashed_line}\\n'\n                \n                #show_number = min(show_number, len(labels))\n                \n                start = random.randint(0,len(labels) - show_number )    \n                for gt, pred, confidence in zip(labels[start:start+show_number], preds[start:start+show_number], confidence_score[start:start+show_number]):\n                    if 'Attn' in opt.Prediction:\n                        gt = gt[:gt.find('[s]')]\n                        pred = pred[:pred.find('[s]')]\n\n                    predicted_result_log += f'{gt:25s} | {pred:25s} | {confidence:0.4f}\\t{str(pred == gt)}\\n'\n                predicted_result_log += f'{dashed_line}'\n                print(predicted_result_log)\n                log.write(predicted_result_log + '\\n')\n                print('validation time: ', time.time()-t1)\n                t1=time.time()\n        # save model per 1e+4 iter.\n        if (i + 1) % 1000 == 0:\n            torch.save(\n                model.state_dict(), f'./saved_models/{opt.experiment_name}/iter_{i+1}.pth')\n        if i == opt.num_iter:\n            print('end the training')\n            break\n        i += 1","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:50.107740Z","iopub.execute_input":"2024-02-14T22:06:50.108150Z","iopub.status.idle":"2024-02-14T22:06:51.377565Z","shell.execute_reply.started":"2024-02-14T22:06:50.108108Z","shell.execute_reply":"2024-02-14T22:06:51.376587Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train(opt, amp=False)","metadata":{"execution":{"iopub.status.busy":"2024-02-14T22:06:51.378764Z","iopub.execute_input":"2024-02-14T22:06:51.379135Z","iopub.status.idle":"2024-02-14T22:19:25.859507Z","shell.execute_reply.started":"2024-02-14T22:06:51.379103Z","shell.execute_reply":"2024-02-14T22:19:25.858150Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}