rehistoGAN.py

"""
 If you find this code useful, please cite our paper:

 Mahmoud Afifi, Marcus A. Brubaker, and Michael S. Brown. "HistoGAN:
 Controlling Colors of GAN-Generated and Real Images via Color Histograms."
 In CVPR, 2021.

 @inproceedings{afifi2021histogan,
  title={Histo{GAN}: Controlling Colors of {GAN}-Generated and Real Images via
  Color Histograms},
  author={Afifi, Mahmoud and Brubaker, Marcus A. and Brown, Michael S.},
  booktitle={CVPR},
  year={2021}
}
"""

from torchvision.transforms.functional import to_tensor
from tqdm import tqdm
from ReHistoGAN import recoloringTrainer
from histoGAN import Trainer, NanException
from datetime import datetime
import torch
import argparse
from retry.api import retry_call
import os
from PIL import Image
from torchvision import transforms
import torchvision
import numpy as np
import copy
from utils.face_preprocessing import face_extraction
import utils.pyramid_upsampling as upsampling
from histogram_classes.RGBuvHistBlock import RGBuvHistBlock


def convert_transparent_to_rgb(image):
  if image.mode == 'RGBA':
    return image.convert('RGB')
  return image


class expand_greyscale(object):
  def __init__(self, num_channels):
    self.num_channels = num_channels

  def __call__(self, tensor):
    return tensor.expand(self.num_channels, -1, -1)


def resize_to_minimum_size(min_size, image):
  if max(*image.size) < min_size:
    return torchvision.transforms.functional.resize(image, min_size)
  return image


def hist_interpolation(hists):
  ratios = torch.abs(torch.rand(hists.shape[0])).to(
    device=torch.cuda.current_device())
  ratios = ratios / torch.sum(ratios)
  out_hist = hists[0, :, :, :, :] * ratios[0]
  for i in range(hists.shape[0] - 1):
    out_hist = out_hist + hists[i + 1, :, :, :, :] * ratios[i + 1]
  return out_hist


