Camouflage detection improvements

.gitignore

@@ -4,3 +4,24 @@ weight/
 wandb/
 *.pth
 *.onnx
+*.mp4
+datasets/
+*.db
+outputs/
+*.engine
+venv/
+
+# Editor
+.vscode/
+.history/
+.codex
+
+# Large binary / data folders
+Traning/
+segmentation_sivert/data/
+
+# Cache
+__pycache__/
+*.egg-info/
+.pytest_cache/
+DETRPose/

configs/dfine_hgnetv2_x_obj2coco.yml

@@ -0,0 +1,61 @@
+__include__: [
+  './dataset/coco_detection.yml',
+  './runtime.yml',
+  './dfine/include/dataloader.yml',
+  './dfine/include/optimizer.yml',
+  './dfine/include/dfine_hgnetv2.yml',
+]
+
+output_dir: ./output/dfine_hgnetv2_x_obj2coco
+
+
+DFINE:
+  backbone: HGNetv2
+
+HGNetv2:
+  name: 'B5'
+  return_idx: [1, 2, 3]
+  freeze_stem_only: True
+  freeze_at: 0
+  freeze_norm: True
+
+HybridEncoder:
+  # intra
+  hidden_dim: 384
+  dim_feedforward: 2048
+
+DFINETransformer:
+  feat_channels: [384, 384, 384]
+  reg_scale: 8
+
+optimizer:
+  type: AdamW
+  params:
+    -
+      params: '^(?=.*backbone)(?!.*norm|bn).*$'
+      lr: 0.0000025
+    -
+      params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
+      weight_decay: 0.
+
+  lr: 0.00025
+  betas: [0.9, 0.999]
+  weight_decay: 0.000125
+
+
+epochs: 36 # Early stop
+train_dataloader:
+  dataset:
+    transforms:
+      policy:
+        epoch: 30
+  collate_fn:
+    stop_epoch: 30
+    ema_restart_decay: 0.9999
+    base_size_repeat: 3
+
+ema:
+  warmups: 0
+
+lr_warmup_scheduler:
+  warmup_duration: 0

detection_training/configs/soldier_finetune.yml

@@ -0,0 +1,62 @@
+__include__: [
+  '../../configs/dataset/coco_detection.yml',
+  '../../configs/runtime.yml',
+  '../../configs/dfine/include/dataloader.yml',
+  '../../configs/dfine/include/optimizer.yml',
+  '../../configs/dfine/include/dfine_hgnetv2.yml',
+]
+
+output_dir: ./outputs/soldier_finetune
+
+num_classes: 2
+remap_mscoco_category: False
+
+HGNetv2:
+  name: 'B5'
+  return_idx: [1, 2, 3]
+  freeze_stem_only: True
+  freeze_at: 0
+  freeze_norm: True
+
+HybridEncoder:
+  hidden_dim: 384
+  dim_feedforward: 2048
+
+DFINETransformer:
+  feat_channels: [384, 384, 384]
+  reg_scale: 8
+  num_classes: 2
+
+epochs: 60
+lr_warmup_scheduler:
+  warmup_duration: 100
+
+optimizer:
+  type: AdamW
+  params:
+    - params: '^(?=.*backbone)(?!.*norm|bn).*$'
+      lr: 0.000005
+    - params: '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$'
+      weight_decay: 0.
+  lr: 0.00005
+  betas: [0.9, 0.999]
+  weight_decay: 0.0001
+
+train_dataloader:
+  dataset:
+    img_folder: ./datasets/soldier_coco/images/train
+    ann_file: ./datasets/soldier_coco/annotations/instances_train.json
+  total_batch_size: 2
+  collate_fn:
+    stop_epoch: 50
+  shuffle: True
+
+val_dataloader:
+  dataset:
+    img_folder: ./datasets/soldier_coco/images/val
+    ann_file: ./datasets/soldier_coco/annotations/instances_val.json
+  total_batch_size: 2
+
+evaluator:
+  type: CocoEvaluator
+  iou_types: ['bbox']

