mirror of https://github.com/kritiksoman/GIMP-ML
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
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
|