Skip to content

Commit fc3b5b9

Browse files
author
Donglai Wei
committed
vesicle seg
1 parent ca70e08 commit fc3b5b9

3 files changed

Lines changed: 229 additions & 25 deletions

File tree

‎connectomics/data/process/distance.py‎

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,17 @@
1-
from __future__ import print_function, division
2-
from typing import Optional, Tuple, Dict
1+
from typing import Dict, Optional, Tuple
32

4-
import numpy as np
53
import kimimaro
6-
from scipy.ndimage import distance_transform_edt, binary_fill_holes
4+
import numpy as np
5+
from scipy.ndimage import binary_fill_holes, distance_transform_edt
6+
from skimage.filters import gaussian
77
from skimage.morphology import (
8-
remove_small_holes,
8+
ball,
99
binary_erosion,
1010
disk,
11-
ball,
11+
remove_small_holes,
1212
)
13-
from skimage.filters import gaussian
1413

15-
from .bbox_processor import BBoxProcessorConfig, BBoxInstanceProcessor
14+
from .bbox_processor import BBoxInstanceProcessor, BBoxProcessorConfig
1615
from .quantize import energy_quantize
1716

1817
__all__ = [
@@ -76,10 +75,11 @@ def _edt_binary_mask(mask, resolution, alpha):
7675
def edt_instance(
7776
label: np.ndarray,
7877
mode: str = "2d",
79-
quantize: bool = True,
78+
quantize: bool = False,
8079
resolution: Tuple[float] = (1.0, 1.0, 1.0),
8180
padding: bool = False,
8281
erosion: int = 0,
82+
bg_value: float = -1.0,
8383
):
8484
assert mode in ["2d", "3d"]
8585
if mode == "3d":
@@ -92,7 +92,7 @@ def edt_instance(
9292
return vol_distance
9393

9494
# Optimized 2D mode: preallocate arrays instead of lists
95-
vol_distance = np.zeros(label.shape, dtype=np.float32)
95+
vol_distance = np.full(label.shape, bg_value, dtype=np.float32)
9696

9797
# Process slices without copying (use view instead)
9898
for i in range(label.shape[0]):
@@ -155,7 +155,7 @@ def compute_instance_edt(
155155
) -> Optional[np.ndarray]:
156156
"""Compute normalized EDT for a single instance within bbox."""
157157
# Extract instance mask
158-
mask = binary_fill_holes((label_crop == instance_id))
158+
mask = binary_fill_holes(label_crop == instance_id)
159159

160160
# Apply erosion if requested
161161
if context["footprint"] is not None:
@@ -198,16 +198,16 @@ def signed_distance_transform(
198198
alpha: float = 8.0,
199199
) -> np.ndarray:
200200
"""Compute smooth signed distance transform for instance segmentation.
201-
201+
202202
This function produces a true signed distance transform where both foreground
203203
and background have meaningful gradient information, solving the class imbalance
204204
problem of traditional EDT approaches.
205-
205+
206206
Returns SDT in range [-1, 1]:
207207
- Positive values: inside instances (distance from boundary)
208208
- Negative values: outside instances (distance to nearest instance)
209209
- Zero: at instance boundaries
210-
210+
211211
Args:
212212
label: Instance segmentation (H, W) or (D, H, W)
213213
resolution: Pixel/voxel resolution for anisotropic data (z, y, x)
@@ -216,21 +216,21 @@ def signed_distance_transform(
216216
alpha: Smoothness parameter for tanh normalization (default: 8.0)
217217
Higher values = sharper transitions at boundaries
218218
Lower values = smoother, more gradual transitions
219-
219+
220220
Returns:
221221
Signed distance transform in range [-1, 1] with same shape as input
222-
222+
223223
Example:
224224
>>> # 2D mitochondria segmentation
225225
>>> label_2d = np.array([[0, 0, 1, 1, 0],
226226
... [0, 1, 1, 1, 0],
227227
... [0, 0, 1, 0, 0]])
228228
>>> sdt_2d = signed_distance_transform(label_2d, resolution=(1.0, 1.0))
229229
>>> # sdt_2d will have positive values inside instance 1, negative outside
230-
230+
231231
>>> # 3D volume with anisotropic resolution
232232
>>> sdt_3d = signed_distance_transform(label_3d, resolution=(40, 16, 16), alpha=8.0)
233-
233+
234234
Notes:
235235
- This approach eliminates class imbalance by ensuring both foreground
236236
and background contribute meaningful gradients
@@ -243,32 +243,32 @@ def signed_distance_transform(
243243
resolution = resolution[-2:] if len(resolution) > 2 else resolution
244244
elif label.ndim == 3:
245245
resolution = resolution if len(resolution) == 3 else (1.0, 1.0, 1.0)
246-
246+
247247
# Create binary masks
248248
foreground_mask = (label > 0).astype(np.uint8)
249249
background_mask = (label == 0).astype(np.uint8)
250-
250+
251251
# Compute EDT for both foreground and background
252252
if foreground_mask.any():
253253
foreground_edt = distance_transform_edt(foreground_mask, sampling=resolution)
254254
else:
255255
# No foreground - return all negative
256256
return -np.ones_like(label, dtype=np.float32)
257-
257+
258258
if background_mask.any():
259259
background_edt = distance_transform_edt(background_mask, sampling=resolution)
260260
else:
261261
# No background - return all positive
262262
return np.ones_like(label, dtype=np.float32)
263-
263+
264264
# Combine into signed distance
265265
# Positive inside instances, negative outside
266266
sdt = np.where(foreground_mask, foreground_edt, -background_edt)
267-
267+
268268
# Normalize to [-1, 1] using tanh
269269
# This provides smooth, bounded values suitable for regression
270270
sdt_normalized = np.tanh(sdt / alpha)
271-
271+
272272
return sdt_normalized.astype(np.float32)
273273

274274

‎justfile‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -111,16 +111,22 @@ tensorboard-run experiment timestamp port='6006':
111111
# just slurm long 8 4 "just train mito_mitoEM_H" vr40g
112112
# just slurm short 8 4 "python scripts/main.py --config tutorials/lucchi.yaml"
113113
# just slurm short 8 4 "just train lucchi++" "" "64G" # override memory
114+
# just slurm medium 8 2 "just train vesicle_xm" vr144g 128G gb001 # pin node
114115
# Time limits: short=12h, medium=2d, long=5d
115116
# CPU-only convenience wrapper for single-task jobs.
116117
# just slurm short 8 0 "python scripts/downsample_nisb.py --splits train"
117-
slurm partition num_cpu num_gpu cmd constraint='' mem='32G':
118+
slurm partition num_cpu num_gpu cmd constraint='' mem='32G' nodelist='':
118119
#!/usr/bin/env bash
119120
constraint_flag=""
120121
if [ -n "{{constraint}}" ]; then
121122
constraint_flag="--constraint={{constraint}}"
122123
fi
123124
125+
nodelist_flag=""
126+
if [ -n "{{nodelist}}" ]; then
127+
nodelist_flag="--nodelist={{nodelist}}"
128+
fi
129+
124130
# Resolve partition time limit (with fallback defaults)
125131
time_limit=$(just _slurm-time-limit {{partition}})
126132
@@ -136,6 +142,7 @@ slurm partition num_cpu num_gpu cmd constraint='' mem='32G':
136142
--mem={{mem}} \
137143
--time=$time_limit \
138144
$constraint_flag \
145+
$nodelist_flag \
139146
--wrap="mkdir -p \$HOME/.just && export JUST_TEMPDIR=\$HOME/.just TMPDIR=\$HOME/.just NCCL_SOCKET_FAMILY=AF_INET && source /projects/weilab/weidf/lib/miniconda3/bin/activate pytc && cd $PWD && srun --ntasks=1 --gpus-per-task={{num_gpu}} --cpus-per-task={{num_cpu}} {{cmd}}"
140147
141148
# Generic CPU-only multi-task launcher (single node, no GPU).

‎tutorials/vesicle_xm.yaml‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
_base_: bases/mednext.yaml
2+
experiment_name: fiber_mednext_bcs
3+
description: Fiber segmentation with MedNeXt and multi-task learning (Binary + Contour + SDT)
4+
system:
5+
training:
6+
num_gpus: -1
7+
num_workers: -1
8+
batch_size: 4
9+
inference:
10+
num_workers: 1
11+
batch_size: 8
12+
seed: 42
13+
model:
14+
out_channels: 3
15+
input_size:
16+
- 32
17+
- 128
18+
- 128
19+
output_size:
20+
- 32
21+
- 128
22+
- 128
23+
mednext_size: S
24+
mednext_kernel_size: 3
25+
deep_supervision: false
26+
loss_functions:
27+
- WeightedBCEWithLogitsLoss
28+
- DiceLoss
29+
- WeightedBCEWithLogitsLoss
30+
- DiceLoss
31+
- WeightedMSELoss
32+
loss_weights:
33+
- 1.0
34+
- 0.5
35+
- 1.0
36+
- 0.5
37+
- 2.0
38+
loss_kwargs:
39+
- reduction: mean
40+
- sigmoid: true
41+
smooth_nr: 1e-5
42+
smooth_dr: 1e-5
43+
- reduction: mean
44+
- sigmoid: true
45+
smooth_nr: 1e-5
46+
smooth_dr: 1e-5
47+
- tanh: true
48+
multi_task_config:
49+
- - 0
50+
- 1
51+
- binary
52+
- - 0
53+
- 1
54+
- - 1
55+
- 2
56+
- instance_boundary
57+
- - 2
58+
- 3
59+
- - 2
60+
- 3
61+
- instance_edt
62+
- - 4
63+
data:
64+
train_image: "datasets/bouton-lv/train/vol0_im.h5"
65+
train_label: "datasets/bouton-lv/train/vol0_lv.h5"
66+
train_mask: "datasets/bouton-lv/train/vol0_mask.h5"
67+
train_resolution: [30, 8, 8]
68+
69+
70+
use_preloaded_cache_train: true
71+
persistent_workers: true
72+
iter_num_per_epoch: 1000
73+
cached_sampling_foreground_threshold: 0.1
74+
cached_sampling_max_attempts: 10
75+
76+
patch_size: [32, 64, 64]
77+
data_transform:
78+
resize: [32, 128, 128]
79+
80+
label_transform:
81+
targets:
82+
- name: binary
83+
kwargs: {}
84+
- name: instance_boundary
85+
kwargs:
86+
thickness: 1
87+
edge_mode: all
88+
mode: 2d
89+
- name: instance_edt # Channel 2: distance transform (bbox-optimized)
90+
kwargs:
91+
mode: "2d" # 2D EDT with per-instance bounding box optimization
92+
93+
augmentation:
94+
preset: some
95+
flip:
96+
enabled: true
97+
rotate:
98+
enabled: true
99+
spatial_axes:
100+
- 1
101+
- 2
102+
intensity:
103+
enabled: true
104+
gaussian_noise_prob: 0.5
105+
gaussian_noise_std: 0.1
106+
shift_intensity_prob: 0.5
107+
shift_intensity_offset: 0.1
108+
contrast_prob: 0.5
109+
contrast_range: [0.9, 1.1]
110+
optimization:
111+
max_epochs: 500
112+
gradient_clip_val: 1
113+
accumulate_grad_batches: 1
114+
precision: "16-mixed"
115+
log_every_n_steps: 100
116+
optimizer:
117+
lr: 0.0003
118+
weight_decay: 0.01
119+
eps: 1.0e-08
120+
scheduler:
121+
name: WarmupCosineLR
122+
warmup_epochs: 3
123+
warmup_start_lr: 1.0e-05
124+
min_lr: 1.0e-06
125+
ema:
126+
enabled: true
127+
decay: 0.999
128+
validate_with_ema: true
129+
monitor:
130+
logging:
131+
scalar:
132+
loss:
133+
- train_loss_total_epoch
134+
loss_every_n_steps: 50
135+
images:
136+
max_images: 2
137+
num_slices: 4
138+
log_every_n_epochs: 50
139+
checkpoint:
140+
save_top_k: 3
141+
save_every_n_epochs: 10
142+
early_stopping:
143+
monitor: train_loss_total_epoch
144+
patience: 100
145+
min_delta: 1.0e-05
146+
threshold: 0.01
147+
divergence_threshold: 100.0
148+
149+
test:
150+
data:
151+
test_image: "datasets/bouton-lv/train/vol0_im.h5"
152+
test_label: "datasets/bouton-lv/train/vol0_lv.h5"
153+
test_mask: "datasets/bouton-lv/train/vol0_mask.h5"
154+
test_resolution: [30, 8, 8]
155+
156+
157+
inference:
158+
sliding_window:
159+
window_size:
160+
- 32
161+
- 128
162+
- 128
163+
overlap: 0.5
164+
blending: gaussian
165+
sigma_scale: 0.25
166+
padding_mode: reflect
167+
test_time_augmentation:
168+
enabled: true
169+
channel_activations:
170+
- - 0
171+
- 1
172+
- sigmoid
173+
- - 1
174+
- 2
175+
- sigmoid
176+
- - 2
177+
- 3
178+
- tanh
179+
decoding:
180+
- name: decode_binary_contour_distance_watershed
181+
kwargs:
182+
binary_threshold:
183+
- 0.9
184+
- 0.85
185+
contour_threshold:
186+
- 0.8
187+
- 1.1
188+
distance_threshold:
189+
- 0.5
190+
- -0.5
191+
min_instance_size: 100
192+
min_seed_size: 20
193+
prediction_scale: 1
194+
evaluation:
195+
enabled: true
196+
metrics:
197+
- adapted_rand

0 commit comments

Comments
 (0)