detection_training/merge_checkpoints.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+"""
+Model Surgery: merge new soldier-finetuned backbone into original combined checkpoint.
+
+Strategy: BACKBONE-ONLY transfer from the new soldier-finetuned model.
+  - backbone   → from new fine-tune (learned to find camouflaged soldiers)
+  - encoder    → from original (keeps encoder features compatible with pose decoder)
+  - det_decoder → from original (80-class COCO, detects 'person')
+  - pose_decoder → from original (unchanged, compatible with original encoder)
+  - seg_head   → from original (unchanged)
+
+Why backbone-only: the pose decoder was trained on the original encoder's feature
+representations. Swapping encoder+decoder breaks pose (encoder features change →
+pose decoder hallucinates keypoints everywhere). Backbone-only surgery preserves
+all downstream compatibility while still improving low-level feature extraction
+for camouflaged persons.
+"""
+
+import argparse
+import torch
+from pathlib import Path
+
+
+def load_state(path: str):
+    ckpt = torch.load(path, map_location="cpu")
+    if isinstance(ckpt, dict):
+        if "ema" in ckpt and isinstance(ckpt["ema"], dict):
+            m = ckpt["ema"].get("module")
+            if m is not None:
+                return m
+        if "model" in ckpt and isinstance(ckpt["model"], dict):
+            return ckpt["model"]
+        if all(isinstance(v, torch.Tensor) for v in ckpt.values()):
+            return ckpt
+    raise RuntimeError(f"Cannot extract state dict from {path}")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--new-det", default="outputs/soldier_finetune/best_stg2.pth",
+                    help="New soldier-finetuned detection checkpoint")
+    ap.add_argument("--orig", default="Traning/best_modelsurgery.pth",
+                    help="Original combined checkpoint (all heads)")
+    ap.add_argument("--out", default="outputs/soldier_finetune/merged_surgery.pth",
+                    help="Output merged checkpoint path")
+    args = ap.parse_args()
+
+    print(f"Loading new detection checkpoint: {args.new_det}")
+    new_state = load_state(args.new_det)
+
+    print(f"Loading original surgery checkpoint: {args.orig}")
+    orig_state = load_state(args.orig)
+
+    merged = {}
+    counts = {"backbone_new": 0, "encoder_orig": 0, "det_decoder_orig": 0,
+              "pose_decoder_orig": 0, "seg_head_orig": 0}
+
+    # 1. Backbone: from NEW model — learns camouflage-aware low-level features
+    for k, v in new_state.items():
+        if k.startswith("backbone."):
+            merged[k] = v
+            counts["backbone_new"] += 1
+
+    # 2-5. Everything else: from ORIGINAL — preserves encoder↔pose_decoder compatibility
+    for k, v in orig_state.items():
+        if k.startswith("encoder."):
+            merged[k] = v
+            counts["encoder_orig"] += 1
+        elif k.startswith("det_decoder."):
+            merged[k] = v
+            counts["det_decoder_orig"] += 1
+        elif k.startswith("pose_decoder."):
+            merged[k] = v
+            counts["pose_decoder_orig"] += 1
+        elif k.startswith("seg_head."):
+            merged[k] = v
+            counts["seg_head_orig"] += 1
+
+    print("\n=== Merge summary ===")
+    for part, n in counts.items():
+        print(f"  {part}: {n} tensors")
+    print(f"  TOTAL: {sum(counts.values())} tensors")
+
+    score_key = "det_decoder.enc_score_head.weight"
+    if score_key in merged:
+        print(f"\nDetection class head shape: {merged[score_key].shape}")
+
+    out_path = Path(args.out)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"model": merged}, str(out_path))
+    print(f"\nSaved merged checkpoint → {out_path}")
+
+
+if __name__ == "__main__":
+    main()