def process_image(model, name, input_image, target_hist, image_size=256,
                  upsampling_output=False,
                  upsampling_method='pyramid',
                  swapping_levels=1,
                  pyramid_levels=5,
                  level_blending=False,
                  post_recoloring=False,
                  sampling=True,
                  target_number=1, results_dir='./results_ReHistoGAN/',
                  hist_insz=150, hist_bin=64,
                  hist_method='inverse-quadratic', hist_resizing='sampling',
                  hist_sigma=0.02):

  img = Image.open(input_image)  # 512x1024, max 255
  original_img = np.array(img) / 255  # np array, 512x1024x3, max 1.0

  # generate reference for pyramid upsampling
  transform = transforms.Compose([transforms.ToTensor()])
  reference = torch.unsqueeze(transform(img), dim=0).to(
    device=torch.cuda.current_device())  # tensor, 1x3x512x1024, 1.0, float32

  if upsampling_output:
    width, height = img.size
    if width > image_size or height > image_size:
      resizing_mode = 'upscaling'
    elif width < image_size or height < image_size:
      resizing_mode = 'downscaling'
    else:
      resizing_mode = 'none'
  else:
    resizing_mode = None
    width = None
    height = None

  if width != image_size or height != image_size:
    img = img.resize((image_size, image_size))  # 256x256, max 255

  now = datetime.now()
  timestamp = now.strftime("%m-%d-%Y_%H-%M-%S")

  postfix = round(np.random.rand() * 1000)
  transform = transforms.Compose([
    transforms.Lambda(convert_transparent_to_rgb),
    transforms.Resize(image_size),
    transforms.ToTensor(),
    transforms.Lambda(expand_greyscale(3))
  ])

  img = torch.unsqueeze(transform(img), dim=0).to(
    device=torch.cuda.current_device())  # 1x3x256x256, max 1.0, float32

  if target_hist is None:
    if sampling:
      target_histograms = np.load('histogram_data/histograms.npy')
      target_histograms = torch.tensor(target_histograms).to(
        device=torch.cuda.current_device())

      for j in range(target_number):
        inds = np.random.randint(0, high=target_histograms.shape[0],
                                 size=5)
        h = hist_interpolation(
          target_histograms[inds, :, :, :, :])
        with torch.no_grad():
          samples_name = f'{j}-output-{timestamp}-{postfix}'
          model.evaluate(samples_name, image_batch=img,
                         hist_batch=h,
                         resizing=resizing_mode,
                         resizing_method=upsampling_method,
                         swapping_levels=swapping_levels,
                         pyramid_levels=pyramid_levels,
                         level_blending=level_blending,
                         original_size=[width, height],
                         input_image_name=input_image,
                         original_image=original_img,
                         save_input=False,
                         post_recoloring=post_recoloring)
        print(f'recolored images generated '
              f'at {results_dir}/{name}/{samples_name}')

    else:
      raise Exception('No target histogram is given.')

  else:
    if os.path.isfile(target_hist):
      files = [target_hist]
    else:
      files = [os.path.join(target_hist, f) for f in os.listdir(target_hist)
               if
               os.path.isfile(os.path.join(target_hist, f))]
    histblock = RGBuvHistBlock(insz=hist_insz, h=hist_bin,
                                resizing=hist_resizing, method=hist_method,
                                sigma=hist_sigma,
                                device=torch.cuda.current_device())
    transform = transforms.Compose([transforms.ToTensor()])
    for f in files:
      extension = os.path.splitext(f)[1]
      if extension == '.npy':
        hist = np.load(f)
        h = torch.from_numpy(hist).to(device=torch.cuda.current_device())
      elif (extension == str.lower(extension) == '.jpg' or str.lower(
          extension) == '.png'):
        img_hist = Image.open(f)  # 640x512
        img_hist = torch.unsqueeze(transform(img_hist), dim=0).to(
          device=torch.cuda.current_device())  # 1x3x512x640, 1.0
        h = histblock(img_hist)  # 1x3x64x64, 1.0
      else:
        print(f'Warning: File extension of {f} is not supported.')
        continue
      with torch.no_grad():
        samples_name = (f'{os.path.basename(os.path.splitext(input_image)[0])}-'
                        f'{os.path.basename(os.path.splitext(f)[0])}')
        generated_images = model.evaluate(samples_name, image_batch=img,
                        hist_batch=h,
                        resizing=resizing_mode,
                        resizing_method=upsampling_method,
                        swapping_levels=swapping_levels,
                        pyramid_levels=pyramid_levels,
                        level_blending=level_blending,
                        original_size=[width, height],
                        original_image=original_img,
                        input_image_name=reference,
                        save_input=False,
                        post_recoloring=post_recoloring)

      generated_images = upsampling.pyramid_upsampling(
        generated_images, reference, levels=pyramid_levels,
        swapping_levels=swapping_levels, blending=level_blending)
      torchvision.utils.save_image(generated_images, os.path.join(results_dir, name, samples_name+'.jpg'), nrow=1)
      print(f'recolored images generated at {results_dir}/{name}/'
            f'{samples_name}')

