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~~