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()