miblab-ssa 0.0.0__py3-none-any.whl

miblab_ssa/pdm.py

@@ -0,0 +1,177 @@
+import numpy as np
+import trimesh
+from skimage.measure import marching_cubes
+from scipy.interpolate import griddata
+from scipy.ndimage import gaussian_filter, binary_dilation, binary_fill_holes
+
+# ---------------------------------------------------------
+# 1. Template & Mapping Utils
+# ---------------------------------------------------------
+
+def get_template_mesh(radius=1.0, subdivisions=3):
+    """
+    Returns the fixed topology icosphere.
+    """
+    mesh = trimesh.creation.icosphere(subdivisions=subdivisions, radius=radius)
+    
+    # Pre-calculate the spherical coordinates (Theta, Phi) for this template ONCE.
+    # We will query these angles on the target kidney.
+    verts = mesh.vertices
+    # Normalize just in case
+    verts /= np.linalg.norm(verts, axis=1)[:, None]
+    
+    # Convert vector -> angles
+    # Mapping (x, y, z) -> (x, z, y) so Z is poles is standard for many libs, 
+    # but let's stick to standard math: Z is up.
+    Xc, Yc, Zc = verts[:,0], verts[:,1], verts[:,2]
+    
+    # Clip z to [-1, 1] to avoid NaNs at poles
+    Theta = np.arccos(np.clip(Zc, -1, 1))
+    Phi = np.arctan2(Yc, Xc) % (2*np.pi)
+    
+    return mesh, Theta, Phi
+
+def get_inflated_mapping(verts, faces, iterations=80):
+    """
+    Iteratively smoothes the target mesh into a sphere.
+    This 'unwraps' the concave kidney so we can map it to the template.
+    """
+    smooth_verts = verts.copy()
+    n_verts = len(verts)
+    
+    # 1. Build Adjacency (Edges)
+    edges = np.concatenate([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]])
+    
+    # 2. Laplacian Smoothing (Inflation)
+    for _ in range(iterations):
+        # Calculate Face centers
+        face_centers = smooth_verts[faces].mean(axis=1)
+        
+        # Accumulate neighbor forces
+        v_sums = np.zeros_like(smooth_verts)
+        v_counts = np.zeros(n_verts)
+        
+        np.add.at(v_sums, faces, face_centers[:, None, :])
+        np.add.at(v_counts, faces, 1)
+        
+        v_counts[v_counts==0] = 1
+        v_avgs = v_sums / v_counts[:, None]
+        
+        # Update (Relax towards average)
+        smooth_verts = 0.5 * smooth_verts + 0.5 * v_avgs
+        
+        # Re-project to unit sphere (The "Inflation" step)
+        smooth_verts -= smooth_verts.mean(axis=0)
+        norms = np.linalg.norm(smooth_verts, axis=1)
+        smooth_verts /= (norms[:, None] + 1e-9)
+
+    # 3. Get Angles of the inflated target
+    Xc, Yc, Zc = smooth_verts[:,0], smooth_verts[:,1], smooth_verts[:,2]
+    Theta = np.arccos(np.clip(Zc, -1, 1))
+    Phi = np.arctan2(Yc, Xc) % (2*np.pi)
+    
+    return Theta, Phi
+
+# ---------------------------------------------------------
+# 2. Main API
+# ---------------------------------------------------------
+
+def features_from_mask(mask: np.ndarray, subdivisions=3):
+    """
+    Extracts PDM features by mapping the Template Mesh to the Target.
+    Preserves Concavities (Hollows).
+    """
+    # 1. Extract Target Mesh
+    # Blur slightly to get smooth normals
+    mask_float = gaussian_filter(mask.astype(float), sigma=1.0)
+    verts_target, faces_target, _, _ = marching_cubes(mask_float, level=0.5)
+    
+    # 2. Center Target (Crucial for PCA)
+    centroid = verts_target.mean(axis=0)
+    verts_centered = verts_target - centroid
+    
+    # 3. Inflate Target (Get its spherical address map)
+    # This finds the unique (Theta, Phi) for every vertex on the kidney
+    Theta_target, Phi_target = get_inflated_mapping(verts_centered, faces_target)
+    
+    # 4. Get Template Angles
+    # We want to move the Template vertices to the corresponding locations
+    template_mesh, Theta_template, Phi_template = get_template_mesh(subdivisions=subdivisions)
+    
+    # 5. Interpolate
+    # We use the Target's (Theta, Phi) -> (X, Y, Z) mapping
+    # to look up the (X, Y, Z) for the Template's (Theta, Phi).
+    
+    # Fix the Seam at 2pi: Augment data
+    T_aug = np.concatenate([Theta_target, Theta_target, Theta_target])
+    P_aug = np.concatenate([Phi_target, Phi_target - 2*np.pi, Phi_target + 2*np.pi])
+    V_aug = np.concatenate([verts_centered, verts_centered, verts_centered])
+    
+    points_in = np.column_stack((T_aug, P_aug))
+    
+    # Look up X, Y, Z coordinates for the template vertices
+    feat_x = griddata(points_in, V_aug[:,0], (Theta_template, Phi_template), method='linear')
+    feat_y = griddata(points_in, V_aug[:,1], (Theta_template, Phi_template), method='linear')
+    feat_z = griddata(points_in, V_aug[:,2], (Theta_template, Phi_template), method='linear')
+    
+    # Fill any interpolation gaps (poles or sparse areas)
+    mask_nan = np.isnan(feat_x)
+    if np.any(mask_nan):
+        feat_x[mask_nan] = griddata(points_in, V_aug[:,0], (Theta_template[mask_nan], Phi_template[mask_nan]), method='nearest')
+        feat_y[mask_nan] = griddata(points_in, V_aug[:,1], (Theta_template[mask_nan], Phi_template[mask_nan]), method='nearest')
+        feat_z[mask_nan] = griddata(points_in, V_aug[:,2], (Theta_template[mask_nan], Phi_template[mask_nan]), method='nearest')
+        
+    # 6. Flatten to Feature Vector
+    # These are the vertices of the TEMPLATE, but moved to the KIDNEY surface
+    return np.column_stack((feat_x, feat_y, feat_z)).flatten()
+
+
+def mask_from_features(features: np.ndarray, shape, subdivisions=3):
+    """
+    Reconstructs the mask. Since features are derived from a Template Mesh,
+    we can use the Template's faces to create a guaranteed watertight surface.
+    """
+    vol = np.zeros(shape, dtype=np.uint8)
+    
+    # 1. Rebuild Mesh Object
+    ref_template, _, _ = get_template_mesh(subdivisions=subdivisions)
+    n_verts = len(ref_template.vertices)
+    
+    if features.size != n_verts * 3:
+         raise ValueError(f"Feature mismatch! Expected {n_verts*3}, got {features.size}. Check subdivision level.")
+
+    vertices = features.reshape((n_verts, 3))
+    
+    # Center back to volume
+    vol_center = np.array(shape) / 2.0
+    vertices += vol_center
+    
+    # Create the mesh using Fixed Topology
+    mesh = trimesh.Trimesh(vertices=vertices, faces=ref_template.faces)
+    
+    # 2. Robust Voxelization (Sample & Fill)
+    # We perform dense sampling on the mesh surface to handle any folding
+    # 50,000 points covers a 300^3 volume well without gaps
+    samples, _ = trimesh.sample.sample_surface(mesh, count=50000)
+    
+    ix = np.round(samples[:, 0]).astype(int)
+    iy = np.round(samples[:, 1]).astype(int)
+    iz = np.round(samples[:, 2]).astype(int)
+    
+    # Bounds check
+    valid = (ix >= 0) & (ix < shape[0]) & \
+            (iy >= 0) & (iy < shape[1]) & \
+            (iz >= 0) & (iz < shape[2])
+            
+    # Paint surface
+    vol[ix[valid], iy[valid], iz[valid]] = 1
+    
+    # 3. Fill
+    vol = binary_dilation(vol, iterations=2)
+    vol = binary_fill_holes(vol)
+    
+    return vol.astype(np.uint8)
+
+def smooth_mask(mask, subdivisions=3):
+    feats = features_from_mask(mask, subdivisions=subdivisions)
+    return mask_from_features(feats, mask.shape, subdivisions=subdivisions)

