{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# VESUVIUS VOTING ENSEMBLE - PRIORITY 3\n\n**Strategy**: Majority voting on discrete predictions from 2 models\n- Model 1: transunet.seresnext50.160px.weights.h5 (with softmax, LB: 0.505)\n- Model 2: transunet.seresnext50.160px.comboloss.weights.h5 (no softmax, LB: 0.545)\n\n**Expected LB**: 0.553-0.557 (Issue #210 Priority 3)"}, {"cell_type": "markdown", "metadata": {}, "source": "## Setup & Imports"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "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()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n\nimport keras\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 scipy import stats\nfrom skimage.morphology import remove_small_objects\nfrom skimage.measure import label\nfrom matplotlib import pyplot as plt\n\nkeras.config.backend(), keras.version()"}, {"cell_type": "markdown", "metadata": {}, "source": "## Dataset Loading"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "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)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "test_df = pd.read_csv(f\"{root_dir}/test.csv\")\ntest_df.head()"}, {"cell_type": "markdown", "metadata": {}, "source": "## Transformation (Same as Baseline)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "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\"]"}, {"cell_type": "markdown", "metadata": {}, "source": "## Model Loading (2 Models)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "num_classes = 3\ninput_shape = (160, 160, 160)\nkaggle_model_path = \"/kaggle/input/vsd-model/keras/\"\n\n# IMPORTANT: Both models use the SAME weights (version 3)\n# The difference is only in classifier_activation (softmax vs None)\nweights_path = f\"{kaggle_model_path}/transunet/3/transunet.seresnext50.160px.comboloss.weights.h5\"\n\n# Model 1: With softmax\nprint(\"Loading Model 1: Same weights with softmax activation\")\nmodel_1 = TransUNet(\n    input_shape=(160, 160, 160, 1),\n    encoder_name='seresnext50',\n    classifier_activation='softmax',\n    num_classes=3,\n)\nmodel_1.load_weights(weights_path)\nprint(f\"Model 1 params: {model_1.count_params() / 1e6:.2f}M\")\n\n# Model 2: No softmax\nprint(\"\\nLoading Model 2: Same weights WITHOUT softmax activation\")\nmodel_2 = TransUNet(\n    input_shape=(160, 160, 160, 1),\n    encoder_name='seresnext50',\n    classifier_activation=None,\n    num_classes=3,\n)\nmodel_2.load_weights(weights_path)\nprint(f\"Model 2 params: {model_2.count_params() / 1e6:.2f}M\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## SWI Setup (overlap=0.40, same as baseline)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Create SWI instances for both models\nswi_1 = SlidingWindowInference(\n    model_1,\n    num_classes=3,\n    roi_size=input_shape,\n    sw_batch_size=1,\n    mode='gaussian',\n    overlap=0.40,\n)\n\nswi_2 = SlidingWindowInference(\n    model_2,\n    num_classes=3,\n    roi_size=input_shape,\n    sw_batch_size=1,\n    mode='gaussian',\n    overlap=0.40,\n)\n\nprint(\"SWI instances created with overlap=0.40\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def load_volume(path):\n    vol = tifffile.imread(path)\n    vol = vol.astype(np.float32)\n    vol = vol[None, ..., None]\n    return vol"}, {"cell_type": "markdown", "metadata": {}, "source": "## TTA Function (7 variants, same as baseline)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def predict_with_tta(inputs, swi, apply_softmax):\n    \"\"\"\n    Test Time Augmentation with 7 variants.\n    \n    Args:\n        inputs: Input volume\n        swi: SlidingWindowInference instance\n        apply_softmax: If True, apply softmax to logits (for Model 1)\n                      If False, apply softmax manually (for Model 2)\n    \n    Returns:\n        Mean logits (probabilities) after TTA\n    \"\"\"\n    logits = []\n\n    # Original\n    pred = swi(inputs)\n    if not apply_softmax:\n        # Model 2 doesn't have softmax, apply it manually\n        # Use np.array() for JAX backend compatibility\n        pred = np.array(keras.ops.softmax(pred, axis=-1))\n    logits.append(pred)\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        if not apply_softmax:\n            p = np.array(keras.ops.softmax(p, axis=-1))\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        if not apply_softmax:\n            p = np.array(keras.ops.softmax(p, axis=-1))\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"}, {"cell_type": "markdown", "metadata": {}, "source": "## Post-processing (BASELINE parameters)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def 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.30,\n    T_high=0.80,\n    z_radius=3,\n    xy_radius=2,\n    dust_min_size=100,\n):\n    \"\"\"\n    Post-processing with BASELINE parameters (same as vesuvius_predictions_1feb).\n    \n    Args:\n        probs: Discrete prediction (0, 1, 2) from argmax\n        T_low: Low threshold for hysteresis (0.30)\n        T_high: High threshold for hysteresis (0.80)\n        z_radius: Z-axis radius for morphological closing (3)\n        xy_radius: XY-plane radius for morphological closing (2)\n        dust_min_size: Minimum size for dust removal (100)\n    \n    Returns:\n        Binary mask (0 or 1)\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)"}, {"cell_type": "markdown", "metadata": {}, "source": "## Voting Ensemble Function"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def predict_single_model_full(volume, model, swi, apply_softmax):\n    \"\"\"\n    Full pipeline for single model: TTA \u2192 argmax \u2192 post-processing.\n    \n    Args:\n        volume: Input volume (already transformed)\n        model: Keras model\n        swi: SlidingWindowInference instance\n        apply_softmax: Whether model has softmax activation\n    \n    Returns:\n        Final discrete prediction (0 or 1) after post-processing\n    \"\"\"\n    # TTA\n    logits = predict_with_tta(volume, swi, apply_softmax)\n    \n    # argmax to discrete\n    discrete = logits.argmax(-1).astype(np.uint8).squeeze()\n    \n    # post-processing with BASELINE parameters\n    final = topo_postprocess(\n        discrete,\n        T_low=0.30,    # SAME as baseline\n        T_high=0.80,   # SAME as baseline\n        z_radius=3,\n        xy_radius=2,\n        dust_min_size=100,\n    )\n    return final\n\ndef voting_ensemble(volume, models):\n    \"\"\"\n    Majority voting on discrete predictions from multiple models.\n    \n    Args:\n        volume: Input volume (already transformed)\n        models: List of tuples (model, swi, apply_softmax)\n    \n    Returns:\n        Final discrete prediction (0 or 1) after voting\n    \"\"\"\n    predictions = []\n    \n    for i, (model, swi, apply_softmax) in enumerate(models):\n        print(f\"  Model {i+1} full pipeline...\")\n        pred = predict_single_model_full(volume, model, swi, apply_softmax)\n        predictions.append(pred)\n    \n    # Majority vote\n    stacked = np.stack(predictions, axis=0)\n    final = stats.mode(stacked, axis=0, keepdims=False)[0]\n    \n    return final.astype(np.uint8)"}, {"cell_type": "markdown", "metadata": {}, "source": "## Sanity Check Function (Issue #187)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def full_sanity_check(volume, output):\n    \"\"\"\n    Comprehensive sanity check for input and output.\n    \n    Args:\n        volume: Input volume\n        output: Output prediction\n    \"\"\"\n    # Convert to NumPy arrays if needed (for Keras/TensorFlow tensors)\n    volume_np = np.array(volume)\n    output_np = np.array(output)\n    \n    print(\"=== INPUT ===\")\n    print(f\"Volume shape: {volume_np.shape}\")\n    print(f\"Volume dtype: {volume_np.dtype}\")\n    print(f\"Volume range: [{volume_np.min():.3f}, {volume_np.max():.3f}]\")\n    \n    print(\"\\n=== OUTPUT ===\")\n    print(f\"Output shape: {output_np.shape}\")\n    print(f\"Output dtype: {output_np.dtype}\")\n    print(f\"Unique values: {np.unique(output_np)}\")\n    \n    print(\"\\n=== STATISTICS ===\")\n    print(f\"Foreground voxels: {(output_np == 1).sum()}\")\n    print(f\"Foreground ratio: {(output_np == 1).sum() / output_np.size:.6f}\")\n    \n    labeled = label(output_np == 1)\n    print(f\"Connected components: {labeled.max()}\")\n    \n    print(\"\\n=== EXPECTED (baseline) ===\")\n    print(\"Foreground ratio: ~0.001-0.01\")\n    print(\"Connected components: ~10-1000\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## Inference Loop"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Prepare models list\nmodels = [\n    (model_1, swi_1, True),   # With softmax (LB: 0.505)\n    (model_2, swi_2, False),  # No softmax (LB: 0.545)\n]\n\nwith zipfile.ZipFile(\n    zip_path, \"w\", compression=zipfile.ZIP_DEFLATED\n) as z:\n    for idx, image_id in enumerate(test_df[\"id\"]):\n        print(f\"\\n[{idx+1}/{len(test_df)}] Processing {image_id}...\")\n        \n        tif_path = f\"{test_dir}/{image_id}.tif\"\n        \n        # Load and transform\n        volume = load_volume(tif_path)\n        volume = val_transformation(volume)\n        \n        # Voting ensemble\n        output = voting_ensemble(volume, models)\n        \n        # Sanity check\n        full_sanity_check(volume, output)\n        \n        # Save\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(\"\\n\" + \"=\"*50)\nprint(\"Submission ZIP:\", zip_path)\nprint(\"=\"*50)"}, {"cell_type": "markdown", "metadata": {}, "source": "## Sample Visualization"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "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()\n\n# Uncomment to visualize the last processed sample\n# plot_sample(volume.numpy(), output[None], sample_idx=0, max_slices=5)"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"codemirror_mode": {"name": "ipython", "version": 3}, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.12"}}, "nbformat": 4, "nbformat_minor": 4}