{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":117682,"databundleVersionId":15062069}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Tensorflow Custom Data Pipeline for Heavy Image datasets\n\nSolution to the problems mentioned in this notebook: [Kaggle notebook for standard Tesnorflow data pipeline](https://www.kaggle.com/code/mohanarc/tensorflow-data-pipeline-demo)  \n\nFeel free to run this notebook in one go.\n\n[GitHub link for notebook](https://github.com/MohanaRC/computefriendly_MLdemo/blob/4c6fdba5a77ffb9056443c306a2f817a3ed70541/tf-data-pipeline-for-heavy-datasets.ipynb)  \n","metadata":{}},{"cell_type":"code","source":"!pip install imagecodecs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T10:34:48.547156Z","iopub.execute_input":"2026-03-12T10:34:48.547610Z","iopub.status.idle":"2026-03-12T10:34:54.051677Z","shell.execute_reply.started":"2026-03-12T10:34:48.547580Z","shell.execute_reply":"2026-03-12T10:34:54.050831Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import warnings\nfrom pathlib import Path\n\n# Declare the paths\nDATA_DIR = Path(\"/kaggle/input/competitions/vesuvius-challenge-surface-detection\")\nTRAIN_IMAGES_DIR = Path(\"/kaggle/input/competitions/vesuvius-challenge-surface-detection/train_images\")\nTRAIN_LABELS_DIR = Path(\"/kaggle/input/competitions/vesuvius-challenge-surface-detection/train_labels\")\nMODEL_INPUT_SIZE = (256, 256, 256)  # (depth, height, width) - standarize volume size to this\n# For training\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T10:34:54.053305Z","iopub.execute_input":"2026-03-12T10:34:54.053552Z","iopub.status.idle":"2026-03-12T10:34:54.058265Z","shell.execute_reply.started":"2026-03-12T10:34:54.053524Z","shell.execute_reply":"2026-03-12T10:34:54.057526Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import tensorflow as tf\nimport os\nimport tifffile as tiff\nimport numpy as np\n\nclass SurfaceDataModuleTF:\n    def __init__(self, images_dir, labels_dir, volume_shape=(256, 256, 256), batch_size=1):\n        self.images_dir = images_dir\n        self.labels_dir = labels_dir\n        self.volume_shape = volume_shape\n        self.batch_size = batch_size\n        self.filenames = sorted([f for f in os.listdir(images_dir) if f.endswith('.tif')])\n\n    def _resize_3d(self, volume):\n        \"\"\"Native TF 3D resize logic\"\"\"\n        # Resizing 3D volumes by 2D slices (D, H, W) -> (256, 256, 256)\n        volume = tf.expand_dims(volume, -1) # (D, H, W, 1)\n        # 1. Resize H, W\n        volume = tf.image.resize(volume, [self.volume_shape[1], self.volume_shape[2]])\n        # 2. Transpose to resize D\n        volume = tf.transpose(volume, perm=[1, 0, 2, 3]) # (H, D, W, 1)\n        volume = tf.image.resize(volume, [self.volume_shape[0], self.volume_shape[2]])\n        volume = tf.transpose(volume, perm=[1, 0, 2, 3]) # (D, H, W, 1)\n        return tf.squeeze(volume, -1)\n\n    def _load_and_process(self, filename):\n        \"\"\"Loads TIF and resizes on the fly\"\"\"\n        # Use py_function to bridge TiffFile into the TF graph\n        def read_tiff(f):\n            img = tiff.imread(os.path.join(self.images_dir, f.numpy().decode()))\n            return img.astype(np.float32) / 255.0\n            \n        # Wrap the Python function\n        vol = tf.py_function(read_tiff, [filename], tf.float32)\n        vol.set_shape([None, 320, 320]) # Assume original dimensions\n        \n        # Apply 3D resize\n        vol = self._resize_3d(vol)\n        return vol\n\n    def get_dataset(self):\n        ds = tf.data.Dataset.from_tensor_slices(self.filenames)\n        ds = ds.map(self._load_and_process, num_parallel_calls=tf.data.AUTOTUNE)\n        \n        # Disk Caching: Prevents crashes by offloading to local storage\n        ds = ds.cache('/tmp/vesuvius_cache')\n        ds = ds.batch(self.batch_size).prefetch(tf.data.AUTOTUNE)\n        return ds\n\n# Usage\ndm = SurfaceDataModuleTF(TRAIN_IMAGES_DIR, TRAIN_LABELS_DIR)\ntrain_ds = dm.get_dataset()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-03-12T10:34:54.059370Z","iopub.execute_input":"2026-03-12T10:34:54.059686Z","iopub.status.idle":"2026-03-12T10:35:19.345686Z","shell.execute_reply.started":"2026-03-12T10:34:54.059658Z","shell.execute_reply":"2026-03-12T10:35:19.344959Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_ds","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T10:35:19.346668Z","iopub.execute_input":"2026-03-12T10:35:19.347229Z","iopub.status.idle":"2026-03-12T10:35:19.352503Z","shell.execute_reply.started":"2026-03-12T10:35:19.347140Z","shell.execute_reply":"2026-03-12T10:35:19.351737Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for data in train_ds.take(1):  ## Visualize\n    print(f\"Data type: {type(data)}\") #data tye\n    vol = data[0] if isinstance(data, (tuple, list)) else data  ### access the tensor\n    slice_to_plot = vol[0, 0, :, :].numpy() ## get a slice from the 3D volume    \n    plt.figure(figsize=(6, 6))\n    plt.imshow(slice_to_plot, cmap=\"gray\")\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T10:36:03.903857Z","iopub.execute_input":"2026-03-12T10:36:03.904669Z","iopub.status.idle":"2026-03-12T10:36:07.511033Z","shell.execute_reply.started":"2026-03-12T10:36:03.904639Z","shell.execute_reply":"2026-03-12T10:36:07.510420Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T10:35:24.269528Z","iopub.status.idle":"2026-03-12T10:35:24.269769Z","shell.execute_reply.started":"2026-03-12T10:35:24.269659Z","shell.execute_reply":"2026-03-12T10:35:24.269672Z"}},"outputs":[],"execution_count":null}]}