dataset_params:
  root_dir: /zlh/VoxCeleb/first-order-256
  frame_shape: [256, 256, 3]
  id_sampling: True
  pairs_list: None
  augmentation_params:
    flip_param:
      horizontal_flip: True
      time_flip: True
    jitter_param:
      brightness: 0.1
      contrast: 0.1
      saturation: 0.1
      hue: 0.1

model_params:
  common_params:
    num_kp: 15 
    image_channel: 3                    
    feature_channel: 32
    estimate_jacobian: False   # True
  kp_detector_params:
     temperature: 0.1
     block_expansion: 32            
     max_features: 1024
     scale_factor: 0.25         # 0.25
     num_blocks: 5
     reshape_channel: 16384  # 16384 = 1024 * 16
     reshape_depth: 16
  he_estimator_params:
     block_expansion: 64            
     max_features: 2048
     num_bins: 66
  generator_params:
    block_expansion: 64
    max_features: 512
    num_down_blocks: 2
    reshape_channel: 32
    reshape_depth: 16         # 512 = 32 * 16
    num_resblocks: 6
    estimate_occlusion_map: True
    dense_motion_params:
      block_expansion: 32
      max_features: 1024
      num_blocks: 5
      # reshape_channel: 32
      reshape_depth: 16
      compress: 4
  discriminator_params:
    scales: [1]
    block_expansion: 32                 
    max_features: 512
    num_blocks: 4
    sn: True

train_params:
  num_epochs: 300
  num_repeats: 75
  epoch_milestones: [180,]
  lr_generator: 2.0e-4
  lr_discriminator: 2.0e-4
  lr_kp_detector: 2.0e-4
  lr_he_estimator: 2.0e-4
  gan_mode: 'hinge'    # hinge or ls
  batch_size: 32
  scales: [1, 0.5, 0.25, 0.125]
  checkpoint_freq: 10
  hopenet_snapshot:  "/mnt/bn/sa-ag-data/yezhenhui/myenv/cache/useful_ckpts/hopenet_robust_alpha1.pkl" # https://drive.google.com/open?id=1m25PrSE7g9D2q2XJVMR6IA7RaCvWSzCR
  transform_params:
    sigma_affine: 0.05
    sigma_tps: 0.005
    points_tps: 5
  loss_weights:
    generator_gan: 1                  
    discriminator_gan: 1
    feature_matching: [10, 10, 10, 10]
    perceptual: [10, 10, 10, 10, 10]
    equivariance_value: 10
    equivariance_jacobian: 0    # 10
    keypoint: 10
    headpose: 20
    expression: 5

visualizer_params:
  kp_size: 5
  draw_border: True
  colormap: 'gist_rainbow'