unet_windowz

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

Model Description

• Architecture: attention_unet
• Input Channels: 3
• Output Classes: 4
• Base Filters: 32
• Window Size: 128

Model-Specific Parameters

Training Configuration

Parameter	Value
Batch Size	24
Learning Rate	0.0006966596027953339
Weight Decay	0.0032033857192098666
Epochs	100
Patience	10
Dataset	GleghornLab/Semi-Automated_LN_Segmentation_10_11_2025

Performance Metrics

Metric	Mean	Class 0	Class 1	Class 2	Class 3
Dice	0.7068	0.7111	0.3810	0.7933	0.9417
IoU	0.6057	0.5602	0.3079	0.6629	0.8919
F1	0.7068	0.7111	0.3810	0.7933	0.9417
MCC	0.7055	0.7222	0.3770	0.7938	0.9289
ROC AUC	0.9239	0.9917	0.7119	0.9955	0.9963
PR AUC	0.7900	0.8115	0.4304	0.9336	0.9844

Usage

import numpy as np
from model import MODEL_REGISTRY, SegmentationConfig

# Load model
config = SegmentationConfig.from_pretrained("lhallee/unet_windowz")
model = MODEL_REGISTRY["attention_unet"].from_pretrained("lhallee/unet_windowz")
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=128,
    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}
}