LN_segmentation_sweep_v3

A unet model for multilabel image segmentation trained with sliding window approach.

Model Description

Wandb Parameters

Parameter	Value
data_path	GleghornLab/Semi-Automated_LN_Segmentation_10_11_2025
img_size	64
downsample_factor	1
num_channels	3
batch_size	8
lr	7.77451918775676e-06
epochs	100
patience	10
weight_decay	8.29726636990404e-05
model_type	unet
n_filts	64
t	3
k	5
augment	False
norm	True
keep	0.7562670422006883
pruning_factor	0.0034113333205044107
output_dir	pooled_metrics_hev_settings
device	None
num_workers	4
prefetch_factor	2
wandb_project	segmentation-sweep
wandb_run_name	hev-only-repro-pooled
wandb_mode	online
push_to_hub	True
hub_model_id	aholk/LN_segmentation_sweep_v3
skip_report	False
sweep_mode	False
num_params	138063876
num_classes	4

Model Parameters

Parameter	Value
num_channels	3
num_classes	4
n_filts	64
t	3
k	5
img_size	64
norm	True
model_arch	unet
transformers_version	5.9.0
architectures	["UNetForSegmentation"]
output_hidden_states	False
return_dict	True
dtype	float32
chunk_size_feed_forward	0
is_encoder_decoder	False
id2label	{"0": "LABEL_0", "1": "LABEL_1"}
label2id	{"LABEL_0": 0, "LABEL_1": 1}
problem_type	None
_name_or_path
batch_size	8
downsample_factor	1.0
model_type	segmentation
output_attentions	False

Performance Metrics

Metric	Mean	Class 0	Class 1	Class 2	Class 3
Dice	0.8095	0.7363	0.7966	0.7977	0.9074
IoU	0.6847	0.5826	0.6620	0.6634	0.8306
F1	0.8095	0.7363	0.7966	0.7977	0.9074
MCC	0.8051	0.7409	0.7990	0.7922	0.8883
ROC AUC	0.9920	0.9897	0.9972	0.9882	0.9928
PR AUC	0.8940	0.8305	0.9001	0.8815	0.9640

Usage

import numpy as np
from model import MODEL_REGISTRY, SegmentationConfig

# Load model
config = SegmentationConfig.from_pretrained("aholk/LN_segmentation_sweep_v3")
model = MODEL_REGISTRY["unet"].from_pretrained("aholk/LN_segmentation_sweep_v3")
model.eval()

# Run inference on a full image with sliding window
image = np.random.rand(2048, 2048, 3).astype(np.float32)  # Your image here
probs = model.predict_full_image(
    image,
    dim=64,
    batch_size=16,
    device="cuda"  # or "cpu"
)
# probs shape: (num_classes, H, W) with values in [0, 1]

# Threshold to get binary masks
masks = (probs > 0.5).astype(np.uint8)

Training Plots

Citation

If you use this model, please cite:

@software{windowz_segmentation,
  title={Multilabel Image Segmentation with Sliding Window U-Net},
  author={Gleghorn Lab},
  year={2025},
  url={https://github.com/GleghornLab/ComputerVision2}
}