ostium_detection.py

import math
from collections import defaultdict, deque
import vtk

ABDOMEN_LABELS = ['Celiac Trunk', 'Superior Mesenteric', 'Renal', 'Renal'] # The 4 abdominal artery labels, assigned top to bottom
FORK_DEDUPLICATION_RADIUS_MM = 3.0     # If 2 bifurcations have distance smaller than this, then they're duplicates of each other
AORTA_CONTINUATION_RATIO = 0.75        # At bifurcation, 2nd segment (outlet) must be at least this fraction of the thickest (inlet) to be considered aorta continuation
ABDOMEN_CLUSTER_MAX_GAP_MM = 50.0      # Maximum gap (mm) within the abdominal cluster
THORACIC_GAP_MM = 80.0                 # Minimum gap (mm) separating the thoracic artery cluster from the abdominal cluster
ILIAC_Z_DROP_RATIO = 1.5               # Z-drop must be at least this times the X-spread

# Helper that sums euclidean distances between consecutive points in a segment
def calculate_length(points):
    length = 0.0
    for i in range(len(points) - 1):
        p1, p2 = points[i], points[i + 1]
        length += math.sqrt((p2['x'] - p1['x']) ** 2 + (p2['y'] - p1['y']) ** 2 + (p2['z'] - p1['z']) ** 2)
    return length

# Returns a unit length 3D vector
def normalize_vector(v):
    norm = math.sqrt(v[0] ** 2 + v[1] ** 2 + v[2] ** 2)
    if norm > 0:
        return (v[0] / norm, v[1] / norm, v[2] / norm)
    return (0.0, 0.0, 0.0)

# Computes average direction vector near a segment endpoint (bifurcation side)
def calculate_average_direction(points, is_start_at_bifurcation, limit_points):
    if not points or len(points) < 2:
        return (0, 0, 0)

    n = min(len(points), limit_points)

    # Walk away from the bifurcation: forward if bifurcation is at start, backward if at end
    if is_start_at_bifurcation:
        pairs = [(points[i], points[i + 1]) for i in range(n - 1)]
    else:
        end = len(points) - 1
        pairs = [(points[i], points[i - 1]) for i in range(end, end - n + 1, -1)]

    total = [0.0, 0.0, 0.0]
    valid = 0
    for a, b in pairs:
        dx, dy, dz = b['x'] - a['x'], b['y'] - a['y'], b['z'] - a['z']
        norm = math.sqrt(dx * dx + dy * dy + dz * dz)
        if norm > 0:
            # Each total will contain the sum of unit vectors
            total[0] += dx / norm 
            total[1] += dy / norm
            total[2] += dz / norm
            valid += 1

    if valid == 0:
        return (0, 0, 0)
    
    # Normalize the total after dividing by thier length to get the mean
    return normalize_vector((total[0] / valid, total[1] / valid, total[2] / valid))

# Helper that sets label on a centerline segment and its points
def apply_label(segment, label):
    segment['label'] = label
    """for pt in segment['points']:
        pt['label'] = label"""
        
# Main Pipeline for extracting centerline data, labeling bifurcations and arteries, computing ostium and reference point positions
def extract_centerline_data(centerlines, branch_junctions=None):
    # Parse VTK centerline polydata into Python list (segments containing points information) and dictionary (mapping each individual point to its parent segment)
    segments, point_to_segments = parse_centerline_polydata(centerlines)

    # Identify bifurcation points of all arteries, supplemented with known branch contact points
    all_bifurcations = find_bifurcation_points(centerlines, branch_junctions)

    # Adds segment fork points (segment diverges) to all_bifurcations, useful for When VMTK produces only few long cells that share the entire aorta which ruins polydata
    fork_pts = find_segment_fork_points(segments, point_to_segments)
    all_bifurcations, bifurcation_ids = merge_bifurcations_with_fork_points(all_bifurcations, fork_pts)

    # Splits diverging segments at all bifurcations into smaller sub segments that can be modified independently
    segments, point_to_segments = split_segments_at_bifurcations(segments, bifurcation_ids)

    # Computes thickest segment branching from each bifurcation point
    all_bifurcation_arms = {}
    for bif in all_bifurcations:
        result = compute_bifurcation_arms(bif, segments, point_to_segments)
        if result is not None:
            all_bifurcation_arms[bif['id']] = result
    print(f"Bifurcations total={len(all_bifurcations)}, real branch points (3+ clusters)={len(all_bifurcation_arms)}")

    # Label centerline segments Aorta or Artery at every bifurcation, returns bifurcation map which maps bifurcation id to its bifurcation point information and connected arms
    bifurcations_map, bifurcations_info = label_aorta_and_arteries(segments, point_to_segments, all_bifurcations, all_bifurcation_arms)

    # Detect if dataset contains thoracic arch or not
    thoracic_info = detect_thoracic_region(centerlines)
    has_thoracic = thoracic_info['detected']
    print(f"has_thoracic={has_thoracic}, reversal_mm={thoracic_info['reversal_mm']}")

    # Classify artery junctions into thoracic, abdomen, iliac zones, determine iliac cutoff, and relabel artery segments to match new classification
    iliac_cutoff = classify_arteries_by_zone(bifurcations_info, segments, has_thoracic)

    # Identify named abdomen arteries (Celiac Trunk / Superior Mesenteric / Renal), return ostium positions
    abdomen_ostia = identify_abdominal_arteries(segments, point_to_segments, bifurcations_info)

    # Place reference point at the most caudal abdomen bifurcation (junction), not at the derived ostium point
    reference_point = compute_reference_point(segments, bifurcations_info, abdomen_ostia)

    response = build_response(centerlines, segments, abdomen_ostia, reference_point, iliac_cutoff, bifurcations_map, bifurcations_info)
    return response

