Method __call__

datasets/augmentv2.py:195–417

Apply augmentations. Parameters ---------- images : torch.Tensor Input tensor of shape ``[N, C, H, W]``. params : dict | None, default = None *If given*, the function will **reuse** these tensors instead of sampling fresh randomnes

(self, images: torch.Tensor, *, params: Optional[Dict[str, torch.Tensor]] = None,
                 return_params: bool = False)

    # Forward pass with optional deterministic parameters.
    # ---------------------------------------------------------------------
    def __call__(self, images: torch.Tensor, *, params: Optional[Dict[str, torch.Tensor]] = None,
                 return_params: bool = False) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict[str, torch.Tensor]]]:
        """Apply augmentations.

        Parameters
        ----------
        images : torch.Tensor
            Input tensor of shape ``[N, C, H, W]``.
        params : dict | None, default = None
            *If given*, the function will **reuse** these tensors instead of
            sampling fresh randomness.  The expected dictionary keys are listed
            in the source; passing fewer keys simply disables the missing
            augmentations.
        return_params : bool, default = False
            If ``True``, the function also returns the dictionary that contains
            *all* random tensors used during this call (both supplied and
            newly‑sampled).  This can be cached by the caller and reused later.

        Returns
        -------
        images : torch.Tensor
            Augmented images with the same shape as the input.
        labels : torch.Tensor
            Per‑image labels recording the effective augmentation parameters
            (identical to the original StyleGAN‑ADA implementation).
        params : dict | None
            Only returned when ``return_params=True``.
        """

        N, C, H, W = images.shape
        device = images.device

        # Internal helper --------------------------------------------------
        rnd: Dict[str, torch.Tensor] = OrderedDict()  # we will fill this

        def _get(name: str, tensor_fn, batch_dim: int = 0):
            """Retrieve a random tensor by name, sampling if necessary."""
            if params is not None and name in params:
                return params[name].transpose(0, batch_dim)
            t = tensor_fn()
            rnd[name] = t.transpose(0, batch_dim)
            return t

        labels: List[torch.Tensor] = [torch.zeros([images.shape[0], 0], device=device)]

        # -------------------------------
        # Pixel blitting : xflip / yflip
        # -------------------------------
        if self.xflip > 0:
            w = _get('xflip', lambda: torch.randint(2, [N, 1, 1, 1], device=device))
            w = _get('xflip_rand', lambda: torch.where(torch.rand_like(w, dtype=torch.float32) < self.xflip * self.p, w, torch.zeros_like(w)))
            images = torch.where(w == 1, images.flip(3), images)
            labels += [w]

        if self.yflip > 0:
            w = _get('yflip', lambda: torch.randint(2, [N, 1, 1, 1], device=device))
            w = _get('yflip_rand', lambda: torch.where(torch.rand_like(w, dtype=torch.float32) < self.yflip * self.p, w, torch.zeros_like(w)))
            images = torch.where(w == 1, images.flip(2), images)