""" Data-loading helpers for the pancreas (PancVMAT) dose-prediction model. Trimmed from the original research toolkit: only the functions the released data loader and inference script actually call are kept. Differs from toolkit_han.py in `combine_optptv`, which subtracts every overlapping OAR from the PTV before building the "optimisation PTV" channel. Keep the two files separate -- swapping them changes model inputs. """ import os import numpy as np import torch import nibabel as nib from monai.transforms import ( Compose, Resized, RandFlipd, RandRotated, SpatialPadd, SpatialCropd, RandSpatialCropd, ) def combine_optptv(tmp_dict, opt_dose_dict, OAR_DICT): ''' Combine the optimal PTV and the dose of the optimal PTV 从PTV mask中扣除与所有OAR重叠的部分 这个修改的目的是解决放疗中的关键问题： - 当PTV与OAR重叠时，重叠区域可能无法达到处方剂量 - 通过从PTV中扣除所有OAR重叠区域，可以更准确地表示实际可达到处方剂量的区域 - 使用OAR_DICT中的所有OAR Args: tmp_dict: 包含所有mask的字典 opt_dose_dict: 包含处方剂量的字典 OAR_DICT: OAR索引字典 Returns: comb_optptv: 组合的优化PTV (扣除OAR重叠后) prs_opt: 处方剂量列表 cat_optptv: 分别的优化PTV通道 ''' comb_optptv = torch.zeros(tmp_dict['img'].shape) cat_optptv = torch.zeros([3] + list(tmp_dict['img'].shape)[1:]) # 假设最多3个PTV prs_opt = [] # 创建组合的所有OAR mask # 使用OAR_DICT中的所有OAR # 注意：OAR_DICT的key带.nii.gz后缀，但tmp_dict中的key可能不带后缀 combined_oar_mask = torch.zeros(tmp_dict['img'].shape) for oar_name in OAR_DICT.keys(): # 尝试多种可能的key格式 oar_key = None for possible_key in [oar_name, oar_name.replace('.nii.gz', ''), f"{oar_name}.nii.gz" if not oar_name.endswith('.nii.gz') else oar_name]: if possible_key in tmp_dict: oar_key = possible_key break if oar_key is not None: combined_oar_mask = torch.maximum(combined_oar_mask, tmp_dict[oar_key]) cnt = 0 for key in opt_dose_dict.keys(): if key in tmp_dict: original_ptv_mask = tmp_dict[key] prescribed_dose = opt_dose_dict[key] # 从PTV mask中扣除与所有OAR重叠的部分 # 只保留不与任何OAR重叠的PTV区域 # 使用 > 0 来创建binary mask，确保OAR区域被完全排除 oar_binary_mask = (combined_oar_mask > 0).float() modified_ptv_mask = original_ptv_mask * (1 - oar_binary_mask) # 创建优化的PTV（应用处方剂量） tmp_optptv = modified_ptv_mask * prescribed_dose # 存储到cat_optptv中（如果有足够的通道） if cnt < cat_optptv.shape[0]: cat_optptv[cnt] = tmp_optptv # 更新组合的optptv comb_optptv = torch.maximum(comb_optptv, tmp_optptv) # 记录处方剂量 prs_opt.append(prescribed_dose) cnt += 1 return comb_optptv, prs_opt, cat_optptv def combine_oar(tmp_dict, need_list, norm_oar=True, OAR_DICT=None): ''' Revised to emphasize Limiting OARs (Stomach, Bowel) over General OARs ''' comb_oar = torch.zeros(tmp_dict['img'].shape) cat_oar = torch.zeros([len(OAR_DICT)] + list(tmp_dict['img'].shape)[1:]) # Define Limiting OARs (need_list不带.nii.gz后缀，所以这里也不带后缀) LIMITING_OARS = [ "Stomach", "Stomach_PRV", "Bowel_Small", "Bowel_Small_PRV", "Bowel_Large", "Bowel_Large_PRV" ] for key in OAR_DICT.keys(): # need_list来自oar_dict.json，不带.nii.gz后缀 # OAR_DICT的key带.nii.gz后缀（除了SpinalCord） # 需要去掉.nii.gz后缀来匹配need_list key_without_suffix = key.replace('.nii.gz', '') key_in_need_list = (key_without_suffix in need_list) if not key_in_need_list: continue # 尝试在tmp_dict中找到数据 # tmp_dict中的key可能是带或不带.nii.gz的格式 single_oar = None # 尝试多种可能的key格式：原始key、去掉后缀、添加后缀 for possible_key in [key, key_without_suffix, f"{key_without_suffix}.nii.gz"]: if possible_key in tmp_dict: single_oar = tmp_dict[possible_key] break if single_oar is None: single_oar = torch.zeros(tmp_dict['img'].shape) cat_oar[OAR_DICT[key]-1: OAR_DICT[key]] = single_oar # --- Special Handling for Limiting OARs Start --- if norm_oar: # Check if current OAR is a Limiting OAR # need_list不带.nii.gz，所以用key_without_suffix来检查 is_limiting = (key_without_suffix in LIMITING_OARS) if is_limiting: # Limiting OARs: base value set to 2.0, with a small gradient to distinguish different bowel segments. # Value range: 2.0 ~ 2.5 (high enough, and significantly different from normal organs) # The index weight is very small, only to distinguish slightly, not to significantly increase the value base_val = 2.0 scale_val = base_val + 0.5 * (OAR_DICT[key] / len(OAR_DICT)) else: # Normal organs: value range 0.5 ~ 0.8 base_val = 0.5 scale_val = base_val + 0.3 * (OAR_DICT[key] / len(OAR_DICT)) # Assign value comb_oar = torch.maximum(comb_oar, single_oar.round() * scale_val) else: # If not normalize, keep the original logic or modify according to needs comb_oar = torch.maximum(comb_oar, single_oar.round() * OAR_DICT[key]) # --- Special Handling for Limiting OARs End --- return comb_oar, cat_oar def combine_ptv(tmp_dict, scaled_dose_dict): prescribed_dose = [] cat_ptv = torch.zeros([3] + list(tmp_dict['img'].shape)[1:]) prescribed_dose = [0] * 3 comb_ptv = torch.zeros(tmp_dict['img'].shape) cnt = 0 for key in scaled_dose_dict.keys(): original_rx = scaled_dose_dict[key] target_val = min(original_rx, 3.0) # tmp_ptv = tmp_dict[key] * scaled_dose_dict[key] tmp_ptv = tmp_dict[key] * target_val # v4: set the prescribed dose to 30 Gy prescribed_dose[cnt] = scaled_dose_dict[key] cat_ptv[cnt] = tmp_ptv comb_ptv = torch.maximum(comb_ptv, tmp_ptv) cnt += 1 # sort the cat_ptv according to the prescribed dose paired = [(cat_ptv[i], prescribed_dose[i]) for i in range(len(prescribed_dose))] paired_sorted = sorted(paired, key=lambda x: x[1], reverse=True) cat_ptv = torch.stack([x[0] for x in paired_sorted]) prescribed_dose = [x[1] for x in paired_sorted] return comb_ptv, prescribed_dose, cat_ptv def tr_augmentation(KEYS, in_size, out_size, crop_center): return Compose([ SpatialCropd(keys = KEYS, roi_center = crop_center, roi_size = [int(in_size[0] * 1.2), int(in_size[1] * 1.2), int(in_size[2] * 1.2)], allow_missing_keys = True), SpatialPadd(keys = KEYS, spatial_size = [int(in_size[0] * 1.2), int(in_size[1] * 1.2), int(in_size[2] * 1.2)], mode = 'constant', allow_missing_keys = True), RandSpatialCropd(keys = KEYS, roi_size = [int(in_size[0] * 0.85), int(in_size[1] * 0.85), int(in_size[2] * 0.85)], max_roi_size = [int(in_size[0] * 1.2), int(in_size[1] * 1.2), int(in_size[2] * 1.2)], random_center = True, random_size = True, allow_missing_keys = True), RandRotated(keys = KEYS, prob=0.8, range_x= 1, range_y = 0.2, range_z = 0.2, allow_missing_keys = True), RandFlipd(keys = KEYS, prob = 0.4, spatial_axis = 0, allow_missing_keys = True), RandFlipd(keys = KEYS, prob = 0.4, spatial_axis = 1, allow_missing_keys = True), RandFlipd(keys = KEYS, prob = 0.4, spatial_axis = 2, allow_missing_keys = True), Resized(keys = KEYS, spatial_size = out_size, allow_missing_keys = True), ]) def tt_augmentation(KEYS, in_size, out_size, crop_center): return Compose([ SpatialCropd(keys = KEYS, roi_center = crop_center, roi_size = in_size, allow_missing_keys = True), SpatialPadd(keys = KEYS, spatial_size = in_size, mode = 'constant', allow_missing_keys = True), Resized(keys = KEYS, spatial_size = out_size, allow_missing_keys = True), ]) def load_nifti(path): """ Load a .nii or .nii.gz file and return a numpy array. torch shape: [B, C, Z, Y, X] where B=1, C=1, Z=depth, Y=height, X=width """ nii = nib.load(path) return np.transpose(nii.get_fdata(), (2, 1, 0)).astype(np.float32)