app-download-image.py

import random
from pathlib import Path

import cv2
import gradio as gr
import numpy as np
import SimpleITK
from huggingface_hub import hf_hub_download, list_repo_files
from scipy import ndimage

IMAGES_REPO = "LMUK-RADONC-PHYS-RES/TrackRAD2025"
DATASET_REPO_TYPE = "dataset"
LABELED_FOLDER = "trackrad2025_labeled_training_data"

OUT_DIR = Path("tmp/videos")


def get_images() -> list[str]:
    images_repo_files = list_repo_files(
        repo_id=IMAGES_REPO,
        repo_type=DATASET_REPO_TYPE,
    )
    image_files = [
        fname
        for fname in images_repo_files
        if fname.startswith(LABELED_FOLDER) and fname.endswith("frames.mha")
    ]
    return image_files


def download_image_files(image_file: str) -> dict[str, str]:
    filename = image_file.split("/")[-1]
    patient = filename.rsplit("_", 1)[0]  # e.g., "A_033_frames.mha" -> "A_033"
    frames_idx = ""
    frames_file = hf_hub_download(
        repo_id=IMAGES_REPO,
        repo_type=DATASET_REPO_TYPE,
        filename=f"{LABELED_FOLDER}/{patient}/images/{patient}_frames{frames_idx}.mha",
    )
    labels_file = hf_hub_download(
        repo_id=IMAGES_REPO,
        repo_type=DATASET_REPO_TYPE,
        filename=f"{LABELED_FOLDER}/{patient}/targets/{patient}_labels{frames_idx}.mha",
    )
    field_strength_file = hf_hub_download(
        repo_id=IMAGES_REPO,
        repo_type=DATASET_REPO_TYPE,
        filename=f"{LABELED_FOLDER}/{patient}/b-field-strength.json",
    )
    scanned_region_file = hf_hub_download(
        repo_id=IMAGES_REPO,
        repo_type=DATASET_REPO_TYPE,
        filename=f"{LABELED_FOLDER}/{patient}/scanned-region{frames_idx}.json",
    )
    frame_rate_file = hf_hub_download(
        repo_id=IMAGES_REPO,
        repo_type=DATASET_REPO_TYPE,
        filename=f"{LABELED_FOLDER}/{patient}/frame-rate{frames_idx}.json",
    )
    return {
        "frames_file": frames_file,
        "labels_file": labels_file,
        "field_strength_file": field_strength_file,
        "scanned_region_file": scanned_region_file,
        "frame_rate_file": frame_rate_file,
    }


def overlay_labels_on_frames(
    frames_array, labels_array, overlay_color="cyan", alpha=1.0
):
    """
    Overlay binary labels on grayscale frames with a bright color.

    Parameters:
    -----------
    frames_array : numpy.ndarray
        Grayscale image sequence of shape [X, Y, T]
    labels_array : numpy.ndarray
        Binary labels of shape [X, Y, T]
    overlay_color : str or tuple
        Color for the overlay ('red', 'green', 'blue', 'yellow', 'cyan', 'magenta')
        or RGB tuple (r, g, b) with values 0-1
    alpha : float
        Transparency of the overlay (0=transparent, 1=opaque)

    Returns:
    --------
    overlaid_frames : numpy.ndarray
        RGB frames with labels overlaid, shape [X, Y, T, 3]
    """

    # Normalize frames to 0-1 range if not already
    frames_norm = frames_array.astype(np.float32)
    if frames_norm.max() > 1.0:
        frames_norm = frames_norm / frames_norm.max()

    # Convert grayscale to RGB by repeating across 3 channels
    rgb_frames = np.stack([frames_norm] * 3, axis=-1)  # Shape: [X, Y, T, 3]

    # Define color mapping
    color_map = {
        "green": (0.0, 1.0, 0.0),
        "blue": (0.0, 0.0, 1.0),
        "yellow": (1.0, 1.0, 0.0),
        "cyan": (0.0, 1.0, 1.0),
        "magenta": (1.0, 0.0, 1.0),
    }

    if overlay_color in color_map:
        r, g, b = color_map[overlay_color]
    else:
        raise ValueError(
            f"Unknown color '{overlay_color}'. Use: {list(color_map.keys())} or RGB tuple"
        )

    # Create the overlaid frames
    overlaid_frames = rgb_frames.copy()

    # Apply overlay where labels are True (assuming binary labels are 0/1 or False/True)
    mask = (ndimage.binary_erosion(labels_array) ^ labels_array).astype(bool)

    # Blend the colors using alpha blending
    overlaid_frames[mask, 0] = (1 - alpha) * rgb_frames[mask, 0] + alpha * r
    overlaid_frames[mask, 1] = (1 - alpha) * rgb_frames[mask, 1] + alpha * g
    overlaid_frames[mask, 2] = (1 - alpha) * rgb_frames[mask, 2] + alpha * b

    return overlaid_frames


def overlay_video(files: dict[str, str]):
    frames = SimpleITK.ReadImage(files["frames_file"])
    frames_array = SimpleITK.GetArrayFromImage(frames)
    # frames_array = [X,Y,T]
    frames_array = np.flip(frames_array, axis=0)

    labels = SimpleITK.ReadImage(files["labels_file"])
    labels_array = SimpleITK.GetArrayFromImage(labels)
    # labels_array = [X,Y,T]
    labels_array = np.flip(labels_array, axis=0)

    overlaid_array = overlay_labels_on_frames(frames_array, labels_array)
    output_path = numpy_to_video_opencv(overlaid_array, "tmp_video", fps=8)
    return output_path