def train_from_folder(
    data='./dataset/',
    results_dir='./results_ReHistoGAN/',
    models_dir='./models/',
    histGAN_models_dir='./models/',
    name='test',
    new=False,
    load_from=-1,
    image_size=128,
    network_capacity=16,
    transparent=False,
    load_histogan_weights=True,
    batch_size=2,
    sampling=True,
    gradient_accumulate_every=8,
    num_train_steps=200000,
    learning_rate=2e-4,
    num_workers=None,
    save_every=10000,
    generate=False,
    trunc_psi=0.75,
    fp16=False,
    skip_conn_to_GAN=False,
    fq_layers=[],
    fq_dict_size=256,
    attn_layers=[],
    hist_method='inverse-quadratic',
    hist_resizing='sampling',
    hist_sigma=0.02,
    hist_bin=64,
    hist_insz=150,
    rec_loss='laplacian',
    alpha=32,
    beta=1.5,
    gamma=4,
    fixed_gan_weights=False,
    initialize_gan=False,
    variance_loss=False,
    target_hist=None,
    internal_hist=False,
    histoGAN_model_name=None,
    input_image=None,
    target_number=None,
    change_hyperparameters=False,
    change_hyperparameters_after=100000,
    upsampling_output=False,
    upsampling_method='pyramid',
    swapping_levels=1,
    pyramid_levels=6,
    level_blending=False,
    post_recoloring=False):

  model = recoloringTrainer(
    name,
    results_dir,
    models_dir,
    batch_size=batch_size,
    gradient_accumulate_every=gradient_accumulate_every,
    image_size=image_size,
    network_capacity=network_capacity,
    transparent=transparent,
    lr=learning_rate,
    num_workers=num_workers,
    save_every=save_every,
    trunc_psi=trunc_psi,
    fp16=fp16,
    fq_layers=fq_layers,
    fq_dict_size=fq_dict_size,
    attn_layers=attn_layers,
    hist_insz=hist_insz,
    hist_bin=hist_bin,
    hist_sigma=hist_sigma,
    hist_resizing=hist_resizing,
    hist_method=hist_method,
    rec_loss=rec_loss,
    fixed_gan_weights=fixed_gan_weights,
    skip_conn_to_GAN=skip_conn_to_GAN,
    initialize_gan=initialize_gan,
    variance_loss=variance_loss,
    internal_hist=internal_hist,
    change_hyperparameters=change_hyperparameters,
    change_hyperparameters_after=change_hyperparameters_after
  )

  if not new:
    status = model.load(load_from)
    if load_histogan_weights and status == -1:
      if histoGAN_model_name is not None:
        GAN_model_name = histoGAN_model_name
      else:
        GAN_model_name = name.replace('_rehistoGAN', '_histoGAN')
      if os.path.exists(os.path.join(models_dir, name)):
        model_HistGAN = Trainer(
          GAN_model_name,
          results_dir,
          histGAN_models_dir,
          batch_size=batch_size,
          gradient_accumulate_every=gradient_accumulate_every,
          image_size=image_size,
          network_capacity=network_capacity,
          transparent=transparent,
          lr=learning_rate,
          num_workers=num_workers,
          save_every=save_every,
          trunc_psi=trunc_psi,
          fp16=fp16,
          fq_layers=fq_layers,
          fq_dict_size=fq_dict_size,
          attn_layers=attn_layers,
          hist_insz=hist_insz,
          hist_bin=hist_bin,
          hist_sigma=hist_sigma,
          hist_resizing=hist_resizing,
          hist_method=hist_method,
        )
        model_HistGAN.load(load_from)
        model.GAN.G.blocks[0] = copy.deepcopy(model_HistGAN.GAN.GE.blocks[-2])
        model.GAN.G.blocks[1] = copy.deepcopy(model_HistGAN.GAN.GE.blocks[-1])
        model.GAN.H = copy.deepcopy(model_HistGAN.GAN.HE)

      else:
        raise Exception('GAN does not exist!')

  else:
    model.clear()
    if load_histogan_weights:
      if os.path.exists(os.path.join(models_dir, name)):
        if histoGAN_model_name is not None:
          GAN_model_name = histoGAN_model_name
        else:
          GAN_model_name = name.replace('_histoGAN', '_rehistoGAN')
        model_HistGAN = Trainer(
          GAN_model_name,
          results_dir,
          histGAN_models_dir,
          batch_size=batch_size,
          gradient_accumulate_every=gradient_accumulate_every,
          image_size=image_size,
          network_capacity=network_capacity,
          transparent=transparent,
          lr=learning_rate,
          num_workers=num_workers,
          save_every=save_every,
          trunc_psi=trunc_psi,
          fp16=fp16,
          fq_layers=fq_layers,
          fq_dict_size=fq_dict_size,
          attn_layers=attn_layers,
          hist_insz=hist_insz,
          hist_bin=hist_bin,
          hist_sigma=hist_sigma,
          hist_resizing=hist_resizing,
          hist_method=hist_method,
        )
        model_HistGAN.load(load_from)
        model.GAN.G.blocks[0] = copy.deepcopy(model_HistGAN.GAN.GE.blocks[-2])
        model.GAN.G.blocks[1] = copy.deepcopy(model_HistGAN.GAN.GE.blocks[-1])
        model.GAN.H = copy.deepcopy(model_HistGAN.GAN.HE)
      else:
        raise Exception('GAN does not exist!')

  if generate:
    if input_image is None:
      raise Exception('No input image is given')

    if os.path.isfile(input_image):
      files = [input_image]
    else:
      files = [os.path.join(input_image, f) for f in
                os.listdir(input_image) if os.path.isfile(os.path.join(
          input_image, f))]
    for f in files:
      extension = os.path.splitext(f)[1]
      if (extension == str.lower(extension) == '.jpg' or str.lower(
          extension) == '.png'):
        process_image(model, name, f, target_hist, image_size=256,
                      upsampling_output=upsampling_output,
                      upsampling_method=upsampling_method,
                      swapping_levels=swapping_levels,
                      pyramid_levels=pyramid_levels,
                      level_blending=level_blending,
                      post_recoloring=post_recoloring,
                      sampling=sampling,
                      target_number=target_number, results_dir=results_dir,
                      hist_insz=hist_insz, hist_bin=hist_bin,
                      hist_method=hist_method, hist_resizing=hist_resizing,
                      hist_sigma=hist_sigma)

  else:
    print('\nStart training....\n')
    print(f'Alpha = {alpha}')
    print(f'Beta = {beta}')
    print(f'Gamma = {gamma}')

    model.set_data_src(data, not fixed_gan_weights)
    for _ in tqdm(range(num_train_steps - model.steps),
                  mininterval=10., desc=f'{name}<{data}>'):
      retry_call(model.train, fargs=[alpha, beta, gamma], tries=3,
                exceptions=NanException)

      if _ % 50 == 0:
        model.print_log()


