# Lunar Crater Detection (NAC only)
# Backbone: NASA-IBM Lunar Foundation Model at patch size 8 (pretrained weights)
#
# RUN:
#   PYTHONPATH=.:$PYTHONPATH \
#   terratorch fit --config terratorch_integration/configs/nac_craters/ni_lfm_ps8.yaml

seed_everything: 44
custom_modules_path: terratorch_integration

trainer:
  accelerator: gpu
  devices: 1
  precision: "32"
  max_epochs: 100
  check_val_every_n_epoch: 1
  log_every_n_steps: 10
  deterministic: true
  benchmark: false
  enable_checkpointing: true
  logger:
    class_path: TensorBoardLogger
    init_args:
      save_dir: tb_logs/nac-crater
      name: ni_lfm_ps8
  callbacks:
    - class_path: LearningRateMonitor
      init_args:
        logging_interval: epoch
    - class_path: terratorch_integration.DeterministicAugmentation
    - class_path: ModelCheckpoint
      init_args:
        dirpath: tb_logs/nac-crater/checkpoints/ni_lfm_ps8
        monitor: val_map
        mode: max
        save_top_k: 1
        filename: best-{epoch:02d}-{val_map:.3f}
        save_last: false

data:
  class_path: terratorch_integration.data_adapter.LunarNACDTMDataModule
  init_args:
    data_dir: data/nac_craters_dataset
    metadata_file: null   # tile paths and splits come from the COCO file below
    annotations_file: data/nac_craters_dataset/annotations_min5px.json
    modalities: [nac]
    metadata_tokenizer_path: ni_lfm/utils/tokenizer/trained/text_tokenizer.json
    output_mode: packed
    batch_size: 4
    num_workers: 4
    image_size: 256
    min_diameter_m: -1    # -1 disables this filter
    drop_no_crater_images: true
    used_for_labeling_only: true
    nac_norm_mean: 0.346548   # train-split stats (289 tiles, DN/255)
    nac_norm_std: 0.159291
    train_transforms:
      class_path: terratorch_integration.data_utils.D4DetectionTransform
      init_args:
        image_keys: [image]
        p: 1.0

model:
  class_path: terratorch_integration.LunarObjectDetectionTask
  init_args:
    model_factory: ObjectDetectionModelFactory
    model_args:
      framework: faster-rcnn
      backbone: ni_lfm_v1_base
      backbone_checkpoint_path: backbone/checkpoint.pt
      backbone_cfg: backbone/config.yaml
      num_classes: 2
      framework_min_size: 256
      framework_max_size: 256

      backbone_modalities:
        - nac

      backbone_patch_size: 8
      backbone_remove_register_tokens: false
      backbone_merge_method: concat

      backbone_drop_path_rate_encoder: 0.1

      framework_rpn_fg_iou_thresh: 0.5
      framework_rpn_bg_iou_thresh: 0.2
      framework_box_fg_iou_thresh: 0.5
      framework_box_bg_iou_thresh: 0.3
      framework_rpn_batch_size_per_image: 512
      framework_box_batch_size_per_image: 1024
      framework_box_positive_fraction: 0.5
      framework_box_detections_per_img: 600
      framework_box_nms_thresh: 0.5
      framework_rpn_post_nms_top_n_train: 3000
      framework_rpn_post_nms_top_n_test: 3000
      framework_rpn_pre_nms_top_n_train: 6000
      framework_rpn_pre_nms_top_n_test: 6000
      in_channels: 1
      necks:
        - name: SelectIndices
          indices: [2, 5, 8, 11]
        - name: ReshapeTokensToImage
          remove_cls_token: false
          h: 32
        - name: MultilayerSimpleFeaturePyramid
          out_channels: 256
        - name: FeaturePyramidNetworkNeck

    freeze_backbone: false
    freeze_decoder: false

    class_names:
      - Background
      - Crater

    score_threshold: 0.05
    iou_threshold: 0.3
    metric_kwargs:
      max_detection_thresholds: [100, 300, 600]
      class_metrics: true

    anchor_sizes: [[4, 6], [8, 12], [16, 24], [48, 96]]
    anchor_aspect_ratios: [0.7, 1.0, 1.4]

    eval_box_detections_per_img: 600
    eval_box_nms_thresh: 0.3
    eval_rpn_post_nms_top_n_test: 3000

    backbone_lr: 5.0e-5
    head_lr: 5.0e-5
    layer_decay: 0.65
    weight_decay: 0.1
    head_weight_decay: 1.0e-3
    betas: [0.9, 0.999]
    warmup_steps: 500
    eta_min: 1.0e-6