{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"competition","sourceId":117682,"databundleVersionId":15062069},{"sourceType":"datasetVersion","sourceId":14453760,"datasetId":9146144,"databundleVersionId":15274059}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install connected-components-3d\n!pip install scikit-fmm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2026-01-11T15:23:53.838216Z","iopub.execute_input":"2026-01-11T15:23:53.838665Z","iopub.status.idle":"2026-01-11T15:24:24.816752Z","shell.execute_reply.started":"2026-01-11T15:23:53.838632Z","shell.execute_reply":"2026-01-11T15:24:24.815574Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport cc3d\nfrom PIL import Image, ImageSequence\nfrom pathlib import Path\nimport torch\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt\nfrom matplotlib.patches import Patch\nfrom matplotlib.lines import Line2D\nimport matplotlib.lines as mlines\nfrom matplotlib import colors\nfrom scipy.spatial.distance import cdist\nfrom scipy.optimize import linear_sum_assignment\nfrom scipy.ndimage import convolve, gaussian_filter, distance_transform_edt, binary_erosion\nimport heapq\nfrom skimage.graph import MCP_Geometric, route_through_array\nfrom skimage.morphology import skeletonize, binary_dilation, square\nfrom skimage.segmentation import find_boundaries\nfrom itertools import product\nimport subprocess\nimport sys\nimport importlib\nfrom tqdm import tqdm\nimport os\nimport cv2\nimport skfmm\nfrom joblib import Parallel, delayed\nfrom scipy.optimize import least_squares\nfrom scipy.optimize import minimize","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T15:24:24.819140Z","iopub.execute_input":"2026-01-11T15:24:24.819785Z","iopub.status.idle":"2026-01-11T15:24:32.366399Z","shell.execute_reply.started":"2026-01-11T15:24:24.819746Z","shell.execute_reply":"2026-01-11T15:24:32.365217Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_volume(path: Path) -> np.ndarray:\n    \"\"\"\n    Load a multi-page TIFF into a 3D NumPy array: (slices, H, W)\n    \"\"\"\n    try:\n        with Image.open(path) as img:\n            frames = [np.array(frame) for frame in ImageSequence.Iterator(img)]\n        volume = np.stack(frames)\n        return volume\n    except Exception as e:\n        raise RuntimeError(f\"Error loading TIFF {path}: {e}\")\n\n\ndef visualize_paths(prob_mask, paths):\n    plt.figure(figsize=(12, 12))\n    \n    # 1. Show the background probability map\n    plt.imshow(prob_mask, cmap='gray', alpha=0.8)\n    \n    # 2. Loop through each path\n    for i, path in enumerate(paths):\n        # Plot the line (Path)\n        line, = plt.plot(path[:, 1], path[:, 0], linewidth=3, label=f'Path {i+1}')\n        \n        # Mark Start (Green) and End (Red)\n        plt.scatter(path[0, 1], path[0, 0], color='green', s=60, edgecolors='white', zorder=10)\n        plt.scatter(path[-1, 1], path[-1, 0], color='red', s=60, edgecolors='white', zorder=10)\n\n    # 3. Create Custom Legend Labels\n    # We create proxy artists for the Start and End points so they only show up once\n    start_handle = mlines.Line2D([], [], color='green', marker='o', linestyle='None',\n                                  markersize=8, label='Start Point')\n    end_handle = mlines.Line2D([], [], color='red', marker='o', linestyle='None',\n                                markersize=8, label='End Point')\n    \n    # Get the handles for the paths automatically\n    path_handles, labels = plt.gca().get_legend_handles_labels()\n    \n    # Combine everything into one legend\n    plt.legend(handles=[start_handle, end_handle] + path_handles, \n               bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=10)\n\n    plt.title(f\"Fast Marching Path Extraction\\n({len(paths)} Components Detected)\", fontsize=14)\n    plt.axis('off')\n    plt.tight_layout()\n    plt.show()\n\n\ndef search_endpoints(binary_img):\n    H, W = binary_img.shape\n    # 2. Skeletonize\n    skeleton = skeletonize(binary_img).astype(np.uint8)\n    # 3. Find Endpoints\n    # Use a kernel to count neighbors in a 3x3 area\n    kernel = np.array([[1, 1, 1],\n                       [1, 10, 1],\n                       [1, 1, 1]], dtype=np.uint8)\n    neighbor_count = convolve(skeleton, kernel, mode='constant', cval=0)\n    # An endpoint will have a value of 11 (10 for itself + 1 for its only neighbor)\n    endpoints = (neighbor_count == 11)\n    # 4. Get Coordinates\n    y, x = np.where(endpoints)\n    endpoint_coords = list(zip(x, y))\n    endpoint_coords = np.array(endpoint_coords)\n    if endpoint_coords.shape[0]==0:\n        return None\n    return endpoint_coords\n\n\ndef extract_path_fmm(prob_mask, start_coords, end_coords):\n    # 1. Define Speed: Wave moves fast in high-prob areas\n    # We use a power to sharpen the preference for the \"center\" of the path\n    speed = np.array(prob_mask**2 + 1e-6) \n    \n    # 2. Setup the \"Phi\" (Initial wavefront)\n    # FMM needs a seed. We set the start point to 0 and everything else to 1\n    phi = np.ones_like(prob_mask)\n    phi[start_coords[0], start_coords[1]] = 0\n    \n    # 3. Compute Travel Time Map\n    t_map = skfmm.travel_time(phi, speed)\n    \n    # 4. Extract Path via Gradient Descent from End to Start\n    path = [end_coords]\n    current = np.array(end_coords)\n    \n    # Simple backtracking loop\n    for _ in range(1000): # Max steps\n        gy, gx = np.gradient(t_map)\n        # Find direction of steepest descent\n        move_y = gy[int(current[0]), int(current[1])]\n        move_x = gx[int(current[0]), int(current[1])]\n        \n        # Update position (step size of 1 pixel)\n        step = np.array([move_y, move_x])\n        step /= (np.linalg.norm(step) + 1e-8)\n        current = current - step\n        \n        path.append(current.copy())\n        \n        # Stop if we are close to start\n        if np.linalg.norm(current - start_coords) < 2:\n            break\n            \n    return np.array(path), t_map\n\n\ndef get_paths_by_component(prob_mask, threshold=0.3):\n    # 1. Threshold and label individual components\n    binary_img = (prob_mask > threshold).astype(np.uint8)\n    labeled_array, num_features = cc3d.connected_components(\n        binary_img,\n        connectivity=8,\n        return_N=True\n    )\n    \n    final_paths = []\n    start_end_points = []\n    output_mask = np.zeros_like(prob_mask, dtype=np.uint8)  # Mask to draw paths\n    \n    # 2. Iterate through each independent path/blob\n    for i in tqdm(range(1, num_features + 1)):\n        component_mask = (labeled_array == i)\n        \n        # Find endpoints ONLY for this component\n        endpoints = search_endpoints(component_mask)\n\n        if endpoints is None:\n            continue\n        \n        if len(endpoints) >= 2:\n            # Use the two points furthest apart in this component\n            start_node = tuple(endpoints[0][::-1])  # (row, col)\n            end_node = tuple(endpoints[-1][::-1])\n            \n            # 3. Execute FMM\n            try:\n                path, _ = extract_path_fmm(prob_mask, start_node, end_node)\n                final_paths.append(path)\n                start_end_points.append((start_node, end_node))\n                \n                # 4. Draw path on mask using cv2.polylines\n                # OpenCV expects int coordinates in (x, y)\n                path_pts = path[:, ::-1].astype(np.int32)  # Convert (row, col) -> (x, y)\n                path_pts = path_pts.reshape((-1, 1, 2))\n                cv2.polylines(output_mask, [path_pts], isClosed=False, color=1, thickness=1)\n            except Exception as e:\n                print(f\"Path extraction failed for component {i}: {e}\")\n                \n    return final_paths, start_end_points, output_mask\n\n\n\ndef viz_report_from_slice(prob_mask, selected_id, selected_slice):\n    valid_mask = gts[selected_id][selected_slice] !=2\n    fix, ax = plt.subplots(ncols = 2, figsize = (16, 16))\n    ax[0].imshow(gts[selected_id][selected_slice] * valid_mask)\n    ax[0].set_title('ground truth')\n    ax[1].imshow(probs[selected_id][selected_slice] * valid_mask)\n    ax[1].set_title('probability map')\n    plt.show()\n    final_paths, start_end_points, output_mask = get_paths_by_component(probs[selected_id][selected_slice])\n    plt.figure(figsize = (12, 12))\n    plt.imshow(output_mask)\n    plt.title('generated mask')\n    plt.show()\n    visualize_paths(probs[selected_id][selected_slice], final_paths)\n    return final_paths, start_end_points, output_mask","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T15:30:40.567271Z","iopub.execute_input":"2026-01-11T15:30:40.567793Z","iopub.status.idle":"2026-01-11T15:30:40.593294Z","shell.execute_reply.started":"2026-01-11T15:30:40.567761Z","shell.execute_reply":"2026-01-11T15:30:40.592181Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"probs = []\ngts = []\nprob_path = Path('/kaggle/input/vesuvius-challenge-test/hard_prob_samples/hard_prob_samples')\ngt_path = Path('/kaggle/input/vesuvius-challenge-surface-detection/train_labels')\nfor prob in tqdm(os.listdir(prob_path)):\n    id_ = prob.split('.')[0]\n    gts.append(load_volume(gt_path/f'{id_}.tif'))\n    probs.append(np.load(prob_path/prob)['prob'])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T15:24:32.395790Z","iopub.execute_input":"2026-01-11T15:24:32.396250Z","iopub.status.idle":"2026-01-11T15:24:58.610941Z","shell.execute_reply.started":"2026-01-11T15:24:32.396212Z","shell.execute_reply":"2026-01-11T15:24:58.609859Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"final_paths, start_end_points, output_mask = viz_report_from_slice(probs, 10, 100)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T15:30:44.025970Z","iopub.execute_input":"2026-01-11T15:30:44.026679Z","iopub.status.idle":"2026-01-11T15:30:53.512515Z","shell.execute_reply.started":"2026-01-11T15:30:44.026648Z","shell.execute_reply":"2026-01-11T15:30:53.511230Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Fitting Part of Spiral","metadata":{}},{"cell_type":"code","source":"def fit_global_spiral(paths):\n    \"\"\"\n    Estimates a single Archimedean spiral template (r = a + b * theta) \n    that best fits multiple local path segments.\n    \"\"\"\n    # Combine all points to find a global center, but keep paths \n    # separate for theta-unwrapping to avoid jumps between segments.\n    all_points = np.vstack(paths)\n    \n    def get_path_polar(xc, yc, path_pts):\n        dx = path_pts[:, 0] - xc\n        dy = path_pts[:, 1] - yc\n        r = np.sqrt(dx**2 + dy**2)\n        theta = np.unwrap(np.arctan2(dy, dx))\n        return r, theta\n\n    def objective(center):\n        xc, yc = center\n        total_residual = 0\n        \n        for path in paths:\n            r, theta = get_path_polar(xc, yc, path)\n            # Linear regression for this specific path segment\n            # r = a + b * theta\n            A = np.vstack([np.ones(len(theta)), theta]).T\n            # Solve for a and b for this segment\n            # Note: 'a' will vary per path if they are concentric but \n            # shifted in phase, but 'b' (growth) is usually constant.\n            sol, resid, _, _ = np.linalg.lstsq(A, r, rcond=None)\n            total_residual += resid[0] if len(resid) > 0 else np.sum((A @ sol - r)**2)\n            \n        return total_residual\n\n    # Initial Guess: Mean of all points\n    init_center = np.mean(all_points, axis=0)\n    \n    # Optimization: Find the common center for all paths\n    res = minimize(objective, init_center, method='Nelder-Mead')\n    xc_opt, yc_opt = res.x\n    \n    # Generate templates for each path based on the optimized center\n    templates = []\n    for path in paths:\n        r, theta = get_path_polar(xc_opt, yc_opt, path)\n        A = np.vstack([np.ones(len(theta)), theta]).T\n        a_opt, b_opt = np.linalg.lstsq(A, r, rcond=None)[0]\n        \n        # Create a smooth \"perfect\" version of the input theta range\n        t_smooth = np.linspace(theta[0], theta[-1], len(theta))\n        r_smooth = a_opt + b_opt * t_smooth\n        \n        template_x = xc_opt + r_smooth * np.cos(t_smooth)\n        template_y = yc_opt + r_smooth * np.sin(t_smooth)\n        templates.append(np.vstack([template_x, template_y]).T)\n        \n    return (xc_opt, yc_opt), templates","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T15:53:49.327815Z","iopub.execute_input":"2026-01-11T15:53:49.328225Z","iopub.status.idle":"2026-01-11T15:53:49.340959Z","shell.execute_reply.started":"2026-01-11T15:53:49.328195Z","shell.execute_reply":"2026-01-11T15:53:49.339778Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"center, templates = fit_global_spiral(final_paths)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T15:55:13.995411Z","iopub.execute_input":"2026-01-11T15:55:13.996402Z","iopub.status.idle":"2026-01-11T15:55:14.670281Z","shell.execute_reply.started":"2026-01-11T15:55:13.996353Z","shell.execute_reply":"2026-01-11T15:55:14.669063Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def visualize_templates_on_image(image, paths, templates):\n    H, W = image.shape\n    plt.figure(figsize=(10, 10))\n\n    # Image\n    plt.imshow(\n        image,\n        cmap='gray',\n        origin='upper',\n        extent=[0, W, H, 0],\n        alpha=0.9\n    )\n\n    for i, tmpl in enumerate(templates):\n        plt.plot(\n            tmpl[:, 1], tmpl[:, 0],\n            color='red',\n            linewidth=2,\n            alpha=0.9,\n            label='Fitted Spiral templates' if i == 0 else None\n        )\n\n    plt.gca().set_aspect('equal')\n    plt.xlim(0, W)\n    plt.ylim(H, 0)\n\n    plt.xlabel(\"x\")\n    plt.ylabel(\"y\")\n    plt.title(\"Mathematical Fitted Partial Spiral from Given Probability Map\")\n\n    # Legend\n    plt.legend(\n        loc='upper right',\n        frameon=True,\n        facecolor='white',\n        framealpha=0.9\n    )\n\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T16:03:17.939791Z","iopub.execute_input":"2026-01-11T16:03:17.940592Z","iopub.status.idle":"2026-01-11T16:03:17.948112Z","shell.execute_reply.started":"2026-01-11T16:03:17.940556Z","shell.execute_reply":"2026-01-11T16:03:17.947092Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"visualize_templates_on_image(probs[10][100], final_paths, templates)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T16:03:18.054446Z","iopub.execute_input":"2026-01-11T16:03:18.055360Z","iopub.status.idle":"2026-01-11T16:03:18.755346Z","shell.execute_reply.started":"2026-01-11T16:03:18.055315Z","shell.execute_reply":"2026-01-11T16:03:18.754374Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}