Method extract_from_batch

diff2flow/trainer_module.py:448–494

Takes batch and extracts data. Returns: x0: Samples from source distribution (can also be None if we start from noise). x0_latent: Source latent codes. If identity first stage, this is the same as x0. x1: Samples from target distribution (must al

(self, batch)

        return x
    
    def extract_from_batch(self, batch):
        """
        Takes batch and extracts data.

        Returns:
            x0: Samples from source distribution (can also be None if we start from noise).
            x0_latent: Source latent codes. If identity first stage, this is the same as x0.
            x1: Samples from target distribution (must always be provided).
            x1_latent: Target latent codes. if identity first stage, this is the same as x1.
            mask: If available, a valid mask for the target.
            mask_latent: If available, latent valid mask for target.
        """
        # source samples (x0 ~ p(x0))
        if "x0" in batch:
            x0 = batch["x0"]
            # check for precomputed latents
            if "x0_latent" in batch:
                x0_latent = batch["x0_latent"]
                x0_latent = x0_latent * self.scale_factor
            else:
                x0_latent = self.encode_first_stage(x0)
        else:
            x0, x0_latent = None, None

        # target samples (x1 ~ p(x1)) - data
        x1 = batch["x1"]
        # check for precomputed latents
        if "x1_latent" in batch:
            x1_latent = batch["x1_latent"]
            x1_latent = x1_latent * self.scale_factor
        else:
            x1_latent = self.encode_first_stage(x1)

        # check for valid mask
        if "mask" in batch or "valid_mask" in batch:
            mask = batch["mask"] if "mask" in batch else batch["valid_mask"]
            # resize valid mask to latent space size
            mask_latent = resize_ims(mask.float(), size=x1_latent.shape[-2:], mode="bilinear")
            mask_latent = mask_latent == 1
        else:
            mask, mask_latent = None, None

        return {
            "x0": x0, "x0_latent": x0_latent,
            "x1": x1, "x1_latent": x1_latent,
            "mask": mask, "mask_latent": mask_latent
        }
    
    def training_step(self, batch, batch_idx):
        """ extract data """