Realcat/skywater_seg

SkyWater-Seg: Sky / Water / Person Segmentation

SegFormer MiT-B2 (24.7M params) fine-tuned on ADE20K for sky, water, and person segmentation. Designed to pre-filter images for Structure from Motion (SfM) and image matching pipelines — mask out sky, water, and person regions to eliminate their interference.

Model Card

• Architecture: SegFormer (MiT-B2 encoder + lightweight MLP decoder)
• Parameters: 24.7M
• Input: RGB image, ImageNet normalization, 384×384
• Output: 4-class logits → 0=background, 1=sky, 2=water, 3=person
• Training data: ADE20K filtered split (11,086 train / 1,111 val)
• License: MIT

Evaluation

ADE20K filtered validation set (1,111 images, 384×384 input):

Class	IoU	Dice	Precision	Recall
Background	96.6%	98.3%	98.9%	97.6%
Sky	92.1%	95.9%	93.3%	98.6%
Water	79.4%	88.5%	89.9%	87.2%
Person	77.8%	87.5%	85.0%	90.1%
Foreground mIoU	88.1%	—	—	—
Overall mIoU	94.6%	—	—	—
Pixel Accuracy	97.2%	—	—	—

Inference Speed (RTX 3060 Laptop, 384×384, batch=1)

Backend	Latency	Platforms
ONNX FP16 GPU	13.3 ms	NVIDIA CUDA
ONNX FP32 GPU	15.4 ms	NVIDIA CUDA
PyTorch FP32	24.1 ms	CUDA / MPS / CPU
ONNX CoreML	~3 ms	Apple Silicon (macOS)
ONNX FP32 CPU	188 ms	Any

Pixel Identity

All ONNX variants produce pixel-identical results to PyTorch. Verified on 50 validation images (12M pixels): only 0.003% pixel difference from FP16 rounding.

Files

File	Format	Size	Description
skywater_segformer_b2.pth	PyTorch	284 MB	Full checkpoint (model + optimizer + scheduler)
skywater_segformer_b2_fp32.onnx	ONNX FP32	95 MB	FP32 — CPU, CUDA, CoreML, TensorRT
skywater_segformer_b2_fp16.onnx	ONNX FP16	48 MB	FP16 — CUDA, CoreML, TensorRT (1.8× faster)
config.yaml	YAML	2 KB	Training config for reproducibility
data/train.txt	Text	11,086 lines	Training split
data/val.txt	Text	1,111 lines	Validation split

Usage

pip install skywater-seg

PyTorch — one-liner from HuggingFace

from skywater_seg import SkyWaterSegModel
import torch

model = SkyWaterSegModel.from_pretrained("Realcat/skywater_seg")
model.eval()

img = ...  # (H, W, 3) RGB, normalized, resized to 384×384
with torch.no_grad():
    logits = model(img)  # (1, 4, H, W)

ONNX Runtime (no PyTorch needed)

from skywater_seg import ONNXRuntimeInference

# NVIDIA GPU
infer = ONNXRuntimeInference("skywater_segformer_b2_fp16.onnx", provider="cuda")

# Apple Silicon
infer = ONNXRuntimeInference("skywater_segformer_b2_fp16.onnx", provider="coreml")

# CPU
infer = ONNXRuntimeInference("skywater_segformer_b2_fp32.onnx", provider="cpu")

result = infer.predict("image.jpg")

CLI

pip install skywater-seg

# PyTorch
skywater-infer --checkpoint skywater_segformer_b2.pth --input image.jpg

# ONNX
skywater-infer --onnx skywater_segformer_b2_fp16.onnx --input image.jpg

Batch Inference

infer = ONNXRuntimeInference("skywater_segformer_b2_fp16.onnx", provider="cuda")
results = infer.predict_batch(["img1.jpg", "img2.jpg", "img3.jpg"], batch_size=4)
for r in results:
    print(r["mask"].shape, r["sky_mask"].mean())

Reproducing Training

1. Download ADE20K (ADEChallengeData2016 subset)
2. Update config.yaml paths to your ADE20K location
3. Use the split files in data/
4. Train: uv run python train.py --config config.yaml

See data/README.md for the class mapping and dataset layout.

Repository

Code & full pipeline: github.com/Vincentqyw/skywater

Citation

@misc{qin2026skywater,
  author       = {Vincent Qin},
  title        = {SkyWater-Seg: Sky, Water, and Person Segmentation
                  for Structure from Motion},
  year         = {2026},
  howpublished = {\url{https://github.com/Vincentqyw/skywater}},
}