kubilaycaglayan/wave-regression

Wave Regression: relative sea-surface waviness from a photograph

A small image-regression model that looks at a photograph of the sea and returns a single score between 0.0 (calm) and 1.0 (rough) describing how wavy the visible sea surface looks.

What the score is not. The score is a relative, visual waviness rating learned from one person's manual labels. It is not a wave height in metres, and it is not calibrated to Douglas, Beaufort, WMO or any other official sea-state scale.

This repository holds one checkpoint and the minimum code to reproduce its preprocessing and run it on a new photograph. You don't need the original development repository.

Task	Single-output image regression
Input	One RGB sea photograph (JPEG, PNG, HEIC, …)
Output	waviness ∈ [0, 1]
Architecture	ResNet18 (ImageNet features) → Linear(512, 1) → Sigmoid
Trainable parameters	513 (of 11,177,025)
Framework	PyTorch (trained with PyTorch Lightning 2.6.5)
Checkpoint	model.ckpt, epoch 45 / global step 276

Problem definition

Given a photograph containing the sea, estimate how rough the water surface looks. Photos contain a lot of irrelevant content (sky, buildings, shoreline), so the model does not see the full photo. A fixed preprocessing pipeline first isolates the lower part of the detected water, where waves were most visible from the camera position, and the regressor only sees that crop.

Architecture

• Backbone: torchvision resnet18, initialized from ResNet18_Weights.DEFAULT (ImageNet-1k). All convolution and BatchNorm affine weights stayed frozen at their ImageNet values. The BatchNorm running statistics were updated during training on the training images.
• Head: the ResNet fc layer is replaced with Linear(512, 1) followed by Sigmoid, so the raw model output is already the score in [0, 1]. No further output transformation is applied.
• The published model.ckpt stores every weight and buffer, so loading it downloads no pretrained ResNet weights.

Dataset

• Source: a limited collection of sea photographs taken with a phone camera from a similar coastal viewpoint. The photographs are not published.

• Labels: one annotator manually rated each preprocessed image on a 0.00–1.00 scale in 0.05 steps using a local labeling tool. There is no written rubric and no second rater. Labels were used as-is, with no further normalization.

