image_pool.py 1.1 KB

12345678910111213141516171819202122232425262728293031
  1. import random
  2. import torch
  3. from torch.autograd import Variable
  4. class ImagePool():
  5. def __init__(self, pool_size):
  6. self.pool_size = pool_size
  7. if self.pool_size > 0:
  8. self.num_imgs = 0
  9. self.images = []
  10. def query(self, images):
  11. if self.pool_size == 0:
  12. return images
  13. return_images = []
  14. for image in images.data:
  15. image = torch.unsqueeze(image, 0)
  16. if self.num_imgs < self.pool_size:
  17. self.num_imgs = self.num_imgs + 1
  18. self.images.append(image)
  19. return_images.append(image)
  20. else:
  21. p = random.uniform(0, 1)
  22. if p > 0.5:
  23. random_id = random.randint(0, self.pool_size-1)
  24. tmp = self.images[random_id].clone()
  25. self.images[random_id] = image
  26. return_images.append(tmp)
  27. else:
  28. return_images.append(image)
  29. return_images = Variable(torch.cat(return_images, 0))
  30. return return_images