@trypag ah man I would stay away from the patch based approach for MRI segmentation… Those methods are incredibly slow at test time and end up taking about as long as (or longer) than traditional methods with much worse accuracy. 2d Conv + Conv transpose on entire slices OR 3d conv +conv-transpose on the entire image works much better. We have ~1 second models for skull stripping (non-DL methods take ~1 min) and ~30s models for tissue segmentation (non-DL takes a few hours)… These are all using conv+Conv transpose ensembles.
Anyways, a very simple approach to sampling 2D patches from 3D MRI is to extract a random slice (see my package torchsample for that transform) then just take a random crop. You can use the “MultiSampler” sampler in torchsample to go through each image more than once in a single epoch. No need to extract EVERY possible patch haha just do it randomly at train time.
It’s still a fun exercise to reimplement them tho… which paper specifically are you interested in? I think I have implementations of most of them and can share or help you develop the right sampler.