GirdSampler for training without backgroud

[kodohoi]

I want to create patches for training without any overlapping. Therefore, although I decided to use GridSmapler, I want to exclude the patches with only the background.
I made my images binary values (0.0 and 1.0, 0.0 is background) and I modified GridSampler as below.

import torch
import torchio as tio

class CustomGridSampler(tio.GridSampler):
    def _generate_patches(self, subject):
        subject = self._pad(subject)
        sizes = subject.spatial_shape, self.patch_size, self.patch_overlap
        self._parse_sizes(*sizes)
        locations = self._get_patches_locations(*sizes)
        
        for location in locations:
            index_ini = location[:3]
            patch = self.extract_patch(subject, index_ini)
            
            # Exclude patches where all values are 0
            if torch.any(patch['img'][tio.DATA]):
                yield patch

Is this an appropriate code?

[fepegar]

Have you run it? Did it seem to work?

[kodohoi]

Thank you for your reply.
I modified SubClass to calculate the number of patches, and it seem to work.

import torch
import torchio as tio
from typing import Generator
from typing import Optional
from typing import Union
from torchio.typing import TypeSpatialShape
from torchio.data.subject import Subject
import numpy as np


class NonZeroGridSampler(tio.GridSampler):
    def __init__(
        self,
        subject: Subject,
        patch_size: TypeSpatialShape,
        patch_overlap: TypeSpatialShape = (0, 0, 0),
        padding_mode: Union[str, float, None] = None,
    ):
        super().__init__(subject, patch_size, patch_overlap, padding_mode)
        self.non_zero_patches = list(self._generate_non_zero_patches())

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

    def __getitem__(self, index):
        return self.non_zero_patches[index]

    def _generate_non_zero_patches(
        self,
    ) -> Generator[Subject, None, None]:
        for patch in super()._generate_patches(self.subject):
            if np.any(patch.img.numpy() != 0):  # I defined the name of the key as "img".
                yield patch

example = torch.zeros([50,50,50])
example[20:30, 10:20, 10:20] = 1
example[30:40, 30:40, 40:50] = 1
example[40:50, 10:20, 40:50] = 1

img_example=tio.ScalarImage(tensor=example.unsqueeze(0))
subject = tio.Subject({"img": img_example})

custom_sample = NonZeroGridSampler(subject, 10)
print(len(custom_sample)) -->3

Are there any more appropriate implementations? Because this SubClass just fit for only name of key "img".