miblab_ssa/sdf_cheby.py

@@ -0,0 +1,153 @@
+import numpy as np
+from scipy.ndimage import distance_transform_edt, label, center_of_mass
+from numpy.polynomial.chebyshev import chebvander
+from sklearn.linear_model import Ridge
+from itertools import product
+
+
+def features_from_mask(mask, order=15, n_samples=50000, random_seed=42):
+    """
+    Extracts PCA-ready shape coefficients using a fixed coordinate system 
+    based on the mask volume dimensions.
+    """
+    # Fix the seed so the random sampling is identical every time
+    np.random.seed(random_seed)
+
+    # 1. Define Fixed Coordinate System (Middle of the Cube)
+    # Since your data is registered, we use the volume's geometrical center.
+    nz, ny, nx = mask.shape
+    center = np.array([nz, ny, nx]) / 2.0
+    
+    # Scale covers the entire volume (radius from center to edge)
+    # For a cube of size N, radius is N/2. 
+    scale = np.max([nz, ny, nx]) / 2.0
+
+    # 2. Compute SDF
+    dist_outside = distance_transform_edt(1 - mask)
+    dist_inside = distance_transform_edt(mask)
+    sdf = dist_outside - dist_inside
+    
+    # 3. Robust Sampling Strategy
+    boundary_mask = np.abs(sdf) < 3.0
+    idx_boundary = np.argwhere(boundary_mask)
+    
+    # Generate random indices globally
+    idx_random = np.random.randint(0, [nz, ny, nx], size=(n_samples, 3))
+    
+    # Combine (50% boundary, 50% random air)
+    n_bound = int(n_samples * 0.5)
+    if len(idx_boundary) > 0:
+        idx_b_select = idx_boundary[np.random.choice(len(idx_boundary), n_bound, replace=True)]
+    else:
+        idx_b_select = idx_boundary # Fallback
+
+    indices = np.vstack([idx_b_select, idx_random])
+    
+    z_idx, y_idx, x_idx = indices[:, 0], indices[:, 1], indices[:, 2]
+    Values = sdf[z_idx, y_idx, x_idx]
+    
+    # 4. Normalize Coords to [-1, 1] using Fixed Center/Scale
+    Z = (z_idx - center[0]) / scale
+    Y = (y_idx - center[1]) / scale
+    X = (x_idx - center[2]) / scale
+    
+    # Filter valid cube (points outside the fitting domain must be ignored)
+    valid_mask = (np.abs(X) <= 1) & (np.abs(Y) <= 1) & (np.abs(Z) <= 1)
+    
+    X, Y, Z = X[valid_mask], Y[valid_mask], Z[valid_mask]
+    Values = Values[valid_mask]
+    
+    # 5. Generate Basis
+    Tx = chebvander(X, order)
+    Ty = chebvander(Y, order)
+    Tz = chebvander(Z, order)
+    
+    basis_cols = []
+    # Deterministic loop order for PCA consistency
+    for i, j, k in product(range(order + 1), repeat=3):
+        if i + j + k <= order:
+            col = Tx[:, i] * Ty[:, j] * Tz[:, k]
+            basis_cols.append(col)
+            
+    A = np.column_stack(basis_cols)
+    
+    # 6. Solve
+    clf = Ridge(alpha=1e-3, fit_intercept=False)
+    clf.fit(A, Values)
+    
+    return clf.coef_
+
+def mask_from_features(coeffs, shape, order):
+    """
+    Reconstructs mask from coefficients using the fixed volume center/scale.
+    """
+    vol = np.zeros(shape, dtype=np.uint8)
+    
+    # 1. Define Fixed Coordinate System (Same as features_from_mask)
+    nz, ny, nx = shape
+    center = np.array([nz, ny, nx]) / 2.0
+    scale = np.max([nz, ny, nx]) / 2.0
+    
+    # 2. ROI optimization (We only loop over the area covered by 'scale')
+    pad = int(scale * 1.0) 
+    z_min, z_max = max(0, int(center[0]-pad)), min(shape[0], int(center[0]+pad))
+    y_min, y_max = max(0, int(center[1]-pad)), min(shape[1], int(center[1]+pad))
+    x_min, x_max = max(0, int(center[2]-pad)), min(shape[2], int(center[2]+pad))
+    
+    roi_nz, roi_ny, roi_nx = z_max - z_min, y_max - y_min, x_max - x_min
+    if roi_nz <= 0 or roi_ny <= 0 or roi_nx <= 0: return vol
+
+    # 3. Coords
+    z_coords = (np.arange(z_min, z_max) - center[0]) / scale
+    y_coords = (np.arange(y_min, y_max) - center[1]) / scale
+    x_coords = (np.arange(x_min, x_max) - center[2]) / scale
+    
+    # 4. Basis
+    Tz = chebvander(z_coords, order).T
+    Ty = chebvander(y_coords, order).T
+    Tx = chebvander(x_coords, order).T
+
+    # 5. Tensor Assembly
+    poly_indices = []
+    for i, j, k in product(range(order + 1), repeat=3):
+        if i + j + k <= order:
+            poly_indices.append((i, j, k))
+
+    C_tensor = np.zeros((order + 1, order + 1, order + 1))
+    for idx, (i, j, k) in enumerate(poly_indices):
+        C_tensor[k, j, i] = coeffs[idx]
+
+    # 6. Contraction
+    recon_roi = np.einsum('kji,kz,jy,ix->zyx', C_tensor, Tz, Ty, Tx, optimize='optimal')
+    
+    # 7. Hard Geometric Crop
+    zz, yy, xx = np.meshgrid(z_coords, y_coords, x_coords, indexing='ij')
+    valid_box = (np.abs(zz) <= 1.0) & (np.abs(yy) <= 1.0) & (np.abs(xx) <= 1.0)
+    
+    mask_roi = (recon_roi < 0) & valid_box
+    
+    # 8. Intelligent Filtering (Center Distance)
+    if np.any(mask_roi):
+        labeled, n_components = label(mask_roi)
+        if n_components > 1:
+            # We assume the kidney is near the center of the volume
+            roi_center = np.array([roi_nz/2, roi_ny/2, roi_nx/2])
+            
+            centers = center_of_mass(mask_roi, labeled, range(1, n_components+1))
+            centers = np.array(centers)
+            
+            if len(centers) > 0:
+                dists = np.linalg.norm(centers - roi_center, axis=1)
+                winner_idx = np.argmin(dists) + 1 
+                mask_roi = (labeled == winner_idx)
+    
+    vol[z_min:z_max, y_min:y_max, x_min:x_max] = mask_roi
+    
+    return vol
+
+def smooth_mask(mask:np.ndarray, order=20):
+    coeffs = features_from_mask(mask, order=order)
+    mask_rec = mask_from_features(coeffs, mask.shape, order)
+    return mask_rec
+
+# N = (L+1) (L+2) (L+3) / 6

