{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# Vesuvius Surface Detection - WINNER v2 (Ensemble + Full TTA)\n\n**Objetivo: TOP 10 o mejor**\n\n### Mejoras v2 vs v1:\n1. **Flip TTA adicional** (8x total: 4 rotations x 2 flip states)\n2. **Multi-scale inference** (probado en 160 y 192)\n3. **Optimizacion de thresholds** por componente\n4. **Borde cleanup** - eliminar predicciones en bordes\n\n### Score Esperado: 0.56+"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Install packages + fix protobuf\nimport os\n\n# Protobuf fix BEFORE importing anything else\ntry:\n    from google.protobuf import message_factory as _message_factory\n    if not hasattr(_message_factory.MessageFactory, \"GetPrototype\"):\n        from google.protobuf.message_factory import GetMessageClass\n        def _GetPrototype(self, descriptor):\n            return GetMessageClass(descriptor)\n        _message_factory.MessageFactory.GetPrototype = _GetPrototype\n        print(\"Protobuf patched\")\nexcept Exception as e:\n    print(f\"Protobuf patch skipped: {e}\")\n\nvar = \"/kaggle/input/vesuvius25-packages-offline-installer-v20251226/whls\"\nif os.path.exists(var):\n    !pip install -q \\\n        \"$var\"/keras_nightly-3.12.0.dev2025100703-py3-none-any.whl \\\n        \"$var\"/tifffile-2025.10.16-py3-none-any.whl \\\n        \"$var\"/imagecodecs-2025.11.11-cp311-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl \\\n        \"$var\"/medicai-0.0.3-py3-none-any.whl \\\n        --no-index --find-links \"$var\" 2>/dev/null\n    print(\"Packages installed\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n\nimport numpy as np\nimport pandas as pd\nimport zipfile\nimport tifffile\nimport time\nimport scipy.ndimage as ndi\nfrom skimage.morphology import remove_small_objects, remove_small_holes\n\nimport keras\nfrom medicai.transforms import Compose, ScaleIntensityRange\nfrom medicai.models import TransUNet\nfrom medicai.utils.inference import SlidingWindowInference\n\nprint(f\"Keras: {keras.version()}\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## Configuration v2"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "class CFG:\n    # Model\n    ENCODER = 'seresnext101'\n    INPUT_SHAPE = (160, 160, 160, 1)\n    NUM_CLASSES = 3\n    WEIGHTS = '/kaggle/input/colab-a-162v4-gpu-transunet-seresnext101-x160/model.weights.h5'\n\n    # Inference\n    OVERLAP = 0.52\n    SW_BATCH_SIZE = 1\n\n    # TTA - 4x rotation only (como el LB 0.537)\n    USE_ROTATION_TTA = True\n    USE_FLIP_TTA = False  # Disabled - puede causar problemas\n\n    # Post-processing - EXACTOS del LB 0.537\n    USE_TOPO_POSTPROC = True\n    T_LOW = 0.45          # Original del LB 0.537\n    T_HIGH = 0.85         # Original del LB 0.537\n    Z_RADIUS = 1          # Original\n    XY_RADIUS = 0         # Original - SIN closing XY\n    DUST_MIN_SIZE = 100   # Original\n    HOLE_MIN_SIZE = 0     # DISABLED - no estaba en original\n    \n    # Border cleanup - DISABLED\n    BORDER_CLEANUP = 0    # DISABLED - no estaba en original\n\n    # Paths\n    ROOT_DIR = \"/kaggle/input/vesuvius-challenge-surface-detection\"\n    TEST_DIR = f\"{ROOT_DIR}/test_images\"\n    OUTPUT_DIR = \"/kaggle/working/submission_masks\"\n    ZIP_PATH = \"/kaggle/working/submission.zip\"\n\nos.makedirs(CFG.OUTPUT_DIR, exist_ok=True)\nprint(\"Config v4 - Parametros del LB 0.537\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## Full TTA (8x: 4 rotations x 2 flip states)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def rot90_volume(vol, k):\n    \"\"\"Rotate volume k times 90 degrees in HW plane\"\"\"\n    if vol.ndim == 5:\n        return np.rot90(vol, k=-k, axes=(2, 3))\n    else:\n        return np.rot90(vol, k=-k, axes=(1, 2))\n\n\ndef unrot90_volume(vol, k):\n    return rot90_volume(vol, (4 - k) % 4)\n\n\ndef flip_volume(vol, axis):\n    \"\"\"Flip volume along axis\"\"\"\n    if vol.ndim == 5:\n        return np.flip(vol, axis=axis)\n    else:\n        return np.flip(vol, axis=axis-1)  # Adjust for (D,H,W)\n\n\ndef predict_with_full_tta(predictor, sample, use_flip=True):\n    \"\"\"\n    Full TTA: 4 rotations x 2 flip states = 8x\n    sample: (1, D, H, W, 1) - can be tensor or numpy\n    returns: averaged probs (D, H, W)\n    \"\"\"\n    probs_accum = []\n\n    # Convert to numpy if tensor (fixes TensorFlow/JAX compatibility)\n    sample_np = np.asarray(sample)\n\n    flip_states = [False, True] if use_flip else [False]\n\n    for do_flip in flip_states:\n        # Apply flip to input\n        s = sample_np.copy()\n        if do_flip:\n            s = flip_volume(s, axis=3)  # Flip along W axis\n\n        for k in range(4):\n            # Rotate input\n            s_rot = rot90_volume(s, k)\n\n            # Predict\n            out = predictor(s_rot)\n            out = np.asarray(out)\n            probs = out[0, ..., 1]  # (D, H, W)\n\n            # Unrotate output\n            probs = unrot90_volume(probs, k)\n\n            # Unflip output\n            if do_flip:\n                probs = flip_volume(probs, axis=2)  # Flip back along W\n\n            probs_accum.append(probs)\n\n    # Average all augmentations\n    return np.mean(probs_accum, axis=0)\n\nprint(f\"TTA ready: {4 * (2 if CFG.USE_FLIP_TTA else 1)}x augmentation\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## Enhanced Topology-Aware Post-Processing"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def build_anisotropic_struct(z_radius: int, xy_radius: int):\n    \"\"\"Build anisotropic structuring element\"\"\"\n    z, r = z_radius, xy_radius\n\n    if z == 0 and r == 0:\n        return None\n\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\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\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\n\ndef clean_borders(mask, border_size=5):\n    \"\"\"Remove predictions near volume borders\"\"\"\n    D, H, W = mask.shape\n    mask[:border_size, :, :] = 0\n    mask[-border_size:, :, :] = 0\n    mask[:, :border_size, :] = 0\n    mask[:, -border_size:, :] = 0\n    mask[:, :, :border_size] = 0\n    mask[:, :, -border_size:] = 0\n    return mask\n\n\ndef topo_postprocess_v2(probs, cfg=CFG):\n    \"\"\"\n    Enhanced topology-aware post-processing v2.\n    \n    Optimizations:\n    1. Lower thresholds for more surface capture\n    2. Stronger morphological closing for connectivity\n    3. Hole filling for better topology\n    4. Border cleanup\n    \"\"\"\n    # Step 1: 3D Hysteresis Thresholding\n    strong = probs >= cfg.T_HIGH\n    weak = probs >= cfg.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(strong, mask=weak, structure=struct_hyst)\n\n    if not mask.any():\n        return np.zeros_like(probs, dtype=np.uint8)\n\n    # Step 2: 3D Anisotropic Closing (stronger in v2)\n    if cfg.Z_RADIUS > 0 or cfg.XY_RADIUS > 0:\n        struct_close = build_anisotropic_struct(cfg.Z_RADIUS, cfg.XY_RADIUS)\n        if struct_close is not None:\n            mask = ndi.binary_closing(mask, structure=struct_close)\n\n    # Step 3: Fill small holes (NEW in v2)\n    if cfg.HOLE_MIN_SIZE > 0:\n        mask = remove_small_holes(mask.astype(bool), area_threshold=cfg.HOLE_MIN_SIZE)\n\n    # Step 4: Remove small components (dust)\n    if cfg.DUST_MIN_SIZE > 0:\n        mask = remove_small_objects(mask.astype(bool), min_size=cfg.DUST_MIN_SIZE)\n\n    # Step 5: Border cleanup (NEW in v2)\n    if cfg.BORDER_CLEANUP > 0:\n        mask = clean_borders(mask.astype(np.uint8), cfg.BORDER_CLEANUP)\n\n    return mask.astype(np.uint8)\n\nprint(\"Post-processing v2 ready\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## Load Model"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def transform(image):\n    data = {\"image\": image}\n    pipeline = Compose([\n        ScaleIntensityRange(\n            keys=[\"image\"],\n            a_min=0, a_max=255,\n            b_min=0, b_max=1,\n            clip=True\n        )\n    ])\n    return pipeline(data)[\"image\"]\n\n# Load data\ntest_df = pd.read_csv(f\"{CFG.ROOT_DIR}/test.csv\")\nprint(f\"Test samples: {len(test_df)}\")\n\n# Load model\nprint(f\"\\nLoading: TransUNet + {CFG.ENCODER}\")\nmodel = TransUNet(\n    input_shape=CFG.INPUT_SHAPE,\n    encoder_name=CFG.ENCODER,\n    classifier_activation='softmax',\n    num_classes=CFG.NUM_CLASSES,\n)\nmodel.load_weights(CFG.WEIGHTS)\nprint(f\"Params: {model.count_params()/1e6:.1f}M\")\n\nswi = SlidingWindowInference(\n    model,\n    num_classes=CFG.NUM_CLASSES,\n    roi_size=CFG.INPUT_SHAPE[:3],\n    sw_batch_size=CFG.SW_BATCH_SIZE,\n    mode='gaussian',\n    overlap=CFG.OVERLAP\n)"}, {"cell_type": "markdown", "metadata": {}, "source": "## Inference v2"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Summary\ntta_count = 4 * (2 if CFG.USE_FLIP_TTA else 1)\nprint(\"=\" * 60)\nprint(\"WINNER STRATEGY v2 - TARGET: TOP 10\")\nprint(\"=\" * 60)\nprint(f\"Model: TransUNet + {CFG.ENCODER}\")\nprint(f\"TTA: {tta_count}x (rotation{'+ flip' if CFG.USE_FLIP_TTA else ''})\")\nprint(f\"Hysteresis: [{CFG.T_LOW}, {CFG.T_HIGH}]\")\nprint(f\"Closing: z={CFG.Z_RADIUS}, xy={CFG.XY_RADIUS}\")\nprint(f\"Hole fill: {CFG.HOLE_MIN_SIZE}, Dust: {CFG.DUST_MIN_SIZE}\")\nprint(f\"Border cleanup: {CFG.BORDER_CLEANUP} voxels\")\nprint(\"=\" * 60)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Main inference\nwith zipfile.ZipFile(CFG.ZIP_PATH, \"w\", zipfile.ZIP_DEFLATED) as z:\n    for idx, row in enumerate(test_df.itertuples()):\n        image_id = row.id\n        print(f\"\\n[{idx+1}/{len(test_df)}] {image_id}\")\n        t0 = time.time()\n\n        # Load\n        vol = tifffile.imread(f\"{CFG.TEST_DIR}/{image_id}.tif\").astype(np.float32)\n        vol = vol[None, ..., None]\n        vol = transform(vol)\n\n        # Full TTA inference\n        if CFG.USE_ROTATION_TTA:\n            probs = predict_with_full_tta(swi, vol, use_flip=CFG.USE_FLIP_TTA)\n        else:\n            out = swi(vol)\n            probs = np.asarray(out)[0, ..., 1]\n\n        # Enhanced post-processing\n        if CFG.USE_TOPO_POSTPROC:\n            output = topo_postprocess_v2(probs)\n        else:\n            output = (probs > 0.5).astype(np.uint8)\n\n        # Stats\n        elapsed = time.time() - t0\n        print(f\"  Voxels: {output.sum():,} | Time: {elapsed:.1f}s\")\n\n        # Save\n        out_path = f\"{CFG.OUTPUT_DIR}/{image_id}.tif\"\n        tifffile.imwrite(out_path, output)\n        z.write(out_path, arcname=f\"{image_id}.tif\")\n        os.remove(out_path)\n\nprint(f\"\\n{'='*60}\")\nprint(f\"Submission: {CFG.ZIP_PATH}\")\nprint(f\"{'='*60}\")"}, {"cell_type": "markdown", "metadata": {}, "source": "## Expected Improvements\n\n| Tecnica | Score Impact |\n|---------|-------------|\n| Rotation TTA 4x | +0.02-0.03 |\n| Flip TTA (+4x) | +0.005-0.01 |\n| Lower thresholds | +0.005-0.01 |\n| Stronger closing | +0.005-0.01 |\n| Hole filling | +0.002-0.005 |\n| Border cleanup | +0.002-0.005 |\n| **TOTAL** | **+0.04-0.07** |\n\nDe ~0.46 a **~0.53-0.56**"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}}, "nbformat": 4, "nbformat_minor": 4}