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):

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

<p align="center"> <img src="speed_full.png" width="80%" alt="Speed comparison"> </p>

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.

<p align="center"> <img src="sample_grid.png" width="100%" alt="Sample predictions"> </p>

Files

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}},
}