def get_args():
  parser = argparse.ArgumentParser(description='Train/Test ReHistoGAN.')
  parser.add_argument('--data', dest='data', default='./dataset/')
  parser.add_argument('--results_dir', dest='results_dir',
                      default='./results_ReHistoGAN')
  parser.add_argument('--models_dir', dest='models_dir', default='./models')
  parser.add_argument('--histGAN_models_dir', dest='histGAN_models_dir',
                      default='./models',
                      help='directory of pre-trained HistoGAN model')
  parser.add_argument('--histoGAN_model_name', dest='histoGAN_model_name',
                      default=None, type=str)
  parser.add_argument('--target_hist', dest='target_hist', default=None)
  parser.add_argument('--input_image', dest='input_image', default=None)
  parser.add_argument('--face_extraction', dest='face_extraction',
                      default=False, type=bool)
  parser.add_argument('--name', dest='name', default='reHistoGAN_model')
  parser.add_argument('--sampling', dest='sampling', default=False, type=bool,
                      help='for testing mode, if no target histogram is '
                           'given, use sampling process')
  parser.add_argument('--target_number', dest='target_number', default=50,
                      type=int,
                      help='number of recolored images; ignore if you specify '
                           'a target histogram')
  parser.add_argument('--new', dest='new', default=False)
  parser.add_argument('--load_from', dest='load_from', default=-1)
  parser.add_argument('--image_size', dest='image_size', default=256,
                      type=int)
  parser.add_argument('--network_capacity', dest='network_capacity',
                      default=16, type=int)
  parser.add_argument('--transparent', dest='transparent', default=False)
  parser.add_argument('--batch_size', dest='batch_size', default=2, type=int)
  parser.add_argument('--gradient_accumulate_every',
                      dest='gradient_accumulate_every', default=8, type=int)
  parser.add_argument('--num_train_steps', dest='num_train_steps',
                      default=200000, type=int)
  parser.add_argument('--learning_rate', dest='learning_rate', default=2e-4,
                      type=float)
  parser.add_argument('--num_workers', dest='num_workers', default=None)
  parser.add_argument('--save_every', dest='save_every', default=10000,
                      type=int)
  parser.add_argument('--trunc_psi', dest='trunc_psi', default=0.75,
                      type=float)
  parser.add_argument('--fp16', dest='fp16', default=False)
  parser.add_argument('--fq_layers', dest='fq_layers', default=[])
  parser.add_argument('--fq_dict_size', dest='fq_dict_size', default=256,
                      type=int)
  parser.add_argument('--attn_layers', dest='attn_layers', default=[])
  parser.add_argument('--gpu', dest='gpu', default=0, type=int)
  parser.add_argument('--hist_bin', dest='hist_bin', default=64, type=int)
  parser.add_argument('--hist_insz', dest='hist_insz', default=150, type=int)
  parser.add_argument('--hist_method', dest='hist_method',
                      default='inverse-quadratic')
  parser.add_argument('--hist_resizing', dest='hist_resizing',
                      default='sampling')
  parser.add_argument('--hist_sigma', dest='hist_sigma', default=0.02,
                      type=float)
  parser.add_argument('--generate', dest='generate', default=False)
  parser.add_argument('--alpha', dest='alpha', default=32, type=float)
  parser.add_argument('--beta', dest='beta', default=1.5, type=float)
  parser.add_argument('--gamma', dest='gamma', default=2, type=float)
  parser.add_argument('--change_hyperparameters',
                      dest='change_hyperparameters', default=False, type=bool)
  parser.add_argument('--change_hyperparameters_after',
                      dest='change_hyperparameters_after', default=100000,
                      type=int)
  parser.add_argument('--rec_loss', dest='rec_loss', default='laplacian',
                      type=str,
                      help='reconstruction loss (sobel or laplacian)')
  parser.add_argument('--internal_hist', dest='internal_hist', default=False,
                      type=bool, help='Internal histogram injection. This was '
                                      'an ablation on a different design; not '
                                      'what we did in the official ReHistoGAN')
  parser.add_argument('--skip_conn_to_GAN', dest='skip_conn_to_GAN',
                      default=True, type=bool,
                      help='See Figures 4 and 6 in the paper.')
  parser.add_argument('--fixed_gan_weights', dest='fixed_gan_weights',
                      default=False,
                      help="To fix weights of the HistoGAN's head")
  parser.add_argument('--load_histoGAN_weights', dest='load_histoGAN_weights',
                      default=False, help="To load weights of HistoGAN's head")
  parser.add_argument('--initialize_gan', dest='initialize_gan',
                      default=True, type=bool)
  parser.add_argument('--variance_loss', dest='variance_loss',
                      default=True, type=bool)
  parser.add_argument('--upsampling_output',
                      dest='upsampling_output',
                      default=False, type=bool,
                      help='TESTING PHASE: Applies guided upsampling. It '
                           'is recommended if input image > 256x256.')
  parser.add_argument('--upsampling_method', dest='upsampling_method',
                      default='pyramid', type=str,
                      help='TESTING PHASE: BGU or pyramid.')
  parser.add_argument('--pyramid_levels', dest='pyramid_levels',
                      default=6, type=int,
                      help='TESTING PHASE: when --upsampling_output True and '
                           '--upsampling_method is pyramid, this controls the '
                           'number of levels in the Laplacian pymraid.')
  parser.add_argument('--swapping_levels', dest='swapping_levels',
                      default=1, type=int,
                      help='TESTING PHASE: when --upsampling_output True and '
                           '--upsampling_method is pyramid, this controls the '
                           'number of levels to swap.')
  parser.add_argument('--level_blending', dest='level_blending',
                      default=False, type=bool,
                      help='TESTING PHASE: when --upsampling_output True and '
                           '--upsampling_method is pyramid, this allows to '
                           'blend between pyramid levels.')
  parser.add_argument('--post_recoloring',
                      dest='post_recoloring',
                      default=False, type=bool,
                      help='TESTING PHASE: Applies post-recoloring to '
                           'reduce artifacts. It is recommended if initial '
                           'results have some color bleeding/artifacts.')

  return parser.parse_args()


