GeoBrain/seismic-interpolation

Seismic Trace Interpolation Benchmark

This repository provides pretrained PyTorch checkpoints for deep-learning-based seismic trace interpolation on pre-stack shot gathers. The models are trained on the SEG C3 synthetic seismic dataset to reconstruct complete shot gathers from inputs with missing traces.

The task is formulated as a supervised reconstruction problem: given a masked shot gather, the network predicts the corresponding fully sampled gather and is trained using mean squared error against the original complete data.

Task Description

Seismic trace interpolation aims to recover missing receiver traces caused by sparse, irregular, or incomplete acquisition. In this benchmark, missing traces are simulated by applying binary trace masks to complete shot gathers.

Given:

• a complete shot gather x
• a binary trace mask m
• a masked input x_masked = x * m

the model learns:

f(x_masked) -> x

where x is the original fully sampled shot gather.

Dataset

• Source: SEG C3 pre-stack synthetic seismic data
• Number of shot gathers: 9 regular shot gathers
• Geometry: 201 traces x 625 time samples per shot
• Sampling interval: dt = 2 ms
• Split: random patch-level 90:10 train/validation split

Preprocessing

The data are preprocessed using:

1. Spherical divergence correction with power 1.2
2. Global max-absolute normalization to [-1, 1]
3. Overlapping 2D patch extraction
4. Patch size: 128 traces x 256 time samples
5. Patch overlap: 50%

Masking Modes

Two masking strategies are used to simulate different acquisition scenarios.

Mode	Description
uniform	Keeps traces at a fixed stride, simulating regular decimation
random	Randomly removes traces according to the specified missing ratio, simulating irregular acquisition gaps

Missing Ratios and Checkpoint Counts

The following table shows the number of available trained checkpoints for each model and missing-ratio setting. Counts include available seed and mask-mode combinations.

Missing Ratio	Attention UNet	UNet Base	ResUNet
30%	6	6	6
50%	6	6	6
70%	6	6	6

Total checkpoint inventory:

Model	Directory	Checkpoints
Attention UNet	unet_atten	18
UNet Base	unet_base	18
ResUNet	unet_res	18

Model Architectures

All models use a U-Net-style encoder-decoder architecture with base channel width 32 and depth 4.

Model	Key	Description
UNet Base	unet_base	Standard U-Net with encoder-decoder skip connections
ResUNet	unet_res	U-Net variant using residual convolutional blocks
Attention UNet	unet_atten	U-Net with additive attention gates on skip connections

Repository Structure

Each leaf directory corresponds to one trained experiment:

models/
├── unet_base/
│   └── uniform/
│       └── miss50/
│           ├── seed42/
│           │   ├── best.pt
│           │   └── config.yaml
│           ├── seed43/
│           └── seed44/
├── unet_res/
│   └── ...
└── unet_atten/
    └── ...

Each best.pt file contains the checkpoint with the lowest validation loss for that experiment. The corresponding config.yaml stores the training and architecture configuration.

Training Details

Hyperparameter	Value
Loss	Mean squared error
Optimizer	AdamW
Learning rate	1e-3
Weight decay	1e-4
Scheduler	Cosine annealing
Minimum learning rate	1e-6
Epochs	200
Gradient clipping	1.0 max norm
Batch size	192
Distributed training	2 x NVIDIA RTX 4090, DDP
Masking modes	uniform, random
Missing ratios	30%, 50%, 70%
Seeds	42, 43, 44

Usage

import torch
from huggingface_hub import hf_hub_download

repo_id = "GeoBrain/seismic-interpolation"

model_key = "unet_res"      # one of: unet_base, unet_res, unet_atten
mask_mode = "uniform"       # one of: uniform, random
ratio = "50"                # one of: 30, 50, 70
seed = "42"

ckpt_path = hf_hub_download(
    repo_id=repo_id,
    filename=f"models/{model_key}/{mask_mode}/miss{ratio}/seed{seed}/best.pt",
)

state_dict = torch.load(
    ckpt_path,
    map_location="cpu",
    weights_only=True,
)

print(state_dict.keys())

Download the corresponding configuration file:

config_path = hf_hub_download(
    repo_id=repo_id,
    filename=f"models/{model_key}/{mask_mode}/miss{ratio}/seed{seed}/config.yaml",
)

The config.yaml file should be used to reconstruct the exact model architecture before loading the checkpoint.

Intended Use

These checkpoints are intended for:

• seismic trace interpolation research
• benchmarking deep-learning reconstruction models
• comparing U-Net, ResUNet, and Attention U-Net architectures
• studying regular versus irregular trace-missing patterns
• transfer learning or fine-tuning on related seismic interpolation tasks

Limitations

• The models are trained on synthetic SEG C3 data and may not generalize directly to field data without fine-tuning.
• The benchmark uses patch-level random train/validation splitting, not a shot-level holdout split.
• The checkpoints are optimized for MSE reconstruction quality and may not preserve all geophysical attributes equally well.
• Performance may vary for acquisition geometries, frequency content, noise levels, or missing-trace patterns that differ from the training setup.

Citation

If you use these checkpoints or the benchmark setup, please cite the relevant architecture papers and the SEG C3 dataset source.

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