# Converts centerline polydata into a python list that contains line segments information (unique segment id, points in that segment, radius of each point, label)
def parse_centerline_polydata(centerlines):
    # Extracts centerline information
    pts = centerlines.GetPoints()
    lines = centerlines.GetLines()
    radius_array = centerlines.GetPointData().GetArray('Radius')

    segments = []
    point_to_segments = defaultdict(list) # Maps each centerline point to the line segment it belongs to, auto creates a default value when a missing key is accessed

    # Initalizes looping over the centerline segments
    lines.InitTraversal()
    id_list = vtk.vtkIdList()
    seg_idx = 0

    while lines.GetNextCell(id_list):
        num_ids = id_list.GetNumberOfIds()
        if num_ids < 2: # If number of points in this line segment is less than 2
            continue

        segment_points = []
        seen_ids_in_seg = set()  # Track unique point IDs within this segment

        for i in range(num_ids):
            point_id = id_list.GetId(i)
            if point_id in seen_ids_in_seg:
                continue  # Skip duplicate point IDs within the same segment
            seen_ids_in_seg.add(point_id)
            pt = pts.GetPoint(point_id)
            radius = radius_array.GetValue(point_id)
            segment_points.append({'id': point_id, 'x': float(pt[0]), 'y': float(pt[1]), 'z': float(pt[2]), 'radius': float(radius), 'label': None})

        # Filter zero-length segments (noise)
        length = calculate_length(segment_points)
        if length > 0.0:
            # Map ALL points to their containing segment (not just start/end), this ensures bifurcation points (interior to cells) can find connected segments
            for pt_entry in segment_points:
                point_to_segments[pt_entry['id']].append(seg_idx)
            segments.append({'points': segment_points, 'length': length, 'label': None})
            seg_idx += 1

    return segments, point_to_segments

# Computes bifurcation points which represent start point of any branching artery
def find_bifurcation_points(centerlines, known_junctions=None):
    # Extracts centerline data
    pts = centerlines.GetPoints()
    lines = centerlines.GetLines()
    radius_array = centerlines.GetPointData().GetArray('Radius')

    # Traverse through each line, then for each point in each line we increment its degree which represents number of line cells touching that point
    degree = defaultdict(int)
    lines.InitTraversal()
    id_list = vtk.vtkIdList()
    while lines.GetNextCell(id_list):
        for i in range(id_list.GetNumberOfIds()):
            degree[id_list.GetId(i)] += 1

    bifurcations = []
    found_ids = set()
    
    # Degree 1 = endpoint, 2 = mid-segment, 3+ = bifurcation
    for point_id, deg in degree.items():
        if deg >= 3:
            pt = pts.GetPoint(point_id)
            radius = radius_array.GetValue(point_id)
            found_ids.add(point_id)
            bifurcations.append({'id': point_id, 'x': float(pt[0]), 'y': float(pt[1]), 'z': float(pt[2]), 'radius': float(radius)})

    # Supplement with known junction coordinates from the pipeline (branch contact points)
    # This ensures bifurcations are detected even if vtkCleanPolyData didn't fully merge junction point IDs
    if known_junctions:
        locator = vtk.vtkPointLocator()
        locator.SetDataSet(centerlines)
        locator.BuildLocator()
        for junction_xyz in known_junctions:
            closest_id = locator.FindClosestPoint(junction_xyz)
            if closest_id < 0 or closest_id in found_ids:
                continue
            if degree.get(closest_id, 0) < 3:
                continue
            pt = pts.GetPoint(closest_id)
            radius = 0.0
            if radius_array:
                radius = radius_array.GetValue(closest_id)
            found_ids.add(closest_id)
            bifurcations.append({'id': closest_id, 'x': float(pt[0]), 'y': float(pt[1]), 'z': float(pt[2]), 'radius': float(radius)})

    return bifurcations

# Detects fork points where overlapping segments diverge
def find_segment_fork_points(segments, point_to_segments):
    # For each segment, walk from start→end and end→start looking for the transition from 'shared with other segments' to 'alone'
    # The last shared point before the unshared tail is a real bifurcation (fork). This catches junctions that degree-based detection misses
    fork_points = []

    for seg_i, seg in enumerate(segments):
        pts = seg['points']
        if len(pts) < 4:
            continue

        # Build per-point "other segment count" (how many OTHER segments also contain this point)
        others = []
        for p in pts:
            segments_at = set(point_to_segments.get(p['id'], []))
            segments_at.discard(seg_i) # Remove CURRENT segment we are at as we are only interested in the count of other segments
            others.append(len(segments_at))

        # Walk start→end: find last shared point before an unshared tail
        for k in range(len(pts) - 1):
            if others[k] > 0 and others[k + 1] == 0:
                fork_points.append(pts[k])
                break  # only save the outermost fork from this side

        # Walk end→start: same check from the other direction
        for k in range(len(pts) - 1, 0, -1):
            if others[k] > 0 and others[k - 1] == 0:
                fork_points.append(pts[k])
                break

    return fork_points

