import argparse
clamp = 2.0
channels_in = 3
log10_lr = -5.2 
lr = 10 ** log10_lr
device_ids = 0

lr2 = 0.0001  #0.0001

epochs = 40  #1000  #200  
weight_decay = 1e-5
init_scale = 0.01

lamda_r =1.0      
lamda_c =1.0        
lamda_l = 1.0  


# Train:
batch_size = 16
batch_size_2 = 1
cropsize = 256
betas = (0.5, 0.999)
weight_step = 1000
gamma = 0.5

# Val:
cropsize_val = 256
batchsize_val = 2
shuffle_val = False
val_freq = 50  

# Display and logging:
loss_display_cutoff = 2.0
loss_names = ['L', 'lr']
silent = False
live_visualization = False
progress_bar = False

# Dateset
# COVER_PATH1 Images for training the INN in AHDeS.
COVER_PATH1 ='image/dataset_coco/train/'

# COVER_PATH2 Images for evuluating the security performance of AHDeS. 
COVER_PATH2 = 'image/dataset_coco/test/'

# Replace the following path with your own dataset.
TRAIN_PATH ='image/dataset_coco/train/'

TEST_PATH = 'image/dataset_coco/test/'

# Saving checkpoints:

MODEL_PATH_1 = 'checkpoint/INN_first_phase_checkpoints/'
MODEL_PATH_2 = 'checkpoint/INN_second_phase_checkpoints/'
MODEL_PATH_3 = 'checkpoint/'
checkpoint_on_error = True
SAVE_freq = 50

#Saving image:
STEGO_PATH2 = 'image/image_deployment/'

# Load:
#train 1
suffix = 'train_first_phase_model.pt'
tain_next = False
trained_epoch = 0

#train 2
tain_next_2 = True
train_first_phase = 'checkpoint/INN_first_phase_checkpoints/train_first_phase_model.pt'

#If you used your own data set, then this CovNet should also be retrained. 
CovNet_model = 'saved_model/steganalyzer.pt'

#deployment
tain_next_3 = True

#If you want to use our trained model directly, you can use the current path directly.
train_second_phase='saved_model/INN_model.pt'