def numpy_to_video_opencv(array: np.ndarray, output_prefix: str, fps: int) -> str:
    limit = 10 * fps
    array_clip = array[:, :, :limit]  # 10s of video
    p99: float = np.percentile(array_clip, 99)  # type: ignore
    array_clip_normalized = cv2.convertScaleAbs(array_clip, alpha=(255.0 / p99))

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    output_path = str((OUT_DIR / output_prefix).with_suffix(".webm"))

    # Define codec and create VideoWriter
    # VP90 is supported by browsers and is available in the pip-installed opencv
    fourcc = cv2.VideoWriter.fourcc(*"VP90")
    X, Y, T, _ = array_clip.shape
    bgr_frames = array_clip_normalized[:, :, :, [2, 1, 0]]
    out = cv2.VideoWriter(output_path, fourcc, fps, (X, Y))
    # Write frames
    for t in range(T):
        frame = bgr_frames[:, :, t, :]
        # OpenCV expects frames in BGR format, but for grayscale we can use as-is
        out.write(frame)

    out.release()
    return output_path


choices = [
    "trackrad2025_labeled_training_data/A_001/images/A_001_frames.mha",
    "trackrad2025_labeled_training_data/A_003/images/A_003_frames.mha",
    "trackrad2025_labeled_training_data/A_004/images/A_004_frames.mha",
    "trackrad2025_labeled_training_data/A_005/images/A_005_frames.mha",
    "trackrad2025_labeled_training_data/A_006/images/A_006_frames.mha",
    "trackrad2025_labeled_training_data/A_007/images/A_007_frames.mha",
    "trackrad2025_labeled_training_data/A_008/images/A_008_frames.mha",
    "trackrad2025_labeled_training_data/A_010/images/A_010_frames.mha",
    "trackrad2025_labeled_training_data/A_011/images/A_011_frames.mha",
    "trackrad2025_labeled_training_data/A_012/images/A_012_frames.mha",
    "trackrad2025_labeled_training_data/A_013/images/A_013_frames.mha",
    "trackrad2025_labeled_training_data/A_014/images/A_014_frames.mha",
    "trackrad2025_labeled_training_data/A_016/images/A_016_frames.mha",
    "trackrad2025_labeled_training_data/A_019/images/A_019_frames.mha",
    "trackrad2025_labeled_training_data/A_020/images/A_020_frames.mha",
    "trackrad2025_labeled_training_data/A_021/images/A_021_frames.mha",
    "trackrad2025_labeled_training_data/A_022/images/A_022_frames.mha",
    "trackrad2025_labeled_training_data/A_023/images/A_023_frames.mha",
    "trackrad2025_labeled_training_data/A_024/images/A_024_frames.mha",
    "trackrad2025_labeled_training_data/A_025/images/A_025_frames.mha",
    "trackrad2025_labeled_training_data/A_026/images/A_026_frames.mha",
    "trackrad2025_labeled_training_data/A_027/images/A_027_frames.mha",
    "trackrad2025_labeled_training_data/A_028/images/A_028_frames.mha",
    "trackrad2025_labeled_training_data/A_029/images/A_029_frames.mha",
    "trackrad2025_labeled_training_data/A_032/images/A_032_frames.mha",
    "trackrad2025_labeled_training_data/B_002/images/B_002_frames.mha",
    "trackrad2025_labeled_training_data/B_003/images/B_003_frames.mha",
    "trackrad2025_labeled_training_data/B_006/images/B_006_frames.mha",
    "trackrad2025_labeled_training_data/B_007/images/B_007_frames.mha",
    "trackrad2025_labeled_training_data/B_008/images/B_008_frames.mha",
    "trackrad2025_labeled_training_data/B_010/images/B_010_frames.mha",
    "trackrad2025_labeled_training_data/B_012/images/B_012_frames.mha",
    "trackrad2025_labeled_training_data/B_017/images/B_017_frames.mha",
    "trackrad2025_labeled_training_data/B_019/images/B_019_frames.mha",
    "trackrad2025_labeled_training_data/B_021/images/B_021_frames.mha",
    "trackrad2025_labeled_training_data/B_022/images/B_022_frames.mha",
    "trackrad2025_labeled_training_data/B_023/images/B_023_frames.mha",
    "trackrad2025_labeled_training_data/B_024/images/B_024_frames.mha",
    "trackrad2025_labeled_training_data/B_025/images/B_025_frames.mha",
    "trackrad2025_labeled_training_data/B_026/images/B_026_frames.mha",
    "trackrad2025_labeled_training_data/C_001/images/C_001_frames.mha",
    "trackrad2025_labeled_training_data/C_004/images/C_004_frames.mha",
    "trackrad2025_labeled_training_data/C_005/images/C_005_frames.mha",
    "trackrad2025_labeled_training_data/C_006/images/C_006_frames.mha",
    "trackrad2025_labeled_training_data/C_008/images/C_008_frames.mha",
    "trackrad2025_labeled_training_data/C_009/images/C_009_frames.mha",
    "trackrad2025_labeled_training_data/C_010/images/C_010_frames.mha",
    "trackrad2025_labeled_training_data/C_011/images/C_011_frames.mha",
    "trackrad2025_labeled_training_data/C_012/images/C_012_frames.mha",
    "trackrad2025_labeled_training_data/C_016/images/C_016_frames.mha",
]


def play_video(fname: str):
    files = download_image_files(fname)
    output_path = overlay_video(files)
    return output_path


demo = gr.Interface(
    play_video,
    [
        gr.Dropdown(
            choices=choices,
            label="Select an MR sequence",
            value=random.choice(choices),
        )
    ],
    gr.Video(
        height=500,
        autoplay=True,
        loop=True,
        label="MR Sequence",
    ),
    live=True,
    title="TrackRAD2025 Labeled Data Viewer",
    examples=[[random.choice(choices)]],
    cache_examples=True,
    preload_example=0,
    flagging_mode="never",
)


demo.launch()