class Inferencer:
  def __init__(self, args) -> None:
    self.args = args

    self.model = recoloringTrainer(
      args.name,
      args.results_dir,
      args.models_dir,
      batch_size=args.batch_size,
      gradient_accumulate_every=args.gradient_accumulate_every,
      image_size=args.image_size,
      network_capacity=args.network_capacity,
      transparent=args.transparent,
      lr=args.learning_rate,
      num_workers=args.num_workers,
      save_every=args.save_every,
      trunc_psi=args.trunc_psi,
      fp16=args.fp16,
      fq_layers=args.fq_layers,
      fq_dict_size=args.fq_dict_size,
      attn_layers=args.attn_layers,
      hist_insz=args.hist_insz,
      hist_bin=args.hist_bin,
      hist_sigma=args.hist_sigma,
      hist_resizing=args.hist_resizing,
      hist_method=args.hist_method,
      rec_loss=args.rec_loss,
      fixed_gan_weights=args.fixed_gan_weights,
      skip_conn_to_GAN=args.skip_conn_to_GAN,
      initialize_gan=args.initialize_gan,
      variance_loss=args.variance_loss,
      internal_hist=args.internal_hist,
      change_hyperparameters=args.change_hyperparameters,
      change_hyperparameters_after=args.change_hyperparameters_after
    )

    self.device = torch.cuda.current_device()
    self.model.load(args.load_from)
    self.resizing_mode = 'upscaling'  # upscaling, downscaling, none
    self.base_transform = transforms.Compose([transforms.ToTensor()])  # used for transforming target
    self.image_size = 256  # resize to 256x256 before feeding to the model
    self.image_transform = transforms.Compose([
      # transforms.Lambda(convert_transparent_to_rgb),
      transforms.ToTensor(),
      transforms.Resize((self.image_size, self.image_size)),  # note slightly different from Image.resize
      # transforms.Lambda(expand_greyscale(3))
    ])
    self.histblock = RGBuvHistBlock(insz=args.hist_insz, h=args.hist_bin,
                                    resizing=args.hist_resizing, method=args.hist_method,
                                    sigma=args.hist_sigma,
                                    device=self.device)

  def process_image(self, image, style):
    # process image
    height, width, _ = image.shape
    reference = image.copy()
    image = torch.unsqueeze(self.image_transform(image), dim=0).to(device=self.device)
    # process hist
    style = torch.unsqueeze(self.base_transform(style), dim=0).to(device=self.device)
    h = self.histblock(style)

    # generate result
    result = self.model.evaluate('', image_batch=image,
                        hist_batch=h,
                        resizing=self.resizing_mode,
                        resizing_method=self.args.upsampling_method,
                        swapping_levels=self.args.swapping_levels,
                        pyramid_levels=self.args.pyramid_levels,
                        level_blending=self.args.level_blending,
                        original_size=[width, height],
                        original_image=reference,  # not used
                        input_image_name=reference,
                        save_input=False,
                        post_recoloring=self.args.post_recoloring)
    # upsample image
    result = upsampling.pyramid_upsampling(
      result, reference, levels=self.args.pyramid_levels,
      swapping_levels=self.args.swapping_levels, blending=self.args.level_blending, to_tensor=False)
    return result