# Merges calculated bifurcations with detected fork points without any duplications
def merge_bifurcations_with_fork_points(all_bifurcations, fork_pts):
    bifurcation_ids = {b['id'] for b in all_bifurcations}

    # Build a spatial locator on the existing degree-based bifurcations so fork points that are physically the same junction are rejected
    existing_bif_polydata = vtk.vtkPolyData()
    existing_bif_pts = vtk.vtkPoints()
    for b in all_bifurcations:
        existing_bif_pts.InsertNextPoint(b['x'], b['y'], b['z'])
    existing_bif_polydata.SetPoints(existing_bif_pts)
    bif_locator = vtk.vtkPointLocator()
    bif_locator.SetDataSet(existing_bif_polydata)
    bif_locator.BuildLocator()

    for fp in fork_pts:
        if fp.get('radius', 0) < 0.5:
            continue
        if fp['id'] in bifurcation_ids:
            continue

        # Reject if any existing bifurcation is within threshold
        if existing_bif_pts.GetNumberOfPoints() > 0:
            nearest_id = bif_locator.FindClosestPoint(fp['x'], fp['y'], fp['z'])
            nearest_pt = existing_bif_pts.GetPoint(nearest_id)
            dist = math.sqrt((fp['x'] - nearest_pt[0]) ** 2 + (fp['y'] - nearest_pt[1]) ** 2 + (fp['z'] - nearest_pt[2]) ** 2)
            if dist < FORK_DEDUPLICATION_RADIUS_MM:
                continue
        
        all_bifurcations.append({'id': fp['id'], 'x': fp['x'], 'y': fp['y'], 'z': fp['z'], 'radius': fp['radius']})
        bifurcation_ids.add(fp['id'])

        # Add the new point to the locator so subsequent fork points also deduplicate against it
        existing_bif_pts.InsertNextPoint(fp['x'], fp['y'], fp['z'])
        bif_locator.BuildLocator()

    return all_bifurcations, bifurcation_ids

# Splits segments that pass through bifurcation points into sub-segments so that each sub-segment can be independently labeled as Aorta or Artery
def split_segments_at_bifurcations(segments, bifurcation_ids):
    new_segments = []
    new_points_to_segments = defaultdict(list)
    seg_idx = 0

    for seg in segments:
        # Find positions of interior bifurcation points (skip endpoints)
        pts = seg['points']
        split_at = []
        for k in range(1, len(pts) - 1):
            if pts[k]['id'] in bifurcation_ids:
                split_at.append(k)

        # If no interior bifurcation is found at this segment, thenk keep segment unchanged
        if not split_at:
            for p in pts:
                new_points_to_segments[p['id']].append(seg_idx)
            new_segments.append(seg)
            seg_idx += 1
            continue

        # Split at each interior bifurcation, the bifurcation point is included in both adjacent sub-segments
        bounds = [0] + split_at + [len(pts) - 1]
        for i in range(len(bounds) - 1):
            sub_segment = pts[bounds[i] : bounds[i + 1] + 1]
            if len(sub_segment) < 2:
                continue

            length = calculate_length(sub_segment)
            if length < 0.5:
                continue
            for p in sub_segment:
                new_points_to_segments[p['id']].append(seg_idx)

            new_segments.append({'points': sub_segment, 'length': length, 'label': None})
            seg_idx += 1

    return new_segments, new_points_to_segments

# For a given bifurcation point, computes all directional arms leaving a bifurcation point
# Arms in a similar direction are clustered together and thickest arm per cluster is kept, returns list of arms sorted by radius (descending)
def compute_bifurcation_arms(bifurcation, segments, point_to_segments):
    connected_seg_indices = point_to_segments.get(bifurcation['id'], [])
    direction_arms = []
    for seg_idx in connected_seg_indices:
        # Get segment points
        seg = segments[seg_idx]
        seg_pts = seg['points']
        if not seg_pts:
            continue
        # Get bifurcation point position on this segment
        bif_pos = None
        for k, p in enumerate(seg_pts):
            if p['id'] == bifurcation['id']:
                bif_pos = k
                break
        if bif_pos is None:
            continue
        
        # Forward arm (bifurcation → end of segment)
        if bif_pos < len(seg_pts) - 1:
            arm = seg_pts[bif_pos:]
            skip = min(3, max(0, len(arm) - 2)) # This avoids sampling the first few points right at the junction where radius is inflated by trunk geometry
            sample = arm[skip:skip + 10] or arm[-min(5, len(arm)):] # Samples up to 10 points further out for a more accurate branch local radius and direction
            local_radius = sum(p['radius'] for p in sample) / len(sample)
            direction = calculate_average_direction(arm, True, 10)
            direction_arms.append({'seg_idx': seg_idx, 'local_radius': local_radius, 'direction': direction, 'arm_points': arm})
        # Backward arm (bifurcation → start of segment)
        if bif_pos > 0:
            arm = seg_pts[:bif_pos + 1][::-1]
            skip = min(3, max(0, len(arm) - 2))
            sample = arm[skip:skip + 10] or arm[-min(5, len(arm)):]
            local_radius = sum(p['radius'] for p in sample) / len(sample)
            direction = calculate_average_direction(arm, True, 10)
            direction_arms.append({'seg_idx': seg_idx, 'local_radius': local_radius, 'direction': direction, 'arm_points': arm})

    # Cluster by direction similarity (dot product > 0.8 = same direction)
    clusters = []
    for arm in direction_arms:
        arm_direction = arm['direction']
        d_norm = math.sqrt(arm_direction[0] ** 2 + arm_direction[1] ** 2 + arm_direction[2] ** 2)
        if d_norm == 0:
            continue

        placed = False
        for cluster in clusters:
            cluster_direction = cluster[0]['direction']
            cd_norm = math.sqrt(cluster_direction[0] ** 2 + cluster_direction[1] ** 2 + cluster_direction[2] ** 2)
            if cd_norm == 0:
                continue
            dot = (arm_direction[0] * cluster_direction[0] + arm_direction[1] * cluster_direction[1] + arm_direction[2] * cluster_direction[2]) / (d_norm * cd_norm)
            if dot > 0.8:
                cluster.append(arm)
                placed = True
                break
        if not placed:
            clusters.append([arm])

    # Keep the thickest arm per cluster
    bifurcation_arms = []
    for cluster in clusters:
        best = max(cluster, key=lambda a: a['local_radius'])
        bifurcation_arms.append(best)
    bifurcation_arms.sort(key=lambda c: c['local_radius'], reverse=True)

    if len(bifurcation_arms) <= 2:
        return None  # Not a true branch point (trunk-overlap)
    return bifurcation_arms

