README.md

---
license: mit
language:
  - en
tags:
  - image-segmentation
  - semantic-segmentation
  - sky-segmentation
  - water-segmentation
  - person-segmentation
  - structure-from-motion
  - sfm
  - ade20k
  - segformer
  - onnx
  - pytorch
datasets:
  - ADE20K
metrics:
  - iou
  - dice
  - accuracy
library_name: segmentation-models-pytorch
pipeline_tag: image-segmentation
---

# 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 |

<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

| 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

```bash
pip install skywater-seg
```

### PyTorch — one-liner from HuggingFace

```python
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)

```python
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

```bash
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

```python
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](https://groups.csail.mit.edu/vision/datasets/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`](data/README.md) for the class mapping and dataset layout.

## Repository

Code & full pipeline: [github.com/Vincentqyw/skywater](https://github.com/Vincentqyw/skywater)

## Citation

```bibtex
@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}},
}
```