text: import random
  import torch
  
  
  class ImagePool():
      """This class implements an image buffer that stores previously generated images.
  
      This buffer enables us to update discriminators using a history of generated images
      rather than the ones produced by the latest generators.
      """
  
      def __init__(self, pool_size):
          """Initialize the ImagePool class
  
          Parameters:
              pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created
          """
          self.pool_size = pool_size
          if self.pool_size > 0:  # create an empty pool
              self.num_imgs = 0
              self.images = []
  
      def query(self, images):
          """Return an image from the pool.
  
          Parameters:
              images: the latest generated images from the generator
  
          Returns images from the buffer.
  
          By 50/100, the buffer will return input images.
          By 50/100, the buffer will return images previously stored in the buffer,
          and insert the current images to the buffer.
          """
          if self.pool_size == 0:  # if the buffer size is 0, do nothing
              return images
          return_images = []
          for image in images:
              image = torch.unsqueeze(image.data, 0)
              if self.num_imgs < self.pool_size:   # if the buffer is not full; keep inserting current images to the buffer
                  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:  # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer
                      random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive
                      tmp = self.images[random_id].clone()
                      self.images[random_id] = image
                      return_images.append(tmp)
                  else:       # by another 50% chance, the buffer will return the current image
                      return_images.append(image)
          return_images = torch.cat(return_images, 0)   # collect all the images and return
          return return_images
  
  
  class MyImagePoolWithSegmentation():
      """This class implements an image buffer that stores previously generated images.
  
      This buffer enables us to update discriminators using a history of generated images
      rather than the ones produced by the latest generators.
      """
  
      def __init__(self, pool_size):
          """Initialize the ImagePool class
  
          Parameters:
              pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created
          """
          self.pool_size = pool_size
          if self.pool_size > 0:  # create an empty pool
              self.num_imgs = 0
              self.images = []
              self.images_segment = []
  
  
      def query(self, images, images_segment):
          """Return an image from the pool.
  
          Parameters:
              images: the latest generated images from the generator
  
          Returns images from the buffer.
  
          By 50/100, the buffer will return input images.
          By 50/100, the buffer will return images previously stored in the buffer,
          and insert the current images to the buffer.
          """
          if self.pool_size == 0:  # if the buffer size is 0, do nothing
              return images, images_segment
          return_images = []
          return_images_segment = []
          for i in range(len(images)):
              image = images[i]
              image_segment = images_segment[i]
              image = torch.unsqueeze(image.data, 0)
              image_segment = torch.unsqueeze(image_segment.data, 0)
              if self.num_imgs < self.pool_size:   # if the buffer is not full; keep inserting current images to the buffer
                  self.num_imgs = self.num_imgs + 1
                  self.images.append(image)
                  return_images.append(image)
                  self.images_segment.append(image_segment)
                  return_images_segment.append(image_segment)
              else:
                  p = random.uniform(0, 1)
                  if p > 0.5:  # by 50% chance, the buffer will return a previously stored image, and insert the current image into the buffer
                      random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive
                      tmp = self.images[random_id].clone()
                      self.images[random_id] = image
                      return_images.append(tmp)
                      
                      tmp_segment = self.images_segment[random_id].clone()
                      self.images_segment[random_id] = image_segment
                      return_images_segment.append(tmp_segment)
                  else:       # by another 50% chance, the buffer will return the current image
                      return_images.append(image)
                      return_images_segment.append(image_segment)
          return_images = torch.cat(return_images, 0)   # collect all the images and return
          return_images_segment = torch.cat(return_images_segment, 0)   # collect all the images and return
          return return_images, return_images_segment