• Split used for this checkpoint (split snapshot of 2026-09-03, grouped by capture session so near-duplicate frames don't cross splits):

  Split	Images	Label range	Label mean
  Train	42	0.00 – 0.80	0.329
  Validation	7	0.05 – 0.90	0.521
  Test (benchmark-v1)	6	0.20 – 0.90	0.567

  More photos were labeled later. They were added to the development repository's validation and test splits, and no image in them was used to train this checkpoint. See Evaluation.

Preprocessing

Implemented in preprocessing.py. It is a direct port of the development repository's Step 1 and Step 2 scripts, in the configuration that produced this checkpoint's training images:

1. Water segmentation. nvidia/segformer-b2-finetuned-ade-512-512 at pinned revision de01bae28967510f9ddd496c60a969357195400c. The whole photograph is resized to 512×512 and segmented in one pass. The class map is resized back with nearest-neighbour interpolation. ADE20K classes water, sea, river and lake count as water.
2. Step 1: standardize. Erode the water mask slightly, black out all non-water pixels, crop to the water bounding box, and fit the crop (aspect preserved, Lanczos) into a centered black 224×224 square. JPEG round trip at quality 95.
3. Step 2: lower-water crop. Segment the Step 1 image again and choose a bottom-anchored rectangle. It must keep ≥ 95 % of the lowest 35 % of the detected water, and among those rectangles the one with the least black padding wins. Fit the rectangle into 224×224 as above. JPEG round trip at quality 95 with 4:4:4 chroma.
4. Tensor. ToTensor() and ImageNet normalization (mean 0.485, 0.456, 0.406, std 0.229, 0.224, 0.225).

The JPEG round trips are deliberate: the model was trained on the decoded JPEG files. Run from the original HEIC photos, this code reproduces the stored training images bit for bit, on both CPU and CUDA.

Input contract. inference.py expects a raw photograph by default. --preprocessed accepts a 224×224 image that this pipeline already produced (for example one written with --save-model-input). A photograph that has just been resized to 224×224 is not a valid preprocessed input.

The SegFormer model is not redistributed here; its license is marked other on the Hub. It is downloaded from the pinned revision on first use (about 200 MB) and then cached by huggingface_hub. Please review its license before use.

Training objective

• Loss: SmoothL1Loss between the sigmoid output and the label
• Optimizer: AdamW on the 513 head parameters, learning rate 1e-3, batch size 8
• Augmentation (training only): RandomHorizontalFlip(p=0.5), ColorJitter(brightness=0.25, contrast=0.25, saturation=0.20, hue=0.02)
• Up to 100 epochs with early stopping on validation MAE (patience 10). Training stopped after 57 epochs.
• The published weights are the best-validation-MAE epoch (epoch index 45, global step 276)

Checkpoint selection

The development repository contained six checkpoints (runs v1, v2, v6, v7, v8, v9). They had been trained and validated on different split snapshots, so the validation MAE in their filenames can't be compared directly. All six were re-evaluated on one common validation set: the 12 current validation images, stored model inputs. None of those images were in any of the six checkpoints' training splits.

Run	Validation MAE (n=12)	Validation RMSE
v1 (published)	0.1437	0.1803
v2 (BatchNorm-frozen rerun)	0.1630	0.2038
v6	0.1625	0.1963
v7	0.1770	0.2143
v8	0.1876	0.2366
v9	0.1876	0.2366

v1 has the lowest validation error, and the development repository had already configured it as its prediction and evaluation checkpoint. The test sets played no part in this choice. Note that each run's own early stopping used a subset of these validation images, so all of these numbers are somewhat optimistic.

Original checkpoint file: wave-regression-baseline-v1-best-val-mae-epoch=45-val_mae=0.1068.ckpt

Evaluation

All numbers below come from running this release's pipeline end to end on the raw photographs, on CPU. Labels are on the 0–1 scale, so an MAE of 0.15 is 15 % of the full range. Baselines always predict the mean (0.329) or median (0.30) of this checkpoint's 42 training labels. No training image appears in any evaluation set.

Held-out test sets (never used for training, early stopping or checkpoint selection):

Test set	n	Model MAE	Model RMSE	Mean-baseline MAE / RMSE	Median-baseline MAE / RMSE
benchmark-v1 (fixed test split of this checkpoint's own split)	6	0.1446	0.1606	0.2810 / 0.3198	0.3000 / 0.3416
Photos labeled after training (current test split, new images only)	6	0.1826	0.2269	0.1357 / 0.1483	0.1167 / 0.1354
Combined current test split	12	0.1636	0.1966	0.2083 / 0.2493	0.2083 / 0.2598

Validation:

Validation set	n	Model MAE	Model RMSE
This checkpoint's own validation split (used for early stopping)	7	0.1068	0.1335
Newer validation photos (not used by this checkpoint)	5	0.2089	0.2662
Combined current validation split	12	0.1493	0.1998

How to read these results:

• On benchmark-v1 the model roughly halves the error of a constant prediction.
• On the six newer test photos (and the five newer validation photos) the model is worse than always predicting the training mean. Most of these photos have low labels (0.1–0.3), and the model over-predicts several of them (for example 0.20 → 0.57). The newer photos may differ from the training photos in conditions or framing. With this few images, the gap can't be attributed to a cause.
• The checkpoint-selection table above was computed on the development repository's stored model inputs. For the five newer validation photos, those inputs had been produced with a later, tiled segmentation setting, which explains why they differ slightly from the end-to-end numbers here.
• Every set has 5–12 images. Treat all figures as rough indications, not benchmarks. No standardized or external evaluation is available.

Usage

hf download kubilaycaglayan/wave-regression --local-dir wave-regression   # or: git clone (needs git-lfs)
cd wave-regression
python -m venv .venv && . .venv/bin/activate
# optional, CPU-only PyTorch wheels:
# pip install torch==2.13.0 torchvision==0.28.0 --index-url https://download.pytorch.org/whl/cpu
pip install -r requirements.txt

python inference.py path/to/sea_photo.jpg
# waviness: 0.5761   (example output)

Useful options:

python inference.py photo.heic --device cpu                        # force CPU (auto: cuda > mps > cpu)
python inference.py photo.jpg --save-model-input model_input.jpg   # keep the 224x224 crop the model saw
python inference.py model_input.jpg --preprocessed                 # skip segmentation for a pipeline output
python inference.py photo.jpg --json                               # {"image": ..., "waviness": ...}

From Python:

from inference import predict

score = predict("sea_photo.jpg", device="cpu")  # float in [0, 1]

If no reliable water region is found, inference stops with an error rather than returning a score. Tested on Python 3.14 with the pinned requirements.txt, on CPU and CUDA. The mps option (Apple Silicon) is supported by the code but has not been tested.

Intended uses

• Learning and experimentation with small-data image regression and segmentation-guided preprocessing.
• Rough, relative comparison of sea-surface conditions between photographs taken from the same or a very similar viewpoint as the training photos, for example a fixed camera feed.

Out-of-scope uses

• Any safety-relevant decision: navigation, swimming, surfing, coastal warnings.
• Estimating physical wave height, wave period or official sea-state categories.
• Photographs from very different viewpoints, lenses, distances, lighting or water bodies, where the score has not been validated.

Known limitations

• Very small dataset. 42 training images, 7 validation images and 6–12 test images. All metrics have wide uncertainty.
• Single viewpoint. All photographs come from a limited collection taken from a similar coastal viewpoint, so generalization to other locations is unknown.
• Subjective single-annotator labels with no rubric. Label noise is probably on the same order as the reported errors.
• Weak generalization to newer photos. On photos labeled after training, the model did worse than a constant baseline (see Evaluation). It also under-predicts very rough water (0.90 → 0.66) and over-predicts several calm scenes.
• Segmentation dependence. If SegFormer misses the sea or labels other surfaces as water, the crop and the score will be wrong. Inspect the crop with --save-model-input when in doubt.
• Sensitivity to encoding. The score depends on exact pixel values. Re-encoding the same photo (for example HEIC → JPEG) shifted one test score from 0.576 to 0.550. Use the original camera file where possible.
• Black padding. Inputs are letterboxed into a black square. With only a linear head, the model may partly respond to the amount of black padding instead of wave texture.
• Speed. Each photo runs SegFormer-B2 twice. On CPU that takes roughly 15–20 s for a 12 MP photo.

Reproducibility

File	Purpose
model.ckpt	Selected checkpoint (inference-only copy, see below)
model.py	Architecture; loads model.ckpt with strict=True
preprocessing.py	Segmentation, Step 1, Step 2 and tensor conversion
inference.py	Command-line and Python inference for one image
config.json	Model, training and preprocessing configuration
requirements.txt	Pinned, tested inference dependencies
LICENSE	MIT License

• model.ckpt SHA-256: fc616d9b797dcba6f1c1cc52abb4e50db978365e6466cf4724d27780cc8e193f
• Original training checkpoint SHA-256: 74b8ca463f46dcd90737d6e4d7fe15291f8731c9f157d0fa4b1ab79454701432

model.ckpt was exported from the original Lightning checkpoint and keeps state_dict, hyper_parameters, epoch, global_step and pytorch-lightning_version. It drops optimizer, loop and callback state, which included absolute paths from the training machine. All 122 tensors are bit-identical to the original. The file loads with torch.load(..., weights_only=True).

Verification before publishing:

• Predictions from this code were compared with the development implementation on the same 24 inputs; they were identical (max absolute difference 0).
• Preprocessing of raw photographs reproduced the stored training images exactly.
• The release was tested in a clean virtual environment with an empty Hugging Face cache and no access to the development repository.

Credits

• Model, data collection, labels and development pipeline: the repository author. Development repository (training, labeling and dataset tooling): github.com/kubilaycaglayan/wave-regression-model.
• Backbone: ResNet18 from torchvision (He et al., 2016), ImageNet weights.
• Water segmentation: SegFormer-B2 fine-tuned on ADE20K by NVIDIA (Xie et al., 2021).

License

The model weights and code in this repository are released under the MIT License. The SegFormer segmentation model is not part of this repository. It is downloaded from its own Hub repository and is subject to NVIDIA's license there.