Skip to content

Move SynthSeg-like GMM sampling to LUT #1486

Description

@vcasellesb

Hi! I was trying to implement in pytorch SynthSeg's approach to generating synthetic scans by sampling from a randomly-parameterized GMM conditioned on a label map, and I came across your code.

I saw that your code iterates for each class value, sampling from a random normal distribution at each iteration, and assigning the values to the area covered by that class to the result.

In contrast, SynthSeg uses a look-up table (LUT) which although it's difficult to understand (at least for me), it is way faster, specially in pytorch with advanced indexing.

Is there a reason why you chose the iterative approach that is currently used in torchio? Would you be interested in an implementation of a LUT-based approach?

I am pasting an example of what I mean (LUT-based approach), done by me with Gemini and taking SynthSeg's original work as a template.

Best,
Vicent

    @staticmethod
    def _sample_uniform(size, a = 0., b = 1.):
        """
        :param a: lower bound
        :param b: upper bound
        """
        return torch.rand(size=size) * (b - a) + a

    def get_parameters(self, label_map: torch.Tensor):
        class_vals = torch.unique(label_map.ravel(), sorted=True)
        max_label = class_vals[-1]
        nlabels = len(class_vals)
        means_lut = torch.zeros(size=(max_label+1,), device=label_map.device, dtype=torch.float32)
        stds_lut = means_lut.clone()

        means = self._sample_uniform((nlabels,), *self.means_bounds).to(device=label_map.device)
        stds = self._sample_uniform((nlabels,), *self.stds_bounds).to(device=label_map.device)

        means_lut[class_vals] = means
        stds_lut[class_vals] = stds
        return means_lut, stds_lut

    def __call__(self,
                 label_map: torch.Tensor
                 ) -> torch.Tensor:
        # coerce to int64 for indexing
        label_map = label_map.long()

        means_lut, stds_lut = self.get_parameters(label_map)
        means_map = means_lut[label_map]
        stds_map = stds_lut[label_map]

        return torch.randn_like(label_map, dtype=torch.float32) * stds_map + means_map

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions