{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpu1vmV38","dataSources":[{"sourceId":61446,"databundleVersionId":6962461,"sourceType":"competition"},{"sourceId":7172666,"sourceType":"datasetVersion","datasetId":4144355},{"sourceId":154467554,"sourceType":"kernelVersion"}],"dockerImageVersionId":30616,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## 1. Load model¶\n","metadata":{"execution":{"iopub.status.busy":"2023-12-11T04:02:49.712423Z","iopub.execute_input":"2023-12-11T04:02:49.713089Z","iopub.status.idle":"2023-12-11T04:02:49.717165Z","shell.execute_reply.started":"2023-12-11T04:02:49.713058Z","shell.execute_reply":"2023-12-11T04:02:49.716096Z"}}},{"cell_type":"markdown","source":"## NOTE: The competition don't accept TPU based submissions!","metadata":{}},{"cell_type":"code","source":"from tensorflow.keras.layers import Conv3D, MaxPool3D, concatenate, Input, Dropout, PReLU, Conv3DTranspose, BatchNormalization, TimeDistributed, Conv2D, ConvLSTM2D\nfrom tensorflow.keras.models import Model\nimport tensorflow as tf\nfrom tqdm import tqdm  # Import tqdm\nimport glob\nimport os\nfrom tensorflow.keras.utils import to_categorical\n\n# Detect and initialize the TPU\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection\n    print('Running on TPU ', tpu.master())\nexcept ValueError:\n    raise BaseException('ERROR: Not connected to a TPU runtime; please see the previous cell in this notebook for instructions!')\n\ntf.config.experimental_connect_to_cluster(tpu)\ntf.tpu.experimental.initialize_tpu_system(tpu)\nstrategy = tf.distribute.experimental.TPUStrategy(tpu)\n\n\ndef unet_core(x, filter_size=8, kernel_size=(3, 3, 3)):\n    x = Conv3D(filters=filter_size,\n               kernel_size=kernel_size,\n               padding='same',\n               kernel_initializer='he_normal',dtype=tf.float32)(x)\n    x = BatchNormalization()(x)\n    x = PReLU()(x)\n    x = Conv3D(filters=filter_size,\n               kernel_size=kernel_size,\n               padding='same',\n               kernel_initializer='he_normal',dtype=tf.float32)(x)\n    x = BatchNormalization()(x)\n    x = PReLU()(x)\n    return x\n\ndef unet3d(patch_size, n_label):\n    # with distributed.scope():\n    input_layer = Input(shape=patch_size,dtype=tf.float32)\n    d1 = unet_core(input_layer, filter_size=96, kernel_size=(3, 3, 3))\n    l = MaxPool3D(strides=(2, 2, 2))(d1)\n    d2 = unet_core(l, filter_size=96*2, kernel_size=(3, 3, 3))\n    l = MaxPool3D(strides=(2, 2, 2))(d2)\n    d3 = unet_core(l, filter_size=96*4, kernel_size=(3, 3, 3))\n    l = MaxPool3D(strides=(2, 2, 2))(d3)\n    d4 = unet_core(l, filter_size=96*8, kernel_size=(3, 3, 3))\n    l = MaxPool3D(strides=(2, 2, 2))(d4)\n\n    b = unet_core(l, filter_size=96*16, kernel_size=(3, 3, 3))\n\n    l = Conv3DTranspose(filters=96*8, kernel_size=(2, 2, 2),  padding='same', strides=2, kernel_initializer='he_normal',dtype=tf.float32)(b)\n    l = concatenate([l, d4])\n    u4 = unet_core(l, filter_size=96*4, kernel_size=(3, 3, 3))\n    l = Conv3DTranspose(filters=192, kernel_size=(2, 2, 2),  padding='same', strides=2, kernel_initializer='he_normal',dtype=tf.float32)(u4)\n    l = concatenate([l, d3])\n    u3 = unet_core(l, filter_size=96*4, kernel_size=(3, 3, 3))\n    l = Conv3DTranspose(filters=192, kernel_size=(2, 2, 2),  padding='same', strides=2, kernel_initializer='he_normal',dtype=tf.float32)(u3)\n    l = concatenate([l, d2])\n    u2 = unet_core(l, filter_size=96*2, kernel_size=(3, 3, 3))\n    l = Conv3DTranspose(filters=96*2, kernel_size=(2, 2, 2),  padding='same', strides=2, kernel_initializer='he_normal',dtype=tf.float32)(u2)\n    l = concatenate([l, d1])\n    u1 = unet_core(l, filter_size=96, kernel_size=(3, 3, 3))\n    output_layer = Conv3D(filters=n_label, kernel_size=(1, 1, 1), activation='sigmoid')(u1)\n    # output_layer = CRF(n_label)\n    model = Model(input_layer, output_layer)\n    return model\nwith strategy.scope():\n    model = unet3d((64,64,64,1),2)\n    model.load_weights('/kaggle/input/training-patch-based-3d-unet-tf-tpu/unet_3d.h5')\n","metadata":{"scrolled":true,"execution":{"iopub.status.busy":"2023-12-11T15:09:08.414935Z","iopub.execute_input":"2023-12-11T15:09:08.415152Z","iopub.status.idle":"2023-12-11T15:10:26.865120Z","shell.execute_reply.started":"2023-12-11T15:09:08.415127Z","shell.execute_reply":"2023-12-11T15:10:26.863847Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2. Inference","metadata":{}},{"cell_type":"code","source":"import os\nimport numpy as np\nfrom PIL import Image\nimport tensorflow as tf\nimport glob\nfrom tqdm import tqdm\nimport pandas as pd\n\n# Parameters\npatch_size = (64, 64, 64)  # Example patch size\nstride = (60, 60, 60)  # Example stride\n\ndef patch_generator_3d(image, patch_size, stride):\n    \"\"\"\n    Generator function for creating 3D patches from an image.\n\n    :param image: 3D numpy array, the original image.\n    :param patch_size: Tuple of 3 integers, the size of the patches to generate.\n    :param stride: Tuple of 3 integers, the stride with which to slide the window\n                   across the image. \n    :return: Yields patches of the specified size.\n    \"\"\"\n    # Calculate the number of patches that can be extracted in each dimension\n    patch_indices = [range(0, image.shape[dim] - patch_size[dim] + 1, stride[dim]) for dim in range(3)]\n\n    for i in patch_indices[0]:\n        for j in patch_indices[1]:\n            for k in patch_indices[2]:\n                yield image[i:i + patch_size[0], j:j + patch_size[1], k:k + patch_size[2]]\n\n                \n# Create a TensorFlow Dataset from generator\ndef create_dataset(image, patch_size, stride, batch_size):\n    def normlize_mean_std(tmp):\n        tmp_std = np.std(tmp) + 0.0001\n        tmp_mean = np.mean(tmp)\n        tmp = (tmp - tmp_mean) / tmp_std\n        return tmp\n    \n    def generator():\n        for patch in patch_generator_3d(image, patch_size, stride):\n            yield normlize_mean_std(patch)\n\n    # Define the output types and shapes of the generator\n    output_types = tf.float32\n    output_shapes = tf.TensorShape(patch_size)\n\n    # Create the dataset\n    dataset = tf.data.Dataset.from_generator(\n        generator, output_types=output_types, output_shapes=output_shapes\n    )\n\n    # Batch the dataset\n    dataset = dataset.batch(batch_size)\n\n    return dataset\n\ndef reconstruct_from_binary_patches(patches, original_shape, patch_size, stride):\n    reconstructed = np.zeros(original_shape, dtype=np.int64)\n    patch_indices = [range(0, i - p + 1, s) for i, p, s in zip(original_shape, patch_size, stride)]\n    counts = np.zeros(original_shape, dtype=np.int64)  # Count the number of patches contributing to each voxel\n    votes_for_one = np.zeros(original_shape, dtype=np.int64)  # Count the number of '1' votes for each voxel\n\n    patch_idx = 0\n    for i in patch_indices[0]:\n        for j in patch_indices[1]:\n            for k in patch_indices[2]:\n                patch = patches[patch_idx]\n                counts[i:i + patch_size[0], j:j + patch_size[1], k:k + patch_size[2]] += 1\n                votes_for_one[i:i + patch_size[0], j:j + patch_size[1], k:k + patch_size[2]] += patch\n                patch_idx += 1\n\n    # Majority voting\n    half_counts = counts // 2\n    reconstructed = (votes_for_one > half_counts).astype(np.int64)\n\n    return reconstructed\n\ndef preprocess_image(image_path):\n    # Open the image, convert to numpy array, and normalize\n    image = np.array(Image.open(image_path))\n    image = image - np.min(image)\n    image = image / np.max(image)\n    image = (image*255).astype(np.uint8)\n    return image\n\n# Load images\ndef read_and_stack_images(folder):\n    file_paths = sorted([os.path.join(folder, f) for f in os.listdir(folder) if f.endswith('.tif')])\n    images = [np.array(preprocess_image(f)) for f in file_paths]\n    return np.stack(images, axis=2)\n\n# Pad image\ndef pad_image(image, patch_size):\n    pad_width = [(0, max(0, ps - sz)) for sz, ps in zip(image.shape, patch_size)]\n    return np.pad(image, pad_width, mode='constant'), pad_width\n\ndef patch_generator_3d(image, patch_size, stride):\n    \"\"\"\n    Generator function for creating 3D patches from an image.\n\n    :param image: 3D numpy array, the original image.\n    :param patch_size: Tuple of 3 integers, the size of the patches to generate.\n    :param stride: Tuple of 3 integers, the stride with which to slide the window\n                   across the image. \n    :return: Yields patches of the specified size.\n    \"\"\"\n    # Calculate the number of patches that can be extracted in each dimension\n    patch_indices = [range(0, image.shape[dim] - patch_size[dim] + 1, stride[dim]) for dim in range(3)]\n\n    for i in patch_indices[0]:\n        for j in patch_indices[1]:\n            for k in patch_indices[2]:\n                yield image[i:i + patch_size[0], j:j + patch_size[1], k:k + patch_size[2]]\n\n# Create a TensorFlow Dataset from generator\ndef create_dataset(image, patch_size, stride, batch_size):\n    def generator():\n        for patch in patch_generator_3d(image, patch_size, stride):\n            yield patch\n\n    # Define the output types and shapes of the generator\n    output_types = tf.float32\n    output_shapes = tf.TensorShape(patch_size)\n\n    # Create the dataset\n    dataset = tf.data.Dataset.from_generator(\n        generator, output_types=output_types, output_shapes=output_shapes\n    )\n\n    # Batch the dataset\n    dataset = dataset.batch(batch_size)\n\n    return dataset\n\n# Predict patches and reconstruct image\ndef predict_and_reconstruct(model, image, patch_size, stride):\n    with strategy.scope():\n        dataset = create_dataset(padded_image, patch_size, stride, 16)\n        predicted_patches = np.argmax(model.predict(dataset), axis=-1)\n    # Reconstruct the image from the predicted patches\n    # Make sure your reconstruction function can handle the new patch size\n    predicted_image = reconstruct_from_binary_patches(predicted_patches, image.shape, patch_size, stride)\n    \n    return predicted_image\n\ndef extract_patches_3d(image, patch_size, stride):\n    # Calculate the number of patches\n    patch_indices = [range(0, i - p + 1, s) for i, p, s in zip(image.shape, patch_size, stride)]\n\n    for i in patch_indices[0]:\n        for j in patch_indices[1]:\n            for k in patch_indices[2]:\n                yield image[i:i + patch_size[0], j:j + patch_size[1], k:k + patch_size[2]]\n\ntest_image_path=[\"/kaggle/input/blood-vessel-segmentation/train/\"]\n\ndef rle_encode(mask):\n    pixel = mask.flatten()\n    pixel = np.concatenate([[0], pixel, [0]])\n    run = np.where(pixel[1:] != pixel[:-1])[0] + 1\n    run[1::2] -= run[::2]\n    rle = ' '.join(str(r) for r in run)\n    if rle == '':\n        rle = '1 0'\n    return rle\n\ndef remove_padding(padded_image, original_shape, pad_width):\n    # Calculate the start and end indices for each dimension\n    start_indices = [pw[0] for pw in pad_width]\n    end_indices = [original_shape[dim] + pw[0] for dim, pw in enumerate(pad_width)]\n\n    # Slice the padded image to get the unpadded image\n    unpadded_image = padded_image[start_indices[0]:end_indices[0],\n                                  start_indices[1]:end_indices[1],\n                                  start_indices[2]:end_indices[2]]\n    return unpadded_image\n\n\nid=[]\nrle=[]\nfor image_path in test_image_path:\n    images = os.listdir(image_path)\n    for image in images[4:5]:\n        # Read and preprocess all slices at once\n        image_3d = read_and_stack_images(os.path.join(image_path, image, 'images'))\n        padded_image, pad_width = pad_image(image_3d, patch_size)\n        predicted_image = predict_and_reconstruct(model, padded_image, patch_size, stride)\n        unpadded_prediction = remove_padding(predicted_image, image_3d.shape, pad_width)\n        \n        # Convert prediction to mask (if applicable)\n        mask = (unpadded_prediction > 0.99).astype(np.uint8)  # Example threshold\n        \n        # Encode each slice (example: RLE encoding)\n        encoded_slices = [rle_encode(mask[..., i]) for i in range(mask.shape[2])]\n        \n        id.extend([image+'_'+img_slice.split('.')[0] for img_slice in os.listdir(os.path.join(image_path, image, 'images'))])\n        rle.extend(encoded_slices)\n        \ndf = pd.DataFrame({'id': id, 'rle': rle})\ndf.to_csv('submission.csv', index=False)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","scrolled":true,"execution":{"iopub.status.busy":"2023-12-11T15:33:20.014487Z","iopub.execute_input":"2023-12-11T15:33:20.014819Z","iopub.status.idle":"2023-12-11T15:39:59.020155Z","shell.execute_reply.started":"2023-12-11T15:33:20.014790Z","shell.execute_reply":"2023-12-11T15:39:59.018928Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n# mask = (unpadded_prediction > 0.5).astype(np.uint8)\nplt.imshow(mask[:,:,900])","metadata":{"execution":{"iopub.status.busy":"2023-12-11T15:44:47.987244Z","iopub.execute_input":"2023-12-11T15:44:47.987825Z","iopub.status.idle":"2023-12-11T15:44:48.399371Z","shell.execute_reply.started":"2023-12-11T15:44:47.987777Z","shell.execute_reply":"2023-12-11T15:44:48.398267Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}