Open AlyssaAmod opened 1 year ago
@Alexandra-Smith I see you already discovered this during inference --> can you check if you happy with my suggested code replace from line 41
As in script under Ask Alex:
image = np.load(self.imgs[idx])
image = torch.from_numpy(image) # 4, 240, 240, 155
if self.mode == "labels":
mask = np.load(self.lbls[idx])
mask = torch.from_numpy(mask) # 240, 240, 155
if self.transform is not None:
if self.mode == "labels":
subject = tio.Subject(
image=tio.ScalarImage(tensor=image),
mask=tio.LabelMap(tensor=mask)
)
tranformed_subject = self.transform(subject)
if self.SSA == False and self.SSAtransform is not None:
tranformed_subject = self.SSAtransform(tranformed_subject)
print("Tranformed_subject: ", tranformed_subject)
image = tranformed_subject["image"].data
mask = tranformed_subject["mask"].data
return image, mask, self.imgs[idx]
else:
subject = tio.Subject(
image=tio.ScalarImage(tensor=image),
)
tranformed_subject = self.transform(subject)
print("Tranformed_subject: ", tranformed_subject)
image = tranformed_subject["image"].data
return image, self.imgs[idx]
UNN_BraTS23/scripts/data_class.py
Should this section not account for when mask does not exist?
see code change suggestion in docstring at end