miblab_ssa/sdf_ft.py

@@ -0,0 +1,78 @@
+import numpy as np
+from scipy.ndimage import distance_transform_edt
+from scipy.fft import dctn, idctn  
+
+
+
+
+def get_spectral_mask(shape, radius):
+    """
+    Generates a boolean mask for spherical truncation.
+    shape: The shape of the coefficient block (e.g., (16, 16, 16))
+    radius: The spectral radius to keep (e.g., 16.0)
+    """
+    # 1. strictly use the 'shape' argument to determine grid dimensions
+    ranges = [np.arange(n) for n in shape]
+    
+    # 2. Create the grid of frequencies (i, j, k)
+    I, J, K = np.meshgrid(*ranges, indexing='ij')
+    
+    # 3. Calculate distance from DC component (0,0,0)
+    # i^2 + j^2 + k^2 <= r^2
+    mask = (I**2 + J**2 + K**2) <= radius**2
+    
+    return mask
+
+def features_from_mask(mask: np.ndarray, order=16, norm="ortho", saturation_threshold=5.0):
+    # 1. Compute SDF & Saturate
+    mask = mask.astype(bool)
+    # Using float32 saves memory/speed for DCT without losing precision relevant for shape
+    sdf = (distance_transform_edt(~mask) - distance_transform_edt(mask)).astype(np.float32)
+    sdf_saturated = saturation_threshold * np.tanh(sdf / saturation_threshold)
+
+    # 2. Compute coefficients
+    full_coeffs = dctn(sdf_saturated, norm=norm, workers=-1)
+    
+    # 3. Cubic Truncation first (Efficiency)
+    # We slice out the corner cube first to avoid generating a mask for the massive 300^3 volume
+    cube_coeffs = full_coeffs[:order, :order, :order]
+    
+    # 4. Spherical Masking
+    # Now we generate the mask specifically for this cube shape
+    mask = get_spectral_mask(cube_coeffs.shape, radius=order)
+    
+    # Flatten ONLY the valid spherical coefficients
+    coeffs_flat = cube_coeffs[mask]
+
+    return coeffs_flat
+
+def mask_from_features(coeffs_flat: np.ndarray, shape, order, norm="ortho"):
+    # 1. Recreate the Mask
+    # We must generate the exact same mask used during encoding
+    # The container is the small corner cube of size (order, order, order)
+    cube_shape = (order, order, order)
+    mask = get_spectral_mask(cube_shape, radius=order)
+    
+    # 2. Rebuild the Cube
+    # Validate that coefficient counts match
+    expected_size = np.sum(mask)
+    if coeffs_flat.size != expected_size:
+        raise ValueError(f"Coefficient mismatch! Expected {expected_size} coeffs for order {order}, got {coeffs_flat.size}.")
+
+    cube_coeffs = np.zeros(cube_shape, dtype=coeffs_flat.dtype)
+    cube_coeffs[mask] = coeffs_flat
+    
+    # 3. Pad to Full Volume
+    full_coeffs = np.zeros(shape, dtype=coeffs_flat.dtype)
+    full_coeffs[:order, :order, :order] = cube_coeffs
+
+    # 4. Inverse DCT
+    sdf_recon = idctn(full_coeffs, norm=norm, workers=-1)
+    
+    return sdf_recon < 0
+
+
+def smooth_mask(mask:np.ndarray, order=16, norm="ortho"):
+    coeffs = features_from_mask(mask, order, norm)
+    mask_recon = mask_from_features(coeffs, mask.shape, order, norm)
+    return mask_recon