if __name__ == "__main__":
  args = get_args()
  torch.cuda.set_device(args.gpu)

  if args.generate and args.face_extraction:
    if args.input_image is None:
      raise Exception('No input image is given')
    extension = os.path.splitext(args.input_image)[1]
    if (extension == str.lower(extension) == '.jpg' or str.lower(
      extension) == '.png'):
      face_extraction(args.input_image)
      input_image = f'./temp-faces/{os.path.split(args.input_image)[-1]}'
    elif extension == '':

      files = [os.path.join('./temp-faces/', f) for f in
               os.listdir('./temp-faces/') if os.path.isfile(os.path.join(
          './temp-faces/', f))]
      for f in files:
        os.remove(f)

      files = [os.path.join(args.input_image, f) for f in
               os.listdir(args.input_image) if os.path.isfile(os.path.join(
          args.input_image, f))]
      for f in files:
        extension = os.path.splitext(f)[1]
        if (extension == str.lower(extension) == '.jpg' or str.lower(
            extension) == '.png'):
          face_extraction(f)
      input_image = './temp-faces/'
    else:
      raise Exception('File extension is not supported!')
  else:
    input_image = args.input_image

  # print(args)

  inferencer = Inferencer(args)
  image = (np.array(Image.open(args.input_image)) / 255.).astype(np.float32)
  style = (np.array(Image.open(args.target_hist)) / 255.).astype(np.float32)
  inferencer.process_image(image, style)

  train_from_folder(
    data=args.data,
    results_dir=args.results_dir,
    models_dir=args.models_dir,
    name=args.name,
    new=args.new,
    histGAN_models_dir=args.histGAN_models_dir,
    load_from=args.load_from,
    load_histogan_weights=args.load_histoGAN_weights,
    image_size=args.image_size,
    network_capacity=args.network_capacity,
    transparent=args.transparent,
    batch_size=args.batch_size,
    gradient_accumulate_every=args.gradient_accumulate_every,
    num_train_steps=args.num_train_steps,
    learning_rate=args.learning_rate,
    num_workers=args.num_workers,
    save_every=args.save_every,
    generate=args.generate,
    trunc_psi=args.trunc_psi,
    fp16=args.fp16,
    fq_layers=args.fq_layers,
    fq_dict_size=args.fq_dict_size,
    attn_layers=args.attn_layers,
    hist_method=args.hist_method,
    hist_resizing=args.hist_resizing,
    hist_sigma=args.hist_sigma,
    hist_bin=args.hist_bin,
    hist_insz=args.hist_insz,
    target_hist=args.target_hist,
    alpha=args.alpha,
    beta=args.beta,
    gamma=args.gamma,
    skip_conn_to_GAN=args.skip_conn_to_GAN,
    fixed_gan_weights=args.fixed_gan_weights,
    sampling=args.sampling,
    rec_loss=args.rec_loss,
    initialize_gan=args.initialize_gan,
    variance_loss=args.variance_loss,
    input_image=input_image,
    internal_hist=args.internal_hist,
    histoGAN_model_name=args.histoGAN_model_name,
    target_number=args.target_number,
    change_hyperparameters=args.change_hyperparameters,
    change_hyperparameters_after=args.change_hyperparameters_after,
    upsampling_output=args.upsampling_output,
    upsampling_method=args.upsampling_method,
    swapping_levels=args.swapping_levels,
    pyramid_levels=args.pyramid_levels,
    level_blending=args.level_blending,
    post_recoloring=args.post_recoloring
  )