from torchvision.utils import save_image
from PIL import Image
import torch


def visualize_patch(viz_patches):
    from IPython.display import display, clear_output
    pix_mean = torch.Tensor([0.485, 0.456, 0.406]).to(viz_patches.device).view(1, -1, 1, 1)
    pix_std = torch.Tensor([0.229, 0.224, 0.225]).to(viz_patches.device).view(1, -1, 1, 1)
    viz_patches = viz_patches * pix_std + pix_mean
    img_size = viz_patches.size(2)
    if img_size < 256:
        viz_patches = torch.nn.functional.interpolate(viz_patches, scale_factor=256 // img_size, mode="nearest")
    save_image(viz_patches, "samples.png", nrow=4, normalize=True, value_range=(0, 1))
    sampled_patches_viz = Image.open("samples.png")
    clear_output(wait=True)
    display(sampled_patches_viz)