Using Rerun

Rerun is both a time-aware multimodal data viewer. That means its designed to display data of all different kinds (text, images, segmentation masks, tensors, meshes, point sets, etc) that may have associated time stamps to them.

In rerun, supported data types are called _archetypes, and each archetype has its own dedicated viewer. For example, numbers can be plotted on a plot-like viewer similar to Plotly. Images are displayed as in Figma, with infinite zoom in/out and lenses for pixel information. Tensors (n-dimensional arrays) are displayed similar to images, but without zooming but the ability to swap directions and dimensions.

All of these different viewers can be displayed simultaneously within a single window, and scrolled simultanously according to their associated timestamps. If no timestamps are given, there is a default one, and objects of a same archetype are assigned sequential timestamps in the order they are sent to the viewer.

The organization of displays within the rerun window is called a _blueprint, which can be saved and used across projects. A project, which we could here call a collection of data, is called a _recording.

Installation

First, install rerun with your favorite package manager:

conda install -c conda-forge rerun-sdk
pixi add rerun-sdk
pip install rerun-sdk

If you plan to use it within jupyter notebooks, you need to install the notebook extension.

conda install -c conda-forge "rerun-sdk[notebook]"
pixi add "rerun-sdk[notebook]" # doesn't seem to work on zsh
pip install "rerun-sdk[notebook]"

Usage

Usage has 2 steps:

1. instantiate the viewer
2. send data to the viewer

Instantiate the viewer

Local use case

We instantiate a rerun viewer upon which to send our data by calling rerun init and giving a name to our _recording.

import rerun as rr

rr.init("Rerun Viewer Example")

This would be analogous to open an empty figure in matplotlib

import matplotlib.pyplot as plt

plt.figure("Matplotlib Figure Example")

If you are using rerun locally, you may wish to spawn a rerun viewer as a separate app by passing the argument spawn=True. You may also with to have rerun not accept malformed commands by passing the argument strict=True. In that case a init call at the top of a script with rather look like.

import rerun as rr

rr.init("Rerun Viewer Example", spawn=True, strict=True)

Remote, or notebook use cases

For within notebook or remote usage, we would not wish to spawn a local viewer, but rather have rerun be served on the web, or placed within the notebook itself. That means, for remote use, we would call

# Remote usage
import rerun as rr

rr.init("Rerun Viewer Example", strict=True)
rr.serve_web(web_port=8081, ws_port=9877) # or whatever

Then on our local computer we would connect to that web_port by calling ssh -NL $LOCAL_PORT:localhost:$WEB_PORT -L WS_PORT:localhost:WS_PORT gpu. and then in point our browser to:

http://localhost:LOCAL_PORT/?url=ws://localhost:WS_PORT

For notebook use, we call notebook_show() at the end of a cell.

# Notebook usage
import rerun as rr

rr.init("Rerun Viewer Example", strict=True)

# do stuff

rr.notebook_show()

Send data to the viewer

Once we have a recording in place we can log our data and associated timestamps to it. Note that timestamps are here a global x-axis, and do not necessarily have to represent time.

If we wish to make use of time or sequence ordering, we first call: rr.set_time_sequence(time). Now, all data logged after this call will be logged associated with time.

Now, the actually send the data to the viewer, we call rr.log:

import rerun as rr
import numpy as np

rr.init("Rerun Viewer Example", strict=True)

