import os import glob import json import torch import pandas as pd import numpy as np from PIL import Image from torch.utils.data import Dataset import torchvision.transforms as transforms from functools import partial import torch.nn.functional as F import nibabel as nib import tqdm import random from monai.transforms import Compose, RandRotate, RandFlip, RandZoom, RandAffine, RandGaussianNoise, RandScaleIntensity,RandAdjustContrast,RandCoarseDropout,RandHistogramShift def resize_array(array, current_spacing, target_spacing): """ Resize the array to match the target spacing. Args: array (torch.Tensor): Input array to be resized. current_spacing (tuple): Current voxel spacing (z_spacing, xy_spacing, xy_spacing). target_spacing (tuple): Target voxel spacing (target_z_spacing, target_x_spacing, target_y_spacing). Returns: np.ndarray: Resized array. """ # Calculate new dimensions original_shape = array.shape[2:] scaling_factors = [ current_spacing[i] / target_spacing[i] for i in range(len(original_shape)) ] new_shape = [ int(original_shape[i] * scaling_factors[i]) for i in range(len(original_shape)) ] # Resize the array resized_array = F.interpolate(array, size=new_shape, mode='trilinear', align_corners=False).cpu().numpy() return resized_array def shuffle_sentences(text): """ 将一个包含多个句子的字符串拆分成单独的句子，然后随机重组这些句子。 Args: text (str): 输入的文本字符串，包含多个由句号分隔的句子 Returns: str: 随机重组后的句子 """ import random import re # 使用正则表达式分割句子，保留句号 # 这种方式可以更好地处理多种句子结束情况 sentences = re.findall(r'[^.!?]+[.!?]', text) # 处理可能的最后一个不带句号的句子 remaining = text[sum(len(s) for s in sentences):] if remaining.strip(): sentences.append(remaining.strip() + '.') # 随机打乱句子顺序，一行代码搞定 random.shuffle(sentences) # 直接用空格连接所有句子 return ' '.join(s.strip() for s in sentences) # 要提前保证数据都是存在的，且数量与对应的josnl文件长度是一致的 注意！ modality_dict = {'3D':0,'2D':1} class ImageReportDataset(Dataset): def __init__(self, jsonl_file,need_aug=True,modality='3D',is_train=True, modal_embedding=False,length_=None ): self.length_ = length_ self.modality = modality self.paths=[] self.is_train = is_train self.samples = self.prepare_samples(jsonl_file) self.modal_embedding = modal_embedding percent = 100 num_files = int((len(self.samples) * percent) / 100) #num_files = 2286 self.samples = self.samples[:num_files] # 是同时包含index以及对应的文本的 print(f'** DATA ** Load {len(self.samples)} samples.') # 假设数据处理过了，不再需要check_integrity print('augmentation:', need_aug) if need_aug: # 不行再回退到原来的版本 # input image must be CDHW # 修改 if '3D' in self.modality: self.augmentator = Compose([ RandRotate(range_x=15, range_y=15, range_z=15, prob=0.3, keep_size=True), # 随机旋转 RandFlip(prob=0.5, spatial_axis=0), # 随机翻转（沿深度轴） RandZoom(min_zoom=0.85, max_zoom=1.15, prob=0.3, keep_size=True), # 随机缩放 RandAffine(prob=0.3, translate_range=(24, 48, 48)), # 随机平移 RandGaussianNoise(prob=0.3, mean=0.0, std=0.05), # 添加高斯噪声 RandScaleIntensity(factors=[-0.1, 0.1], prob=0.3), RandAdjustContrast(prob=0.3, gamma=(0.9, 1.1)), # 增加对比度调整 RandCoarseDropout(prob=0.2, holes=8, spatial_size=3) # 模拟局部缺失 ]) elif '2D' in self.modality: self.augmentator = Compose([ RandRotate(range_x=15, range_y=15, prob=0.5, keep_size=True), # 随机旋转（仅x和y轴） RandFlip(prob=0.5, spatial_axis=(0,1)), # 随机翻转（仅水平和垂直轴） RandZoom(min_zoom=0.8, max_zoom=1.1, prob=0.5, keep_size=True), # 随机缩放 RandAffine(prob=0.5, translate_range=(32, 32)), # 随机平移（仅2D平移） RandGaussianNoise(prob=0.5, mean=0.0, std=0.1), # 添加高斯噪声 RandScaleIntensity(factors=[-0.15, 0.15], prob=0.5), # 随机调整像素强度 RandAdjustContrast(prob=0.4, gamma=(0.85, 1.15)) # 增加对比度变化 ]) else: self.augmentator = None if int(os.environ.get("RANK", 0)) == 0: print(f'** DATA ** Load {len(self.samples)} images.') def prepare_samples(self, jsonl_file): samples = [] with open(jsonl_file, 'r') as f: data = f.readlines() data = [json.loads(l) for l in data] for d in data: samples.append((d['img_path'], d['text'])) self.paths.append(d['img_path']) return samples def __len__(self): if self.length_ is not None: return self.length_ else: return len(self.samples) def MLS_nii_img_to_tensor(self, path): nii_img = safe_load_image(sample["image"]) img_data = nii_img.get_fdata() img_data = np.flip(img_data, axis=0) img_data = np.flip(img_data, axis=1) # WARNING Respacing img_data = img_data.transpose(2, 0, 1) img_data = np.copy(img_data) tensor = torch.tensor(img_data) tensor = tensor.unsqueeze(0).unsqueeze(0) target_x_spacing = 0.75 target_y_spacing = 0.75 target_z_spacing = 1.5 current = (3, 1, 1) # this is all set to 3 1 1 target = (target_z_spacing, target_x_spacing, target_y_spacing) img_data = resize_array(tensor, current, target) img_data = img_data[0][0] img_data= np.transpose(img_data, (1, 2, 0)) # WARNING Normalization 2 hu_min, hu_max = -1000, 1000 img_data = np.clip(img_data, hu_min, hu_max) img_data = (((img_data ) / 1000)).astype(np.float32) tensor = torch.tensor(img_data) # WARNING Padding or Crop # Get the dimensions of the input tensor target_shape = (224,224,240) # h w d # Extract dimensions h, w, d = tensor.shape # Calculate cropping/padding values for height, width, and depth dh, dw, dd = target_shape h_start = max((h - dh) // 2, 0) h_end = min(h_start + dh, h) w_start = max((w - dw) // 2, 0) w_end = min(w_start + dw, w) d_start = max((d - dd) // 2, 0) d_end = min(d_start + dd, d) # Crop or pad the tensor tensor = tensor[h_start:h_end, w_start:w_end, d_start:d_end] pad_h_before = (dh - tensor.size(0)) // 2 pad_h_after = dh - tensor.size(0) - pad_h_before pad_w_before = (dw - tensor.size(1)) // 2 pad_w_after = dw - tensor.size(1) - pad_w_before pad_d_before = (dd - tensor.size(2)) // 2 pad_d_after = dd - tensor.size(2) - pad_d_before tensor = torch.nn.functional.pad(tensor, (pad_d_before, pad_d_after, pad_w_before, pad_w_after, pad_h_before, pad_h_after), value=-1) tensor = tensor.permute(2, 0, 1) # d h w tensor = tensor.unsqueeze(0) # 1 d h w return tensor def nii_img_to_tensor(self, path): nii_img = safe_load_image(sample["image"]) img_data = nii_img.get_fdata() img_data = np.flip(img_data, axis=0) img_data = np.flip(img_data, axis=1) # WARNING Respacing img_data = img_data.transpose(2, 0, 1) img_data = np.copy(img_data) tensor = torch.tensor(img_data) tensor = tensor.unsqueeze(0).unsqueeze(0) target_x_spacing = 0.75 target_y_spacing = 0.75 target_z_spacing = 1.5 current = (3, 1, 1) # this is all set to 3 1 1 target = (target_z_spacing, target_x_spacing, target_y_spacing) img_data = resize_array(tensor, current, target) img_data = img_data[0][0] img_data= np.transpose(img_data, (1, 2, 0)) # WARNING Normalization 2 hu_min, hu_max = -1000, 1000 img_data = np.clip(img_data, hu_min, hu_max) img_data = (((img_data ) / 1000)).astype(np.float32) tensor = torch.tensor(img_data) target_shape = (480,480,240) # h w d # Extract dimensions h, w, d = tensor.shape # Calculate cropping/padding values for height, width, and depth dh, dw, dd = target_shape h_start = max((h - dh) // 2, 0) h_end = min(h_start + dh, h) w_start = max((w - dw) // 2, 0) w_end = min(w_start + dw, w) d_start = max((d - dd) // 2, 0) d_end = min(d_start + dd, d) # # Add random shift to the crop region # hw_max_shift = 128 # z_max_shift = 64 # if z_max_shift > 0 or hw_max_shift > 0: # # Generate random shifts within [-max_shift, max_shift] # h_shift = np.random.randint(-hw_max_shift, hw_max_shift + 1) # w_shift = np.random.randint(-hw_max_shift, hw_max_shift + 1) # d_shift = np.random.randint(-z_max_shift, z_max_shift + 1) # # Apply shifts, ensuring the crop region remains within the image bounds # h_start = max(h_start + h_shift, 0) # h_end = min(h_start + dh, h) # w_start = max(w_start + w_shift, 0) # w_end = min(w_start + dw, w) # d_start = max(d_start + d_shift, 0) # d_end = min(d_start + dd, d) # Crop or pad the tensor tensor = tensor[h_start:h_end, w_start:w_end, d_start:d_end] pad_h_before = (dh - tensor.size(0)) // 2 pad_h_after = dh - tensor.size(0) - pad_h_before pad_w_before = (dw - tensor.size(1)) // 2 pad_w_after = dw - tensor.size(1) - pad_w_before pad_d_before = (dd - tensor.size(2)) // 2 pad_d_after = dd - tensor.size(2) - pad_d_before tensor = torch.nn.functional.pad(tensor, (pad_d_before, pad_d_after, pad_w_before, pad_w_after, pad_h_before, pad_h_after), value=-1) tensor = tensor.permute(2, 0, 1) # d h w tensor = tensor.unsqueeze(0) # 1 d h w return tensor def load_2d_image_to_tensor(self, image_path,resize=True): """ Load a 2D grayscale image and convert it to a tensor with dimensions [1, 1, 480, 480]. Args: image_path (str): Path to the image file. Returns: tensor (torch.Tensor): Tensor with shape [1, 1, 224, 24]. """ # Load the image as grayscale image = Image.open(image_path).convert('L') # Resize to 224*224 if needed if resize: if image.size != (224, 224): image = image.resize((224, 224), Image.BILINEAR) else: if image.size != (224, 224): # 计算可以开始裁剪的最大左上角坐标 width, height = image.size crop_size = 224 if width < crop_size or height < crop_size: raise ValueError(f"Image size {image.size} is smaller than the crop size {crop_size}.") max_left = width - crop_size max_top = height - crop_size # 随机选择裁剪的左上角坐标 left = random.randint(0, max(0, max_left)) top = random.randint(0, max(0, max_top)) # 裁剪图像 image = image.crop((left, top, left + crop_size, top + crop_size)) # Convert to numpy array and normalize to [0, 1] img_array = (np.array(image, dtype=np.float32) / 255.0) * 2 -1 # Convert to tensor tensor = torch.from_numpy(img_array) # Add batch and channel dimensions tensor = tensor.unsqueeze(0).unsqueeze(0) # [1, 1, 480, 480] return tensor def __getitem__(self, index): nii_file, input_text = self.samples[index] if '3D' in self.modality: try: # video_tensor = self.MLS_nii_img_to_tensor(nii_file) if self.is_train: video_tensor = self.nii_img_to_tensor(nii_file) else: video_tensor = self.MLS_nii_img_to_tensor(nii_file) assert video_tensor.dim() == 4, f"Expected 4D tensor, got {video_tensor.dim()}D tensor for {index}" if self.augmentator is not None and self.is_train: video_tensor = self.augmentator(video_tensor) # video_tensor = torch.zeros((1, 240, 224, 224), dtype=torch.float32) except Exception as e: print(f"Error processing {nii_file} at index {index}: {e}. Skipping this sample.") video_tensor = torch.zeros((1, 240, 224, 224), dtype=torch.float32) elif '2D' in self.modality: try: video_tensor = self.load_2d_image_to_tensor(nii_file, resize=True) video_tensor = video_tensor.squeeze().unsqueeze(0) if self.augmentator is not None and self.is_train: video_tensor = self.augmentator(video_tensor) video_tensor = video_tensor.unsqueeze(0) # Add batch dimension except Exception as e: print(f"Error processing {nii_file} at index {index}: {e}. Skipping this sample.") video_tensor = torch.zeros((1, 1, 224, 224), dtype=torch.float32) input_text = str(input_text) input_text = input_text.replace('"', '') input_text = input_text.replace('\'', '') input_text = input_text.replace('(', '') input_text = input_text.replace(')', '') # 不再打乱句子顺序了 # if self.is_train: # # 训练时打乱句子顺序 # input_text = shuffle_sentences(input_text) if not self.modal_embedding: return video_tensor, input_text else: return video_tensor, input_text, modality_dict[self.modality] if __name__ == '__main__': from torch.utils.data import DataLoader import random import time import itertools """测试 CTReportDataset 类的功能""" # # 初始化数据集 dataset_path = '/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/CTRATE-IR/test.jsonl' dataset = CTReportDataset(dataset_path, need_aug=False, modality='2D', is_train=False) print(f"数据集大小: {len(dataset)} 样本") # 测试数据集的基本信息 print("\n1. 检查数据集基本信息...") print(f"数据集路径: {dataset_path}") print(f"文件是否存在: {os.path.exists(dataset_path)}") # 随机选择几个样本进行测试 print("\n2. 测试随机样本...") # indices = random.sample(range(len(dataset)), min(5, len(dataset))) indices = [0, 1, 2, 3, 4] # 测试前5个样本 for idx in indices: print(f"\n样本 #{idx}:") try: # 获取样本 video_tensor, input_text, index = dataset[idx] print(f" - 视频张量形状: {video_tensor.shape}") print(f" - 输入文本: {input_text}") print(f" - 索引: {index}") print(f' - 视频张量数据: {video_tensor[0, 0, 0, 0:10]}') # 打印前10个元素 except Exception as e: print(f"错误: {e}") # 测试 DataLoader print("\n3. 测试 DataLoader...") try: dataloader = DataLoader(dataset, batch_size=5, shuffle=True, num_workers=0) batch = next(iter(dataloader)) print(f"成功加载批次数据") if isinstance(batch, tuple): for i, item in enumerate(batch): if isinstance(item, torch.Tensor): print(f" - 批次项目 {i}: Tensor, 形状: {item.shape}") else: print(f" - 批次项目 {i}: {type(item)}") except Exception as e: print(f"DataLoader 错误: {e}") print("\n测试完成!") def safe_load_image(path): """ Robust image loader that supports both .nii.gz and .npy files. If .nii.gz is missing, automatically falls back to _compressed.npy. """ try: if os.path.exists(path): if path.endswith('.nii.gz'): data = nib.load(path).get_fdata() return data elif path.endswith('.npy'): data = np.load(path) return data # fallback to compressed npy npy_fallback = path.replace('.nii.gz', '_compressed.npy') if os.path.exists(npy_fallback): data = np.load(npy_fallback) print(f"[Info] Fallback to compressed file: {npy_fallback}") return data # all failed raise FileNotFoundError(f"{path} and fallback {npy_fallback} not found") except Exception as e: print(f"[Warning] Error loading {path}: {e}") return None