| 48 | |
| 49 | |
| 50 | class Synapse_dataset(Dataset): |
| 51 | def __init__(self, base_dir, list_dir, split, nclass=9, transform=None): |
| 52 | self.transform = transform # using transform in torch! |
| 53 | self.split = split |
| 54 | self.sample_list = open(os.path.join(list_dir, self.split+'.txt')).readlines() |
| 55 | self.data_dir = base_dir |
| 56 | self.nclass = nclass |
| 57 | |
| 58 | def __len__(self): |
| 59 | return len(self.sample_list) |
| 60 | |
| 61 | def __getitem__(self, idx): |
| 62 | if self.split == "train": |
| 63 | slice_name = self.sample_list[idx].strip('\n') |
| 64 | data_path = os.path.join(self.data_dir, slice_name+'.npz') |
| 65 | data = np.load(data_path) |
| 66 | image, label = data['image'], data['label'] |
| 67 | #print(image.shape) |
| 68 | #image = np.reshape(image, (512, 512)) |
| 69 | #image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) |
| 70 | |
| 71 | #label = np.reshape(label, (512, 512)) |
| 72 | |
| 73 | |
| 74 | else: |
| 75 | vol_name = self.sample_list[idx].strip('\n') |
| 76 | filepath = self.data_dir + "/{}.npy.h5".format(vol_name) |
| 77 | data = h5py.File(filepath) |
| 78 | image, label = data['image'][:], data['label'][:] |
| 79 | #image = np.reshape(image, (image.shape[2], 512, 512)) |
| 80 | #label = np.reshape(label, (label.shape[2], 512, 512)) |
| 81 | #label[label==5]= 0 |
| 82 | #label[label==9]= 0 |
| 83 | #label[label==10]= 0 |
| 84 | #label[label==12]= 0 |
| 85 | #label[label==13]= 0 |
| 86 | #label[label==11]= 5 |
| 87 | |
| 88 | if self.nclass == 9: |
| 89 | label[label==5]= 0 |
| 90 | label[label==9]= 0 |
| 91 | label[label==10]= 0 |
| 92 | label[label==12]= 0 |
| 93 | label[label==13]= 0 |
| 94 | label[label==11]= 5 |
| 95 | |
| 96 | sample = {'image': image, 'label': label} |
| 97 | if self.transform: |
| 98 | sample = self.transform(sample) |
| 99 | sample['case_name'] = self.sample_list[idx].strip('\n') |
| 100 | return sample |
no outgoing calls
no test coverage detected