| 9 | self.images = [] |
| 10 | |
| 11 | def query(self, images): |
| 12 | if self.pool_size == 0: |
| 13 | return images |
| 14 | return_images = [] |
| 15 | for image in images.data: |
| 16 | image = torch.unsqueeze(image, 0) |
| 17 | if self.num_imgs < self.pool_size: |
| 18 | self.num_imgs = self.num_imgs + 1 |
| 19 | self.images.append(image) |
| 20 | return_images.append(image) |
| 21 | else: |
| 22 | p = random.uniform(0, 1) |
| 23 | if p > 0.5: |
| 24 | random_id = random.randint(0, self.pool_size-1) |
| 25 | tmp = self.images[random_id].clone() |
| 26 | self.images[random_id] = image |
| 27 | return_images.append(tmp) |
| 28 | else: |
| 29 | return_images.append(image) |
| 30 | return_images = Variable(torch.cat(return_images, 0)) |
| 31 | return return_images |