Spaces:
Running
Running
File size: 3,887 Bytes
caf6ee7 906fcb9 caf6ee7 906fcb9 caf6ee7 1baebae 906fcb9 caf6ee7 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import argparse
import logging
import os
import nrrd
import numpy as np
from tqdm import tqdm
def get_heatmap(args: argparse.Namespace) -> argparse.Namespace:
"""
Generate heatmaps from DWI (Diffusion Weighted Imaging) and ADC (Apparent Diffusion Coefficient) medical imaging data.
This function processes medical imaging files (DWI and ADC) along with their corresponding
segmentation masks to create normalized heatmaps. It combines the DWI
and ADC heatmaps through element-wise multiplication.
Args:
args: An object containing the following attributes:
- t2_dir (str): Directory path containing T2 image files.
- dwi_dir (str): Directory path containing DWI image files.
- adc_dir (str): Directory path containing ADC image files.
- seg_dir (str): Directory path containing segmentation mask files.
- output_dir (str): Base output directory where 'heatmaps/' subdirectory will be created.
- heatmapdir (str): Output directory for heatmap files (created by function).
Returns:
args: The modified args object with heatmapdir attribute set.
Raises:
FileNotFoundError: If input directories or files do not exist.
ValueError: If NRRD files cannot be read properly.
Notes:
- DWI heatmap is normalized as (dwi - min) / (max - min)
- ADC heatmap is normalized as (max - adc) / (max - min) (inverted)
- Final heatmap is re-normalized to [0, 1] range
- If all values in a mask region are identical, the heatmap is skipped for that modality
- Output files are written in NRRD format with the same header as the input DWI file
"""
files = os.listdir(args.t2_dir)
args.heatmapdir = os.path.join(args.output_dir, "heatmaps/")
os.makedirs(args.heatmapdir, exist_ok=True)
logging.info("Starting heatmap generation")
for file in tqdm(files):
bool_dwi = False
bool_adc = False
mask, _ = nrrd.read(os.path.join(args.seg_dir, file))
dwi, header_dwi = nrrd.read(os.path.join(args.dwi_dir, file))
adc, header_adc = nrrd.read(os.path.join(args.adc_dir, file))
nonzero_vals_dwi = dwi[mask > 0]
if len(nonzero_vals_dwi) > 0:
min_val = nonzero_vals_dwi.min()
max_val = nonzero_vals_dwi.max()
heatmap_dwi = np.zeros_like(dwi, dtype=np.float32)
if min_val != max_val:
heatmap_dwi = (dwi - min_val) / (max_val - min_val)
masked_heatmap_dwi = np.where(mask > 0, heatmap_dwi, heatmap_dwi[mask > 0].min())
else:
bool_dwi = True
else:
bool_dwi = True
nonzero_vals_adc = adc[mask > 0]
if len(nonzero_vals_adc) > 0:
min_val = nonzero_vals_adc.min()
max_val = nonzero_vals_adc.max()
heatmap_adc = np.zeros_like(adc, dtype=np.float32)
if min_val != max_val:
heatmap_adc = (max_val - adc) / (max_val - min_val)
masked_heatmap_adc = np.where(mask > 0, heatmap_adc, heatmap_adc[mask > 0].min())
else:
bool_adc = True
else:
bool_adc = True
if not bool_dwi and not bool_adc:
mix_mask = masked_heatmap_dwi * masked_heatmap_adc
write_header = header_dwi
elif bool_dwi:
mix_mask = masked_heatmap_adc
write_header = header_adc
elif bool_adc:
mix_mask = masked_heatmap_dwi
write_header = header_dwi
else:
mix_mask = np.ones_like(adc, dtype=np.float32)
write_header = header_dwi
mix_mask = (mix_mask - mix_mask.min()) / (mix_mask.max() - mix_mask.min())
nrrd.write(os.path.join(args.heatmapdir, file), mix_mask, write_header)
return args
|