| 3 | import torch |
| 4 | from torch.autograd import Variable |
| 5 | class ImagePool(): |
| 6 | def __init__(self, pool_size): |
| 7 | self.pool_size = pool_size |
| 8 | if self.pool_size > 0: |
| 9 | self.num_imgs = 0 |
| 10 | self.images = [] |
| 11 | |
| 12 | def query(self, images): |
| 13 | if self.pool_size == 0: |
| 14 | return images |
| 15 | return_images = [] |
| 16 | for image in images.data: |
| 17 | image = torch.unsqueeze(image, 0) |
| 18 | if self.num_imgs < self.pool_size: |
| 19 | self.num_imgs = self.num_imgs + 1 |
| 20 | 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) |
| 26 | tmp = self.images[random_id].clone() |
| 27 | self.images[random_id] = image |
| 28 | return_images.append(tmp) |
| 29 | else: |
| 30 | return_images.append(image) |
| 31 | return_images = Variable(torch.cat(return_images, 0)) |
| 32 | return return_images |
nothing calls this directly
no outgoing calls
no test coverage detected