MLBench/ReaLens

1import random2import torch3 4 5class ImagePool:6    """This class implements an image buffer that stores previously generated images.7 8    This buffer enables us to update discriminators using a history of generated images9    rather than the ones produced by the latest generators.10    """11 12    def __init__(self, pool_size):13        """Initialize the ImagePool class14 15        Parameters:16            pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created17        """18        self.pool_size = pool_size19        if self.pool_size > 0:  # create an empty pool20            self.num_imgs = 021            self.images = []22 23    def query(self, images):24        """Return an image from the pool.25 26        Parameters:27            images: the latest generated images from the generator28 29        Returns images from the buffer.30 31        By 50/100, the buffer will return input images.32        By 50/100, the buffer will return images previously stored in the buffer,33        and insert the current images to the buffer.34        """35        if self.pool_size == 0:  # if the buffer size is 0, do nothing36            return images37        return_images = []38        for image in images:39            image = torch.unsqueeze(image.data, 0)40            if self.num_imgs < self.pool_size:  # if the buffer is not full; keep inserting current images to the buffer41                self.num_imgs = self.num_imgs + 142                self.images.append(image)43                return_images.append(image)44            else:45                p = random.uniform(0, 1)46                if p > 0.5:  # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer47                    random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive48                    tmp = self.images[random_id].clone()49                    self.images[random_id] = image50                    return_images.append(tmp)51                else:  # by another 50% chance, the buffer will return the current image52                    return_images.append(image)53        return_images = torch.cat(return_images, 0)  # collect all the images and return54        return return_images55