12345678910111213141516171819202122232425262728293031 |
- 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
|