{"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":"gpu","dataSources":[{"sourceId":117682,"databundleVersionId":14443416,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":281390877,"sourceType":"kernelVersion"},{"sourceId":143395,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":121490,"modelId":144634}],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Alternative YOLO-based Approach\n\nUsing a YOLO (You Only Look Once) model for this volumetric segmentation task offers several distinct advantages over traditional 3D CNNs or patch-based U-Net approaches:\n\n1.  **Inference Speed & Efficiency**: YOLO models are state-of-the-art in terms of speed. By treating the 3D volume as a stack of 2D slices, we can process vast amounts of data much faster than heavy 3D convolutional networks.\n\n2.  **Reduced Computational Cost**: 3D models require significant VRAM and computation power. This slice-by-slice 2D approach allows us to work with high-resolution inputs on standard GPUs (like those in Colab) without running out of memory.\n\n3.  **Leveraging Pre-trained Weights**: YOLOv11 comes pre-trained on massive datasets (COCO), providing a robust feature extractor out-of-the-box. This transfer learning can be more effective than training a 3D model from scratch, especially when labeled data is sparse.\n\n4.  **Simplified Pipeline**: Handling 2D images is generally less complex than managing 3D tensors. This simplifies data augmentation, debugging, and visualization, allowing for faster iteration cycles.","metadata":{"_kg_hide-input":true}},{"cell_type":"code","source":"!pip install --find-links=\"/kaggle/input/surface-package-scraper\" --no-index --no-deps imagecodecs ultralytics \"opencv-python<4.11\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T16:55:49.893729Z","iopub.execute_input":"2025-11-24T16:55:49.893936Z","iopub.status.idle":"2025-11-24T16:55:59.063852Z","shell.execute_reply.started":"2025-11-24T16:55:49.893918Z","shell.execute_reply":"2025-11-24T16:55:59.062950Z"},"_kg_hide-input":true,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport shutil\nimport imagecodecs\nimport numpy as np\nimport cv2\nimport tifffile\nimport torch\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\n\n# Define input and output directories\nINPUT_DIR = '/kaggle/input/vesuvius-challenge-surface-detection'\nDATASET_DIR = './vesuvius_yolo_dataset'\n# Device\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-24T16:55:59.065995Z","iopub.execute_input":"2025-11-24T16:55:59.066345Z","iopub.status.idle":"2025-11-24T16:56:03.854593Z","shell.execute_reply.started":"2025-11-24T16:55:59.066304Z","shell.execute_reply":"2025-11-24T16:56:03.853713Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 🔎 Sample","metadata":{}},{"cell_type":"code","source":"# Define paths based on the directory structure inspection\ntrain_images_dir = os.path.join(INPUT_DIR, 'train_images')\ntrain_labels_dir = os.path.join(INPUT_DIR, 'train_labels')\n\n# Get all .tif files\nall_files = [f for f in os.listdir(train_images_dir) if f.endswith('.tif')]\nif not all_files:\n    raise FileNotFoundError(f\"No .tif files found in {train_images_dir}\")\n\nprint(f\"Found {len(all_files)} files to inspect.\")\n\nfile_path = all_files[0]\nimg_path = os.path.join(train_images_dir, file_path)\nlabel_path = os.path.join(train_labels_dir, file_path)\n\nprint(f\"\\nInspecting file: {file_path}\")\n\n# Load data\nimage = tifffile.imread(img_path)\nlabel = tifffile.imread(label_path)\n\n# Inspect\nfor name, data in [('Image', image), ('Label', label)]:\n    print(f\"\\n{name} Stats:\")\n    print(f\"  Shape: {data.shape}\")\n    print(f\"  Dtype: {data.dtype}\")\n    print(f\"  Min: {data.min()}, Max: {data.max()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T16:56:03.855532Z","iopub.execute_input":"2025-11-24T16:56:03.856515Z","iopub.status.idle":"2025-11-24T16:56:05.033781Z","shell.execute_reply.started":"2025-11-24T16:56:03.856481Z","shell.execute_reply":"2025-11-24T16:56:05.032838Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 📊 Dataset","metadata":{}},{"cell_type":"code","source":"def mask_to_yolo(mask, width, height):\n    \"\"\"\n    Converts a 2D label mask to YOLO polygon format.\n    The mask is expected to be 0 (background) and 1 (foreground).\n    Uses connected components to separate distinct objects.\n    \"\"\"\n    yolo_strings = []\n    \n    # Ensure mask is binary and uint8\n    mask = (mask > 0).astype(np.uint8)\n    \n    # Identify connected components (distinct objects)\n    num_labels, labels = cv2.connectedComponents(mask, connectivity=8)\n    \n    # Iterate through each detected object (label 0 is background)\n    for i in range(1, num_labels):\n        # Isolate the specific object\n        object_mask = (labels == i).astype(np.uint8)\n        \n        # Find the contour of this object\n        contours, _ = cv2.findContours(object_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n        \n        for contour in contours:\n            # Normalize coordinates to [0, 1]\n            contour = contour.squeeze()\n            \n            # Filter out invalid or too small contours\n            if contour.ndim < 2 or len(contour) < 3:\n                continue\n\n            normalized_coords = []\n            for point in contour:\n                x = point[0] / width\n                y = point[1] / height\n                normalized_coords.extend([x, y])\n\n            # We assume a single class (ID 0) for these objects\n            class_id = 0\n            \n            # Create YOLO string: class_id x1 y1 x2 y2 ...\n            coords_str = \" \".join([f\"{c:.6f}\" for c in normalized_coords])\n            yolo_strings.append(f\"{class_id} {coords_str}\")\n\n    return yolo_strings","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T16:56:05.034801Z","iopub.execute_input":"2025-11-24T16:56:05.035241Z","iopub.status.idle":"2025-11-24T16:56:05.042585Z","shell.execute_reply.started":"2025-11-24T16:56:05.035220Z","shell.execute_reply":"2025-11-24T16:56:05.041691Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def process_volume(vol_image, vol_label, base_name, out_img_dir, out_lbl_dir):\n    \"\"\"Iterates through slices of a 3D volume, saves valid images and YOLO annotations.\"\"\"\n    depth, height, width = vol_image.shape\n    saved_count = 0\n\n    for z in range(depth):\n        # Extract slice\n        img_slice = vol_image[z]\n        lbl_slice = vol_label[z]\n\n        # Skip slice if it contains label 2\n        if np.any(lbl_slice == 2):\n            continue\n\n        # Check if label contains any objects\n        if not np.any(lbl_slice > 0):\n            continue\n\n        # Generate annotations\n        annotations = mask_to_yolo(lbl_slice, width, height)\n        if not annotations:\n            continue\n\n        # Define filenames\n        file_id = f\"{base_name}_{z:04d}\"\n        img_path = os.path.join(out_img_dir, f\"{file_id}.png\")\n        txt_path = os.path.join(out_lbl_dir, f\"{file_id}.txt\")\n\n        # Save image (assuming vol_image is already uint8)\n        cv2.imwrite(img_path, img_slice)\n\n        # Save annotations\n        with open(txt_path, 'w') as f:\n            f.write(\"\\n\".join(annotations))\n\n        saved_count += 1\n\n    if saved_count:\n        print(f\"Processed {base_name}: Saved {saved_count} slices with objects.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T16:56:05.043495Z","iopub.execute_input":"2025-11-24T16:56:05.043911Z","iopub.status.idle":"2025-11-24T16:56:05.063904Z","shell.execute_reply.started":"2025-11-24T16:56:05.043888Z","shell.execute_reply":"2025-11-24T16:56:05.063042Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import random\nimport multiprocessing\nfrom functools import partial\n\n# 1. Define Directories\nTEMP_DIR = './temp_vesuvius_processing' # Temp storage for processing\n\n# Clean/Create Directories\nif os.path.exists(DATASET_DIR): shutil.rmtree(DATASET_DIR)\nif os.path.exists(TEMP_DIR): shutil.rmtree(TEMP_DIR)\n\n# Create Final Structure\nfor split in ['train', 'val']:\n    os.makedirs(os.path.join(DATASET_DIR, 'images', split), exist_ok=True)\n    os.makedirs(os.path.join(DATASET_DIR, 'labels', split), exist_ok=True)\n\n# Create Temp Structure\ntemp_images_dir = os.path.join(TEMP_DIR, 'images')\ntemp_labels_dir = os.path.join(TEMP_DIR, 'labels')\nos.makedirs(temp_images_dir, exist_ok=True)\nos.makedirs(temp_labels_dir, exist_ok=True)\n\n# 2. Process ALL Volumes\ntrain_images_dir = os.path.join(INPUT_DIR, 'train_images')\ntrain_labels_dir = os.path.join(INPUT_DIR, 'train_labels')\n\nall_files = [f for f in os.listdir(train_images_dir) if f.endswith('.tif')]\nprint(f\"Processing all {len(all_files)} volumes to identify valid data...\")\n\ndef process_wrapper(filename, train_images_dir, train_labels_dir, temp_img_dir, temp_lbl_dir):\n    \"\"\"Wrapper to load volume and process it.\"\"\"\n    img_path = os.path.join(train_images_dir, filename)\n    lbl_path = os.path.join(train_labels_dir, filename)\n\n    vol_img = tifffile.imread(img_path)\n    vol_lbl = tifffile.imread(lbl_path)\n    base_name = os.path.splitext(filename)[0]\n\n    # process_volume must be defined in previous cell\n    process_volume(vol_img, vol_lbl, base_name, temp_img_dir, temp_lbl_dir)\n\n# Prepare the partial function with fixed arguments\nprocess_func = partial(\n    process_wrapper,\n    train_images_dir=train_images_dir,\n    train_labels_dir=train_labels_dir,\n    temp_img_dir=temp_images_dir,\n    temp_lbl_dir=temp_labels_dir\n)\n\n# Run in parallel\nnum_workers = min(os.cpu_count(), 3)\nprint(f\"Using {num_workers} workers for multiprocessing.\")\n\nwith multiprocessing.Pool(num_workers) as pool:\n    list(tqdm(pool.imap_unordered(process_func, all_files), total=len(all_files)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T16:56:05.065032Z","iopub.execute_input":"2025-11-24T16:56:05.065379Z","iopub.status.idle":"2025-11-24T17:03:00.778933Z","shell.execute_reply.started":"2025-11-24T16:56:05.065345Z","shell.execute_reply":"2025-11-24T17:03:00.777875Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 3. Identify Usable Volumes and Split\n# We only care about volumes that produced at least one image\ngenerated_files = os.listdir(os.path.join(TEMP_DIR, 'images'))\nvolumes_with_objects = set()\n\nfor f in generated_files:\n    # Extract base_name (remove the last _sliceID segment)\n    # Filename format from process_volume: {base_name}_{slice_id}.jpg\n    # As noted, for this dataset base_name is just the first part\n    base_name = f.split('_')[0]\n    volumes_with_objects.add(base_name)\n\nvolumes_with_objects = list(volumes_with_objects)\nrandom.shuffle(volumes_with_objects)\n\nprint(f\"\\nFound {len(volumes_with_objects)} volumes containing objects.\")\n\nsplit_ratio = 0.8\nsplit_idx = int(len(volumes_with_objects) * split_ratio)\n\n# 4. Move Files to Final Destinations\nvol_to_split = {v: 'train' for v in volumes_with_objects[:split_idx]}\nvol_to_split.update({v: 'val' for v in volumes_with_objects[split_idx:]})\n\nprint(\"\\nDistributing files to train/val folders...\")\nfor img_file in tqdm(generated_files):\n    name, _ = os.path.splitext(img_file)\n    lbl_file = f\"{name}.txt\"\n    base_name = img_file.split('_')[0]\n    \n    if base_name not in vol_to_split:\n        continue\n        \n    split = vol_to_split[base_name]\n    \n    src_img = os.path.join(TEMP_DIR, 'images', img_file)\n    src_lbl = os.path.join(TEMP_DIR, 'labels', lbl_file)\n    dst_img = os.path.join(DATASET_DIR, 'images', split, img_file)\n    dst_lbl = os.path.join(DATASET_DIR, 'labels', split, lbl_file)\n    shutil.move(src_img, dst_img)\n    shutil.move(src_lbl, dst_lbl)\n\n# 5. Cleanup\nshutil.rmtree(TEMP_DIR)\nprint(\"Dataset creation complete.\")\n\nfor split in (\"train\", \"val\"):\n    for tp in (\"images\", \"labels\"):\n        print(f\"{tp} / {split}: {len(os.listdir(f'/kaggle/working/vesuvius_yolo_dataset/{tp}/{split}'))}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T17:03:00.781994Z","iopub.execute_input":"2025-11-24T17:03:00.782326Z","iopub.status.idle":"2025-11-24T17:03:01.300464Z","shell.execute_reply.started":"2025-11-24T17:03:00.782300Z","shell.execute_reply":"2025-11-24T17:03:01.299708Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import random\nimport matplotlib.pyplot as plt\n\ndef visualize_yolo_sample(img_path, lbl_path):\n    \"\"\"\n    Loads the image and YOLO labels, draws the polygons, and displays the result.\n    Takes direct paths to image and label files.\n    \"\"\"\n    # Load image\n    image = cv2.imread(img_path)\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n    h, w, _ = image.shape\n\n    filename = os.path.basename(img_path)\n    print(f\"Annotation: {lbl_path}\")\n\n    with open(lbl_path, 'r') as f:\n        lines = f.readlines()\n\n    for line in lines:\n        parts = line.strip().split()\n        if len(parts) < 3:\n            continue\n        coords = list(map(float, parts[1:]))\n        # Rescale coordinates\n        points = []\n        for i in range(0, len(coords), 2):\n            px = int(coords[i] * w)\n            py = int(coords[i+1] * h)\n            points.append([px, py])\n        points = np.array(points, dtype=np.int32).reshape((-1, 1, 2))\n        # Draw polygon\n        cv2.polylines(image, [points], isClosed=True, color=(255, 0, 0), thickness=1)\n\n    plt.figure(figsize=(10, 10))\n    plt.imshow(image)\n    plt.title(f\"Sample: {filename}\")\n    plt.axis('off')\n    plt.show()\n\n# Logic to pick a random file\nimg_dir = os.path.join(DATASET_DIR, 'images', 'train')\nlbl_dir = os.path.join(DATASET_DIR, 'labels', 'train')\n\nfiles = [f for f in os.listdir(img_dir) if f.endswith('.jpg') or f.endswith('.png')]\nrandom_file = random.choice(files)\nimg_path_sample = os.path.join(img_dir, random_file)\nlbl_file_sample = os.path.splitext(random_file)[0] + '.txt'\nlbl_path_sample = os.path.join(lbl_dir, lbl_file_sample)\n\nvisualize_yolo_sample(img_path_sample, lbl_path_sample)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T17:03:01.301213Z","iopub.execute_input":"2025-11-24T17:03:01.301420Z","iopub.status.idle":"2025-11-24T17:03:01.857566Z","shell.execute_reply.started":"2025-11-24T17:03:01.301404Z","shell.execute_reply":"2025-11-24T17:03:01.856646Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import yaml\n\n# Define the configuration for the YOLO dataset\nyaml_content = {\n    'path': os.path.abspath(DATASET_DIR), # Absolute path to the dataset root\n    'train': 'images/train',              # Relative path to training images\n    'val': 'images/val',                  # Relative path to validation images\n    'names': {\n        0: 'surface'                      # Class mapping\n    }\n}\n\n# Save to dataset.yaml\nyaml_path = 'dataset.yaml' # Save in current directory for easy access\n\nwith open(yaml_path, 'w') as f:\n    yaml.dump(yaml_content, f, sort_keys=False)\n\nprint(f\"Dataset YAML created at: {os.path.abspath(yaml_path)}\")\nprint(\"\\nFile Content:\")\nprint(yaml.dump(yaml_content, sort_keys=False))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T17:03:01.858495Z","iopub.execute_input":"2025-11-24T17:03:01.858742Z","iopub.status.idle":"2025-11-24T17:03:01.890215Z","shell.execute_reply.started":"2025-11-24T17:03:01.858715Z","shell.execute_reply":"2025-11-24T17:03:01.889413Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 🏋️ Training\n\nMain training pipeline...","metadata":{}},{"cell_type":"code","source":"from ultralytics import YOLO\n\n# Load a model\nmodel = YOLO(\"/kaggle/input/yolo11-seg/pytorch/default/1/yolo11m-seg.pt\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T17:03:01.891045Z","iopub.execute_input":"2025-11-24T17:03:01.891336Z","iopub.status.idle":"2025-11-24T17:04:06.408133Z","shell.execute_reply.started":"2025-11-24T17:03:01.891309Z","shell.execute_reply":"2025-11-24T17:04:06.407402Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train the model\nprint(\"Starting training...\")\nresults = model.train(\n    data=\"dataset.yaml\",  # path to dataset YAML created earlier\n    epochs=50,            # number of training epochs\n    imgsz=320,            # image size (matches our slice size)\n    batch=12,             # batch size\n    name=\"vesuvius_yolo11\" # save results to runs/detect/vesuvius_yolo11\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T17:04:06.408986Z","iopub.execute_input":"2025-11-24T17:04:06.409254Z","iopub.status.idle":"2025-11-24T17:16:07.646579Z","shell.execute_reply.started":"2025-11-24T17:04:06.409234Z","shell.execute_reply":"2025-11-24T17:16:07.645660Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import glob\nfrom IPython.display import Image, display\n\n# Find the latest training run directory\nruns_dir = 'runs/segment'\n\n# Path to the results plot\nfor run_img in sorted(glob.glob(f'{runs_dir}/*/*.png')):\n    print(run_img)\n    display(Image(filename=run_img))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T17:19:17.378699Z","iopub.execute_input":"2025-11-24T17:19:17.379352Z","iopub.status.idle":"2025-11-24T17:19:17.439959Z","shell.execute_reply.started":"2025-11-24T17:19:17.379304Z","shell.execute_reply":"2025-11-24T17:19:17.439100Z"},"_kg_hide-output":false,"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 🔮 Inference\n\nPredict on new 3D volumes...","metadata":{}},{"cell_type":"code","source":"from skimage.morphology import remove_small_objects, ball\nimport torch.nn.functional as F\n\ndef get_spherical_kernel(radius):\n    \"\"\"Generates a spherical kernel (structuring element) for 3D morphological operations.\"\"\"\n    # Generate boolean ball on CPU\n    kernel_np = ball(radius)\n    # Convert to float tensor: (1, 1, D, H, W)\n    kernel = torch.from_numpy(kernel_np.astype(np.float32)).unsqueeze(0).unsqueeze(0)\n    return kernel\n\ndef post_process_3d(\n    volume: np.ndarray,\n    min_size: int = 1000,\n    closing_radius: int = 5,\n    device: str = DEVICE,\n) -> np.ndarray:\n    \"\"\"Applies 3D morphological operations to clean up a segmentation volume using a spherical element.\n\n    Steps:\n    1. Performs morphological closing (Dilation -> Erosion) with a spherical kernel.\n    2. Removes small connected components.\n    \"\"\"\n    # Ensure input is boolean\n    binary_vol = volume > 0\n    clean_vol_np = binary_vol\n\n    # 1. Close gaps with Spherical Element (GPU accelerated)\n    if closing_radius > 0:\n        # print(f\"Closing gaps with spherical radius {closing_radius} (GPU accelerated)...\")\n        # Prepare Input: (1, 1, D, H, W)\n        input_tensor = torch.from_numpy(clean_vol_np.astype(np.float32)).unsqueeze(0).unsqueeze(0).to(device)\n        # Prepare Kernel\n        kernel = get_spherical_kernel(closing_radius).to(device)\n        # Dilation: (Input * Kernel) > 0\n        # We use padding=closing_radius to maintain the same spatial dimensions (same as 'same' padding)\n        dilated = (F.conv3d(input_tensor, kernel, padding=closing_radius) > 0).float()\n        # Erosion: 1 - ((1 - Dilated) * Kernel > 0)\n        # This relies on the duality: Erosion(A) = ~Dilation(~A)\n        eroded = 1.0 - (F.conv3d(1.0 - dilated, kernel, padding=closing_radius) > 0).float()\n        # Retrieve result\n        clean_vol_np = eroded.squeeze().cpu().numpy().astype(bool)\n\n    # 2. Remove small objects (CPU-bound)\n    # print(f\"Removing small objects < {min_size} voxels (CPU-bound)...\")\n    clean_vol_np = remove_small_objects(clean_vol_np, min_size=min_size)\n\n    return clean_vol_np.astype(np.uint8)","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-11-24T18:06:26.174370Z","iopub.execute_input":"2025-11-24T18:06:26.175253Z","iopub.status.idle":"2025-11-24T18:06:26.183633Z","shell.execute_reply.started":"2025-11-24T18:06:26.175228Z","shell.execute_reply":"2025-11-24T18:06:26.182708Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict_and_save_volume(volume_path, model, output_dir='.'):\n    \"\"\"Loads a volume, runs YOLO inference slice-by-slice, and saves the predicted mask volume.\"\"\"\n    base_name = Path(volume_path).stem\n    # print(f\"Processing volume: {base_name}...\")\n    volume = tifffile.imread(volume_path)\n    depth, height, width = volume.shape\n    predictions_list_2d = []\n\n    # Iterate over all slices\n    for i in tqdm(range(depth), desc=f\"Slices for {base_name}\", leave=False):\n        img_slice = volume[i]\n        # Preprocess: Normalize to uint8 (0-255) if needed\n        if img_slice.dtype != np.uint8:\n            img_slice = cv2.normalize(img_slice, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)\n        # Convert to BGR for YOLO\n        img_bgr = cv2.cvtColor(img_slice, cv2.COLOR_GRAY2BGR)\n        # Run inference\n        results = model.predict(img_bgr, verbose=False, retina_masks=True)\n        result = results[0]\n\n        # Initialize empty mask (background)\n        mask_slice = np.zeros((height, width), dtype=np.uint8)\n        if result.masks is not None:\n            masks_tensor = result.masks.data\n            if masks_tensor is not None:\n                # Combine all detected masks\n                combined_mask = torch.any(masks_tensor, dim=0).int().cpu().numpy()\n                # Resize if needed\n                combined_mask = cv2.resize(\n                    combined_mask.astype(np.uint8), (width, height), interpolation=cv2.INTER_NEAREST)\n                mask_slice = combined_mask.astype(np.uint8)\n        predictions_list_2d.append(mask_slice)\n\n    # Stack slices\n    predictions_volume = np.stack(predictions_list_2d, axis=0).astype(np.uint8)\n    predictions_volume = post_process_3d(\n        predictions_volume, min_size=20*20*20, closing_radius=7)\n\n    # Save volume\n    output_filename = os.path.join(output_dir, f\"{base_name}.tif\")\n    tifffile.imwrite(output_filename, predictions_volume)\n    print(f\"Saved prediction volume: {output_filename} (Shape: {predictions_volume.shape})\")\n    return output_filename","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T18:07:38.545175Z","iopub.execute_input":"2025-11-24T18:07:38.546219Z","iopub.status.idle":"2025-11-24T18:07:38.555163Z","shell.execute_reply.started":"2025-11-24T18:07:38.546179Z","shell.execute_reply":"2025-11-24T18:07:38.554309Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Main Execution Block\ntest_images_dir = os.path.join(INPUT_DIR, 'test_images')\ntest_volume_files = [f for f in os.listdir(test_images_dir) if f.endswith('.tif')]\nprint(f\"Found {len(test_volume_files)} test volumes.\")\n\ntest_filenames = []\nfor test_volume_file in tqdm(test_volume_files, desc=\"Total Progress\"):\n    full_path = os.path.join(test_images_dir, test_volume_file)\n    saved_file = predict_and_save_volume(full_path, model)\n    test_filenames.append(saved_file)\n\nprint(\"All volumes processed.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T18:07:40.238880Z","iopub.execute_input":"2025-11-24T18:07:40.239251Z","iopub.status.idle":"2025-11-24T18:08:01.085578Z","shell.execute_reply.started":"2025-11-24T18:07:40.239225Z","shell.execute_reply":"2025-11-24T18:08:01.084563Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_three_axis_cuts(image_vol_path, mask_vol_path):\n    \"\"\"Plots the middle slice of the XY, XZ, and YZ planes for both the image volume and the predicted mask.\"\"\"\n    print(f\"Visualizing cuts for: {os.path.basename(image_vol_path)}\")\n    # Load volumes\n    image_vol = tifffile.imread(image_vol_path)\n    mask_vol = tifffile.imread(mask_vol_path)\n    \n    # Get dimensions\n    d, h, w = image_vol.shape\n    z_mid, y_mid, x_mid = d // 2, h // 2, w // 2\n    \n    # Extract slices\n    slices = {\n        'XY Plane (Z-axis)': (image_vol[z_mid, :, :], mask_vol[z_mid, :, :]),\n        'XZ Plane (Y-axis)': (image_vol[:, y_mid, :], mask_vol[:, y_mid, :]),\n        'YZ Plane (X-axis)': (image_vol[:, :, x_mid], mask_vol[:, :, x_mid])\n    }\n    \n    fig, axes = plt.subplots(3, 2, figsize=(12, 15))\n    for i, (plane_name, (img_slice, mask_slice)) in enumerate(slices.items()):\n        # Image Volume\n        axes[i, 0].imshow(img_slice, cmap='gray')\n        axes[i, 0].set_title(f\"{plane_name} - Image Volume\")\n        axes[i, 0].axis('off')\n        \n        # Mask\n        axes[i, 1].imshow(mask_slice, cmap='gray')\n        axes[i, 1].set_title(f\"{plane_name} - Predicted Mask\")\n        axes[i, 1].axis('off')\n        \n    plt.tight_layout()\n    plt.show()\n\n# Visualize the first processed volume\nif test_filenames:\n    # We assume test_volume_files and test_filenames are aligned from the previous loop\n    # test_filenames contains the output paths we just saved\n    # test_volume_files contains the original filenames in test_images_dir\n    \n    image_path = os.path.join(test_images_dir, test_volume_files[0])\n    mask_path = test_filenames[0]\n    \n    plot_three_axis_cuts(image_path, mask_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T18:08:01.086910Z","iopub.execute_input":"2025-11-24T18:08:01.087230Z","iopub.status.idle":"2025-11-24T18:08:02.242789Z","shell.execute_reply.started":"2025-11-24T18:08:01.087211Z","shell.execute_reply":"2025-11-24T18:08:02.241768Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Export submission","metadata":{}},{"cell_type":"code","source":"import zipfile\n\nwith zipfile.ZipFile('submission.zip', 'w', zipfile.ZIP_DEFLATED) as zipf:\n    for filename in tqdm(test_filenames, desc=\"Zipping files\"):\n        if not os.path.exists(filename):\n            print(f\"Missing <> {filename}\")\n            continue\n        zipf.write(filename)\n        os.remove(filename)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-24T17:16:28.533268Z","iopub.execute_input":"2025-11-24T17:16:28.533524Z","iopub.status.idle":"2025-11-24T17:16:28.955749Z","shell.execute_reply.started":"2025-11-24T17:16:28.533504Z","shell.execute_reply":"2025-11-24T17:16:28.954737Z"}},"outputs":[],"execution_count":null}]}