miblab_ssa/sdf_ft_simple.py

@@ -0,0 +1,47 @@
+import numpy as np
+from scipy.ndimage import distance_transform_edt
+from scipy.fftpack import dctn, idctn
+
+
+def features_from_mask(mask:np.ndarray, order=16, norm="ortho"):
+
+    # Compute sdf
+    mask = mask.astype(bool)
+    dist_out = distance_transform_edt(~mask)
+    dist_in  = distance_transform_edt(mask)
+    sdf = dist_out - dist_in 
+
+    # Compute coefficients, truncate and flatten
+    coeffs = dctn(sdf, norm=norm)
+    coeffs = coeffs[:order, :order, :order]
+    coeffs = coeffs.flatten()
+
+    return coeffs
+
+def mask_from_features(coeffs_trunc:np.ndarray, shape, norm="ortho"):
+
+    # Rebuild coefficients at required shape
+    order = int(np.cbrt(coeffs_trunc.size))
+    coeffs_trunc = coeffs_trunc.reshape(3 * [order])
+    coeffs = np.zeros(shape, dtype=coeffs_trunc.dtype)
+    coeffs[:order, :order, :order] = coeffs_trunc
+
+    # Compute mask from sdf
+    sdf = idctn(coeffs, norm=norm)
+    mask = sdf < 0
+
+    return mask
+
+def smooth_mask(mask:np.ndarray, order=16, norm="ortho"):
+    coeffs = features_from_mask(mask, order, norm)
+    mask_recon = mask_from_features(coeffs, mask.shape, norm)
+    return mask_recon
+
+def features_from_dataset(masks, order=16, norm="ortho"):
+    features = [features_from_mask(mask, order, norm) for mask in masks]
+    return np.array(features)
+
+
+
+
+