# Label every segment as 'Aorta' or 'Artery', further artery classification will build upon these labels
def label_aorta_and_arteries(segments, point_to_segments, bifurcations, bifurcation_arms):
    if not bifurcations:
        return {}, []
    
    # Reset labels
    for seg in segments:
        seg['label'] = None
        seg['confirmed_aorta'] = False # True means the bifurcation loop explicitly kept it as trunk, False means BFS propagated the label into a branch stub

    # Sort bifurcations radius wise descendingly
    sorted_bifs = sorted(bifurcations, key=lambda b: b['radius'], reverse=True)
    bifurcations_map = {}   # Maps bifurcation_id → (bifurcation_dict, connected_segments_list)
    bifurcations_info = []

    # Classify every bifurcation
    for bif in sorted_bifs:
        segment_connections = bifurcation_arms.get(bif['id'])
        if segment_connections is None:
            continue
        bifurcations_map[bif['id']] = (bif, segment_connections)
        # Retrieve information about connected segments to bifurcation (average radius and aorta vs artery segments count)
        arm_info = [(c, calculate_distal_radius(c['arm_points'])) for c in segment_connections]
        confirmed_count = sum(1 for c, _ in arm_info if segments[c['seg_idx']].get('confirmed_aorta'))
        artery_count = sum(1 for c, _ in arm_info if segments[c['seg_idx']]['label'] == 'Artery')

        # If bifurcation is inside an artery sub-tree (no aorta arms), we label all segments connected to this bifurcation as artery
        if confirmed_count == 0 and artery_count > 0:
            for c, _ in arm_info:
                if segments[c['seg_idx']]['label'] is None:
                    apply_label(segments[c['seg_idx']], 'Artery')
            bifurcations_info.append({
                'id': int(bif['id']), 'x': float(bif['x']), 'y': float(bif['y']), 'z': float(bif['z']), 'radius': float(bif['radius']),
                'artery_type': None, 'gap_to_next_bifurcation': None, 'is_iliac_cutoff': False, # Needed when classifying arteries into different regions
                'connected_arms_labels': [segments[c['seg_idx']]['label'] for c, _ in arm_info],
                'arm_endpoint_z': [float(c['arm_points'][-1]['z']) if c['arm_points'] else 0.0 for c, _ in arm_info], # Needed when checking if this is iliac cutoff
                'arm_endpoint_x': [float(c['arm_points'][-1]['x']) if c['arm_points'] else 0.0 for c, _ in arm_info], # Needed when checking if this is iliac cutoff
            })
            continue

        # Break if there are no arms that are unlabeled or already aorta labeled
        aorta_arms = [(c, r) for c, r in arm_info if segments[c['seg_idx']]['label'] in (None, 'Aorta')]
        if not aorta_arms:
            continue
        # Sort connected aorta/None arms descending by radius
        aorta_arms_candidates = sorted(aorta_arms, key=lambda x: x[1], reverse=True)

        # Decide number of arms that belong to aorta trunk
        aorta_arms_count = compute_aorta_arms_count(aorta_arms_candidates, segments)
        largest_arm_radius = aorta_arms_candidates[0][1] # Largest radius of aorta arms

        # Apply labels to the connected arms to this bifurcation
        for i, (c, _) in enumerate(aorta_arms_candidates):
            seg = segments[c['seg_idx']]
            if i < aorta_arms_count:
                apply_label(seg, 'Aorta')
                seg['confirmed_aorta'] = True
            else:
                if not seg.get('confirmed_aorta'):
                    apply_label(seg, 'Artery')

        # Handle direction-clustered hidden segments (At noisy bifurcations, multiple connected segments may exist at the same bifurcation point)
        connection_segments = set(c['seg_idx'] for c, _ in arm_info)
        all_connected = set(point_to_segments.get(bif['id'], []))
        for seg_idx in (all_connected - connection_segments):
            seg = segments[seg_idx]
            if seg.get('confirmed_aorta'):
                continue
            seg_r = 0.0
            pts = segments[seg_idx]['points']
            if pts:
                seg_r = sum(p['radius'] for p in pts) / len(pts)
                
            if largest_arm_radius > 0 and seg_r / largest_arm_radius >= AORTA_CONTINUATION_RATIO:
                apply_label(seg, 'Aorta')
                seg['confirmed_aorta'] = True
            elif seg['label'] is None:
                apply_label(seg, 'Artery')

        bifurcations_info.append({
            'id': int(bif['id']), 'x': float(bif['x']), 'y': float(bif['y']), 'z': float(bif['z']), 'radius': float(bif['radius']), 
            'artery_type': None, 'gap_to_next_bifurcation': None, 'is_iliac_cutoff': False,
            'connected_arms_labels': [segments[c['seg_idx']]['label'] for c, _ in arm_info],
            'arm_endpoint_z': [float(c['arm_points'][-1]['z']) if c['arm_points'] else 0.0 for c, _ in arm_info],
            'arm_endpoint_x': [float(c['arm_points'][-1]['x']) if c['arm_points'] else 0.0 for c, _ in arm_info],
        })

    # Propagate the labeled segments to their unlabeled neighbors
    propagate_labeled_segments(segments, point_to_segments)

    # Safety net: if no segment ended up labeled 'Aorta' (happens when every junction failed the 3-cluster gate, like a dataset with only simple T-junctions)
    # So we re-seed the longest segment as 'Aorta' and re-flood the labels propagation
    if not any(seg.get('label') == 'Aorta' for seg in segments):
        seeded = max((s for s in segments if s.get('points')), key=lambda s: len(s['points']), default=None)
        if seeded is not None:
            for s in segments:
                s['label'] = None
            apply_label(seeded, 'Aorta')
            propagate_labeled_segments(segments, point_to_segments)

    # Safety-net relabel may change final connected-arm labels
    for bif_info in bifurcations_info:
        seg_connections = bifurcation_arms.get(bif_info['id'])
        if not seg_connections:
            continue
        bif_info['connected_arms_labels'] = [segments[c['seg_idx']]['label'] for c in seg_connections]

    return bifurcations_map, bifurcations_info

