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 tqdm import nibabel as nib import copy def resize_array(array, current_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:] new_shape = [original_shape[0], 256, 256] scaling_factors = [new_shape[i] / original_shape[i] for i in range(len(original_shape))] resized_spacing = [current_spacing[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() resized_array = F.interpolate(array, size=new_shape, mode='nearest').cpu().numpy() # breakpoint() return resized_array, resized_spacing def resize_mask(array, current_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:] new_shape = [original_shape[0], 256, 256] resized_array = F.interpolate(array, size=new_shape, mode='nearest').cpu().numpy() # breakpoint() return resized_array class CTReportDatasetinfer(Dataset): def __init__(self, data_folder, csv_file, min_slices=20, resize_dim=500, force_num_frames=True, labels = "labels.csv"): self.data_folder = data_folder self.min_slices = min_slices self.labels = labels # self.accession_to_text = self.load_accession_text(csv_file) self.paths=[] self.samples = self.prepare_samples() self.transform = transforms.Compose([ transforms.Resize((resize_dim,resize_dim)), transforms.ToTensor() ]) self.nii_to_tensor = partial(self.nii_img_to_tensor, transform = self.transform) self.sample_length=64 def load_accession_text(self, csv_file): df = pd.read_csv(csv_file) accession_to_text = {} for index, row in df.iterrows(): accession_to_text[row['VolumeName']] = row["Findings_EN"],row['Impressions_EN'] # accession_to_text[row['VolumeName']] = row["disease_findings"] return accession_to_text def prepare_samples(self): samples = [] # patient_folders = glob.glob(os.path.join(self.data_folder, '*')) # Read labels once outside the loop test_df = pd.read_csv(self.labels) test_label_cols = list(test_df.columns[1:]) test_df['one_hot_labels'] = list(test_df[test_label_cols].values) import json with open('/sd/shuhan/CT-RATE/effusion_mask_json/disease_valid_single_prompt_opacity.json', 'r') as f: items = [json.loads(line) for line in f] # 2. 提取所有 volume_path effusion_mask_paths = [item['disease_mask'] for item in items if 'disease_mask' in item] organ_mask_paths = [item['organ_mask'] for item in items if 'organ_mask' in item] disease_findings_list = [item['disease_findings'] for item in items if 'disease_findings' in item] disease_mask_channels = [item['disease_mask_channel'] for item in items if 'disease_mask_channel' in item] disease_labels = [item['disease_label'] for item in items if 'disease_label' in item] disease_classes = [item['disease_class'] for item in items if 'disease_class' in item] # breakpoint() for (organ_mask_file, effusion_mask_file, disease_findings, disease_mask_channel, disease_label, disease_class) in tqdm.tqdm(zip(organ_mask_paths, effusion_mask_paths, disease_findings_list, disease_mask_channels, disease_labels, disease_classes)): if effusion_mask_file=='seg_rxg/valid_851_a_2.nii.gz': continue # for patient_folder in tqdm.tqdm(patient_folders): # accession_folders = glob.glob(os.path.join(patient_folder, '*')) # for accession_folder in accession_folders: # nii_files = glob.glob(os.path.join(accession_folder, '*.npz')) # for nii_file in nii_files: accession_number = organ_mask_file.split("/")[-1] seg_file = '/sd/shuhan/CT-RATE/'+effusion_mask_file nii_file = '/sd/shuhan/CT-RATE/'+organ_mask_file # accession_number = accession_number.replace(".npz", ".nii.gz") # if accession_number not in self.accession_to_text: # continue text_final = disease_findings # text_final = "" # for text in list(impression_text): # text = str(text) # if text == "Not given.": # text = "" # text_final = text_final + text onehotlabels = test_df[test_df["VolumeName"] == accession_number]["one_hot_labels"].values if len(onehotlabels) > 0: samples.append((nii_file, seg_file, text_final, onehotlabels[0], disease_mask_channel, disease_label, disease_class)) self.paths.append(nii_file) return samples def __len__(self): return len(self.samples) def nii_img_to_tensor(self, path, seg_file, disease_mask_channel, disease_label, disease_class, transform): nii_img = nib.load(str(path)) img_data = nii_img.get_fdata() df = pd.read_csv("/sd/shuhan/CT-RATE/metadata/all_metadata.csv") #select the metadata file_name = path.split("/")[-1] row = df[df['VolumeName'] == file_name] slope = float(row["RescaleSlope"].iloc[0]) intercept = float(row["RescaleIntercept"].iloc[0]) xy_spacing = float(row["XYSpacing"].iloc[0][1:][:-2].split(",")[0]) z_spacing = float(row["ZSpacing"].iloc[0]) nii_seg = nib.load(str(seg_file)) # breakpoint() # xy_spacing = nib.affines.voxel_sizes(nii_seg.affine) # z_spacing = nib.affines.voxel_sizes(nii_seg.affine) mask_data = nii_seg.get_fdata()[int(disease_mask_channel)] current = (z_spacing, xy_spacing, xy_spacing) img_data = img_data.transpose(2, 0, 1) tensor = torch.tensor(img_data) tensor = tensor.unsqueeze(0).unsqueeze(0) img_data, target_spacing = resize_array(tensor, current) img_data = img_data[0][0] # breakpoint() mask_data = mask_data.transpose(2, 0, 1) tensor = torch.tensor(mask_data) tensor = tensor.unsqueeze(0).unsqueeze(0) mask_data = resize_mask(tensor, current) mask_data = mask_data[0][0] # breakpoint() assert mask_data.shape == img_data.shape num_slices = img_data.shape[0] # 找到 mask>0 的切片索引 valid_indices = np.where(mask_data.sum(axis=(1, 2)) > 0)[0] if len(valid_indices) > 0: # 只保留能完整取 sample_length 的起点 valid_start_ids = valid_indices[valid_indices <= num_slices - self.sample_length] if len(valid_start_ids) > 0: start_id = np.random.choice(valid_start_ids) else: # 如果没有合法起点就退化到随机，但保证能取 sample_length start_id = np.random.randint(0, num_slices - self.sample_length + 1) else: # 如果 mask 全部为0，则完全随机 start_id = np.random.randint(0, num_slices - self.sample_length + 1) start_id = np.random.randint(0, img_data.shape[0]-self.sample_length) end_id = start_id + self.sample_length # 保证长度严格等于 sample_length # img_data = img_data[start_id:end_id] # mask_data = mask_data[start_id:end_id] # mask_all = copy.deepcopy(img_data) mask_all = np.zeros_like(img_data) mask_all[mask_data>0] = 280 mask_data = (((mask_all ) / 300)).astype(np.float32) * 2 -1 img_data = (((img_data ) / 300)).astype(np.float32) * 2 -1 # breakpoint() # slices=[] img_data = torch.tensor(img_data) mask_data = torch.tensor(mask_data) # # Get the dimensions of the input tensor # target_shape = (480,480,240) img_data = img_data.unsqueeze(0) mask_data = mask_data.unsqueeze(0) img_data=img_data.repeat(3,1,1,1) mask_data=mask_data.repeat(3,1,1,1) # example = {} # example['name'] = file_name # example['volume_data'] = tensor # # example['organ_mask'] = volume_seg # example['spacing'] = target_spacing return img_data, mask_data, target_spacing, file_name def __getitem__(self, index): nii_file, seg_file, input_text, onehotlabels, disease_mask_channel, disease_label, disease_class = self.samples[index] video_tensor, volume_seg, spacing, file_name = self.nii_to_tensor(nii_file, seg_file, disease_mask_channel, disease_label, disease_class) input_text = input_text.replace('"', '') input_text = input_text.replace('\'', '') input_text = input_text.replace('(', '') input_text = input_text.replace(')', '') name_acc = nii_file.split("/")[-2] # return dict(name=file_name, disease_mask_channel=disease_mask_channel,edited=torch.cat([video_tensor.float(), volume_seg.float()],dim=-1), edit=dict(c_concat=torch.cat([video_tensor.float(), torch.ones_like(video_tensor)*-1],dim=-1), c_crossattn=input_text)) return dict(name=file_name, edited=volume_seg.float(), c_concat=video_tensor.float(), c_crossattn=input_text, disease_mask_channel=disease_mask_channel)