| import random |
| import torch |
| from torch.autograd import Variable |
| class ImagePool(): |
| def __init__(self, pool_size): |
| self.pool_size = pool_size |
| if self.pool_size > 0: |
| self.num_imgs = 0 |
| self.images = [] |
|
|
| def query(self, images): |
| if self.pool_size == 0: |
| return images |
| return_images = [] |
| for image in images.data: |
| image = torch.unsqueeze(image, 0) |
| if self.num_imgs < self.pool_size: |
| self.num_imgs = self.num_imgs + 1 |
| self.images.append(image) |
| return_images.append(image) |
| else: |
| p = random.uniform(0, 1) |
| if p > 0.5: |
| random_id = random.randint(0, self.pool_size-1) |
| tmp = self.images[random_id].clone() |
| self.images[random_id] = image |
| return_images.append(tmp) |
| else: |
| return_images.append(image) |
| return_images = Variable(torch.cat(return_images, 0)) |
| return return_images |
|
|