# Helper that calculates average inscribed-sphere radius over the distal half of an arm
# Skipping the proximal half avoids radius inflation that occurs when the iliac artery origin lies inside the aortic sac
def calculate_distal_radius(arm_points):
    # arm_points[0] is the junction (proximal) end; arm_points[-1] is the distal tip.
    n = len(arm_points)
    if n == 0:
        return 0
    start = 0
    if n > 4:
        start = n // 2
    pts = arm_points[start:]
    return sum(p['radius'] for p in pts) / len(pts)

# Decide how many of the labeled Aorta/None arms connected to a bifurcation to label/keep as 'Aorta'
def compute_aorta_arms_count(aorta_arms_candidates, segments):
    largest_arm_radius = 0
    aorta_arms_count = 0
    if aorta_arms_candidates:
        largest_arm_radius = aorta_arms_candidates[0][1]

    # If candidates are 3 or more for a bifurcations (maximum arms count is usually 4) then compute radius differences, where largest gap indicates branch separation
    if len(aorta_arms_candidates) >= 3:
        radii = [r for _, r in aorta_arms_candidates]
        gaps  = [radii[i] - radii[i + 1] for i in range(len(radii) - 1)]
        aorta_arms_count = gaps.index(max(gaps)) + 1

    # If candidates are 2 and they have similar radii, then they both are probably the aorta inlet and outlet of the bifurcation so keep both, otherwise keep only one
    elif len(aorta_arms_candidates) == 2:
        second_largest_arm_radius = aorta_arms_candidates[1][1]
        if (largest_arm_radius > 0 and second_largest_arm_radius / largest_arm_radius >= AORTA_CONTINUATION_RATIO):
            aorta_arms_count = 2
        else:
            aorta_arms_count = 1
    else:
        aorta_arms_count = 1

    if aorta_arms_count <= 1:
        return aorta_arms_count

    # Add guards only if a confirmed-aorta arm is already in the keep group
    confirmed_aorta_max_radius = max((r for c, r in aorta_arms_candidates[:aorta_arms_count] if segments[c['seg_idx']].get('confirmed_aorta')), default=0)
    if confirmed_aorta_max_radius == 0:
        return aorta_arms_count

    # Ratio threshold: any non-confirmed kept arm must be ≥ 75% of aorta trunk
    for i in range(aorta_arms_count - 1, 0, -1):
        curr_connection, curr_radius = aorta_arms_candidates[i]
        if (not segments[curr_connection['seg_idx']].get('confirmed_aorta') and curr_radius < AORTA_CONTINUATION_RATIO * confirmed_aorta_max_radius):
            aorta_arms_count = i
            break
    aorta_arms_count = max(1, aorta_arms_count)

    # Check symmetric branches where two thinner arms are closer to each other than to aorta trunk
    if aorta_arms_count == 2 and len(aorta_arms_candidates) == 3:
        ra = aorta_arms_candidates[0][1]
        conn2, rb = aorta_arms_candidates[1]
        rc = aorta_arms_candidates[2][1]
        if (rb > 0 and ra > 0 and (rc / rb) > (rb / ra) and not segments[conn2['seg_idx']].get('confirmed_aorta')):
            aorta_arms_count = 1

    return aorta_arms_count

# Fill unlabeled segments with labels from already-labeled neighbors
def propagate_labeled_segments(segments, point_to_segments):
    # Add all labeled segments to our queue and set them as visited
    queue = deque()
    visited = set()
    for seg_idx, seg in enumerate(segments):
        if seg['label'] is not None:
            queue.append(seg_idx)
            visited.add(seg_idx)

    # Loop over each labeled segment, filling its unlabeled neighbors with the same label
    while queue:
        # Retrieve current labeled segment info
        seg_idx = queue.popleft()
        label = segments[seg_idx]['label']
        pts = segments[seg_idx]['points']
        if not pts:
            continue
        
        # Loop over each point in the labeled segment
        for point_id in {pts[0]['id'], pts[-1]['id']}:
            for neighbor_seg_idx in point_to_segments.get(point_id, []):
                # Exit if neighbor was already visited or labeled
                if neighbor_seg_idx in visited:
                    continue
                if segments[neighbor_seg_idx]['label'] is not None:
                    visited.add(neighbor_seg_idx)
                    continue
                # Apply label only to unlabeled segments
                visited.add(neighbor_seg_idx)
                apply_label(segments[neighbor_seg_idx], label)
                queue.append(neighbor_seg_idx) # Add this neighbor to the queue as after labeling, it may still have unlabeled neighbors

    # Fallback that labels unseen segments as artery
    for seg in segments:
        if seg['label'] is None:
            apply_label(seg, 'Artery')

