{"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":"Hello fellow Kagglers,\n\nIn this notebook I will demonstrate how to clean up the images and create TFRecords.\nThis is a first approach and updates will follow as the competition progresses.\n\nCurrent image cleaning techniques applied are\n* Blob removal\n* Filling missing pixels in lines\n* Crop image by detecting contours\n\nIf you have any questions, don't hesitate to leave a comment.\n\nConsider upvoting this work if you find it helpful :D","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\n\ntqdm.pandas()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Choosing an appropriate image size is difficult, as complex molecules will need a high resolution to preserve details, but training on 2.4 million in high resolution is unfeasable. The chosen resolution of 256\\*448 should preserve enough detail and allow for training on a TPU within the 3 hours limit","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 = 200 # maximum InChI length to prevent to much padding","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train/Test DataFrames","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":"print(train.info())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"display(train.head(3))","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":"print(test.info())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"display(test.head(3))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Vocabulary\nUsing sets, which don't allow for duplicate elements, all unique characters are determined","metadata":{}},{"cell_type":"code","source":"def get_vocabulary():\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\nvocabulary = get_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\ntrain['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(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        \ntrain['InChI_int'] = train['InChIClean'].progress_apply(inchi_str2int)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Mean Image Ratio\n\nThe mean image ratio on the first 10K images is computed here.\nUsed image height/width are slightly higher than the mean to preserve details in complex molecule structures.","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":"markdown","source":"Many images contain random pixels as noise. The following function detects small blobs of images and remove them.","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":"markdown","source":"Crop image after removing blobs, this will reduce computational cost when filling missing pixels","metadata":{}},{"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":"markdown","source":"This function tries to fill missing pixels. It uses kernels applied to the image to locate horizontal, vertical and diagonal lines with missing pixels. Both lines of one pixel wide and multiple pixels wide are detected.\n\nKernels detect if a pixel is missing in a line by checking if there is a single pixel mising. In the kernel *-1* indicates no colored pixels are allowed there and and *a* indicates a colored pixel is expected. Detecting an unwanted colored pixel will prevent the pixel from being filled and thresholds are used to detect a minimum amount of pixels in a line. Unwanted colored pixels are needed to prevent characters such as \"OH\" to be filled.","metadata":{}},{"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":"markdown","source":"Add padding to preserve the aspect ratio and resize to target resolution","metadata":{}},{"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":"markdown","source":"The next function processes an image_id to a png encoded cleaned image.","metadata":{}},{"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\nsplits the train, val and test images into chunk of ~100MB TFRecords as is [recommended](https://www.tensorflow.org/guide/tpu#input_datasets)","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\nThis processes the 2.3M training images, 100K val images and 1.6M test images, this process will take several hours (~8)","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\n\nCheck if TFRecords are correctly created, just to be sure","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":{"trusted":true},"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":{"trusted":true},"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":{"trusted":true},"execution_count":null,"outputs":[]}]}