pytorch_connectomics/tutorials/neuron_liconn_mit.yaml at master · PytorchConnectomics/pytorch_connectomics · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
experiment_name: liconn_affinity_sdt_instance_seg
description: >
  LiConn neuron instance segmentation using long-range affinity + SDT model
  (9-channel affinity + 1-channel skeleton-aware distance transform).

  Data: liconn_18nm_labelled_cropped.zarr
    img:   shape (415, 4290, 3345), dtype uint8,  chunks (64,64,64), ZYX order
    label: shape (415, 4290, 3345), dtype uint64, chunks (64,64,64), ZYX order
    resolution: 25 nm (z) x 18 nm (y) x 18 nm (x)

  Pipeline: aff9_sdt (same as neuron_snemi_sdt.yaml)
    Ch 0-2: short-range affinity (offset 1) z,y,x
    Ch 3-5: mid-range   affinity (offset 3) z,y,x
    Ch 6-8: long-range  affinity (offset 9) z,y,x
    Ch 9:   skeleton-aware EDT (distance to skeleton)

  Loss (from loss_bd profile):
    BCE + Dice on affinity channels [0:9]
    MSE on SDT channel [9:]

# ── Inherit all shared profile libraries ──────────────────────────────────────
_base_:
- ../connectomics/config/all_profiles.yaml

# ── Profile selectors ─────────────────────────────────────────────────────────
default:
  pipeline_profile: aff9_sdt
  system:
    profile: all-gpu-cpu
  model:
    arch:
      profile: mednext_m
    input_size: [64, 256, 256]
    output_size: [64, 256, 256]
    mednext:
      size: M
      kernel_size: 3
      dim: 3d
      checkpoint_style: outside_block

  data:
    root_path: ""   # override per site: /projects/weilab/dataset/liconn/ffn/mansour
    label_transform:
      resolution: [25, 18, 18]
      erosion: 1
      targets:
      - name: affinity
        kwargs:
          offsets: ["1-0-0", "0-1-0", "0-0-1", "3-0-0", "0-3-0", "0-0-3", "9-0-0", "0-9-0", "0-0-9"]
          affinity_mode: deepem
      - name: skeleton_aware_edt
        kwargs:
          alpha: 0.8
          bg_value: -1.0
    image_transform:
      normalize: "0-1"
      clip_percentile_low: 0.0
      clip_percentile_high: 1.0
    data_transform:
      pad_size: [32, 128, 128]
    dataloader:
      profile: cached
      use_cache: false
      use_lazy_zarr: true
      use_preloaded_cache_train: false
      use_preloaded_cache_val: false
      cached_sampling_foreground_threshold: 0.10
      patch_size: [64, 256, 256]
    augmentation:
      profile: aug_em_neuron

  inference:
    sliding_window:
      window_size: [64, 256, 256]
      sw_batch_size: 1
      overlap: 0.5
      blending: gaussian
      keep_input_on_cpu: true
      sw_device: cuda
      output_device: cpu
    test_time_augmentation:
      enabled: false
    crop_pad: [23, 32, 119, 128, 119, 128]
    save_prediction:
      enabled: true

  decoding:
    steps:
    - template: decoding_waterz
      kwargs:
        thresholds: 0.4
        merge_function: aff85_his256
        aff_threshold: [0.001, 0.999]

    postprocessing:
      enabled: true
  evaluation:
    enabled: true
    metrics: [adapted_rand]

# ── Training stage ────────────────────────────────────────────────────────────
train:
  data:
    train:
      image: liconn_18nm_labelled_cropped.zarr/img
      label: liconn_18nm_labelled_cropped.zarr/label
      # label_aux: auto-derived from label if null (skeletonization on-the-fly)
      resolution: [25, 18, 18]
    val:
      image: liconn_18nm_labelled_cropped.zarr/img
      label: liconn_18nm_labelled_cropped.zarr/label
      resolution: [25, 18, 18]
    dataloader:
      batch_size: 1

  optimization:
    profile: warmup_cosine_lr
    max_epochs: 500
    n_steps_per_epoch: 200
    accumulate_grad_batches: 4
    val_check_interval: 10

  system:
    num_gpus: 2
    num_workers: 16
    seed: 42

  monitor:
    logging:
      scalar:
        loss: [train_loss_total_epoch, train_loss_affinity_total, train_loss_sdt_total]
        loss_every_n_steps: 100
      images:
        enabled: true
        max_images: 8
        num_slices: 4
        log_every_n_epochs: 5
        channel_mode: all
    checkpoint:
      save_top_k: 3
      monitor: val_loss_total
      mode: min

# ── Test / inference stage ────────────────────────────────────────────────────
test:
  data:
    test:
      image: DL288B_251222S_cond5_40x_12tiles_round1_fused_488_crop512x1024x1024.tif
      # label: omitted — no ground truth
      resolution: [25, 18, 18]

  inference:
    sliding_window:
      window_size: [32, 256, 256]

  evaluation:
    enabled: false