<reponame>chenhsiu48/HiDDeN
import os
import time


class TrainingOptions:
    """
    Configuration options for the training
    """

    def __init__(self, batch_size: int, number_of_epochs: int, data_ratio: float,
                 data_dir: str, runs_folder: str, tb_logger_folder: str,
                 start_epoch: int, experiment_name: str, best_epoch=-1, best_cond=None):
        self.batch_size = batch_size
        self.number_of_epochs = number_of_epochs
        self.data_dir = data_dir
        self.runs_folder = runs_folder
        self.tb_logger_folder = tb_logger_folder
        self.start_epoch = start_epoch
        self.experiment_name = experiment_name
        self.best_epoch = best_epoch
        self.best_cond = best_cond
        self.data_ratio = data_ratio

        self.run_label = f'{experiment_name}_{time.strftime("%Y%m%d-%H%M%S")}'
        self.this_run_folder = os.path.join(self.runs_folder, self.run_label)
        self.checkpoint_dir = os.path.join(self.this_run_folder, 'checkpoints')
        self.image_dir = os.path.join(self.this_run_folder, 'images')
        self.tb_logger_dir = os.path.join(self.tb_logger_folder, self.run_label)

        self.options_file = os.path.join(self.this_run_folder, 'options-and-config.pickle')


class HiDDenConfiguration():
    """
    The HiDDeN network configuration.
    """

    def __init__(self, H: int, W: int, input_channels: int, message_length: int,
                 encoder_blocks: int, encoder_channels: int,
                 decoder_blocks: int, decoder_channels: int,
                 use_discriminator: bool,
                 use_vgg: bool,
                 discriminator_blocks: int, discriminator_channels: int,
                 decoder_loss: float,
                 encoder_loss: float,
                 adversarial_loss: float,
                 enable_fp16: bool = False):
        self.H = H
        self.W = W
        self.input_channels = input_channels
        self.message_length = message_length
        self.encoder_blocks = encoder_blocks
        self.encoder_channels = encoder_channels
        self.use_discriminator = use_discriminator
        self.use_vgg = use_vgg
        self.decoder_blocks = decoder_blocks
        self.decoder_channels = decoder_channels
        self.discriminator_blocks = discriminator_blocks
        self.discriminator_channels = discriminator_channels
        self.decoder_loss = decoder_loss
        self.encoder_loss = encoder_loss
        self.adversarial_loss = adversarial_loss
        self.enable_fp16 = enable_fp16

import os
import time


class TrainingOptions:
    """
    Configuration options for the training
    """

    def __init__(self, batch_size: int, number_of_epochs: int, data_ratio: float,
                 data_dir: str, runs_folder: str, tb_logger_folder: str,
                 start_epoch: int, experiment_name: str, best_epoch=-1, best_cond=None):
        self.batch_size = batch_size
        self.number_of_epochs = number_of_epochs
        self.data_dir = data_dir
        self.runs_folder = runs_folder
        self.tb_logger_folder = tb_logger_folder
        self.start_epoch = start_epoch
        self.experiment_name = experiment_name
        self.best_epoch = best_epoch
        self.best_cond = best_cond
        self.data_ratio = data_ratio

        self.run_label = f'{experiment_name}_{time.strftime("%Y%m%d-%H%M%S")}'
        self.this_run_folder = os.path.join(self.runs_folder, self.run_label)
        self.checkpoint_dir = os.path.join(self.this_run_folder, 'checkpoints')
        self.image_dir = os.path.join(self.this_run_folder, 'images')
        self.tb_logger_dir = os.path.join(self.tb_logger_folder, self.run_label)

        self.options_file = os.path.join(self.this_run_folder, 'options-and-config.pickle')


class HiDDenConfiguration():
    """
    The HiDDeN network configuration.
    """

    def __init__(self, H: int, W: int, input_channels: int, message_length: int,
                 encoder_blocks: int, encoder_channels: int,
                 decoder_blocks: int, decoder_channels: int,
                 use_discriminator: bool,
                 use_vgg: bool,
                 discriminator_blocks: int, discriminator_channels: int,
                 decoder_loss: float,
                 encoder_loss: float,
                 adversarial_loss: float,
                 enable_fp16: bool = False):
        self.H = H
        self.W = W
        self.input_channels = input_channels
        self.message_length = message_length
        self.encoder_blocks = encoder_blocks
        self.encoder_channels = encoder_channels
        self.use_discriminator = use_discriminator
        self.use_vgg = use_vgg
        self.decoder_blocks = decoder_blocks
        self.decoder_channels = decoder_channels
        self.discriminator_blocks = discriminator_blocks
        self.discriminator_channels = discriminator_channels
        self.decoder_loss = decoder_loss
        self.encoder_loss = encoder_loss
        self.adversarial_loss = adversarial_loss
        self.enable_fp16 = enable_fp16
Configuration options for the training The HiDDeN network configuration.