303 lines
11 KiB
Python
303 lines
11 KiB
Python
|
# import nibabel as nib
|
||
|
|
# import matplotlib.pyplot as plt
|
||
|
|
# import numpy as np
|
||
|
|
# import os
|
||
|
|
# from typing import Tuple, Set, Dict, Optional
|
||
|
|
|
||
|
|
# def load_nifti(file_path: str) -> np.ndarray:
|
||
|
|
# """Load a NIfTI file and return its data array."""
|
||
|
|
# return nib.load(file_path).get_fdata()
|
||
|
|
|
||
|
|
# def create_colored_mask(data: np.ndarray) -> np.ndarray:
|
||
|
|
# """Create a colored mask for segmentation visualization.
|
||
|
|
|
||
|
|
# Args:
|
||
|
|
# data: Input segmentation array
|
||
|
|
# Returns:
|
||
|
|
# RGB mask array
|
||
|
|
# """
|
||
|
|
# color_dict = {
|
||
|
|
# 1: [144/255, 238/255, 144/255], # Brainstem - Light green
|
||
|
|
# 2: [255/255, 218/255, 150/255], # Right_Eye - Light yellow
|
||
|
|
# 3: [205/255, 170/255, 125/255], # Left_Eye - Brown
|
||
|
|
# 4: [135/255, 206/255, 235/255], # Optic_Chiasm - Sky blue
|
||
|
|
# 5: [255/255, 99/255, 71/255], # Right_Optic_Nerve - Bright coral
|
||
|
|
# 6: [255/255, 160/255, 122/255], # Left_Optic_Nerve - Light coral
|
||
|
|
# 7: [147/255, 112/255, 219/255] # TV - Medium purple
|
||
|
|
# }
|
||
|
|
|
||
|
|
# mask = np.zeros((*data.shape, 3))
|
||
|
|
|
||
|
|
# for label, color in color_dict.items():
|
||
|
|
# mask_3d = np.dstack([data == label] * 3)
|
||
|
|
# mask = np.where(mask_3d, color, mask)
|
||
|
|
|
||
|
|
# return mask
|
||
|
|
|
||
|
|
# def get_brain_bounds(image: np.ndarray, padding_factor: float = 0.15) -> Tuple[int, int, int, int]:
|
||
|
|
# """Get the boundaries of the brain region with adjustable padding.
|
||
|
|
|
||
|
|
# Args:
|
||
|
|
# image: Input image array
|
||
|
|
# padding_factor: Factor to determine padding around the region of interest
|
||
|
|
# Returns:
|
||
|
|
# Tuple of (rmin, rmax, cmin, cmax) coordinates
|
||
|
|
# """
|
||
|
|
# mask = image > np.percentile(image, 1)
|
||
|
|
# rows = np.any(mask, axis=1)
|
||
|
|
# cols = np.any(mask, axis=0)
|
||
|
|
# rmin, rmax = np.where(rows)[0][[0, -1]]
|
||
|
|
# cmin, cmax = np.where(cols)[0][[0, -1]]
|
||
|
|
|
||
|
|
# center_r = (rmin + rmax) // 2
|
||
|
|
# center_c = (cmin + cmax) // 2
|
||
|
|
|
||
|
|
# radius = int(max(rmax - rmin, cmax - cmin) * 0.55)
|
||
|
|
# margin = int(radius * padding_factor)
|
||
|
|
|
||
|
|
# return (
|
||
|
|
# max(center_r - radius - margin, 0),
|
||
|
|
# min(center_r + radius + margin, image.shape[0]),
|
||
|
|
# max(center_c - radius - margin, 0),
|
||
|
|
# min(center_c + radius + margin, image.shape[1])
|
||
|
|
# )
|
||
|
|
|
||
|
|
# def check_organs_in_slice(slice_data: np.ndarray) -> Set[str]:
|
||
|
|
# """Check which organs are present in a given slice.
|
||
|
|
|
||
|
|
# Args:
|
||
|
|
# slice_data: 2D array of the slice
|
||
|
|
# Returns:
|
||
|
|
# Set of organ names present in the slice
|
||
|
|
# """
|
||
|
|
# organ_dict = {
|
||
|
|
# 1: "Brainstem",
|
||
|
|
# 2: "Right_Eye",
|
||
|
|
# 3: "Left_Eye",
|
||
|
|
# 4: "Optic_Chiasm",
|
||
|
|
# 5: "Right_Optic_Nerve",
|
||
|
|
# 6: "Left_Optic_Nerve",
|
||
|
|
# 7: "TV"
|
||
|
|
# }
|
||
|
|
|
||
|
|
# unique_labels = set(np.unique(slice_data)) - {0}
|
||
|
|
# return {organ_dict[label] for label in unique_labels if label in organ_dict}
|
||
|
|
|
||
|
|
# def save_slice(image: np.ndarray, ground_truth: np.ndarray,
|
||
|
|
# prediction: np.ndarray, slice_num: int,
|
||
|
|
# save_path: str) -> None:
|
||
|
|
# """Save comparison visualization of a specific slice.
|
||
|
|
|
||
|
|
# Args:
|
||
|
|
# image: MRI volume
|
||
|
|
# ground_truth: Ground truth segmentation volume
|
||
|
|
# prediction: Predicted segmentation volume
|
||
|
|
# slice_num: Slice number to visualize
|
||
|
|
# save_path: Path to save the visualization
|
||
|
|
# """
|
||
|
|
# plt.clf()
|
||
|
|
# fig, axes = plt.subplots(1, 2, figsize=(25, 11))
|
||
|
|
# plt.subplots_adjust(wspace=0.01)
|
||
|
|
|
||
|
|
# # Process image
|
||
|
|
# img_slice = np.rot90(image[:, :, slice_num])
|
||
|
|
|
||
|
|
# # Get bounds
|
||
|
|
# rmin, rmax, cmin, cmax = get_brain_bounds(img_slice)
|
||
|
|
|
||
|
|
# # Normalize image for display
|
||
|
|
# img_norm = (img_slice - np.min(img_slice)) / (np.max(img_slice) - np.min(img_slice))
|
||
|
|
|
||
|
|
# # Ground Truth
|
||
|
|
# gt_slice = np.rot90(ground_truth[:, :, slice_num])
|
||
|
|
# gt_mask = create_colored_mask(gt_slice)
|
||
|
|
|
||
|
|
# axes[0].imshow(img_norm[rmin:rmax, cmin:cmax], cmap='gray')
|
||
|
|
# axes[0].imshow(gt_mask[rmin:rmax, cmin:cmax], alpha=0.5)
|
||
|
|
# axes[0].set_title('Ground Truth', fontsize=32, pad=20, weight='bold')
|
||
|
|
# axes[0].axis('off')
|
||
|
|
|
||
|
|
# # Prediction
|
||
|
|
# pred_slice = np.rot90(prediction[:, :, slice_num])
|
||
|
|
# pred_mask = create_colored_mask(pred_slice)
|
||
|
|
|
||
|
|
# axes[1].imshow(img_norm[rmin:rmax, cmin:cmax], cmap='gray')
|
||
|
|
# axes[1].imshow(pred_mask[rmin:rmax, cmin:cmax], alpha=0.5)
|
||
|
|
# axes[1].set_title('Prediction', fontsize=32, pad=20, weight='bold')
|
||
|
|
# axes[1].axis('off')
|
||
|
|
|
||
|
|
# plt.savefig(save_path, bbox_inches='tight', dpi=300, pad_inches=0.05)
|
||
|
|
# plt.close()
|
||
|
|
# print(f"Saved slice {slice_num} to: {save_path}")
|
||
|
|
|
||
|
|
# def main():
|
||
|
|
# """Main function to run the visualization pipeline."""
|
||
|
|
# # Load images
|
||
|
|
# base_path = '/mnt/1248/onlylian/nnUNet/nnUNet_raw/Dataset999'
|
||
|
|
# image_mri = load_nifti(os.path.join(base_path, 'imagesTs/T1C_043_0000.nii.gz'))
|
||
|
|
# ground_truth = load_nifti(os.path.join(base_path, 'labelsTs/T1C_043.nii.gz'))
|
||
|
|
# prediction = load_nifti('/mnt/1248/onlylian/T1C/output_predictions/2d/fold_2/T1C_043.nii.gz')
|
||
|
|
|
||
|
|
# save_dir = '/mnt/1248/onlylian/T1C'
|
||
|
|
|
||
|
|
# # Analyze slices
|
||
|
|
# total_slices = image_mri.shape[2]
|
||
|
|
# for slice_num in range(total_slices):
|
||
|
|
# gt_organs = check_organs_in_slice(ground_truth[:, :, slice_num])
|
||
|
|
# pred_organs = check_organs_in_slice(prediction[:, :, slice_num])
|
||
|
|
|
||
|
|
# if gt_organs or pred_organs:
|
||
|
|
# print(f"\nSlice {slice_num}:")
|
||
|
|
# print(f"Ground Truth contains: {', '.join(sorted(gt_organs))}")
|
||
|
|
# print(f"Prediction contains: {', '.join(sorted(pred_organs))}")
|
||
|
|
|
||
|
|
# # Interactive slice selection
|
||
|
|
# while True:
|
||
|
|
# try:
|
||
|
|
# slice_num = int(input("\nEnter slice number to save (-1 to exit): "))
|
||
|
|
# if slice_num == -1:
|
||
|
|
# break
|
||
|
|
# if 0 <= slice_num < total_slices:
|
||
|
|
# save_path = os.path.join(save_dir, f'comparison_{slice_num:03d}.png')
|
||
|
|
# save_slice(image_mri, ground_truth, prediction,
|
||
|
|
# slice_num, save_path)
|
||
|
|
# else:
|
||
|
|
# print(f"Slice number must be between 0 and {total_slices-1}")
|
||
|
|
# except ValueError:
|
||
|
|
# print("Please enter a valid number")
|
||
|
|
# except Exception as e:
|
||
|
|
# print(f"Error occurred: {str(e)}")
|
||
|
|
|
||
|
|
# if __name__ == "__main__":
|
||
|
|
# main()
|
||
|
|
|
||
|
|
import nibabel as nib
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
import numpy as np
|
||
|
|
import os
|
||
|
|
from PIL import Image
|
||
|
|
from typing import Tuple, Set, Dict, Optional
|
||
|
|
import io
|
||
|
|
|
||
|
|
def load_nifti(file_path: str) -> np.ndarray:
|
||
|
|
"""Load a NIfTI file and return its data array."""
|
||
|
|
return nib.load(file_path).get_fdata()
|
||
|
|
|
||
|
|
def create_colored_mask(data: np.ndarray) -> np.ndarray:
|
||
|
|
"""Create a colored mask for segmentation visualization."""
|
||
|
|
color_dict = {
|
||
|
|
1: [144/255, 238/255, 144/255], # Brainstem - Light green
|
||
|
|
2: [255/255, 218/255, 150/255], # Right_Eye - Light yellow
|
||
|
|
3: [205/255, 170/255, 125/255], # Left_Eye - Brown
|
||
|
|
4: [135/255, 206/255, 235/255], # Optic_Chiasm - Sky blue
|
||
|
|
5: [255/255, 99/255, 71/255], # Right_Optic_Nerve - Bright coral
|
||
|
|
6: [255/255, 160/255, 122/255], # Left_Optic_Nerve - Light coral
|
||
|
|
7: [147/255, 112/255, 219/255] # TV - Medium purple
|
||
|
|
}
|
||
|
|
|
||
|
|
mask = np.zeros((*data.shape, 3))
|
||
|
|
for label, color in color_dict.items():
|
||
|
|
mask_3d = np.dstack([data == label] * 3)
|
||
|
|
mask = np.where(mask_3d, color, mask)
|
||
|
|
return mask
|
||
|
|
|
||
|
|
def get_brain_bounds(image: np.ndarray, padding_factor: float = 0.15) -> Tuple[int, int, int, int]:
|
||
|
|
"""Get the boundaries of the brain region with adjustable padding."""
|
||
|
|
mask = image > np.percentile(image, 1)
|
||
|
|
rows = np.any(mask, axis=1)
|
||
|
|
cols = np.any(mask, axis=0)
|
||
|
|
rmin, rmax = np.where(rows)[0][[0, -1]]
|
||
|
|
cmin, cmax = np.where(cols)[0][[0, -1]]
|
||
|
|
|
||
|
|
center_r = (rmin + rmax) // 2
|
||
|
|
center_c = (cmin + cmax) // 2
|
||
|
|
radius = int(max(rmax - rmin, cmax - cmin) * 0.55)
|
||
|
|
margin = int(radius * padding_factor)
|
||
|
|
|
||
|
|
return (
|
||
|
|
max(center_r - radius - margin, 0),
|
||
|
|
min(center_r + radius + margin, image.shape[0]),
|
||
|
|
max(center_c - radius - margin, 0),
|
||
|
|
min(center_c + radius + margin, image.shape[1])
|
||
|
|
)
|
||
|
|
|
||
|
|
def create_slice_image(image: np.ndarray, ground_truth: np.ndarray,
|
||
|
|
prediction: np.ndarray, slice_num: int) -> Image.Image:
|
||
|
|
"""Create image for a specific slice."""
|
||
|
|
plt.clf()
|
||
|
|
fig, axes = plt.subplots(1, 2, figsize=(25, 11))
|
||
|
|
plt.subplots_adjust(wspace=0.01)
|
||
|
|
|
||
|
|
img_slice = np.rot90(image[:, :, slice_num])
|
||
|
|
rmin, rmax, cmin, cmax = get_brain_bounds(img_slice)
|
||
|
|
img_norm = (img_slice - np.min(img_slice)) / (np.max(img_slice) - np.min(img_slice))
|
||
|
|
|
||
|
|
# Ground Truth
|
||
|
|
gt_slice = np.rot90(ground_truth[:, :, slice_num])
|
||
|
|
gt_mask = create_colored_mask(gt_slice)
|
||
|
|
|
||
|
|
axes[0].imshow(img_norm[rmin:rmax, cmin:cmax], cmap='gray')
|
||
|
|
axes[0].imshow(gt_mask[rmin:rmax, cmin:cmax], alpha=0.5)
|
||
|
|
axes[0].set_title('Ground Truth', fontsize=32, pad=20, weight='bold')
|
||
|
|
axes[0].axis('off')
|
||
|
|
|
||
|
|
# Prediction
|
||
|
|
pred_slice = np.rot90(prediction[:, :, slice_num])
|
||
|
|
pred_mask = create_colored_mask(pred_slice)
|
||
|
|
|
||
|
|
axes[1].imshow(img_norm[rmin:rmax, cmin:cmax], cmap='gray')
|
||
|
|
axes[1].imshow(pred_mask[rmin:rmax, cmin:cmax], alpha=0.5)
|
||
|
|
axes[1].set_title('Prediction', fontsize=32, pad=20, weight='bold')
|
||
|
|
axes[1].axis('off')
|
||
|
|
|
||
|
|
# Convert plot to PIL Image
|
||
|
|
buf = io.BytesIO()
|
||
|
|
plt.savefig(buf, format='png', bbox_inches='tight', dpi=300, pad_inches=0.05)
|
||
|
|
plt.close()
|
||
|
|
buf.seek(0)
|
||
|
|
return Image.open(buf)
|
||
|
|
|
||
|
|
def create_animated_gif(image: np.ndarray, ground_truth: np.ndarray,
|
||
|
|
prediction: np.ndarray, start_slice: int,
|
||
|
|
end_slice: int, output_path: str,
|
||
|
|
duration: int = 500) -> None:
|
||
|
|
"""Create an animated GIF from a range of slices."""
|
||
|
|
frames = []
|
||
|
|
|
||
|
|
for slice_num in range(start_slice, end_slice + 1):
|
||
|
|
print(f"Processing slice {slice_num}...")
|
||
|
|
frame = create_slice_image(image, ground_truth, prediction, slice_num)
|
||
|
|
frames.append(frame)
|
||
|
|
|
||
|
|
# Save the animated GIF
|
||
|
|
frames[0].save(
|
||
|
|
output_path,
|
||
|
|
save_all=True,
|
||
|
|
append_images=frames[1:],
|
||
|
|
duration=duration,
|
||
|
|
loop=0
|
||
|
|
)
|
||
|
|
print(f"Saved animated GIF to: {output_path}")
|
||
|
|
|
||
|
|
def main():
|
||
|
|
"""Main function to create the animated visualization."""
|
||
|
|
# Load images with new paths
|
||
|
|
base_path = '/mnt/1248/onlylian/nnUNet/nnUNet_raw/Dataset999'
|
||
|
|
image_mri = load_nifti(os.path.join(base_path, 'imagesTs/T1C_043_0000.nii.gz'))
|
||
|
|
ground_truth = load_nifti(os.path.join(base_path, 'labelsTs/T1C_043.nii.gz'))
|
||
|
|
prediction = load_nifti('/mnt/1248/onlylian/T1C/output_predictions/2d/fold_2/T1C_043.nii.gz')
|
||
|
|
|
||
|
|
# Create animated GIF for slices 111-121
|
||
|
|
output_path = '/mnt/1248/onlylian/T1C/slice_animation.gif'
|
||
|
|
create_animated_gif(
|
||
|
|
image_mri,
|
||
|
|
ground_truth,
|
||
|
|
prediction,
|
||
|
|
start_slice=111,
|
||
|
|
end_slice=121,
|
||
|
|
output_path=output_path,
|
||
|
|
duration=500 # 500ms per frame
|
||
|
|
)
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|