GoodWin/Deep-Multi-scale

1import random2import torch3 4 5class ImagePool():6    def __init__(self, pool_size):7        self.pool_size = pool_size8        if self.pool_size > 0:9            self.num_imgs = 010            self.images = []11 12    def query(self, images):13        if self.pool_size == 0:14            return images15        return_images = []16        for image in images:17            image = torch.unsqueeze(image.data, 0)18            if self.num_imgs < self.pool_size:19                self.num_imgs = self.num_imgs + 120                self.images.append(image)21                return_images.append(image)22            else:23                p = random.uniform(0, 1)24                if p > 0.5:25                    random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive26                    tmp = self.images[random_id].clone()27                    self.images[random_id] = image28                    return_images.append(tmp)29                else:30                    return_images.append(image)31        return_images = torch.cat(return_images, 0)32        return return_images33