# Detects if the dataset contains the thoracic aorta (aortic arch)
def detect_thoracic_region(centerlines_vtk):
    thoracic_info = {'detected': False, 'reversal_mm': 0.0, 'arch_zone_z': 0.0}
    pts = centerlines_vtk.GetPoints()
    if pts is None or pts.GetNumberOfPoints() == 0:
        return thoracic_info

    n_pts = pts.GetNumberOfPoints()
    z_all = [pts.GetPoint(i)[2] for i in range(n_pts)]
    z_min = min(z_all)
    z_max = max(z_all)
    z_range = z_max - z_min

    # The indicator of the aortic arch is a near-180° U-turn in the superior portion of the scan.
    # This appears as a Z-reversal: along a single centerline cell the Z-coordinate first increases (ascending aorta) then decreases(descending aorta) or vice-versa 
    TOP_FRACTION = 0.30   # Inspect top 30% of Z extent
    MIN_ARCH_SPAN_MM = 40.0   # Reversal region must span >= 40 mm
    z_arch_start = z_max - TOP_FRACTION * z_range
    biggest_reversal = 0.0

    n_cells = centerlines_vtk.GetNumberOfCells()
    for cell_idx in range(n_cells):
        cell = centerlines_vtk.GetCell(cell_idx)
        cell_pts = cell.GetPoints()
        if cell_pts is None:
            continue
        n_cell = cell_pts.GetNumberOfPoints()

        # Extract Z of points that lie inside the arch zone
        arch_z = [cell_pts.GetPoint(j)[2] for j in range(n_cell) if cell_pts.GetPoint(j)[2] >= z_arch_start]

        if len(arch_z) < 4:
            continue

        peak_z = max(arch_z)
        trough_z = min(arch_z)

        # Case 1: dome (ascending → arch peak → descending)
        dome_span = min(peak_z - arch_z[0], peak_z - arch_z[-1])
        # Case 2: inverse dome (descending → bottom → ascending)
        inv_span  = min(arch_z[0] - trough_z, arch_z[-1] - trough_z)

        reversal = max(dome_span, inv_span)
        if reversal > biggest_reversal:
            biggest_reversal = reversal

        if reversal >= MIN_ARCH_SPAN_MM:
            return {'detected': True, 'reversal_mm': round(reversal, 1), 'arch_zone_z': round(z_arch_start, 1)}
        
    return {'detected': False, 'reversal_mm': round(biggest_reversal, 1), 'arch_zone_z': round(z_arch_start, 1)}

# Classify artery bifurcations into thoracic / abdomen / iliac zones, returns iliac cutoff bifurcation
def classify_arteries_by_zone(bifurcations_info, segments, has_thoracic):
    # Keep only bifurcations where at least one arm ended up labeled 'Artery'
    bifurcations_indices = [i for i, bif in enumerate(bifurcations_info) if 'Artery' in (bif.get('connected_arms_labels') or [])]
    if not bifurcations_indices:
        return None

    # Sort descending by Z (highest = most superior)
    bifurcations_indices.sort(key=lambda i: bifurcations_info[i]['z'], reverse=True)

    # Compute Z-gap to the next artery bifurcation below each one
    for idx, bif_idx in enumerate(bifurcations_indices):
        if idx + 1 < len(bifurcations_indices):
            nxt_bif_idx = bifurcations_indices[idx + 1]
            gap = bifurcations_info[bif_idx]['z'] - bifurcations_info[nxt_bif_idx]['z']
            bifurcations_info[bif_idx]['gap_to_next_bifurcation'] = round(gap, 1)

    # Zone classification
    thoracic_zone = []
    abdomen_zone = []
    ptr = 0
    n = len(bifurcations_indices)

    # Collect thoracic bifurcations until thoracic gap is found (or end of bifurcations list), has_thoracic indicates if dataset contains thoracic part or not
    if has_thoracic:
        while ptr < n:
            idx = bifurcations_indices[ptr]
            thoracic_zone.append(idx)
            gap = bifurcations_info[idx]['gap_to_next_bifurcation']
            ptr += 1
            if gap is None or gap >= THORACIC_GAP_MM:
                break

    # Collect abdomen bifurcations after the thoracic cluster (or its the first bifurcations in datasets that start at abdomen region having no thoracic arch)
    while ptr < n and len(abdomen_zone) < len(ABDOMEN_LABELS):
        idx = bifurcations_indices[ptr]
        abdomen_zone.append(idx)
        gap = bifurcations_info[idx]['gap_to_next_bifurcation']
        ptr += 1
        if gap is None or gap >= ABDOMEN_CLUSTER_MAX_GAP_MM:
            break 

    for idx in thoracic_zone:
        bifurcations_info[idx]['artery_type'] = 'thoracic'
    for idx in abdomen_zone:
        bifurcations_info[idx]['artery_type'] = 'abdomen'

    # Relabel artery segments which appear in thoracic region as thoraic arteries
    if has_thoracic and thoracic_zone:
        # Compute boundary between thoracic bifurcations and abdomen bifurcations
        if abdomen_zone:
            z_min_thoracic = min(bifurcations_info[i]['z'] for i in thoracic_zone)
            z_max_abdomen = max(bifurcations_info[i]['z'] for i in abdomen_zone)
            z_boundary = (z_min_thoracic + z_max_abdomen) / 2.0
        else:
            z_min_thoracic = min(bifurcations_info[i]['z'] for i in thoracic_zone)
            z_boundary = z_min_thoracic - 20.0  # fallback: 20 mm below lowest thoracic junction

        for seg in segments:
            if seg.get('label') == 'Artery' and seg.get('points'):
                avg_z = sum(pt['z'] for pt in seg['points']) / len(seg['points'])
                if avg_z > z_boundary:
                    apply_label(seg, 'Thoracic Artery')

    # Walk the remaining bifurcations (below the abdomen cluster, Z descending), detect iliac cutoff by scoring how strongly an arm trends inferiorly versus laterally
    reference_zone = abdomen_zone if abdomen_zone else thoracic_zone
    iliac_cutoff = None
    cutoff_pos = None
    cutoff_idx = None
    if reference_zone:
        for k in range(ptr, n):
            idx = bifurcations_indices[k]
            bif = bifurcations_info[idx]
            bif_z = bif['z']
            bif_x = bif['x']
            arm_labels = bif.get('connected_arms_labels', [])
            arm_endpoint_z = bif.get('arm_endpoint_z', [])
            arm_endpoint_x = bif.get('arm_endpoint_x', [])

            # Iliac like means any Artery arm drops more in Z than it spreads in X (lateral noise branches like lumbars move far in X but barely drop in Z)
            iliac_like = any(((bif_z - ez) > ILIAC_Z_DROP_RATIO * abs(ex - bif_x)) for ez, ex, lbl in zip(arm_endpoint_z, arm_endpoint_x, arm_labels) if lbl == 'Artery')
            if iliac_like:
                cutoff_pos = k
                cutoff_idx = idx
                break

        # Fallback: no iliac-like junction found — use first bifurcation below the reference zone if one exists
        # otherwise fall back to the last bifurcation in the reference zone
        if cutoff_idx is None:
            if ptr < n:
                cutoff_pos = ptr
                cutoff_idx = bifurcations_indices[ptr]
            else:
                cutoff_pos = n - 1
                cutoff_idx = reference_zone[-1]

        iliac_bif = bifurcations_info[cutoff_idx]
        bifurcations_info[cutoff_idx]['is_iliac_cutoff'] = True
        iliac_cutoff = {'x': float(iliac_bif['x']), 'y': float(iliac_bif['y']), 'z': float(iliac_bif['z']), 'radius': float(iliac_bif['radius'])}

    # Mark the cutoff bifurcation and everything below it as 'iliac'
    if cutoff_pos is not None:
        for k in range(cutoff_pos, n):
            bifurcations_info[bifurcations_indices[k]]['artery_type'] = 'iliac'

    # Relabel segments below iliac cutoff as iliac arteries
    if iliac_cutoff:
        cutoff_z = iliac_cutoff['z']
        cutoff_x = iliac_cutoff['x']
        for seg in segments:
            if seg.get('label') in ('Aorta', 'Artery') and seg.get('points'):
                pts = seg['points']
                # Use centroid Z: avoids mislabeling the aortic trunk whose last point sits exactly at the bifurcation
                # Iliac arms extend well below cutoff_z, so their centroid is clearly < cutoff_z
                centroid_z = sum(pt['z'] for pt in pts) / len(pts)
                if centroid_z < cutoff_z:
                    avg_x = sum(pt['x'] for pt in pts) / len(pts)
                    if avg_x < cutoff_x:
                        apply_label(seg, 'Left Iliac Artery')
                    else:
                        apply_label(seg, 'Right Iliac Artery')

    return iliac_cutoff

