README.md

---
license: mit
tags:
- seismic
- interpolation
- trace-interpolation
- reconstruction
- unet
- resunet
- attention-unet
- pytorch
- geophysics
library_name: pytorch
#datasets:
#- SEG C3 synthetic seismic dataset
metrics:
- mse
- snr
- psnr
- ssim
---

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

```text
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:

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

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

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