README.md

metadata

tags:
  - seismic
  - ground-roll
  - denoising
  - unet
  - resunet
  - dncnn
  - attention-unet
  - pytorch
library_name: pytorch

Ground-Roll Attenuation Benchmark

Deep-learning-based coherent noise (ground roll) suppression on pre-stack seismic shot gathers, using the SEG C3 synthetic dataset.

Task

Given a noisy shot gather contaminated by dispersive ground-roll noise, the model predicts the additive noise component. The denoised signal is obtained by:

denoised = noisy_input - predicted_noise

This is a paired regression task trained with a noise-label objective (the ground-truth noise component). The supervision target is the residual between the noisy input and the clean reference.

Dataset

• Source: SEG C3 pre-stack synthetic data, 9 regular shot gathers
• Geometry: 201 traces × 625 time samples per shot, dt = 2 ms
• Noise modeling: Reflection signals modeled with the acoustic wave equation; ground roll modeled with the elastic wave equation to capture its dispersive, low-velocity character
• Split: Shot-level (FFID) sequential 7:1:1 — 7 training shots, 1 validation, 1 held-out test

Noise Intensity Levels

Five ground-roll intensity levels produce paired noisy / noise-label records:

Level	SNR (dB)	PSNR (dB)	SSIM	MAE	MSE	RMSE
1.0	2.71	15.81	0.9480	0.030624	0.026232	0.161964
3.0	-6.83	6.27	0.9418	0.091871	0.236091	0.485892
5.0	-11.27	1.83	0.9402	0.153118	0.655809	0.809820
7.0	-14.19	-1.09	0.9395	0.214366	1.285386	1.133749
9.0	-16.37	-3.27	0.9390	0.275613	2.124822	1.457677

Metrics computed on the full dataset in the original amplitude domain before normalization.

Model Architectures

• Attention UNet (atten_unet) — U-Net with additive attention gates on skip connections (Oktay et al., 2018, MIDL). Base channels: 32, depth: 4.
• DnCNN (dncnn) — Flat 17-layer Conv-BN-ReLU stack with residual learning (Zhang et al., 2017, IEEE TIP). Base channels: 64.

Preprocessing

• Normalization: max_abs, global scope — the entire dataset scaled to [-1, 1]
• Patching: Overlapping 2D patches (128 × 256) with 50% overlap, channel-last format (1, H, W)

Repository Structure

models/
├── unet/
│   ├── level1.0_seed42/
│   │   ├── best.pt          # Best checkpoint (minimum validation loss)
│   │   └── config.yaml      # Full training configuration
│   ├── level1.0_seed43/
│   ├── level1.0_seed44/
│   ├── level3.0_seed42/
│   └── ...
└── res_unet/
    └── ...

Each subdirectory corresponds to one experiment: a model architecture trained at a specific noise level with a specific random seed.

Training Details

Hyperparameter	Value
Loss	MSE (predicted noise vs. label noise)
Optimizer	AdamW (lr=1e-3, weight_decay=1e-4)
Scheduler	Cosine annealing (min_lr=1e-6)
Epochs	200
Gradient clipping	1.0 (max norm)
Batch size	196
Distributed training	2 × NVIDIA RTX 4090, DDP
Seeds	42, 43, 44 per experiment

Usage

import torch
from huggingface_hub import hf_hub_download

# Download a checkpoint
repo = "GeoBrain/coherent-noise-attenuation"
model_key = "res_unet"
level = "3.0"
seed = "42"

ckpt_path = hf_hub_download(
    repo_id=repo,
    filename=f"models/{model_key}/level{level}_seed{seed}/best.pt",
)

# Load state dict
state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True)

# For full model loading, instantiate the corresponding architecture
# and load the state dict (see config.yaml for exact architecture params).

See the companion benchmark documentation for detailed experimental setup and full evaluation results.

Results (SNR dB)

Model	Params (M)	Level 1.0	Level 3.0	Level 5.0	Level 7.0	Level 9.0
Raw (noisy)	—	2.71	−6.83	−11.27	−14.19	−16.37
UNet	7.76	28.39±1.09	23.07±0.10	20.22±0.32	17.90±0.28	16.60±0.46
ResUNet	8.11	31.62±0.56	22.65±0.93	18.11±2.12	16.60±1.44	14.61±0.01
DnCNN	0.56	31.67±0.11	30.02±0.43	22.91±3.74	—	—
Attention UNet	7.85	28.79±0.46	23.10±1.00	19.66±0.22	17.57±0.11	16.61±0.19

Mean ± std over 3 seeds. All models achieve 15–31 dB SNR improvement. DnCNN delivers the best performance at low-to-mid noise levels with the smallest footprint (0.56 M parameters). See the benchmark documentation for per-level detailed metrics (PSNR, SSIM, MAE, MSE, RMSE).

References

• Ronneberger et al., U-Net: Convolutional Networks for Biomedical Image Segmentation, MICCAI 2015
• He et al., Deep Residual Learning for Image Recognition, CVPR 2016
• Zhang et al., Image Denoising via Deep CNN (DnCNN), IEEE TIP 2017
• Oktay et al., Attention U-Net: Learning Where to Look for the Pancreas, MIDL 2018
• SEG C3 Velocity Model: https://wiki.seg.org/wiki/C3