# Finds and labels 4 main abdominal arteries (celiac, SMA, renals) using already-classified abdomen-zone bifurcations
def identify_abdominal_arteries(segments, point_to_segments, bifurcations_info):
    if not bifurcations_info:
        return []

    # Take only abdomen-zone bifurcations and order from superior -> inferior
    abdomen_candidates = [b for b in bifurcations_info if b.get('artery_type') == 'abdomen']
    if not abdomen_candidates:
        return []
    abdomen_candidates.sort(key=lambda b: b['z'], reverse=True)

    abdomen_ostia = []
    label_idx = 0
    labeled_seg_indices = set()

    for bif in abdomen_candidates:
        if label_idx >= len(ABDOMEN_LABELS):
            break

        bif_id = bif.get('id')
        if bif_id is None:
            continue

        # Recompute directional arms for this bifurcation on current segments
        connections = compute_bifurcation_arms(bif, segments, point_to_segments)
        if not connections:
            continue

        # Skip the first two arms (aorta inlet/outlet), label remaining branch arms
        # Sort by arm endpoint Z descending: superior branch always gets the earlier label
        branch_arms = sorted(connections[2:], key=lambda b: b['arm_points'][-1]['z'] if b['arm_points'] else 0.0, reverse=True)
        for branch in branch_arms:
            if label_idx >= len(ABDOMEN_LABELS):
                break
            if branch['seg_idx'] in labeled_seg_indices:
                continue
            labeled_seg_indices.add(branch['seg_idx'])

            label = ABDOMEN_LABELS[label_idx]
            branch_dir = branch['direction']
            arm_pts = branch['arm_points']

            branch_label = label
            if label == 'Renal':
                if branch_dir[0] > 0:
                    branch_label = 'Right Renal' 
                else:
                    branch_label = 'Left Renal'

            apply_label(segments[branch['seg_idx']], branch_label)

            target = (bif['x'] + bif['radius'] * branch_dir[0], bif['y'] + bif['radius'] * branch_dir[1], bif['z'] + bif['radius'] * branch_dir[2])
            best_pt = min(arm_pts, key=lambda p: (p['x'] - target[0]) ** 2 + (p['y'] - target[1]) ** 2 + (p['z'] - target[2]) ** 2)

            abdomen_ostia.append({
                'x': best_pt['x'], 'y': best_pt['y'], 'z': best_pt['z'],
                'nx': branch_dir[0], 'ny': branch_dir[1], 'nz': branch_dir[2],
                'label': branch_label, 'radius': best_pt['radius'],
            })
            label_idx += 1

    return abdomen_ostia

