# Visualize Anatomical Groups
# ===========================

# This example visualizes different anatomical groups using the segmentation
# labels available from the downloadable datasets
# `download_whole_body_ct_female()` and
# `download_whole_body_ct_male()`.
# These datasets include labels for 117 anatomical structures. In this example,
# the labels are grouped by filtering the list of labels and coloring the
# labels with the recommended RGB values used by the 3DSlicer
# TotalSegmentator Extension.
import pyvista as pv
from pyvista import examples

# Load Dataset
# ------------

# Load a TotalSegmentator dataset. Can be the
# `male` or
# `female` female
# subjects included with PyVista, or any other subject from the
# TotalSegmentator dataset.
dataset = examples.download_whole_body_ct_female()

# Get the label map used for plotting different anatomical groups.
label_map = dataset['label_map']

# Get metadata associated with the dataset
# Get a list of all label names. This list will be filtered by group.
label_names = dataset['segmentations'].keys()

# Get color and id mappings included with the dataset. These are used to filter and
# color the contours.
names_to_colors = dataset.user_dict['names_to_colors']
names_to_ids = dataset.user_dict['names_to_ids']
ids_to_colors = dataset.user_dict['ids_to_colors']

# Color Mapping
# ~~~~~~~~~~~~~

# Show the color mapping included with the dataset. Print the dictionary and format
# it to visually align the RGB values. The formatted dictionary is valid python code.
print('{')
for name, (R, G, B) in names_to_colors.items():
    print(f'{name!r:<32}: ({R:>3}, {G:>3}, {B:>3}),')
print('}')

# Utility Functions
# -----------------

# Define helper functions to visualize the data.

# filter_labels
# ~~~~~~~~~~~~~

# To visualize a particular group of anatomic regions, we first define a function
# to filter the labels by name. Given a list of terms, the function returns
# any label which contains any of the search terms.
def filter_labels(label_names: list[str], search_terms: list[str]):
    def include_label(label_name: str):
        return any(target in label_name for target in search_terms)

    return [label for label in label_names if include_label(label)]

# plot_anatomy
# ~~~~~~~~~~~~

# Define a function which, given a list of terms, will look up labels associated
# with those terms, generate contours for the labels, and plot the result. The function
# uses `contour_labels()` for generating contours and
# `color_labels()` for coloring them.
def plot_anatomy(search_terms: list[str]):
    # Get a list of labels which contain any of the listed terms.
    group_names = filter_labels(label_names, search_terms)

    # Get the label ids corresponding to the matched labels.
    group_ids = [names_to_ids[name] for name in group_names]

    # Selectively generate surfaces for the specified labels
    group_surface = dataset['label_map'].contour_labels(select_inputs=group_ids)

    # Color the labels with the color mapping
    colored_surface = group_surface.color_labels(colors=ids_to_colors)

    # Plot the label map.
    pl = pv.Plotter()
    pl.add_mesh(colored_surface)
    pl.view_zx()
    pl.camera.up = (0, 0, 1)
    pl.show()

# Anatomical Groups
# -----------------

# Group the labels and visualize the result.

# Cardiovascular System
# ~~~~~~~~~~~~~~~~~~~~~

# Show segments of the cardiovascular system.
# Define terms which describe all relevant segments.
cardio = [
    'heart',
    'aorta',
    'artery',
    'brachiocephalic_trunk',
    'vein',
    'atrial_appendage',
    'vena_cava',
]

# Plot the labels associated with these terms.
plot_anatomy(cardio)

# Gastrointestinal System
# ~~~~~~~~~~~~~~~~~~~~~~~

# Show segments of the gastrointestinal system.
# Define terms which describe all relevant segments.
gastro = [
    'esophagus',
    'stomach',
    'duodenum',
    'small_bowel',
    'colon',
    'urinary_bladder',
]

# Plot the labels associated with these terms.
plot_anatomy(gastro)

# Spine
# ~~~~~

# Show segments of the spinal.
# Define terms which describe all relevant segments.
spine = [
    'spinal_cord',
    'vertebrae',
    'sacrum',
]

# Plot the labels associated with these terms.
plot_anatomy(spine)

# Other Organs
# ~~~~~~~~~~~~

# Show other organs not included in the cardiovascular or gastrointestinal
# systems.
# Define terms which describe all relevant segments.
other_organs = [
    'brain',
    'spinal_cord',
    'thyroid_gland',
    'trachea',
    'lung',
    'adrenal_gland',
    'spleen',
    'liver',
    'gallbladder',
    'kidney',
    'pancreas',
    'prostate',
]

# Plot the labels associated with these terms.
plot_anatomy(other_organs)

# Muscles
# ~~~~~~~

# Show the muscles.
# Define terms which describe all relevant segments.
muscles = [
    'gluteus',
    'autochthon',
    'iliopsoas',
]

# Plot the labels associated with these terms.
plot_anatomy(muscles)

# Ribs
# ~~~~

# Show the ribs.
# Define terms which describe all relevant segments.
ribs = [
    'rib',
    'sternum',
    'costal_cartilages',
]

# Plot the labels associated with these terms.
plot_anatomy(ribs)

# Skeleton
# ~~~~~~~~

# Show the skeleton.
# Define terms which describe all relevant segments.
skeleton = [
    'skull',
    'clavicula',
    'scapula',
    'humerus',
    'vertebrae',
    'sternum',
    'rib',
    'costal_cartilages',
    'hip',
    'sacrum',
    'femur',
]

# Plot the labels associated with these terms.
plot_anatomy(skeleton)

# ----------------------------------------------------------------------
# Generated by sphinx-examples-as-code https://github.com/pyvista/sphinx-examples-as-code