{"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":"This notebook is baselined by [this notebook](https://www.kaggle.com/markwijkhuizen/advanced-image-cleaning-and-tfrecord-generation). Only difference between this and baseline notebook is going to be tokenizer.","metadata":{}},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport tensorflow as tf\nfrom multiprocessing import cpu_count\nfrom tqdm.notebook import tqdm\n\nimport cv2\nimport imageio\nimport joblib\nimport os\nimport pickle\nimport re\n\ntqdm.pandas()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Tokenizer Helper Functions","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Preprocess functions\n# ====================================================\ndef split_form(form):\n    string = ''\n    for i in re.findall(r\"[A-Z][^A-Z]*\", form):\n        elem = re.match(r\"\\D+\", i).group()\n        num = i.replace(elem, \"\")\n        if num == \"\":\n            string += f\"{elem} \"\n        else:\n            string += f\"{elem} {str(num)} \"\n    return string.rstrip(' ')\n\ndef split_form2(form):\n    string = ''\n    for i in re.findall(r\"[a-z][^a-z]*\", form):\n        elem = i[0]\n        num = i.replace(elem, \"\").replace('/', \"\")\n        num_string = ''\n        for j in re.findall(r\"[0-9]+[^0-9]*\", num):\n            num_list = list(re.findall(r'\\d+', j))\n            assert len(num_list) == 1, f\"len(num_list) != 1\"\n            _num = num_list[0]\n            if j == _num:\n                num_string += f\"{_num} \"\n            else:\n                extra = j.replace(_num, \"\")\n                num_string += f\"{_num} {' '.join(list(extra))} \"\n        string += f\"/{elem} {num_string}\"\n    return string.rstrip(' ')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Global Variables","metadata":{}},{"cell_type":"code","source":"# Mean ratio of first 10K images is 1.73, this image width/height result in a ratio of 1.75\nDEBUG = False\nIMG_HEIGHT = 256\nIMG_WIDTH = 448\nVAL_SIZE = int(100) if DEBUG else int(100e3) # 100K validation molecules\nCHUNK_SIZE = 40000 # to get ~100MB TFRecords\n\nMAX_INCHI_LEN_OLD = 200 # maximum InChI length to prevent to much paddin\nMAX_INCHI_LEN = 275 # maximum InChI length to prevent to much paddin","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Loading CSV","metadata":{}},{"cell_type":"code","source":"if DEBUG:\n    train = pd.read_csv('/kaggle/input/bms-molecular-translation/train_labels.csv', dtype={ 'image_id': 'string', 'InChI': 'string' }).head(int(1e3))\nelse:\n    train = pd.read_csv('/kaggle/input/bms-molecular-translation/train_labels.csv', dtype={ 'image_id': 'string', 'InChI': 'string' })\n\n# Drop all InChI longer than MAX_INCHI_LEN - 2,  <start>InChI <end>, remove 'InChI=1S/' at start\ntrain['InChI_len'] = train['InChI'].apply(len).astype(np.uint16)\ntrain = train.loc[train['InChI_len'] <= MAX_INCHI_LEN - 2 + 9].reset_index(drop=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if DEBUG:\n    test = pd.read_csv('/kaggle/input/bms-molecular-translation/sample_submission.csv', usecols=['image_id'], dtype={ 'image_id': 'string' }).head(int(1e3))\nelse:\n    test = pd.read_csv('/kaggle/input/bms-molecular-translation/sample_submission.csv', usecols=['image_id'], dtype={ 'image_id': 'string' })","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Parse InChi","metadata":{}},{"cell_type":"code","source":"train['InChI_1'] = train['InChI'].progress_apply(lambda x: x.split('/')[1])\ntrain['InChI_text'] = train['InChI_1'].progress_apply(split_form) + ' ' + train['InChI'].apply(lambda x: '/'.join(x.split('/')[2:])).progress_apply(split_form2).values","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_vocabulary_OLD():\n    tokens = ['<start>', '<end>', '<pad>']\n    vocabulary = set()\n    for s in tqdm(train['InChI'].values):\n        vocabulary.update(s)\n    return tokens + list(vocabulary)\n\ndef get_vocabulary():\n    tokens = ['<start>', '<end>', '<pad>']\n    vocabulary = []\n    for s in tqdm(train['InChI_text'].values):\n        vocabulary += s.split(' ')\n        \n        #vocabulary.update(s)\n    return tokens + list(set(vocabulary))\n\nvocabulary = get_vocabulary()\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(vocabulary)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Save vocabulary mappings\n# , character -> integer\nvocabulary_to_int = dict(zip(vocabulary, np.arange(len(vocabulary), dtype=np.int8)))\nwith open('vocabulary_to_int.pkl', 'wb') as handle:\n    pickle.dump(vocabulary_to_int, handle)\n\n#  integer -> character\nint_to_vocabulary = dict(zip(np.arange(len(vocabulary), dtype=np.int8), vocabulary))\nwith open('int_to_vocabulary.pkl', 'wb') as handle:\n    pickle.dump(int_to_vocabulary, handle)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Remove the \"InChI=1S/\" part from the InChI strings\n# It is equal for all InChI's, thus redundant\n#train['InChIClean'] = train['InChI'].apply(lambda InChI: '/'.join(InChI.split('=')[1].split('/')[1:]))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# convert the InChI strings to integer lists\n# start/end/pad tokens are used\ndef inchi_str2int_OLD(InChI):\n    res = []\n    res.append(vocabulary_to_int.get('<start>'))\n    for c in InChI:\n        res.append(vocabulary_to_int.get(c))\n    \n    res.append(vocabulary_to_int.get('<end>'))\n    while len(res) < MAX_INCHI_LEN: \n        res.append(vocabulary_to_int.get('<pad>'))\n        \n    return np.array(res, dtype=np.uint8)\n\ndef inchi_str2int(InChI):\n    res = []\n    res.append(vocabulary_to_int.get('<start>'))\n    for c in InChI.split(' '):\n        res.append(vocabulary_to_int.get(c))\n    \n    res.append(vocabulary_to_int.get('<end>'))\n    while len(res) < MAX_INCHI_LEN: \n        res.append(vocabulary_to_int.get('<pad>'))\n        \n    return np.array(res, dtype=np.uint8)\n        \ntrain['InChI_int'] = train['InChI_text'].progress_apply(inchi_str2int)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Mean Image Ratio","metadata":{}},{"cell_type":"code","source":"ws = []\nhs = []\nfor image_id in tqdm(train.loc[:int(10e3), 'image_id'].values):\n    file_path =  f'/kaggle/input/bms-molecular-translation/train/{image_id[0]}/{image_id[1]}/{image_id[2]}/{image_id}.png'\n    h, w = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE).shape\n    ws.append(w)\n    hs.append(h)\n    \nws_mean = int(np.array(ws).mean())\nhs_mean = int(np.array(hs).mean())\nprint(f'mean width: {ws_mean}, mean height: {hs_mean}, mean ratio: {round(ws_mean / hs_mean, 2)}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Val Split","metadata":{}},{"cell_type":"code","source":"val = train.iloc[-VAL_SIZE:].reset_index(drop=True)\ntrain = train.iloc[:-VAL_SIZE].reset_index(drop=True)\nN_IMGS = len(train)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pd.options.display.max_colwidth = 100\ndisplay(train.head(3))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(train.info())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(val.info())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Image Cleaning","metadata":{}},{"cell_type":"code","source":"def remove_blobs(img, min_size=10, debug=False):\n    if debug:\n        fig, ax = plt.subplots(1,2, figsize=(30,8))\n        ax[0].imshow(img)\n        ax[0].set_title('original image', size=16)\n    \n    height, width = img.shape\n\n    # find all the connected components (white blobs in your image)\n    nb_components, output, stats, centroids = cv2.connectedComponentsWithStats(img, connectivity=8)\n    # Removes background, which is seen as a big component\n    sizes = stats[1:, -1]\n  \n    blob_idxs = []    \n    for idx, s in enumerate(sizes):\n        if s < min_size:\n            blob_idxs.append(idx+1)\n    \n    img[np.isin(output, blob_idxs)] = 0\n    \n    if debug:\n        ax[1].imshow(img)\n        ax[1].set_title('image with removed blobs', size=16)\n        plt.show()\n    \n    return img","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def crop(img, debug=False):\n    if debug:\n        fig, ax = plt.subplots(1,2, figsize=(30,8))\n        ax[0].imshow(img)\n        ax[0].set_title(f'original image, shape: {img.shape}', size=16)\n        \n    _, thresh = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)\n    contours, _ = cv2.findContours(thresh,cv2.RETR_LIST,cv2.CHAIN_APPROX_SIMPLE)[-2:]\n    \n    x_min, y_min, x_max, y_max = np.inf, np.inf, 0, 0\n    for cnt in contours:\n        x, y, w, h = cv2.boundingRect(cnt)\n        x_min = min(x_min, x)\n        y_min = min(y_min, y)\n        x_max = max(x_max, x + w)\n        y_max = max(y_max, y + h)\n\n    img_cropped = img[y_min:y_max, x_min:x_max]\n    \n    if debug:\n        ax[1].imshow(img_cropped)\n        ax[1].set_title(f'cropped image, shape: {img_cropped.shape}', size=16)\n        plt.show()\n    \n    return img_cropped","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# pad the kernels to create squares, max_pad can be used to create rectangles\ndef pad_kernel(kernel, max_pad=np.inf):\n    kernel = np.array(kernel)\n    h, w = kernel.shape\n    pad_h = min((max(h, w) - h) // 2, max_pad)\n    pad_w = min((max(h, w) - w) // 2, max_pad)\n    return np.pad(kernel, ([pad_h, pad_h], [pad_w, pad_w]), 'constant', constant_values=-1)\n\n# creates a mask of missing pixels to be filled using\ndef create_mask(kernel, img_b):\n    mask = cv2.filter2D(img_b, -1, kernel)\n    kernel_flat_sum = (kernel == a).flatten().sum()\n    threshold_min = kernel_flat_sum * threshold_ratio\n    threshold_max = kernel_flat_sum + 1\n    return (mask > threshold_min) & (mask < threshold_max)\n\n# make kernels\na = np.float32(1.0 / 255.0)\nthreshold_ratio = 0.50\n# single pixel width horizontal line with 1 pixel missing\nkernel_h_single_mono = pad_kernel([\n    [ a, a,  a, -1,  a,  a, a ]\n], max_pad=1)\n# single pixel width horizontal line with 3 pixels missing\nkernel_h_single_triple = pad_kernel([\n    [ a, a, a, -1, -1, -1, a, a, a ]\n], max_pad=1)\n\nkernel_h_multi = pad_kernel([\n    [ a, a, a, a, a, a, a ],\n    [ a, a, a,-1, a, a, a ],\n    [ a, a, a, a, a, a, a ],\n], max_pad=1)\n\nkernel_v_single = pad_kernel([\n    [ a],\n    [ a],\n    [ a],\n    [-1],\n    [ a],\n    [ a],\n    [ a],\n], max_pad=1)\n\nkernel_v_multi = pad_kernel([\n    [ a, a, a ],\n    [ a, a, a ],\n    [ a, a, a ],\n    [ a,-1, a ],\n    [ a, a, a ],\n    [ a, a, a ],\n    [ a, a, a ],\n], max_pad=1)\n\nkernel_lr_single = pad_kernel([\n    [ -1,-1,-1,-1, a ],\n    [ -1,-1,-1, a,-1 ],\n    [ -1,-1,-1,-1,-1 ],\n    [ -1, a,-1,-1,-1 ],\n    [  a,-1,-1,-1,-1 ],\n])\n\nkernel_lr_multi = pad_kernel([\n    [ -1,-1,-1, a, a ],\n    [ -1,-1, a, a, a ],\n    [ -1, a,-1, a,-1 ],\n    [  a, a, a,-1,-1 ],\n    [  a, a,-1,-1,-1 ],\n])\n\nkernel_rl_single = pad_kernel([\n    [  a,-1,-1,-1,-1 ],\n    [ -1, a,-1,-1,-1 ],\n    [ -1,-1,-1,-1,-1 ],\n    [ -1,-1,-1, a,-1 ],\n    [ -1,-1,-1,-1, a ],\n])\n\nkernel_rl_multi = pad_kernel([\n    [ a, a,-1,-1,-1],\n    [ a, a, a,-1,-1],\n    [-1, a,-1, a,-1],\n    [-1,-1, a, a, a],\n    [-1,-1,-1, a, a],\n])\n\ndef fill_missing_pixels(img, debug):\n    img_b = img.astype(np.float32)\n    img_b[img_b > 0] = 255\n\n    mask_h_single_mono = create_mask(kernel_h_single_mono, img_b)\n\n    mask_h_single_triple = create_mask(kernel_h_single_triple, img_b)\n\n    mask_h_single = mask_h_single_mono | mask_h_single_triple\n\n    mask_h_multi = create_mask(kernel_h_multi, img_b)\n\n\n    mask_v_single = create_mask(kernel_v_single, img_b)\n\n\n    mask_v_multi = create_mask(kernel_v_multi, img_b)\n\n\n    mask_lr_single = create_mask(kernel_lr_single, img_b)\n\n\n    mask_lr_multi = create_mask(kernel_lr_multi, img_b)\n\n\n    mask_rl_single = create_mask(kernel_lr_single, img_b)\n\n\n    mask_rl_multi = create_mask(kernel_rl_multi, img_b)\n\n    mask_single = mask_h_single | mask_v_single | mask_lr_single | mask_rl_single\n    mask_multi = mask_h_multi  | mask_v_multi |mask_lr_multi | mask_rl_multi\n    mask = mask_single | mask_multi\n\n    if debug:\n        fig, ax = plt.subplots(2, 2 ,figsize=(35,20))\n        ax[0,0].imshow(mask_h_single)\n        ax[0,0].set_title('mask_h_single', size=16)\n        ax[0,1].imshow(mask_v_single)\n        ax[0,1].set_title('mask_v_single', size=16)\n        ax[1,0].imshow(mask_lr_single)\n        ax[1,0].set_title('mask_lr_single', size=16)\n        ax[1,1].imshow(mask_lr_single)\n        ax[1,1].set_title('mask_lr_single', size=16)\n        plt.show()\n\n        fig, ax = plt.subplots(2, 2, figsize=(35,20))\n        ax[0,0].imshow(mask_h_multi)\n        ax[0,0].set_title('mask_h_multi', size=16)\n        ax[0,1].imshow(mask_v_multi)\n        ax[0,1].set_title('mask_v_multi', size=16)\n        ax[1,0].imshow(mask_lr_multi)\n        ax[1,0].set_title('mask_lr_multi', size=16)\n        ax[1,1].imshow(mask_rl_multi)\n        ax[1,1].set_title('mask_rl_multi', size=16)\n        plt.show()\n\n        fig, ax = plt.subplots(2, 1 ,figsize=(15,20))\n        ax[0].imshow(img)\n        ax[0].set_title('original image', size=16)\n\n        img_fill = mask.copy()\n        img_fill[img_fill > 0] = 255\n\n        img_rgb = np.stack([\n            img_fill,\n            img_b,\n            np.zeros(img.shape),\n        ], axis=2)\n\n        ax[1].imshow(img_rgb)\n        ax[1].set_title('image with filled missing pixels (red)', size=16)\n        plt.show()    \n\n    # all pixels in the mask are filled up\n    img[mask] = 255\n\n    return img","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def pad_resize(img):\n    h, w = img.shape\n    s = max(w, h)\n    pad_h, pad_v = 0, 0\n    hw_ratio = (h / w) - (IMG_HEIGHT / IMG_WIDTH)\n    if hw_ratio < 0:\n        pad_h = int(abs(hw_ratio) * w / 2)\n    else:\n        wh_ratio = (w / h) - (IMG_WIDTH / IMG_HEIGHT)\n        pad_v = int(abs(wh_ratio) * h // 2)\n\n    img = np.pad(img, [(pad_h, pad_h), (pad_v, pad_v)], mode='constant')\n    img = cv2.resize(img,(IMG_WIDTH, IMG_HEIGHT), interpolation=cv2.INTER_NEAREST)\n\n    return img","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def process_img(image_id, folder='train', debug=False):\n    # read image and invert colors to get black background and white molecule\n    file_path =  f'/kaggle/input/bms-molecular-translation/{folder}/{image_id[0]}/{image_id[1]}/{image_id[2]}/{image_id}.png'\n    img0 = 255 - cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)\n    \n    # rotate counter clockwise to get horizontal images\n    h, w = img0.shape\n    if h > w:\n        img0 = np.rot90(img0)\n    \n    # remove blobs, crop, fill missing pixels, pad and resize\n    img = remove_blobs(img0, min_size=2, debug=debug)\n    img = crop(img, debug=debug)\n    img = fill_missing_pixels(img, debug=debug)\n    img = pad_resize(img)\n    \n    if debug:\n        fig, ax = plt.subplots(1, 2, figsize=(20,10))\n        ax[0].imshow(img0)\n        ax[0].set_title('Original image', size=16)\n        ax[1].imshow(img)\n        ax[1].set_title('Fully processed image', size=16)\n    \n    # normalize to range 0-255 and encode as png\n    img = (img / img.max() * 255).astype(np.uint8)\n    img = cv2.imencode('.png', img)[1].tobytes()\n\n    return img","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# example of image cleaning process\nprocess_img(train.loc[9, 'image_id'], debug=True)\npass","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Split data in chunks for TFRecords","metadata":{}},{"cell_type":"code","source":"def split_in_chunks(data):\n    return [data[i:i + CHUNK_SIZE] for i in range(0, len(data), CHUNK_SIZE)]\n\ntrain_data_chunks = {\n    'train': {\n        'image_id': split_in_chunks(train['image_id'].values),\n        'InChI': split_in_chunks(train['InChI_int'].values),\n    },\n    'val': {\n        'image_id': split_in_chunks(val['image_id'].values),\n        'InChI': split_in_chunks(val['InChI_int'].values),\n    }\n}\n\ntest_data_chunks = {\n    'test': {\n        'image_id': split_in_chunks(test['image_id'].values),\n    }\n}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# TFRecords","metadata":{}},{"cell_type":"code","source":"def make_tfrecords(data_chunks, folder='train'):\n    # Try to make output folder\n    try:\n        os.makedirs(f'./train')\n        os.makedirs(f'./val')\n        os.makedirs(f'./test')\n    except:\n        print(f'folders already created')\n\n    for k, v in data_chunks.items():\n        for chunk_idx, image_id_chunk in tqdm(enumerate(v['image_id']), total=len(v['image_id'])):\n            # process images in parallel\n            jobs = [joblib.delayed(process_img)(fp, folder) for fp in image_id_chunk]\n            bs = 10\n            processed_images_chunk = joblib.Parallel(\n                n_jobs=cpu_count(),\n                verbose=0,\n                require='sharedmem',\n                batch_size=bs,\n                backend='threading',\n            )(jobs)\n\n            # Create the TFRecords from the processed images\n            with tf.io.TFRecordWriter(f'./{k}/batch_{chunk_idx}.tfrecords') as file_writer:\n                if 'InChI' in v.keys(): # TRAIN/VAL, InChI included\n                    for image, InChI in zip(processed_images_chunk, v['InChI'][chunk_idx]):\n                        record_bytes = tf.train.Example(features=tf.train.Features(feature={\n                            'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])),\n                            'InChI': tf.train.Feature(int64_list=tf.train.Int64List(value=InChI)),\n                        })).SerializeToString()\n                        file_writer.write(record_bytes)\n                else: # TEST, image_id included for submission file\n                    for image, image_id in zip(processed_images_chunk, image_id_chunk):\n                        record_bytes = tf.train.Example(features=tf.train.Features(feature={\n                            'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])),\n                            'image_id': tf.train.Feature(bytes_list=tf.train.BytesList(value=[str.encode(image_id)])),\n                        })).SerializeToString()\n                        file_writer.write(record_bytes)\n\nmake_tfrecords(train_data_chunks)\nmake_tfrecords(test_data_chunks, 'test')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Check TFRecords","metadata":{}},{"cell_type":"code","source":"# convert in int encoded InChI to string\ndef inchi_int2char(InChI):\n    res = []\n    for i in InChI:\n        c = int_to_vocabulary.get(i)\n        if c not in ['<start>', '<end>', '<pad>']:\n            res.append(c)\n    return ''.join(res)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Check train TFRecords\ndef decode_tfrecord(record_bytes):\n    features = tf.io.parse_single_example(record_bytes, {\n        'image': tf.io.FixedLenFeature([], tf.string),\n        'InChI': tf.io.FixedLenFeature([MAX_INCHI_LEN], tf.int64),\n    })\n\n    image = tf.io.decode_jpeg(features['image'])    \n    image = tf.reshape(image, [IMG_HEIGHT, IMG_WIDTH, 1])\n    image = tf.cast(image, tf.float32)  / 255.0\n    \n    InChI = features['InChI']\n    InChI = tf.reshape(InChI, [MAX_INCHI_LEN])\n    \n    return image, InChI\n\ndef show_tfrecords(file_path, rows=3, cols=2):\n    fig, ax = plt.subplots(rows, cols, figsize=(cols*7, rows*4))\n    tfrecord = tf.data.TFRecordDataset(file_path)\n    for idx, (image, InChI) in enumerate(tfrecord.map(decode_tfrecord).take(rows*cols)):\n        if idx == 0:\n            print(f'first InChI int: {InChI}')\n            print(f'first InChI char {inchi_int2char(InChI.numpy())}')\n        image = tf.cast(image * 255, tf.uint8)\n        image = tf.squeeze(image)\n        row, col = idx // cols, idx % cols\n        ax[row, col].imshow(image)\n\n    plt.show()\n\nprint('TRAIN BATCH')\nshow_tfrecords(f'./train/batch_0.tfrecords')\nprint('VAL BATCH')\nshow_tfrecords(f'./val/batch_0.tfrecords')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Check test TFRecords\ndef decode_test_tfrecord(record_bytes):\n    features = tf.io.parse_single_example(record_bytes, {\n        'image': tf.io.FixedLenFeature([], tf.string),\n        'image_id': tf.io.FixedLenFeature([], tf.string),\n    })\n\n    image = tf.io.decode_jpeg(features['image'])\n    image = tf.reshape(image, [IMG_HEIGHT, IMG_WIDTH, 1])\n    image = tf.cast(image, tf.float32)  / 255.0\n    \n    image_id = features['image_id']\n    \n    return image, image_id\n\ndef show_test_tfrecords(file_path, rows=3, cols=2):\n    fig, ax = plt.subplots(rows, cols, figsize=(cols*7, rows*4))\n    tfrecord = tf.data.TFRecordDataset(file_path)\n    for idx, (image, image_id) in enumerate(tfrecord.map(decode_test_tfrecord).take(rows*cols)):\n        image = tf.cast(image * 255, tf.uint8)\n        image = tf.squeeze(image)\n        row, col = idx // cols, idx % cols\n        ax[row, col].imshow(image)\n        ax[row, col].set_title(image_id.numpy().decode(), size=16)\n\n    plt.show()\n    \nprint('TEST BATCH')\nshow_test_tfrecords(f'./test/batch_0.tfrecords')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}