{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":117682,"databundleVersionId":15062069,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":765836,"databundleVersionId":15833078,"modelInstanceId":585137,"modelId":510647,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":732880,"databundleVersionId":15477237,"modelInstanceId":516822,"modelId":510647,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":660383,"databundleVersionId":14626934,"modelInstanceId":499479,"modelId":510647,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":672178,"databundleVersionId":14761443,"modelInstanceId":495238,"modelId":510647,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":673516,"databundleVersionId":14789938,"modelInstanceId":499479,"modelId":510647,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":655294,"databundleVersionId":14570138,"modelInstanceId":495238,"modelId":510647,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":665924,"databundleVersionId":14691066,"modelInstanceId":504051,"modelId":510647,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":765835,"databundleVersionId":15833018,"modelInstanceId":516822,"modelId":510647,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":681152,"databundleVersionId":14910215,"modelInstanceId":516822,"modelId":510647,"isSourceIdPinned":false},{"sourceType":"modelInstanceVersion","sourceId":674747,"databundleVersionId":14811492,"modelInstanceId":503784,"modelId":510647,"isSourceIdPinned":false},{"sourceType":"kernelVersion","sourceId":290917305,"isSourceIdPinned":false}],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Training Notebooks\n\n- [Vesuvius Surface 3D Detection in Keras-JAX](https://www.kaggle.com/code/ipythonx/vesuvius-surface-3d-detection-in-jax)\n- [Vesuvius Surface 3D Detection in PyTorch](https://www.kaggle.com/code/ipythonx/vesuvius-surface-3d-detection-in-pytorch)\n- [Vesuvius Surface 3D Detection in PyTorch Lightning](https://www.kaggle.com/code/ipythonx/train-vesuvius-surface-3d-detection-in-lightning)\n- [[WIP] Vesuvius Surface 2.5D Detection](https://www.kaggle.com/code/ipythonx/wip-vesuvius-surface-2-5d-detection)\n\n**Note**\n1. The inference code below is adapted from the **Keras-JAX** version. The PyTorch and Lightning implementations follow the same workflow. Training was performed on a single Tesla T4 (16 GB VRAM) with extended epochs.\n2. Both the training and inference pipelines are implemented using [`medicai`](https://github.com/innat/medic-ai), a **Keras 3** based multi-backend medical ML library designed for 2D and 3D classification and segmentation tasks. However, please note, `medicai` project is still new and actively evolving.","metadata":{}},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"from IPython.display import clear_output\n\nvar=\"/kaggle/input/vsdetection-packages-offline-installer-only/whls\"\n!pip install \\\n  \"$var\"/keras_nightly-*.whl \\\n  \"$var\"/tifffile-*.whl \\\n  \"$var\"/imagecodecs-*.whl \\\n  \"$var\"/medicai-*.whl \\\n  --no-index \\\n  --find-links \"$var\"\n\nclear_output()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:01:51.872683Z","iopub.execute_input":"2026-02-25T20:01:51.872968Z","iopub.status.idle":"2026-02-25T20:02:00.368932Z","shell.execute_reply.started":"2026-02-25T20:01:51.872946Z","shell.execute_reply":"2026-02-25T20:02:00.368064Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n\nimport keras\nfrom keras import ops\nfrom medicai.transforms import (\n    Compose,\n    ScaleIntensityRange,\n    NormalizeIntensity\n)\nfrom medicai.models import SegFormer, TransUNet\nfrom medicai.utils.inference import SlidingWindowInference\n\nimport numpy as np\nimport pandas as pd\nimport zipfile\nimport tifffile\nimport scipy.ndimage as ndi\nfrom skimage.morphology import remove_small_objects\nfrom matplotlib import pyplot as plt\n\nkeras.config.backend(), keras.version()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:25:27.713396Z","iopub.execute_input":"2026-02-25T20:25:27.713856Z","iopub.status.idle":"2026-02-25T20:25:27.720983Z","shell.execute_reply.started":"2026-02-25T20:25:27.713829Z","shell.execute_reply":"2026-02-25T20:25:27.720285Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Dataset**","metadata":{}},{"cell_type":"code","source":"root_dir = \"/kaggle/input/vesuvius-challenge-surface-detection\"\ntest_dir = f\"{root_dir}/test_images\"\noutput_dir = \"/kaggle/working/submission_masks\"\nzip_path = \"/kaggle/working/submission.zip\"\nos.makedirs(output_dir, exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:02:25.047561Z","iopub.execute_input":"2026-02-25T20:02:25.048015Z","iopub.status.idle":"2026-02-25T20:02:25.051857Z","shell.execute_reply.started":"2026-02-25T20:02:25.047995Z","shell.execute_reply":"2026-02-25T20:02:25.051185Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df = pd.read_csv(f\"{root_dir}/test.csv\")\ntest_df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:02:25.052615Z","iopub.execute_input":"2026-02-25T20:02:25.052848Z","iopub.status.idle":"2026-02-25T20:02:25.102190Z","shell.execute_reply.started":"2026-02-25T20:02:25.052832Z","shell.execute_reply":"2026-02-25T20:02:25.101478Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Transformation**","metadata":{}},{"cell_type":"code","source":"def val_transformation(image):\n    data = {\"image\": image}\n    pipeline = Compose([\n        NormalizeIntensity(\n            keys=[\"image\"], \n            nonzero=True,\n            channel_wise=False\n        ),\n    ])\n    result = pipeline(data)\n    return result[\"image\"]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:02:25.103730Z","iopub.execute_input":"2026-02-25T20:02:25.103985Z","iopub.status.idle":"2026-02-25T20:02:25.108129Z","shell.execute_reply.started":"2026-02-25T20:02:25.103968Z","shell.execute_reply":"2026-02-25T20:02:25.107326Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Model**","metadata":{}},{"cell_type":"code","source":"tta=1\nnum_classes=3\ninput_shape=(160, 160, 160)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:03:19.274882Z","iopub.execute_input":"2026-02-25T20:03:19.275186Z","iopub.status.idle":"2026-02-25T20:03:19.279568Z","shell.execute_reply.started":"2026-02-25T20:03:19.275164Z","shell.execute_reply":"2026-02-25T20:03:19.278829Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_model_a():\n    # 0.551 (tta+pp)\n    model = TransUNet(\n        input_shape=(160, 160, 160, 1),\n        encoder_name='seresnext50',\n        classifier_activation=None, # For tta, do softmax later.\n        num_classes=3,\n    )\n    model.load_weights(\n        \"/kaggle/input/models/ipythonx/vsd-model/keras/transunet/4/transunet.seresnext50.160px.duel.weights.h5\"\n    )\n    return model\n\ndef get_model_b():\n    # 0.522 (tta+pp)\n    model = SegFormer(\n        input_shape=input_shape + (1,),\n        encoder_name='mit_b4',\n        classifier_activation=None, # For tta, do softmax later.\n        num_classes=num_classes,\n    )\n    model.load_weights(\n        \"/kaggle/input/models/ipythonx/vsd-model/keras/segformer/1/segformer.mit.b4.weights.h5\"\n    )\n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:05:52.801254Z","iopub.execute_input":"2026-02-25T20:05:52.801568Z","iopub.status.idle":"2026-02-25T20:05:52.806776Z","shell.execute_reply.started":"2026-02-25T20:05:52.801547Z","shell.execute_reply":"2026-02-25T20:05:52.806182Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"models = [get_model_a(), get_model_b()]\n[model.count_params() / 1e6 for model in models]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:05:53.201421Z","iopub.execute_input":"2026-02-25T20:05:53.201724Z","iopub.status.idle":"2026-02-25T20:06:01.975941Z","shell.execute_reply.started":"2026-02-25T20:05:53.201704Z","shell.execute_reply":"2026-02-25T20:06:01.975343Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for model in models:\n    print(model.instance_describe())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:07:13.169864Z","iopub.execute_input":"2026-02-25T20:07:13.170351Z","iopub.status.idle":"2026-02-25T20:07:13.230239Z","shell.execute_reply.started":"2026-02-25T20:07:13.170331Z","shell.execute_reply":"2026-02-25T20:07:13.229685Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Sliding Window Inference**","metadata":{}},{"cell_type":"code","source":"swis = [\n    SlidingWindowInference(\n        model=m,\n        num_classes=num_classes,\n        roi_size=input_shape,\n        sw_batch_size=1,\n        overlap=0.5,\n        mode=\"gaussian\",\n    )\n    for m in models\n]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:07:22.433187Z","iopub.execute_input":"2026-02-25T20:07:22.433686Z","iopub.status.idle":"2026-02-25T20:07:22.437514Z","shell.execute_reply.started":"2026-02-25T20:07:22.433665Z","shell.execute_reply":"2026-02-25T20:07:22.436828Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_volume(path):\n    vol = tifffile.imread(path)\n    vol = vol.astype(np.float32)\n    vol = vol[None, ..., None]\n    return vol","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:07:23.869107Z","iopub.execute_input":"2026-02-25T20:07:23.869394Z","iopub.status.idle":"2026-02-25T20:07:23.873256Z","shell.execute_reply.started":"2026-02-25T20:07:23.869374Z","shell.execute_reply":"2026-02-25T20:07:23.872619Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Test Time Augmentation (TTA)**","metadata":{}},{"cell_type":"code","source":"def predict_with_tta(inputs, swi):\n    logits = []\n\n    # Original\n    logits.append(swi(inputs))\n\n    # Flips (spatial only)\n    for axis in [1, 2, 3]:\n        img_f = np.flip(inputs, axis=axis)\n        p = swi(img_f)\n        p = np.flip(p, axis=axis)\n        logits.append(p)\n\n    # Axial rotations (H, W)\n    for k in [1, 2, 3]:\n        img_r = np.rot90(inputs, k=k, axes=(2, 3))\n        p = swi(img_r)\n        p = np.rot90(p, k=-k, axes=(2, 3))\n        logits.append(p)\n\n    mean_logits = np.mean(logits, axis=0)\n    return mean_logits","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:14:28.428526Z","iopub.execute_input":"2026-02-25T20:14:28.429189Z","iopub.status.idle":"2026-02-25T20:14:28.434243Z","shell.execute_reply.started":"2026-02-25T20:14:28.429161Z","shell.execute_reply":"2026-02-25T20:14:28.433670Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Post Processing**","metadata":{}},{"cell_type":"code","source":"# https://www.kaggle.com/code/choudharymanas/inference-baseline-transunet-lb-0-537\ndef build_anisotropic_struct(z_radius: int, xy_radius: int):\n    z, r = z_radius, xy_radius\n    if z == 0 and r == 0:\n        return None\n    if z == 0 and r > 0:\n        size = 2 * r + 1\n        struct = np.zeros((1, size, size), dtype=bool)\n        cy, cx = r, r\n        for dy in range(-r, r + 1):\n            for dx in range(-r, r + 1):\n                if dy * dy + dx * dx <= r * r:\n                    struct[0, cy + dy, cx + dx] = True\n        return struct\n    if z > 0 and r == 0:\n        struct = np.zeros((2 * z + 1, 1, 1), dtype=bool)\n        struct[:, 0, 0] = True\n        return struct\n    depth = 2 * z + 1\n    size = 2 * r + 1\n    struct = np.zeros((depth, size, size), dtype=bool)\n    cz, cy, cx = z, r, r\n    for dz in range(-z, z + 1):\n        for dy in range(-r, r + 1):\n            for dx in range(-r, r + 1):\n                if dy * dy + dx * dx <= r * r:\n                    struct[cz + dz, cy + dy, cx + dx] = True\n    return struct\n\ndef topo_postprocess(\n    probs,\n    T_low=0.90,\n    T_high=0.90,\n    z_radius=1,\n    xy_radius=0,\n    dust_min_size=100,\n):\n    # Step 1: 3D Hysteresis\n    strong = probs >= T_high\n    weak   = probs >= T_low\n\n    if not strong.any():\n        return np.zeros_like(probs, dtype=np.uint8)\n\n    struct_hyst = ndi.generate_binary_structure(3, 3)\n    mask = ndi.binary_propagation(\n        strong, mask=weak, structure=struct_hyst\n    )\n\n    if not mask.any():\n        return np.zeros_like(probs, dtype=np.uint8)\n\n    # Step 2: 3D Anisotropic Closing\n    if z_radius > 0 or xy_radius > 0:\n        struct_close = build_anisotropic_struct(z_radius, xy_radius)\n        if struct_close is not None:\n            mask = ndi.binary_closing(mask, structure=struct_close)\n\n    # Step 3: Dust Removal\n    if dust_min_size > 0:\n        mask = remove_small_objects(\n            mask.astype(bool), min_size=dust_min_size\n        )\n\n    return mask.astype(np.uint8)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:14:43.789367Z","iopub.execute_input":"2026-02-25T20:14:43.789677Z","iopub.status.idle":"2026-02-25T20:14:43.799053Z","shell.execute_reply.started":"2026-02-25T20:14:43.789656Z","shell.execute_reply":"2026-02-25T20:14:43.798305Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Weighted Ensembling**","metadata":{}},{"cell_type":"code","source":"weights = np.array([0.8, 0.2], dtype=np.float32)\ndef predict_weighted_ensemble(volume):\n    logits = []\n    for swi in swis:\n        logits.append(\n            predict_with_tta(volume, swi) if tta else swi(volume) \n        )\n    logits = np.stack(logits, axis=0)\n    weighted_logits = logits * weights[:, None, None, None, None, None]\n    mean_logits = weighted_logits.sum(axis=0) / weights.sum()\n    probs = ops.softmax(mean_logits, axis=-1)\n    return probs.argmax(axis=-1).astype(np.uint8).squeeze()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:14:49.644763Z","iopub.execute_input":"2026-02-25T20:14:49.645455Z","iopub.status.idle":"2026-02-25T20:14:49.650084Z","shell.execute_reply.started":"2026-02-25T20:14:49.645430Z","shell.execute_reply":"2026-02-25T20:14:49.649485Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Prediction and Zip Submission**","metadata":{}},{"cell_type":"code","source":"def inference_pipelines(\n    volume,\n    T_low=0.50,\n    T_high=0.90,\n    z_radius=1,\n    xy_radius=0,\n    dust_min_size=100,\n):\n    probs = predict_weighted_ensemble(volume)\n    final = topo_postprocess(\n        probs,\n        T_low=T_low,\n        T_high=T_high,\n        z_radius=z_radius,\n        xy_radius=xy_radius,\n        dust_min_size=dust_min_size,\n    )\n    return final","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:14:50.759193Z","iopub.execute_input":"2026-02-25T20:14:50.759879Z","iopub.status.idle":"2026-02-25T20:14:50.763940Z","shell.execute_reply.started":"2026-02-25T20:14:50.759852Z","shell.execute_reply":"2026-02-25T20:14:50.763322Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"with zipfile.ZipFile(\n    zip_path, \"w\", compression=zipfile.ZIP_DEFLATED\n) as z:\n    for image_id in test_df[\"id\"]:\n        tif_path = f\"{test_dir}/{image_id}.tif\"\n        \n        volume = load_volume(tif_path)\n        volume = val_transformation(volume)\n        output = inference_pipelines(\n            volume, T_low=0.35, T_high=0.85, z_radius=3, xy_radius=2\n        ) \n        out_path = f\"{output_dir}/{image_id}.tif\"\n        tifffile.imwrite(out_path, output.astype(np.uint8))\n\n        z.write(out_path, arcname=f\"{image_id}.tif\")\n        os.remove(out_path)\n\nprint(\"Submission ZIP:\", zip_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:14:51.711027Z","iopub.execute_input":"2026-02-25T20:14:51.711664Z","iopub.status.idle":"2026-02-25T20:18:28.996439Z","shell.execute_reply.started":"2026-02-25T20:14:51.711616Z","shell.execute_reply":"2026-02-25T20:18:28.995796Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Sample View**","metadata":{}},{"cell_type":"code","source":"def plot_sample(x, y, sample_idx=0, max_slices=16):\n    img = np.squeeze(x[sample_idx])  # make (D, H, W)\n    mask = np.squeeze(y[sample_idx])  # make (D, H, W)\n    D = img.shape[0]\n\n    # Decide which slices to plot\n    step = max(1, D // max_slices)\n    slices = range(0, D, step)\n\n    n_slices = len(slices)\n    fig, axes = plt.subplots(2, n_slices, figsize=(3*n_slices, 6))\n\n    for i, s in enumerate(slices):\n        axes[0, i].imshow(img[s], cmap='gray')\n        axes[0, i].set_title(f\"Slice {s}\")\n        axes[0, i].axis('off')\n\n        axes[1, i].imshow(mask[s], cmap='gray')\n        axes[1, i].set_title(f\"Mask {s}\")\n        axes[1, i].axis('off')\n\n    plt.suptitle(f\"Sample {sample_idx}\")\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:18:32.012658Z","iopub.execute_input":"2026-02-25T20:18:32.013198Z","iopub.status.idle":"2026-02-25T20:18:32.019234Z","shell.execute_reply.started":"2026-02-25T20:18:32.013172Z","shell.execute_reply":"2026-02-25T20:18:32.018669Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plot_sample(\n    volume.numpy(), output[None], sample_idx=0, max_slices=5\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-25T20:18:32.639404Z","iopub.execute_input":"2026-02-25T20:18:32.639700Z","iopub.status.idle":"2026-02-25T20:18:33.408343Z","shell.execute_reply.started":"2026-02-25T20:18:32.639677Z","shell.execute_reply":"2026-02-25T20:18:33.407514Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}