\nThis notebook presents the complete inference pipeline of a model trained for\nsurface detection (3D segmentation) using the MONAI UNETR architecture.\nIt includes the necessary preprocessing steps, model loading, and inference\nprocedures applied to 3D volumetric data, as well as 3D visualizations of the\nresulting prediction outputs to qualitatively assess the model’s performance.\n
\n","metadata":{}},{"cell_type":"markdown","source":"## 1. Dataset Analysis","metadata":{}},{"cell_type":"code","source":"!pip install /kaggle/input/unetr-vit/*.whl --no-index","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:38:31.619797Z","iopub.execute_input":"2026-01-29T14:38:31.620402Z","iopub.status.idle":"2026-01-29T14:38:39.666880Z","shell.execute_reply.started":"2026-01-29T14:38:31.620371Z","shell.execute_reply":"2026-01-29T14:38:39.666149Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport zipfile\nimport numpy as np\nimport torch\nimport tifffile as tiff\nfrom glob import glob\nfrom tqdm import tqdm\n\nfrom monai.networks.nets import UNETR\nimport matplotlib.pyplot as plt\nimport cv2\n\nimport imageio\nfrom mpl_toolkits.mplot3d import Axes3D","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:38:39.668431Z","iopub.execute_input":"2026-01-29T14:38:39.668678Z","iopub.status.idle":"2026-01-29T14:39:17.624812Z","shell.execute_reply.started":"2026-01-29T14:38:39.668649Z","shell.execute_reply":"2026-01-29T14:39:17.624154Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"image_path = \"/kaggle/input/vesuvius-challenge-surface-detection/train_images/1004283650.tif\"\nmask_path = \"/kaggle/input/vesuvius-challenge-surface-detection/train_labels/1004283650.tif\"\n\nimages = tiff.imread(image_path) \nmasks = tiff.imread(mask_path) \n\nprint(\"Image shape:\", images.shape)\nprint(\"Mask shape :\", masks.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:31:21.833660Z","iopub.execute_input":"2026-01-29T14:31:21.835226Z","iopub.status.idle":"2026-01-29T14:31:22.325870Z","shell.execute_reply.started":"2026-01-29T14:31:21.835169Z","shell.execute_reply":"2026-01-29T14:31:22.324580Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"images = images[:, ::2, ::2]\nmasks = masks[:, ::2, ::2]\n\nnum_slices = 25\nalpha_img = 0.15\nalpha_mask = 0.9\n\nD, H, W = images.shape\n\nslice_indices = np.linspace(\n D // 4,\n 3 * D // 4 - 1,\n num_slices\n).astype(int)\n\nslice_indices = np.clip(slice_indices, 0, D - 1)\n\nview_angles = [0, 90, 180, 270]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:31:22.327652Z","iopub.execute_input":"2026-01-29T14:31:22.329104Z","iopub.status.idle":"2026-01-29T14:31:22.336488Z","shell.execute_reply.started":"2026-01-29T14:31:22.329044Z","shell.execute_reply":"2026-01-29T14:31:22.335100Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig = plt.figure(figsize=(25, 6))\n\nfor i, azim in enumerate(view_angles):\n ax = fig.add_subplot(1, len(view_angles), i + 1, projection=\"3d\")\n\n for z in slice_indices:\n Y, X = np.mgrid[0:H, 0:W]\n Z = np.ones_like(X) * z\n\n img_slice = images[z]\n img_norm = (img_slice - img_slice.min()) / (np.ptp(img_slice) + 1e-6)\n\n colors_img = plt.cm.gray(img_norm)\n colors_img[..., -1] = alpha_img\n\n ax.plot_surface(\n X, Y, Z,\n facecolors=colors_img,\n rstride=4,\n cstride=4,\n shade=False\n )\n papyrus_mask = masks[z] == 1\n\n if papyrus_mask.any():\n ax.scatter(\n X[papyrus_mask],\n Y[papyrus_mask],\n Z[papyrus_mask],\n c=\"red\",\n s=3,\n alpha=alpha_mask\n )\n\n ax.set_xlim(0, 256)\n ax.set_ylim(0, 256)\n ax.set_zlim(0, D)\n\n ax.set_xlabel(\"X\")\n ax.set_ylabel(\"Y\")\n ax.set_zlabel(\"Slice\")\n\n ax.view_init(elev=25, azim=azim)\n ax.set_title(f\"Z rotation {azim}°\")\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:31:22.770100Z","iopub.execute_input":"2026-01-29T14:31:22.771108Z","iopub.status.idle":"2026-01-29T14:32:10.355017Z","shell.execute_reply.started":"2026-01-29T14:31:22.771067Z","shell.execute_reply":"2026-01-29T14:32:10.353633Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Inference\n\n\nThe model was trained using full-resolution 320×320×320 volumetric inputs\nin FP32 precision. However, during inference in the Kaggle environment,\nGPU memory limitations may prevent direct execution at full precision.\nTherefore, FP16 (half-precision) inference is employed to reduce memory\nconsumption while preserving the trained model structure and enabling\nsuccessful inference under the available GPU constraints.\n
\n","metadata":{}},{"cell_type":"code","source":"IMG_SIZE = (320, 320, 320)\nCKPT_PATH = \"/kaggle/input/unetr-vit/unetr_papyrus.pth\"\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nmodel = UNETR(\n in_channels=1,\n out_channels=2,\n img_size=IMG_SIZE,\n feature_size=16,\n hidden_size=768,\n mlp_dim=3072,\n num_heads=12,\n norm_name=\"instance\",\n res_block=True,\n).to(device)\n\ncheckpoint = torch.load(CKPT_PATH, map_location=device)\nmodel.load_state_dict(checkpoint[\"model_state_dict\"])\nmodel.eval()\n\n# FP16 (half-precision)\nmodel = model.half()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:40:22.941136Z","iopub.execute_input":"2026-01-29T14:40:22.941779Z","iopub.status.idle":"2026-01-29T14:40:33.704352Z","shell.execute_reply.started":"2026-01-29T14:40:22.941745Z","shell.execute_reply":"2026-01-29T14:40:33.703504Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TEST_DIR = \"/kaggle/input/vesuvius-challenge-surface-detection/test_images\"\nOUT_DIR = \"predictions\"\nos.makedirs(OUT_DIR, exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:40:33.706149Z","iopub.execute_input":"2026-01-29T14:40:33.706510Z","iopub.status.idle":"2026-01-29T14:40:33.710292Z","shell.execute_reply.started":"2026-01-29T14:40:33.706485Z","shell.execute_reply":"2026-01-29T14:40:33.709518Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.backends.cudnn.benchmark = True\ntorch.backends.cuda.matmul.allow_tf32 = True\ntorch.set_float32_matmul_precision(\"high\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:40:33.711340Z","iopub.execute_input":"2026-01-29T14:40:33.711602Z","iopub.status.idle":"2026-01-29T14:40:33.726165Z","shell.execute_reply.started":"2026-01-29T14:40:33.711577Z","shell.execute_reply":"2026-01-29T14:40:33.725275Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict_3d_volume(model, tif_path):\n vol = tiff.imread(tif_path).astype(np.float32)\n\n # normalize\n vol = (vol - vol.min()) / (vol.max() - vol.min() + 1e-6)\n\n # (1, 1, D, H, W)\n vol = torch.from_numpy(vol).unsqueeze(0).unsqueeze(0).to(device).half()\n\n with torch.inference_mode():\n with torch.cuda.amp.autocast(dtype=torch.float16):\n logits = model(vol)\n pred = torch.argmax(logits, dim=1)\n\n return pred.squeeze(0).cpu().numpy().astype(np.uint8)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:40:33.727924Z","iopub.execute_input":"2026-01-29T14:40:33.728230Z","iopub.status.idle":"2026-01-29T14:40:33.741556Z","shell.execute_reply.started":"2026-01-29T14:40:33.728206Z","shell.execute_reply":"2026-01-29T14:40:33.740826Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_images = sorted(glob(os.path.join(TEST_DIR, \"*.tif\")))\nprint(f\" Found {len(test_images)} test volumes\")\n\nfor img_path in tqdm(test_images):\n image_id = os.path.splitext(os.path.basename(img_path))[0]\n\n pred_mask = predict_3d_volume(model, img_path)\n\n save_path = os.path.join(OUT_DIR, f\"predicted_{image_id}.tif\")\n tiff.imwrite(save_path, pred_mask)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:40:33.742595Z","iopub.execute_input":"2026-01-29T14:40:33.742982Z","iopub.status.idle":"2026-01-29T14:40:52.018640Z","shell.execute_reply.started":"2026-01-29T14:40:33.742947Z","shell.execute_reply":"2026-01-29T14:40:52.017851Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# visualize results\nimport tifffile as tiff\nimport matplotlib.pyplot as plt\nimport numpy as np\n\npred = tiff.imread(\"/kaggle/working/predictions/predicted_1407735.tif\") # (Z, H, W)\n\nnum_slices = pred.shape[0]\n\nindices = np.linspace(0, num_slices - 1, 25, dtype=int)\n\nfig, axes = plt.subplots(5, 5, figsize=(10, 10))\n\nfor ax, idx in zip(axes.flat, indices):\n ax.imshow(pred[idx])\n ax.set_title(f\"Slice {idx}\", fontsize=8)\n ax.axis(\"off\")\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-29T14:41:13.517535Z","iopub.execute_input":"2026-01-29T14:41:13.518333Z","iopub.status.idle":"2026-01-29T14:41:14.856010Z","shell.execute_reply.started":"2026-01-29T14:41:13.518299Z","shell.execute_reply":"2026-01-29T14:41:14.855260Z"}},"outputs":[],"execution_count":null}]}