From 3D datas to 2D slices

1

Hi there I have faced a problem and then I have read some forums but I couldn`t get the correct answer. I would like to ask again, in Dataset part I would like to read my 3D data which has first 2 dimension same shape ([512,512,600] or [512,512,550]) and I would like to create like 600 slices or 550 slices which has just [512,512] dimensions as 2d array.

class TrainSet(Dataset):
    def __init__(self, image_dir, mask_dir):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.images = os.listdir(image_dir)
        self.masks = os.listdir(mask_dir)

  

    def __getitem__(self, index):
        img_path = os.path.join(self.image_dir, self.images[index])
        mask_path = os.path.join(self.mask_dir, self.masks[index])
        image = nibabel.load(img_path)
        img = image.get_fdata()
        mask, _ = nrrd.read(mask_path)
        
        for k in range(img.shape[2]):
            img = img[:, :, k]
            mask = mask[:, :, k]
            
            return img, mask

    def __len__(self):
        return len(self.images)

I have tried something like that but as you guess it couldnt apply for whole images. And I didnt prefer to make list by slice by actually because it causes other problems. At that point can you suggest me some solutions please? How can I create my 2D dataset from 3D arrays? Thank you very much in advance~~

2

I don’t know how these slices should be used, i.e. if you want to push each slice separately to the model or if you would like to create a batch of slices from different files.
In the former case it might be easier to just return all slices from the current file and iterate them in the training loop.
In the latter case, a custom sampler (e.g. using two indices for the file and slice index) might work.

3

Thank you very much for fast reply~
I think both idea looks okay but the first one is muc more close. From your suggestion I understood that I should return 3d arrays and seperate on the training loop… I am sorry I couldn`t understand clearly how can I iterate them in training loop to 2D arrays?

4

Yes, you could directly return the 3D tensor, which would have the shape [batch_size=1, 512, 512, 600] inside the training loop.
Afterwards, you could index each slice via:

for data, target in loader:
    # data will have the shape [1, 512, 512, 600]
    for idx in data.size(3):
        x = data[:, :, :, i]
        # x will have the shape [1, 512, 512]

and could then pass x to the model.

5

I would like to try that way but in advance why I am having this problem: RuntimeError: stack expects each tensor to be equal size, but got [512, 512, 533] at entry 0 and [512, 512, 661] at entry 1
I didn`t change the code I just deleted the for part…

class TrainSet(Dataset):
    def __init__(self, image_dir, mask_dir):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.images = os.listdir(image_dir)
        self.masks = os.listdir(mask_dir)

    def __getitem__(self, index):
        img_path = os.path.join(self.image_dir, self.images[index])
        mask_path = os.path.join(self.mask_dir, self.masks[index])
        image = nibabel.load(img_path)
        img = image.get_fdata()
        mask, _ = nrrd.read(mask_path)



        return img, mask

    def __len__(self):
        return len(self.images)
6

This error is raised in the collate_fn which tries to create a batch of multiple samples.
Since each sample has a different shape, torch.stack will fail.
If you want to iterate each slice of a single sample, you could thus set batch_size=1 and it should work.

7

I am really thankful for fast replying I really need to finish this code so I am also sorry for asking so question~~
After I finished my dataset part I was started to train part like that basically:

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = UNet()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)


for epoch in range(epochs):
    for i, (images, masks) in enumerate(train_loader, 1):
        print(images.size())
        for idx in range(images.size(3)):
            images = images[:, :, :, idx].unsqueeze(1).to(device)
            masks = masks[:, :, :, idx].unsqueeze(1).to(device)
            images = images.type(torch.FloatTensor)
            masks = images.type(torch.FloatTensor)
            print(images.shape)
            print(masks.shape)
            outputs = model(images)
            loss = criterion(outputs, masks)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if (i) % 1 == 0:
                print(f"Epoch [{epoch + 1}/{epochs}], Step [{i}/{total_steps}], Loss: {loss.item():4f}")

But I couldn`t understand why it was connected to CrossEntropyLoss because it is giving this error from loss part : RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4

8

I guess you are working on a multi-class segmentation?
If so, the model output should have the shape [batch_size, nb_classes, height, width], while the target should have the shape [batch_size, height, width] and container class indices in the range [0, nb_classes-1] as shown here:

batch_size, nb_classes, height, width = 2, 3, 4, 4
output = torch.randn(batch_size, nb_classes, height, width)
target = torch.randint(0, nb_classes, (batch_size, height, width))

criterion = nn.CrossEntropyLoss()
loss = criterion(output, target)

Based on your error message I guess your target has 4 dimensions, which is wrong. If you are using a one-hot encoded target, use target = torch.argmax(target, dim=1) to create the expected shape (assuming the class dimension is in dim1).