miblab_ssa/sdf_mono.py

@@ -0,0 +1,214 @@
+import numpy as np
+from scipy.ndimage import distance_transform_edt
+from itertools import product
+from sklearn.linear_model import Ridge
+
+def compute_implicit_coeffs(mask, order=6, n_samples=50000):
+    """
+    Fits an implicit surface using narrow-band sampling for high accuracy.
+    Lower order (4-8) is often more stable than high order (10+).
+    """
+    # 1. Compute SDF
+    dist_outside = distance_transform_edt(1 - mask)
+    dist_inside = distance_transform_edt(mask)
+    sdf = dist_outside - dist_inside
+
+    # 2. Define Coordinate System (Relative to Object)
+    coords = np.argwhere(mask)
+    min_c = coords.min(axis=0)
+    max_c = coords.max(axis=0)
+    center = (min_c + max_c) / 2.0
+    # Use max extent for uniform scaling (preserves aspect ratio)
+    scale = (max_c - min_c).max() / 2.0 * 1.1  # 10% padding
+
+    # 3. Intelligent Sampling (The Key Fix)
+    # Instead of zooming the grid, we pick specific points to learn from.
+    
+    # Grid of all voxel indices
+    nz, ny, nx = mask.shape
+    
+    # Strategy: Pick mostly points near the boundary (narrow band)
+    # plus some points inside and outside to define the bulk.
+    boundary_mask = np.abs(sdf) < 2.0  # Points within 2 pixels of surface
+    inside_mask = (sdf < -2.0)         # Deep inside
+    outside_mask = (sdf > 2.0)         # Far outside
+
+    # Get indices
+    idx_boundary = np.argwhere(boundary_mask)
+    idx_inside = np.argwhere(inside_mask)
+    idx_outside = np.argwhere(outside_mask)
+
+    # Sampling counts
+    n_bound = int(n_samples * 0.5)  # 50% samples on boundary
+    n_in = int(n_samples * 0.25)    # 25% inside
+    n_out = int(n_samples * 0.25)   # 25% outside
+
+    # Helper to safely sample
+    def sample_indices(indices, n):
+        if len(indices) == 0: return indices
+        # If we have fewer points than requested, take them all
+        if len(indices) < n: return indices 
+        choice = np.random.choice(len(indices), n, replace=False)
+        return indices[choice]
+
+    s_bound = sample_indices(idx_boundary, n_bound)
+    s_in = sample_indices(idx_inside, n_in)
+    s_out = sample_indices(idx_outside, n_out)
+
+    # Combine samples
+    all_indices = np.vstack([s_bound, s_in, s_out])
+    
+    # 4. Extract Coordinates and Values
+    # Z, Y, X (indices)
+    z_idx, y_idx, x_idx = all_indices[:, 0], all_indices[:, 1], all_indices[:, 2]
+    
+    # Normalize to [-1, 1]
+    Z_norm = (z_idx - center[0]) / scale
+    Y_norm = (y_idx - center[1]) / scale
+    X_norm = (x_idx - center[2]) / scale
+    
+    Values = sdf[z_idx, y_idx, x_idx]
+
+    # 5. Build Design Matrix
+    # Basis: x^i * y^j * z^k
+    # Pre-calculating powers is faster than looping
+    basis_cols = []
+    for i, j, k in product(range(order + 1), repeat=3):
+        if i + j + k <= order:
+            col = (X_norm**i) * (Y_norm**j) * (Z_norm**k)
+            basis_cols.append(col)
+    
+    A = np.column_stack(basis_cols)
+
+    # 6. Fit
+    clf = Ridge(alpha=1e-4) # Small regularization
+    clf.fit(A, Values)
+
+    return clf.coef_, center, scale, order
+
+def reconstruct_implicit(coeffs, shape, center, scale, order):
+    """
+    Reconstructs the volume. 
+    Optimization: Only evaluate within the bounding box of the object.
+    """
+    vol = np.zeros(shape, dtype=np.uint8)
+    
+    # Determine Bounding Box in indices (where normalized coords are roughly [-1, 1])
+    # Evaluating the polynomial over the entire 301^3 volume is wasteful and slow.
+    pad = int(scale * 1.1)
+    z_min, z_max = max(0, int(center[0]-pad)), min(shape[0], int(center[0]+pad))
+    y_min, y_max = max(0, int(center[1]-pad)), min(shape[1], int(center[1]+pad))
+    x_min, x_max = max(0, int(center[2]-pad)), min(shape[2], int(center[2]+pad))
+
+    # Generate grid only for the ROI
+    z_idx, y_idx, x_idx = np.meshgrid(
+        np.arange(z_min, z_max),
+        np.arange(y_min, y_max),
+        np.arange(x_min, x_max),
+        indexing='ij'
+    )
+    
+    # Normalize
+    Z = (z_idx - center[0]) / scale
+    Y = (y_idx - center[1]) / scale
+    X = (x_idx - center[2]) / scale
+    
+    # Flatten ROI
+    Xf, Yf, Zf = X.flatten(), Y.flatten(), Z.flatten()
+    
+    # Evaluate
+    recon_sdf = np.zeros_like(Xf)
+    col_idx = 0
+    for i, j, k in product(range(order + 1), repeat=3):
+        if i + j + k <= order:
+            weight = coeffs[col_idx]
+            recon_sdf += weight * (Xf**i) * (Yf**j) * (Zf**k)
+            col_idx += 1
+            
+    # Threshold
+    # SDF < 0 is inside
+    mask_roi = (recon_sdf < 0)
+    
+    # Place back into full volume
+    vol[z_min:z_max, y_min:y_max, x_min:x_max] = mask_roi.reshape(z_idx.shape)
+    
+    return vol
+
+import numpy as np
+from itertools import product
+
+def reconstruct_implicit_fast(coeffs, shape, center, scale, order):
+    """
+    Optimized reconstruction using precomputed powers and broadcasting.
+    """
+    vol = np.zeros(shape, dtype=np.uint8)
+    
+    # 1. Determine ROI (Bounding Box)
+    pad = int(scale * 1.1)
+    z_min, z_max = max(0, int(center[0]-pad)), min(shape[0], int(center[0]+pad))
+    y_min, y_max = max(0, int(center[1]-pad)), min(shape[1], int(center[1]+pad))
+    x_min, x_max = max(0, int(center[2]-pad)), min(shape[2], int(center[2]+pad))
+    
+    # Dimensions of the ROI
+    nz, ny, nx = z_max - z_min, y_max - y_min, x_max - x_min
+    if nz <= 0 or ny <= 0 or nx <= 0: return vol
+
+    # 2. Create Normalized Coordinate Vectors (1D)
+    # Instead of making a meshgrid immediately, we keep them 1D first.
+    z_coords = (np.arange(z_min, z_max) - center[0]) / scale
+    y_coords = (np.arange(y_min, y_max) - center[1]) / scale
+    x_coords = (np.arange(x_min, x_max) - center[2]) / scale
+    
+    # 3. Precompute Powers (Vandermonde-ish)
+    # Shape: (order+1, N)
+    # P_x[p, :] contains x^p for all x in the ROI
+    P_z = np.array([z_coords**p for p in range(order + 1)])
+    P_y = np.array([y_coords**p for p in range(order + 1)])
+    P_x = np.array([x_coords**p for p in range(order + 1)])
+
+    # 4. Re-assemble the Polynomial via Tensor Contraction
+    # The polynomial is: Sum_{i,j,k} ( C_{ijk} * Z^i * Y^j * X^k )
+    # This is a tensor dot product.
+    
+    # First, we need to unpack the flat 'coeffs' list back into a 3D tensor C[i,j,k].
+    # Since our flat list skipped terms where i+j+k > order, we must be careful.
+    C_tensor = np.zeros((order + 1, order + 1, order + 1))
+    
+    col_idx = 0
+    # Use the exact same iteration order as the training step
+    for i in range(order + 1):
+        for j in range(order + 1):
+            for k in range(order + 1):
+                if i + j + k <= order:
+                    # Note: Your training loop was (X^i, Y^j, Z^k), 
+                    # so map indices carefully: C[k, j, i] if Z is axis 0, etc.
+                    # Your generic loop: (X**i) * (Y**j) * (Z**k)
+                    # Let's align C_tensor indices to (k, j, i) -> (Z, Y, X)
+                    C_tensor[k, j, i] = coeffs[col_idx]
+                    col_idx += 1
+
+    # 5. Perform Contraction (Einstein Summation)
+    # We want result[z, y, x] = Sum_{k,j,i} C[k,j,i] * P_z[k,z] * P_y[j,y] * P_x[i,x]
+    # 'kji,kz,jy,ix->zyx'
+    #   k,j,i are power indices
+    #   z,y,x are spatial indices
+    
+    # einsum is efficient but can be memory hungry for large chunks.
+    # If memory fails, we can do it in two steps.
+    
+    # Step A: Contract X axis -> Temp[k, j, y, x]
+    # T1 = np.einsum('kji,ix->kjx', C_tensor, P_x) 
+    # But contracting 3 axes at once is cleanest if it fits in RAM:
+    
+    recon_roi = np.einsum('kji,kz,jy,ix->zyx', C_tensor, P_z, P_y, P_x, optimize='optimal')
+    
+    # 6. Threshold and Place
+    mask_roi = (recon_roi < 0)
+    vol[z_min:z_max, y_min:y_max, x_min:x_max] = mask_roi
+    
+    return vol
+
+def reconstruct_shape(mask, order=8): # Order 8 is a good balance for kidneys
+    coeffs, center, scale, order = compute_implicit_coeffs(mask, order=order)
+    mask_rec = reconstruct_implicit_fast(coeffs, mask.shape, center, scale, order)
+    return mask_rec