Add unpatchify model utils operation

randombenj commented 1 month ago

This utility function is usefull to visualize the predicted images by the MAE model for example:

def validation_step(self, batch: torch.Tensor, batch_idx: int):
    images, _ = batch
    loss, preds, idx_mask = self(images)

    idx_mask = idx_mask - 1  # must adjust idx_mask for missing class token

    # visualize predictions
    patches = utils.patchify(images, self.patch_size)

    pred_patches = utils.set_at_index(patches, idx_mask, preds)
    pred_img = utils.unpatchify(pred_patches, self.patch_size)

    masked_patches = utils.set_at_index(patches, idx_mask, torch.zeros_like(preds))
    masked_img = utils.unpatchify(masked_patches, self.patch_size)

    img = torch.stack([zip(images, pred_img, masked_img)])
    grid = torchvision.utils.make_grid(img, nrow=6)
    self.logger.experiment.add_image("generated_images", grid, 0)

randombenj commented 1 month ago

@guarin Thanks, I guess the CI failure in the tests is unrelated?

guarin commented 1 month ago

I guess the CI failure in the tests is unrelated?

Yes it is unrelated, we are looking into it :)

codecov[bot] commented 1 month ago

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 82.23%. Comparing base (a3de571) to head (f6f6071).

Additional details and impacted files

```diff @@ Coverage Diff @@ ## master #1544 +/- ## ========================================== + Coverage 82.21% 82.23% +0.02% ========================================== Files 144 144 Lines 6116 6124 +8 ========================================== + Hits 5028 5036 +8 Misses 1088 1088 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

lightly-ai / lightly

Add unpatchify model utils operation #1544

Codecov Report