# Places the default reference point at the most caudal abdomen bifurcation(actual junction), otherwise falls back to abdomen ostia
def compute_reference_point(segments, bifurcations_info, abdomen_ostia=None):
    abdomen_bifs = [b for b in (bifurcations_info or []) if b.get('artery_type') == 'abdomen']
    if abdomen_bifs:
        last_abdomen_bif = min(abdomen_bifs, key=lambda b: b['z'])
        return {'x': last_abdomen_bif['x'], 'y': last_abdomen_bif['y'], 'z': last_abdomen_bif['z']}

    if abdomen_ostia:
        last_abdomen_ostium = min(abdomen_ostia, key=lambda b: b['z'])
        return {'x': last_abdomen_ostium['x'], 'y': last_abdomen_ostium['y'], 'z': last_abdomen_ostium['z']}

    aorta_points = [pt for seg in segments if seg['label'] == 'Aorta' for pt in seg['points']]
    if not aorta_points:
        return {'x': 0, 'y': 0, 'z': 0}

    return {
        'x': sum(p['x'] for p in aorta_points) / len(aorta_points),
        'y': sum(p['y'] for p in aorta_points) / len(aorta_points),
        'z': sum(p['z'] for p in aorta_points) / len(aorta_points),
    }

# Build Frontend response JSON
def build_response(centerlines_vtk, segments, abdomen_ostia, reference_point, iliac_cutoff, bifurcations_map=None, bifurcations_info=None):
    bounds = centerlines_vtk.GetBounds()

    # Orient artery segments so points run bifurcation → endpoint
    orient_artery_segments(segments, bifurcations_map)

    # Flatten segments into a single points list + segment index ranges
    flat_points = []
    segment_ranges = []

    for seg in segments:
        if not seg['points']:
            continue  # Skip empty segments
        start_idx = len(flat_points) # Currently it's empty so start_idx is 0, second iteration it will be the length of first segment and so on
        for pt in seg['points']:
            flat_points.append({'x': pt['x'], 'y': pt['y'], 'z': pt['z'], 'radius': pt['radius'], 'label': seg['label']})
        segment_ranges.append({'start': start_idx, 'end': len(flat_points) - 1, 'label': seg['label']})

    return {
        'points': flat_points,
        'segments': segment_ranges,
        'abdomen_ostia': abdomen_ostia,
        'bifurcations': [
            {
                'id': int(b['id']),
                'x': float(b['x']),
                'y': float(b['y']),
                'z': float(b['z']),
                'radius': float(b['radius']),
                'artery_type': b.get('artery_type'),
                'is_iliac_cutoff': bool(b.get('is_iliac_cutoff', False)),
            }
            for b in (bifurcations_info or [])
        ],
        'reference_point': reference_point,
        'iliac_cutoff': iliac_cutoff,
        'bounds': {
            'x_min': float(bounds[0]), 'x_max': float(bounds[1]),
            'y_min': float(bounds[2]), 'y_max': float(bounds[3]),
            'z_min': float(bounds[4]), 'z_max': float(bounds[5]),
        },
    }

# Ensures artery (non-Aorta) segments are ordered bifurcation → endpoint
def orient_artery_segments(segments, bifurcations_map=None):
    # Fast path: use bifurcation-map arm directions to orient covered segments
    oriented_from_map = set()
    if bifurcations_map:
        for _, (_, connections) in bifurcations_map.items():
            for c in connections:
                seg_idx = c['seg_idx']
                if seg_idx < 0 or seg_idx >= len(segments):
                    continue
                seg = segments[seg_idx]
                if seg.get('label') == 'Aorta':
                    continue
                pts = seg.get('points', [])
                arm_pts = c.get('arm_points', [])
                if len(pts) < 2 or len(arm_pts) < 2:
                    continue
                # arm_points are built bifurcation -> distal. Align segment to that
                if pts[0]['id'] != arm_pts[0]['id'] and pts[-1]['id'] == arm_pts[0]['id']:
                    pts.reverse()
                oriented_from_map.add(seg_idx)

    # Collect all Aorta segment endpoints (first and last point of each Aorta seg)
    aorta_endpoints = []
    aorta_endpoint_ids = set()
    for seg in segments:
        if seg.get('label') != 'Aorta':
            continue
        pts = seg.get('points', [])
        if not pts:
            continue
        aorta_endpoints.append(pts[0])
        aorta_endpoint_ids.add(pts[0]['id'])
        if len(pts) > 1:
            aorta_endpoints.append(pts[-1])
            aorta_endpoint_ids.add(pts[-1]['id'])

    for seg_idx, seg in enumerate(segments):
        if seg.get('label') == 'Aorta':
            continue
        if seg_idx in oriented_from_map:
            continue
        pts = seg.get('points', [])
        if len(pts) < 2:
            continue

        if aorta_endpoints:
            # O(1) endpoint-id check first
            if pts[0]['id'] in aorta_endpoint_ids:
                continue
            if pts[-1]['id'] in aorta_endpoint_ids:
                pts.reverse()
                continue

            # Proximity check: which end of artery segment is closest to any Aorta endpoint?
            p0, p_last = pts[0], pts[-1]
            min_dist_start = min((p0['x'] - end_pt['x']) ** 2 + (p0['y'] - end_pt['y']) ** 2 + (p0['z'] - end_pt['z']) ** 2 for end_pt in aorta_endpoints)
            min_dist_end = min((p_last['x'] - end_pt['x']) ** 2 + (p_last['y'] - end_pt['y']) ** 2 + (p_last['z'] - end_pt['z']) ** 2 for end_pt in aorta_endpoints)
            # If the END is closer to the aorta, the segment is backwards → reverse it
            if min_dist_end < min_dist_start:
                pts.reverse()
        else:
            # Fallback: radius heuristic (bifurcation end is thicker)
            r_first = pts[0]['radius']
            r_last = pts[-1]['radius']
            if r_first < r_last:
                pts.reverse()