{"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":"### Stratified KFold tils and TFRecords  \n\nIn this notebook, my goal is to generate Stratified tile images, tile masks and corresponding TFRecords from original images and rle encodings.  \n\n* **Why tile?**  \nCutting images into small tiles may help to better preserve the features that exist in high resolution images. Tiles with the same size can easily be processed by various pretrained backbones.\n\n* **Why Stratified?**  \nWhen we split the whole dataset into training set and validation set, we can do a total random sampling. But there might be a problem. From [this data exploration notebook](https://www.kaggle.com/code/electro/hubmap-hpa-data-exploration/notebook), we know FTUs of different organs present different features, both graphically and statistically. So it may be a good idea to keep the percentages of different organs the same both in training set and validation set. Stratified sampling can do this. Otherwise, you might find all lung images crowding in validation set by chance, then the validation metrics would not be representative.  \n\n* **Why Kfold?**  \nFor model and hyperparameters tuning, we need to split the whole dataset into K folds. Do the training process K times, each time we use one fold as validation set, other folds as training set. Then we aggregate all metrics from K training processes to evaluate models and hyperparameters. It's cross validation.\n\n* **Why TFRecord?**  \nTFRecord is a binary format dataset file that can be processed fast on [TPU](https://www.kaggle.com/docs/tpu). TPU can help us accelerate experiments.\n","metadata":{}},{"cell_type":"code","source":"import numpy as np\nfrom PIL import Image\nimport pandas as pd\nimport os\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nfrom sklearn.model_selection import StratifiedKFold\nimport tensorflow as tf","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:52:45.805442Z","iopub.execute_input":"2022-08-07T05:52:45.806110Z","iopub.status.idle":"2022-08-07T05:52:52.481856Z","shell.execute_reply.started":"2022-08-07T05:52:45.806022Z","shell.execute_reply":"2022-08-07T05:52:52.480829Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Tile size and fold number can be set here. If you don't want TFRecords, set it to be false.","metadata":{}},{"cell_type":"code","source":"RESIZE_DIM = 1024\nTILE_SIZE = 256\nN_FOLD = 10\nTFRECORD = True","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:52:52.484319Z","iopub.execute_input":"2022-08-07T05:52:52.485179Z","iopub.status.idle":"2022-08-07T05:52:52.490781Z","shell.execute_reply.started":"2022-08-07T05:52:52.485136Z","shell.execute_reply":"2022-08-07T05:52:52.489652Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BASE_DIR = \"/kaggle/input/hubmap-organ-segmentation\"\nTRAIN_IMAGES_DIR = os.path.join(BASE_DIR, \"train_images\")\ntrain_df = pd.read_csv(os.path.join(BASE_DIR, \"train.csv\"))\ndisplay(train_df)","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:52:52.492258Z","iopub.execute_input":"2022-08-07T05:52:52.492923Z","iopub.status.idle":"2022-08-07T05:52:52.894167Z","shell.execute_reply.started":"2022-08-07T05:52:52.492893Z","shell.execute_reply":"2022-08-07T05:52:52.892872Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"First, we need to generate masks from rle encoding for later use.","metadata":{}},{"cell_type":"code","source":"# https://www.kaggle.com/paulorzp/rle-functions-run-lenght-encode-decode\ndef rle2mask(mask_rle, shape):\n    '''\n    mask_rle: run-length as string formated (start length)\n    shape: (width,height) of array to return \n    Returns numpy array, 1 - mask, 0 - background\n\n    '''\n    s = mask_rle.split()\n    starts, lengths = [\n        np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])\n    ]\n    starts -= 1\n    ends = starts + lengths\n    img = np.zeros(shape[0] * shape[1], dtype=np.uint8)\n    for lo, hi in zip(starts, ends):\n        img[lo : hi] = 1\n    return img.reshape(shape).T","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:52:52.896939Z","iopub.execute_input":"2022-08-07T05:52:52.897391Z","iopub.status.idle":"2022-08-07T05:52:52.905026Z","shell.execute_reply.started":"2022-08-07T05:52:52.897347Z","shell.execute_reply":"2022-08-07T05:52:52.903986Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"OUTPUT_DIR = f\"/kaggle/working/\"\nMASK_IMG_DIR = os.path.join(OUTPUT_DIR, \"mask_img\")\nif not os.path.isdir(MASK_IMG_DIR): os.makedirs(MASK_IMG_DIR, exist_ok=True)\nfor index, row in tqdm(train_df.iterrows()):\n    mask = rle2mask(row[\"rle\"],(row[\"img_width\"], row[\"img_height\"]))\n    m_path = os.path.join(MASK_IMG_DIR, str(row[\"id\"])+\"_mask.png\")\n    Image.fromarray(mask).save(m_path)","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:52:52.906922Z","iopub.execute_input":"2022-08-07T05:52:52.907459Z","iopub.status.idle":"2022-08-07T05:54:12.517417Z","shell.execute_reply.started":"2022-08-07T05:52:52.907393Z","shell.execute_reply":"2022-08-07T05:54:12.516649Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Utility functions** to split images and masks into tiles and to assembly them back are from [this excellent notebook](https://www.kaggle.com/code/jirkaborovec/ftus-segm-decompose-large-images-to-tiles), please check it out!  \n\nThe below **tile_image** function will apply to both images(TIFF format, original) and masks(PNG format, we generate last step).   \n\nIt reads an image or a mask, resize to a smaller shape, divide it into tiles with (size, size) shape, then save tiles in PNG format in the folder. It returns two lists. One contains all saved tile PNG images' paths. Another one contains corresponding tuples of indices. One tuple of indices indicates which part of the resized image or mask this tile belongs to. For example, (0, 512, 0, 512) indicates this tile is at the top left of the resized image or mask.   \n\nAll tiles, file paths and indices generated by this function will be used to make TFRecords and to reconstruct the resized images later.","metadata":{}},{"cell_type":"code","source":"# https://www.kaggle.com/code/jirkaborovec/ftus-segm-decompose-large-images-to-tiles\ndef tile_image(p_img, folder, resize_shape=(RESIZE_DIM, RESIZE_DIM),size: int = 512) -> list:\n    w = h = size\n    im = np.array(Image.open(p_img).resize(resize_shape))\n    # https://stackoverflow.com/a/47581978/4521646\n    tiles = [im[i:(i + h), j:(j + w), ...] for i in range(0, im.shape[0], h) for j in range(0, im.shape[1], w)]\n    idxs = [(i, (i + h), j, (j + w)) for i in range(0, im.shape[0], h) for j in range(0, im.shape[1], w)]\n    name, _ = os.path.splitext(os.path.basename(p_img))\n    files = []\n    for k, tile in enumerate(tiles):\n        if tile.shape[:2] != (h, w):\n            tile_ = tile\n            tile = np.zeros_like(tiles[0])\n            tile[:tile_.shape[0], :tile_.shape[1], ...] = tile_\n        p_img = os.path.join(folder, f\"{name}_{k:03}.png\")\n        Image.fromarray(tile).save(p_img)\n        files.append(p_img)\n    return files, idxs","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:54:36.714833Z","iopub.execute_input":"2022-08-07T05:54:36.715235Z","iopub.status.idle":"2022-08-07T05:54:36.738650Z","shell.execute_reply.started":"2022-08-07T05:54:36.715201Z","shell.execute_reply":"2022-08-07T05:54:36.737520Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Utility functions** to make examples and to write [TFRecords](https://www.tensorflow.org/tutorials/load_data/tfrecord).   \n\nTo write one example of TFRecord, we have to:\n1. Read the raw data.\n2. Convert raw data into types tf.train.Feature can handle.\n3. Organize tf.train.Feature in tf.train.Example.\n4. Call SerializeToString method to turn one example into a binary-string. \n5. Write the binary-string to a TFRecords file.\n  \nHere I make one example as this: {tile image, tile mask, tile id, tile position(indices), label, class name}.  \nTile images and tile masks are dataset for training. Tile IDs are tile identities. Tile positions are where they sit on resized images. Labels and class names are refer to the types of organs.\n\nYou can make your own TFRecord struct, put into what ever information you see necessary.","metadata":{}},{"cell_type":"markdown","source":"**Step 1 and 2**","metadata":{}},{"cell_type":"code","source":"def encode_features(image_file, mask_file, id_str, class_name):\n    image = tf.image.convert_image_dtype(Image.open(image_file), dtype=tf.uint8)\n    image = tf.io.encode_png(image)\n    mask = tf.image.convert_image_dtype(Image.open(mask_file), dtype=tf.uint8)\n    mask = tf.expand_dims(mask, -1)\n    mask = tf.io.encode_png(mask)\n    tile_id = id_str.encode('utf-8')\n    class_name = class_name.encode('utf-8')\n    return image, mask, tile_id, class_name","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:54:45.411053Z","iopub.execute_input":"2022-08-07T05:54:45.411467Z","iopub.status.idle":"2022-08-07T05:54:45.418546Z","shell.execute_reply.started":"2022-08-07T05:54:45.411418Z","shell.execute_reply":"2022-08-07T05:54:45.417334Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# The following functions can be used to convert a value to a type compatible\n# with tf.train.Example.\n\ndef _bytes_feature(value):\n  \"\"\"Returns a bytes_list from a string / byte.\"\"\"\n  if isinstance(value, type(tf.constant(0))):\n    value = value.numpy() # BytesList won't unpack a string from an EagerTensor.\n  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))\n\ndef _float_feature(value):\n  \"\"\"Returns a float_list from a float / double.\"\"\"\n  return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))\n\ndef _int64_feature(value):\n  \"\"\"Returns an int64_list from a bool / enum / int / uint.\"\"\"\n  return tf.train.Feature(int64_list=tf.train.Int64List(value=value))","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:54:48.161627Z","iopub.execute_input":"2022-08-07T05:54:48.162133Z","iopub.status.idle":"2022-08-07T05:54:48.171670Z","shell.execute_reply.started":"2022-08-07T05:54:48.162087Z","shell.execute_reply":"2022-08-07T05:54:48.170508Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Step 3 and 4**\n  \nTile images, tile masks and tile ids are converted into tf.train.BytesList. Tile positions are converted into tf.train.Int64List.","metadata":{}},{"cell_type":"code","source":"def make_example(encoded_image, encoded_mask, tile_id, tile_pos, label, class_name):\n    features = tf.train.Features(feature={\n        'image': _bytes_feature(encoded_image),\n        'mask': _bytes_feature(encoded_mask),\n        'tile_id': _bytes_feature(tile_id),\n        'tile_position': _int64_feature(tile_pos),\n        'label': _int64_feature([label]),\n        'class_name': _bytes_feature(class_name)\n    })\n    \n    example = tf.train.Example(features=features)\n    \n    return example.SerializeToString()","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:54:51.695326Z","iopub.execute_input":"2022-08-07T05:54:51.695846Z","iopub.status.idle":"2022-08-07T05:54:51.705413Z","shell.execute_reply.started":"2022-08-07T05:54:51.695802Z","shell.execute_reply":"2022-08-07T05:54:51.704063Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now we do the actual process.   \n\n1. We split 10 fold datasets, stratified by organ types: kidney, prostate, large intestine, spleen and lung.  \n  \n2. For each fold, we create tile files, and then write image tiles, mask tiles, tile ids, tile position to TFRecords(**step 5**).    \n\n3. We plot organ counts of each fold.","metadata":{}},{"cell_type":"code","source":"skf = StratifiedKFold(n_splits=N_FOLD, shuffle=True, random_state=42)\nCLASS_NAMES = ['kidney', 'largeintestine', 'lung', 'prostate', 'spleen']\nplt.figure(figsize=(20, 8))\nfor fold_ind, (train_index, val_index) in enumerate(skf.split(train_df, train_df[\"organ\"])):\n    strat_val_set = train_df.loc[val_index]\n    tile_img_files = []\n    tile_msk_files = []\n    tile_idxs_list = []\n    print(f\"FOLD {fold_ind} has the following IDs:\", list(strat_val_set[\"id\"].values))\n    print(\"Creating tiles...\")\n    \n    #Create one fold image tiles and mask tiles\n    TILE_IMG_DIR = os.path.join(OUTPUT_DIR, f\"FOLD_{fold_ind}\", \"img_tile\")\n    if not os.path.isdir(TILE_IMG_DIR): os.makedirs(TILE_IMG_DIR, exist_ok=True)\n\n    TILE_MSK_DIR = os.path.join(OUTPUT_DIR, f\"FOLD_{fold_ind}\", \"mask_tile\")\n    if not os.path.isdir(TILE_MSK_DIR): os.makedirs(TILE_MSK_DIR, exist_ok=True)\n    \n    for _, row in tqdm(strat_val_set.iterrows()):\n        tiles_img, _ = tile_image(os.path.join(TRAIN_IMAGES_DIR, str(row[\"id\"])+\".tiff\"), \n                                  TILE_IMG_DIR, \n                                  size=TILE_SIZE)\n        tiles_msk, tile_idxs = tile_image(os.path.join(MASK_IMG_DIR, str(row[\"id\"])+\"_mask.png\"), \n                                     TILE_MSK_DIR, \n                                     size=TILE_SIZE)\n        tile_img_files.extend(tiles_img)\n        tile_msk_files.extend(tiles_msk)\n        tile_idxs_list.extend(tile_idxs)\n#         break\n    print(f\"FOLD {fold_ind} finished.\\n\\n\")\n    \n    # Write TFRecord of this fold\n    if TFRECORD:\n        print(f\"FOLD {fold_ind} TFRecord creating...\")\n        NUM_EXAMPLE = len(tile_img_files)\n        record_file = f'fold_{fold_ind}-{NUM_EXAMPLE}.tfrec'\n        with tf.io.TFRecordWriter(record_file) as writer:\n            for p_img, p_msk, _tile_idxs in zip(tile_img_files, tile_msk_files, tile_idxs_list):\n                tile_id_str = os.path.splitext(os.path.basename(p_img))[0]\n                img_id = tile_id_str.split(\"_\")[0]\n                class_name = train_df[train_df['id']==int(img_id)]['organ'].values[0]\n                label = CLASS_NAMES.index(class_name)\n                encoded_image, encoded_mask, encoded_tile_id, encoded_class_name = encode_features(p_img, p_msk, tile_id_str, class_name)\n                writer.write(make_example(encoded_image, encoded_mask, encoded_tile_id, _tile_idxs, label, encoded_class_name))\n        print(f\"TFRecord finished.\\n\\n\")      \n    \n    #Plot organ counts to check the result of stratified split\n    plt.subplot(2, 5, fold_ind + 1)\n    plt.title(f\"FOLD {fold_ind}\")\n    strat_val_set[\"organ\"].value_counts().loc[[\"kidney\", \"prostate\", \"largeintestine\", \"spleen\", \"lung\"]].plot.bar()\n#     break \nplt.tight_layout()","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:55:09.284729Z","iopub.execute_input":"2022-08-07T05:55:09.285159Z","iopub.status.idle":"2022-08-07T05:55:11.941019Z","shell.execute_reply.started":"2022-08-07T05:55:09.285122Z","shell.execute_reply":"2022-08-07T05:55:11.940263Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"From the charts above, we can see that after the stratified split, each fold remains the same percentages of different organ types as the whole dataset.","metadata":{}},{"cell_type":"markdown","source":"TFRecords writing finished. Now we read from TFRecords to check if we can reconstruct resized images.   \n\nThe next cell is the function to read one example from TFRecord and parse it into data we can understand.","metadata":{}},{"cell_type":"code","source":"tile_dataset = tf.data.TFRecordDataset(tf.io.gfile.glob('fold_0-*.tfrec'))\n\n# Create a dictionary describing the features.\nfeature_description = {\n    'image': tf.io.FixedLenFeature([], tf.string),\n    'mask': tf.io.FixedLenFeature([], tf.string),\n    'tile_id': tf.io.FixedLenFeature([], tf.string),\n    'tile_position': tf.io.RaggedFeature(dtype=tf.int64),\n    'label': tf.io.FixedLenFeature([], tf.int64),\n    'class_name': tf.io.FixedLenFeature([], tf.string),\n}\n\ndef _parse_image_function(example_proto):\n  # Parse the input tf.train.Example proto using the dictionary above.\n  return tf.io.parse_single_example(example_proto, feature_description)\n\nparsed_tile_dataset = tile_dataset.map(_parse_image_function)","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:55:28.613532Z","iopub.execute_input":"2022-08-07T05:55:28.614619Z","iopub.status.idle":"2022-08-07T05:55:28.721417Z","shell.execute_reply.started":"2022-08-07T05:55:28.614564Z","shell.execute_reply":"2022-08-07T05:55:28.720499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"As printed before, Fold 0 has ID 10611. We parse all examples with tile id start with \"10611\", decode their tile images, mask images, get tile positions.","metadata":{}},{"cell_type":"code","source":"IMAGE_ID = \"10611\"\ntile_arrs = []\nmsk_tile_arrs = []\ntile_pos_list = []\nlabel_list = []\nclass_name_list = []\nfor record_features in parsed_tile_dataset:\n    if record_features['tile_id'].numpy().decode(encoding='utf-8').split(\"_\")[0] == IMAGE_ID:\n        tile_arrs.append(tf.io.decode_png(record_features['image']).numpy())\n        msk_tile_arrs.append(tf.io.decode_png(record_features['mask']).numpy())\n        tile_pos_list.append(record_features['tile_position'].numpy())\n        label_list.append(record_features['label'].numpy())\n        class_name_list.append(record_features['class_name'].numpy().decode(encoding='utf-8'))","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-08-07T05:55:43.718535Z","iopub.execute_input":"2022-08-07T05:55:43.718970Z","iopub.status.idle":"2022-08-07T05:55:43.859328Z","shell.execute_reply.started":"2022-08-07T05:55:43.718922Z","shell.execute_reply":"2022-08-07T05:55:43.857135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Plot tiles**","metadata":{}},{"cell_type":"code","source":"nb_tiles_sqrt = int(np.sqrt(len(tile_arrs)))\n\nfig, axes = plt.subplots(nrows=nb_tiles_sqrt, ncols=nb_tiles_sqrt, figsize=(9, 9))\nfig.suptitle(f'{label_list[0]} - {class_name_list[0]}')\nfor i, (tile_arr, msk_tile_arr) in enumerate(zip(tile_arrs, msk_tile_arrs)):\n    axes[i // nb_tiles_sqrt, i % nb_tiles_sqrt].imshow(tile_arr)\n    axes[i // nb_tiles_sqrt, i % nb_tiles_sqrt].imshow(msk_tile_arr, cmap=\"hot\", alpha=0.5)\n    axes[i // nb_tiles_sqrt, i % nb_tiles_sqrt].set_axis_off()\nfig.tight_layout()","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:55:58.869523Z","iopub.execute_input":"2022-08-07T05:55:58.869938Z","iopub.status.idle":"2022-08-07T05:56:00.060083Z","shell.execute_reply.started":"2022-08-07T05:55:58.869903Z","shell.execute_reply":"2022-08-07T05:56:00.059222Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Assemble tiles by tile positions(indices)**   ","metadata":{}},{"cell_type":"code","source":"assembled_image = np.zeros((RESIZE_DIM, RESIZE_DIM, 3), dtype=np.uint8)\nassembled_mask = np.zeros((RESIZE_DIM, RESIZE_DIM, 1), dtype=np.uint8)\nfor tile_i, tile_m, (i1, i2, j1, j2) in zip(tile_arrs, msk_tile_arrs, tile_pos_list):\n    assembled_image[i1:i2, j1:j2,:] = tile_i[:(i2 - i1), :(j2 - j1),:]\n    assembled_mask[i1:i2, j1:j2,:] = tile_m[:(i2 - i1), :(j2 - j1),:]\n\nplt.figure(figsize=(16, 8))\nplt.subplot(1, 2, 1)\nplt.imshow(assembled_image)\nplt.subplot(1, 2, 2)\nplt.imshow(assembled_mask)\nplt.tight_layout()\n\nprint(f\"Assembled image has shape {assembled_image.shape}.\\n\\n\")\nprint(f\"Assembled mask has shape {assembled_mask.shape}.\\n\\n\")","metadata":{"execution":{"iopub.status.busy":"2022-08-07T05:56:16.401966Z","iopub.execute_input":"2022-08-07T05:56:16.402357Z","iopub.status.idle":"2022-08-07T05:56:17.171549Z","shell.execute_reply.started":"2022-08-07T05:56:16.402324Z","shell.execute_reply":"2022-08-07T05:56:17.170486Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#comment this cell to keep png files\n!rm -r ./mask_img/\n!rm -r ./FOLD_*/","metadata":{"execution":{"iopub.status.busy":"2022-08-02T08:44:47.178838Z","iopub.execute_input":"2022-08-02T08:44:47.179222Z","iopub.status.idle":"2022-08-02T08:44:49.552603Z","shell.execute_reply.started":"2022-08-02T08:44:47.179191Z","shell.execute_reply":"2022-08-02T08:44:49.550971Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"If you find this notebook helpful, please give it a upvote. Thank you.","metadata":{}}]}