Spaces:
Running
Running
| import io | |
| from PIL import Image as im | |
| from tqdm import tqdm | |
| import matplotlib.pyplot as plt | |
| import torch | |
| import numpy as np | |
| class Painter(object): | |
| def __init__(self) : | |
| pass | |
| def show_images(self, images, title : str = '', index : bool = False, cmap = None, show = True): | |
| images = images.permute(0, 2, 3, 1) | |
| if type(images) is torch.Tensor: | |
| images = images.detach().cpu().numpy() | |
| images = np.clip(images / 2 + 0.5, 0, 1) | |
| fig = plt.figure(figsize=(8, 8)) | |
| rows = int(len(images) ** (1 / 2)) | |
| cols = round(len(images) / rows) | |
| idx = 0 | |
| for _ in range(rows): | |
| for _ in range(cols): | |
| fig.add_subplot(rows, cols, idx + 1) | |
| if idx < len(images): | |
| plt.imshow(images[idx], cmap = cmap) | |
| if index : | |
| plt.title(idx + 1) | |
| plt.axis('off') | |
| idx += 1 | |
| fig.suptitle(title, fontsize=30) | |
| if show: | |
| plt.show() | |
| def show_first_batch(self, loader): | |
| for batch in loader: | |
| self.show_images(images = batch, title = "First Batch") | |
| break | |
| def make_gif(self, images, file_name): | |
| imgs = [] | |
| for i in tqdm(range(len(images))): | |
| img_buf = io.BytesIO() | |
| self.show_images(images[i], title = 't = ' + str(i), show=False) | |
| plt.savefig(img_buf, format='png') | |
| imgs.append(im.open(img_buf)) | |
| imgs[0].save(file_name + '.gif', format='GIF', append_images=imgs, save_all=True, duration=1, loop=0) | |
| plt.close('all') | |