tensor = np.random.rand(32, 32, 16)
rr.log("Example tensor", rr.Tensor(tensor, dim_names=("x", "y", "z"))

image2d = np.random.rand(32, 32)
rr.log("Example image", rr.Image(image2d))

Tensor API. Image API.

For more info and archetypes, refer to their docs.

Two examples

Here are two examples showing how to view images.

Viewing all images and segmentation masks from the Bamberg dataset.

The following example will open a remote server with all of the T2 and anatomy segmentation masks from the Bamberg dataset. It should work as is.

"""
Example Rerun to log all of the segmentation masks overlaid on the anatomical images.
"""

import os
import rerun as rr
from nilearn import image
import numpy as np
import einops
import glob


# Start rerun, but don't spawn the app.
# strict: raise errors instead of warnings if we misuse the API
rr.init("Bamberg anatomy explorer", spawn=False, strict=True)
# since we are not spawning, serve on the web instead.
rr.serve_web(web_port=8081, ws_port=9877)


def pad(img, target_nslice=35):
    """
    Pad the image on the z direction so that all mosaics have the same size.
    """
    padval = target_nslice - img.shape[2]
    return np.pad(img, [(0, 0), (0, 0), (0, padval)], mode="constant")

def clamp(img):
    """
    Clamp the image to a certain range.
    """
    # vmin, vmax are the 2 and 98 percentiles
    vmin, vmax = np.percentile(img, (2, 98))
    return np.clip(img, vmin, vmax)

def mosaic(img):
    """
    Make a 3D image into 2D.
    Put it in RAS orientation.

    Note this could be wrong if data is acquired in a different orientation than Bamberg TRA.
    x: R->L
    y: A->P
    z: Top left corner is bottom slice
    """
    return einops.rearrange(
        pad(clamp(img))[::-1, ::-1, :], "x y (z1 z2) -> (z1 y) (z2 x)", z1=5, z2=7
    )


# Create an annotation context (each SegmentationImage will have idx assigned to colors below)
rr.log(
    "/",  # Applies to all entities.
    rr.AnnotationContext(
        [
            rr.AnnotationInfo(id=0, label="Background", color=(0, 0, 0, 0)),
            rr.AnnotationInfo(id=1, label="Bladder Lumen", color=(255, 0, 0)),
            rr.AnnotationInfo(id=2, label="Bladder Wall", color=(0, 128, 255)),
            rr.AnnotationInfo(id=3, label="Urethra", color=(0, 200, 0)),
            rr.AnnotationInfo(id=4, label="Rectum", color=(255, 0, 255)),
            rr.AnnotationInfo(
                id=5, label="Prostate (Peripheral Zone)", color=(255, 165, 0)
            ),
            rr.AnnotationInfo(
                id=6, label="Prostate (Transition Zone)", color=(128, 0, 128)
            ),
            rr.AnnotationInfo(
                id=7, label="Prostate (Anterior Stromae)", color=(255, 255, 0)
            ),
            rr.AnnotationInfo(
                id=8, label="Prostate (Central Zone)", color=(0, 128, 128)
            ),
            rr.AnnotationInfo(id=9, label="Neurovascular Bundles", color=(165, 42, 42)),
            rr.AnnotationInfo(id=10, label="Seminal Vesicles", color=(70, 130, 180)),
        ],
    ),
    static=True,  # Make it valid across all time. (we will use "time" to log different subjects)
)

# Lazy hack to not have to deal with directory.
os.chdir("/mnt/storage/projects/mismo/bamberg/vify/data/train")


for subject in glob.glob("BB_*"):
    imgf = glob.glob(f"{subject}/*t2_tse_tra_384*.nii.gz")[0]
    anatf = glob.glob(f"{subject}/*t2_anatomy_segmentation_*.nii.gz")[0]
    img = image.load_img(imgf).get_fdata()
    anat = image.load_img(anatf).get_fdata()
    # Skip if for whatever reason these images don't match in size.
    if not img.shape == anat.shape:
        continue
    # Skip if image is too large since that'd mess up our mosaic. This could be adapted.
    if img.shape[2] > 35:
        continue
    # Set the subject number on the timeline.
    rr.set_time_sequence("subject", int(subject[3:]))
    # Log our T2 and the segmentation mask.
    rr.log("T2", rr.Image(mosaic(img)))
    rr.log("segmentation", rr.SegmentationImage(mosaic(anat)))

Viewing all of the diffusion images from the Bamberg dataset.

This example brings all DWI and TRACE images from all patients together into two large ND tensors. The images are reshaped to be viewed as mosaics, and additional dimensions are left for b value, and whether images have diffusion of not.

This function needs to be modified to use vicom. Currently it will not work because of file naming conventions. Only functions get_subject_array and get_subject_trace need adaption, though.

"""
Visualize the different images of the Bamberg dataset.

on the server:

Images are located here: /mnt/storage/data/clinical_trials/bamberg/v3/ and viseg


Error processing JSON BB_015/study_clinical.json: list index out of range
Error processing JSON BB_067/study_clinical.json: list index out of range
Error processing JSON BB_070/study_clinical.json: list index out of range
Error processing JSON BB_093/study_clinical.json: list index out of range
"""
import os
import rerun as rr
from nilearn import image
from glob import glob
import einops
import numpy as np
import json
os.chdir("/Users/dangom/virdx/data/bamberg/train/")


def extract_clinical_values(file_path):
    try:
        with open(file_path, 'r') as file:
            json_data = json.load(file)

        # Extract age and score directly
        age = json_data.get("age")
        score = json_data.get("pi_rads_score")
        psa_value = json_data.get("psa_values")[0]['psa_value']

        return {"age": age, "score": score, "psa_value": psa_value}

    except Exception as e:
        print(f"Error processing JSON {file_path}: {str(e)}")
        return None


def get_subject_array(subject):
    # These 3 first lines need to be modified
    resolves_x = list(sorted(glob(f"{subject}/b*dir-x*dwi.nii.gz")))
    resolves_y = list(sorted(glob(f"{subject}/b*dir-y*dwi.nii.gz")))
    resolves_z = list(sorted(glob(f"{subject}/b*dir-z*dwi.nii.gz")))
    if len(resolves_x) != len(resolves_y) or len(resolves_x) != len(resolves_z):
        raise ValueError("Different number of images for x, y, and z directions.")
    if len(resolves_x) + len(resolves_y) + len(resolves_z) < 21:
        raise ValueError("Not enough images for all directions.")
    imgs_x = []
    for img in resolves_x:
        img_x = image.load_img(img).get_fdata()
        imgs_x.append(img_x)
    img_arr_x = np.stack(imgs_x, axis=-1)
    imgs_y = []
    for img in resolves_y:
        img_y = image.load_img(img).get_fdata()
        imgs_y.append(img_y)
    img_arr_y = np.stack(imgs_y, axis=-1)
    imgs_z = []
    for img in resolves_z:
        img_z = image.load_img(img).get_fdata()
        imgs_z.append(img_z)
    img_arr_z = np.stack(imgs_z, axis=-1)

    # Now stack them all into a new dimension
    img_arr = np.stack([img_arr_x, img_arr_y, img_arr_z], axis=-1)
    return img_arr

def get_subject_trace(subject):
    traces_b0 = list(sorted(glob(f"{subject}/*b0*trace.nii.gz")))
    traces_bval = list(sorted(glob(f"{subject}/*bval*trace.nii.gz")))
    # append the first traces_b0 to bval, since the b0 does not have another b0
    traces_b0.insert(0, traces_bval[0])
    if len(traces_b0) != len(traces_bval):
        print(len(traces_b0))
        print(len(traces_bval))
        raise ValueError("Different number of images for trace b0 and bval.")
    if len(traces_b0) + len(traces_bval) < 16:
        raise ValueError("Not enough trace images.")
    imgs_x = []
    for img in traces_b0:
        img_x = image.load_img(img).get_fdata()
        imgs_x.append(img_x)
    img_arr_x = np.stack(imgs_x, axis=-1)
    imgs_y = []
    for img in traces_bval:
        img_y = image.load_img(img).get_fdata()
        imgs_y.append(img_y)
    img_arr_y = np.stack(imgs_y, axis=-1)

    # Now stack them all into a new dimension
    img_arr = np.stack([img_arr_y, img_arr_x], axis=-1)
    return img_arr


# img_arr = get_subject_array("BB_003")
# img_arr_3mm = einops.rearrange(
#     img_arr, "x y (z1 z2) bval dir -> (z1 x) (z2 y) bval dir", z1=5, z2=5
# )
# rr.init("Bamberg MRI", spawn=True, strict=True)
# rr.log("tensor", rr.Tensor(img_arr_3mm, dim_names=("x", "y", "bval", "dir")))


# View all subjects.
rr.init("Bamberg MRI viewer", spawn=True, strict=True)
for subject in list(sorted(glob("BB_*"))):
    try:
        age, pirads, psa_values = extract_clinical_values(f"{subject}/study_clinical.json").values()
        img_arr = get_subject_array(subject)
        img_arr_3mm = einops.rearrange(
            img_arr, "x y (z1 z2) bval dir -> (z1 x) (z2 y) bval dir", z1=5, z2=5
        )
        trace_arr = get_subject_trace(subject)
        trace_arr_3mm = einops.rearrange(
            trace_arr, "x y (z1 z2) bval b0 -> (z1 x) (z2 y) bval b0", z1=5, z2=5
        )
    except:
        continue
    rr.set_time_sequence("subject", int(subject[3:]))
    rr.log("dwi", rr.Tensor(img_arr_3mm, dim_names=("x", "y", "bval", "dir")))
    rr.log("trace", rr.Tensor(trace_arr_3mm, dim_names=("x", "y", "bval", "b0")))
    rr.log("PIRADS T2", rr.Scalar(int(pirads)))

To adapt get_subject_array(subject) above to use VICOM, we can use Moritz' code from his dwi_interp notebook as an example, and then adapt the other calls to nilearn to use vicom instead.

dataset_path = Path("/mnt/storage/data/clinical_trials/bamberg/v3/train")
study_paths = list(sorted(list(Path(dataset_path).glob("*"))))

vols = load_volumes(path=study_path, file_type="VOLUME")
resolves_x = [
    v
    for v in vols
    if v.volume_type == "DWI"
    and "TRACEW" not in v.image_type
    and "ADC" not in v.image_type
    and (
        v.diffusion_gradient_orientation_0 > 0.9
        or v.diffusion_gradient_orientation_0 < -0.9
    )
]
resolves_y = [
    v
    for v in vols
    if v.volume_type == "DWI"
    and "TRACEW" not in v.image_type
    and "ADC" not in v.image_type
    and (
        v.diffusion_gradient_orientation_1 > 0.9
        or v.diffusion_gradient_orientation_1 < -0.9
    )
]
resolves_z = [
    v
    for v in vols
    if v.volume_type == "DWI"
    and "TRACEW" not in v.image_type
    and "ADC" not in v.image_type
    and (
        v.diffusion_gradient_orientation_2 > 0.9
        or v.diffusion_gradient_orientation_2 < -0.9
    )
]

resolves_x = sorted(resolves_x, key=lambda v: v.b_value)
resolves_y = sorted(resolves_y, key=lambda v: v.b_value)
resolves_z = sorted(resolves_z, key=lambda v: v.b_value)