# DeepMAC meta architecture from the "The surprising impact of mask-head
# architecture on novel class segmentation" [1] paper with an Hourglass-52[2]
# mask head. This config is only trained on masks from the VOC classes in COCO
# and achieves a mask mAP of 32.5% on non-VOC classes.
# [1]: https://arxiv.org/abs/2104.00613
# [2]: https://arxiv.org/abs/1904.07850

# Train on TPU-32

model {
  center_net {
    num_classes: 90
    feature_extractor {
      type: "hourglass_104"
      bgr_ordering: true
      channel_means: [104.01362025, 114.03422265, 119.9165958 ]
      channel_stds: [73.6027665 , 69.89082075, 70.9150767 ]
    }
    image_resizer {
      keep_aspect_ratio_resizer {
        min_dimension: 512
        max_dimension: 512
        pad_to_max_dimension: true
      }
    }
    object_detection_task {
      task_loss_weight: 1.0
      offset_loss_weight: 1.0
      scale_loss_weight: 0.1
      localization_loss {
        l1_localization_loss {
        }
      }
    }
    object_center_params {
      object_center_loss_weight: 1.0
      min_box_overlap_iou: 0.7
      max_box_predictions: 100
      classification_loss {
        penalty_reduced_logistic_focal_loss {
          alpha: 2.0
          beta: 4.0
        }
      }
    }

    deepmac_mask_estimation {
      dim: 32
      task_loss_weight: 5.0
      pixel_embedding_dim: 16
      mask_size: 32
      use_xy: true
      use_instance_embedding: true
      network_type: "hourglass52"
      classification_loss {
        weighted_sigmoid {}
      }

      allowed_masked_classes_ids: [
        1,   # person
        2,   # bicycle
        3,   # car
        4,   # motorcycle/motorbike
        5,   # airplane/aeroplane,
        6,   # bus
        7,   # train
        9,   # boat
        16,  # bird
        17,  # cat
        18,  # dog
        19,  # horse
        20,  # sheep
        21,  # cow
        44,  # bottle
        62,  # chair
        63,  # couch/sofa
        64,  # potted plant
        67,  # dining table
        72  # tvmonitor
      ]
    }
  }
}

train_config: {

  batch_size: 128
  num_steps: 50000

  data_augmentation_options {
    random_horizontal_flip {
    }
  }

  data_augmentation_options {
    random_adjust_hue {
    }
  }

  data_augmentation_options {
    random_adjust_contrast {
    }
  }

  data_augmentation_options {
    random_adjust_saturation {
    }
  }

  data_augmentation_options {
    random_adjust_brightness {
    }
  }

   data_augmentation_options {
     random_square_crop_by_scale {
      scale_min: 0.6
      scale_max: 1.3
    }
  }

  optimizer {
    adam_optimizer: {
      epsilon: 1e-7  # Match tf.keras.optimizers.Adam's default.
      learning_rate: {
        cosine_decay_learning_rate {
          learning_rate_base: 1e-3
          total_steps: 50000
          warmup_learning_rate: 2.5e-4
          warmup_steps: 5000
        }
      }
    }
    use_moving_average: false
  }
  max_number_of_boxes: 100
  unpad_groundtruth_tensors: false

  fine_tune_checkpoint_version: V2
  fine_tune_checkpoint: "PATH_TO_BE_CONFIGURED/ckpt-1"
  fine_tune_checkpoint_type: "detection"
}

train_input_reader: {
  load_instance_masks: true
  label_map_path: "PATH_TO_BE_CONFIGURED/label_map.txt"
  mask_type: PNG_MASKS
  tf_record_input_reader {
    input_path: "PATH_TO_BE_CONFIGURED/train2017-?????-of-00256.tfrecord"
  }
}

eval_config: {
  metrics_set: "coco_detection_metrics"
  metrics_set: "coco_mask_metrics"
  include_metrics_per_category: true
  use_moving_averages: false
  batch_size: 1;
  super_categories {
    key: "VOC"
    value: "person,bicycle,car,motorcycle,airplane,bus,train,boat,bird,cat,"
           "dog,horse,sheep,cow,bottle,chair,couch,potted plant,dining table,tv"
  }
  super_categories {
    key: "NonVOC"
    value: "truck,traffic light,fire hydrant,stop sign,parking meter,bench,"
    "elephant,bear,zebra,giraffe,backpack,umbrella,handbag,tie,suitcase,"
    "frisbee,skis,snowboard,sports ball,kite,baseball bat,baseball glove,"
    "skateboard,surfboard,tennis racket,wine glass,cup,fork,knife,spoon,bowl,"
    "banana,apple,sandwich,orange,broccoli,carrot,hot dog,pizza,donut,cake,bed,"
    "toilet,laptop,mouse,remote,keyboard,cell phone,microwave,oven,toaster,"
    "sink,refrigerator,book,clock,vase,scissors,teddy bear,hair drier,"
    "toothbrush"
  }
  super_categories {
    key: "person"
    value: "person"
  }
  super_categories {
    key: "vehicle"
    value: "bicycle,car,motorcycle,airplane,bus,train,truck,boat"
  }
  super_categories {
    key: "outdoor"
    value: "traffic light,fire hydrant,stop sign,parking meter,bench"
  }
  super_categories {
    key: "animal"
    value: "bird,cat,dog,horse,sheep,cow,elephant,bear,zebra,giraffe"
  }
  super_categories {
    key: "accessory"
    value: "backpack,umbrella,handbag,tie,suitcase"
  }
  super_categories {
    key: "sports"
    value: "frisbee,skis,snowboard,sports ball,kite,baseball bat,"
           "baseball glove,skateboard,surfboard,tennis racket"
  }
  super_categories {
    key: "kitchen"
    value: "bottle,wine glass,cup,fork,knife,spoon,bowl"
  }
  super_categories {
    key: "food"
    value: "banana,apple,sandwich,orange,broccoli,carrot,hot dog,pizza,donut,"
           "cake"
  }
  super_categories {
    key: "furniture"
    value: "chair,couch,potted plant,bed,dining table,toilet"
  }
  super_categories {
    key: "electronic"
    value: "tv,laptop,mouse,remote,keyboard,cell phone,microwave,oven,toaster,"
           "sink,refrigerator"
  }
  super_categories {
    key: "indoor"
    value: "book,clock,vase,scissors,teddy bear,hair drier,toothbrush"
  }
}

eval_input_reader: {
  load_instance_masks: true
  mask_type: PNG_MASKS
  label_map_path: "PATH_TO_BE_CONFIGURED/label_map.txt"
  shuffle: false
  num_epochs: 1
  tf_record_input_reader {
    input_path: "PATH_TO_BE_CONFIGURED/val2017-?????-of-00032.tfrecord"
  }
}