detection_training/train.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+"""
+Detection finetuning for camouflaged soldier detection.
+Finetunes the DFINE detection backbone on soldier/civilian dataset.
+
+Usage:
+  python detection_training/train.py \
+    --config detection_training/configs/soldier_finetune.yml \
+    --resume outputs/phase2_run/best_modelsurgery.pth \
+    --finetune
+"""
+
+import argparse
+import os
+import sys
+
+# Ensure repo root is on path
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, REPO_ROOT)
+
+import torch
+from src.core import YAMLConfig
+from src.solver import DetSolver
+
+
+def get_args():
+    parser = argparse.ArgumentParser('DFINE soldier detection finetuning')
+    parser.add_argument('--config', '-c', type=str,
+                        default='detection_training/configs/soldier_finetune.yml')
+    parser.add_argument('--resume', '-r', type=str,
+                        default='outputs/phase2_run/best_modelsurgery.pth',
+                        help='Checkpoint to resume or finetune from')
+    parser.add_argument('--finetune', action='store_true',
+                        help='Load weights only (ignore optimizer/scheduler state)')
+    parser.add_argument('--test-only', action='store_true')
+    parser.add_argument('--amp', action='store_true', default=True)
+    return parser.parse_args()
+
+
+def main():
+    args = get_args()
+
+    cfg = YAMLConfig(args.config, resume=args.resume if not args.finetune else None)
+
+    if args.finetune and args.resume:
+        print(f'Finetuning from: {args.resume}')
+        ckpt = torch.load(args.resume, map_location='cpu')
+        # The surgery checkpoint stores state under 'model' key
+        state = ckpt.get('model', ckpt)
+        # Strip 'module.' prefix if present
+        state = {k.replace('module.', ''): v for k, v in state.items()}
+        # Load only matching keys (detection head; ignore pose/seg heads)
+        missing, unexpected = cfg.model.load_state_dict(state, strict=False)
+        print(f'  Loaded weights — missing: {len(missing)}, unexpected: {len(unexpected)}')
+        if missing:
+            print(f'  Missing keys (first 5): {missing[:5]}')
+
+    if args.test_only:
+        solver = DetSolver(cfg)
+        solver.val()
+        return
+
+    solver = DetSolver(cfg)
+    solver.fit()
+
+
+if __name__ == '__main__':
+    main()

pose_estimation_berna/base_dfine/dfine_hgnetv2_x_obj2coco_detrpose_paper.yml

@@ -0,0 +1,55 @@
+__include__: [
+  '../../configs/dataset/coco_detection.yml',
+  '../../configs/runtime.yml',
+  '../../configs/dfine/include/dataloader.yml',
+  '../../configs/dfine/include/optimizer.yml',
+  '../../configs/dfine/include/dfine_hgnetv2.yml',
+]
+
+output_dir: ./output/dfine_hgnetv2_x_obj2coco_detrpose_paper
+
+# IMPORTANT:
+# coco_detection.yml sets num_classes=80 globally. DETRPoseTransformer shares `num_classes`,
+# so without overriding it here, the decoder will silently become 80-class.
+num_classes: 2
+
+DFINE:
+  backbone: HGNetv2
+  encoder: HybridEncoder
+  decoder: DETRPoseTransformer
+
+HGNetv2:
+  name: 'B5'
+  return_idx: [1, 2, 3]
+  freeze_stem_only: True
+  freeze_at: 0
+  freeze_norm: True
+
+# X-size encoder: hidden_dim=384 to match D-FINE-X detection config
+HybridEncoder:
+  hidden_dim: 384
+  dim_feedforward: 2048
+
+# Paper-like DETRPose defaults — hidden_dim=384 to match encoder output
+DETRPoseTransformer:
+  num_classes: 2
+  num_queries: 60
+  hidden_dim: 384
+  num_decoder_layers: 6
+  dim_feedforward: 2048
+  dropout: 0.0
+  activation: relu
+  num_feature_levels: 3
+  dec_n_points: 4
+  nhead: 8
+  aux_loss: True
+  num_body_points: 17
+  feat_strides: [8, 16, 32]
+  eval_spatial_size: [640, 640]
+  reg_max: 32
+  reg_scale: 8.0
+  dn_number: 20
+  dn_label_noise_ratio: 0.5
+
+DFINEPostProcessor:
+  num_top_queries: 60