GeoSeg3D — Promptable 3D Segmentation of Seismic Geo-Bodies

A 4-channel 3D-UNet (4.77 M parameters) that extracts a single, topologically-connected geological body (salt dome, paleo-channel, reef) from a seismic probability cube, refinable by + / − user clicks at interactive rate (~17 ms warm per click on RTX 5090).

The model is the AI core of the GeoSeg3D prototype; the full pipeline (SEG-Y I/O, prompt heatmaps, single-CC + border-safe morphology post-processing, Gradio web viewer, multi-seed evaluation harness) lives in the GitHub repository:

🔗 Code: https://github.com/Cooler5D/geoseg3d 📓 Walk-through notebook: notebooks/demo.ipynb 📄 Spec-format report: REPORT.md

Model details


Architecture	3D-UNet, base width 32, three down-levels (/1 /2 /4 /8)
Input channels	4 — `[seismic, probability, pos_heatmap, neg_heatmap]`
Bottleneck	identity-init FiLM gate (γ=1, β=0 at start)
Regularisation	Dropout3d (p=0.1) — also used at inference for MC Dropout uncertainty
Parameters	4,766,497
Training data	200 procedural synthetic scenes (cube 80³), 24 val
Targets	3 classes: salt dome / paleo-channel / reef
Loss	BCE + Dice + 3.0·prompt-consistency (margin=10) + 2.0·click-region BCE
Augmentation	RandConv (p=0.3), σ jitter U[4, 12], D4 axial flips
Optimiser	AdamW, lr 1e-3, cosine schedule
Epochs	30 (~16 min on RTX 5090, bf16 mixed precision)
Best ckpt selection	prompted Dice on val (selecting on no-prompt Dice removes prompt incentive)

Provenance: the file promptable.pt ships with a sidecar promptable.pt.geoseg3d.json carrying SHA-256, file size, training timestamp, exact torch.__version__, Python version, platform, and parameter count. The loader recomputes the hash at load and warns on mismatch.

Intended use

Primary: research prototype for the AI R&D task "promptable 3D segmentation of one geo-body from a seismic probability cube". The model is meant to be the AI core of a SAM-like interactive segmentation pipeline for seismic interpretation.

Inputs the model expects (see geoseg3d.PromptableSegmenter.from_cubes):

seismic: 3D float array, axis order (Z, Y, X) = (samples, crosslines, inlines)
probability: 3D float array, same shape, values in [0, 1]
pos_clicks, neg_clicks: lists of (z, y, x) voxel coordinates

Out of scope: standalone segmentation from raw seismic without a probability cube; bodies that span the cube boundary (predict_logits_tiled is the production path for cubes >256³ but is not validated end-to-end on real surveys).

How to use

import torch
from huggingface_hub import hf_hub_download
from geoseg3d.segmenter_v2 import PromptableSegmenter

# Download the canonical checkpoint + provenance sidecar
ckpt_path = hf_hub_download(repo_id="Cooler3D/geoseg3d", filename="promptable.pt")
hf_hub_download(repo_id="Cooler3D/geoseg3d", filename="promptable.pt.geoseg3d.json")

# Build the segmenter (FiLM is auto-detected from the state-dict keys)
seg = PromptableSegmenter.from_cubes(
    seismic=my_seismic_cube,        # (Z, Y, X) float32
    probability=my_probability_cube, # (Z, Y, X) float32, in [0, 1]
    checkpoint=ckpt_path,
)

# Add positive / negative seeds and read the mask back
seg.add_positive((40, 40, 40))
seg.add_negative((10, 10, 10))
mask = seg.get_mask()  # (Z, Y, X) uint8, single connected component

# Optional: MC Dropout uncertainty
mean, std = seg.predict_uncertainty(n_samples=10)

The full Gradio web viewer (click directly on a slice to drop a seed) ships in the GitHub repository:

python -m geoseg3d webdemo --encoder /path/to/promptable.pt
# → http://localhost:7860

Evaluation

Multi-seed synthetic (30 OOD seeds, cube 80³, three target classes, mean ± std):

Configuration	salt	channel	reef
Threshold(prob) + largest CC, no AI	0.988 ± 0.002	0.171 ± 0.295	0.440 ± 0.143
Promptable, no prompts (model alone)	0.995 ± 0.001	0.658 ± 0.346	0.727 ± 0.177
Promptable, +/− clicks, blend OFF	0.995 ± 0.001	0.721 ± 0.229	0.724 ± 0.187
Promptable, +/− clicks, blend ON (UX path)	0.992 ± 0.001	0.384 ± 0.273	0.659 ± 0.187

Headline number: on the hard channel class, threshold-only median IoU is 0.000 (thin sinuous bodies break into disjoint connected components at the 0.5 cut). The trained network alone — no prompts, no inference blend — lifts median to 0.814. This is the project's strongest "AI as core" evidence.

F3 sim-to-real check (Yalaudah 2019, Zechstein salt sub-cube, probability computed from real seismic via class-independent local-std heuristic — labels never enter the probability):

Path	IoU	Dice
Promptable, no prompts (AI alone)	0.709	0.830

Honest scope: real-data validation is salt only. Channel/reef on real seismic (TGS Salt / SEAM / Volve) is on the future-work list.

For paired Wilcoxon p-values, full distribution stats, and ablation discussion see README.md.

Limitations

Real-data validation covers Zechstein salt only; channels and reefs are validated on synthetic only.
The concat-at-input architecture doesn't scale beyond ~256³ cubes on 32 GB VRAM; predict_logits_tiled is the production path but is not validated end-to-end.
The click effect on channel is directional but not statistically established at n=30 (paired Wilcoxon p=0.56, mean Δ +0.063 IoU, median Δ −0.001).
click_region_loss uses GT as target — generalising the synthetic prior locally rather than just enforcing "respect the click". A constraint-style region loss is on the future-work list.

Reproducing

git clone https://github.com/Cooler5D/geoseg3d
cd geoseg3d
uv venv --python 3.12 .venv && .venv\Scripts\activate
uv pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128
uv pip install -e .[viz,dev,web]

python scripts/cache_scenes.py --out data/cache_train --n 200 --cube 80
python scripts/cache_scenes.py --out data/cache_val   --n 24  --cube 80
python -m geoseg3d.train_promptable \
    --cache-train data/cache_train --cache-val data/cache_val \
    --out checkpoints/promptable.pt --epochs 30

Training takes ~16 min on RTX 5090 (Blackwell, sm_120).

Citation

If you use this model in research, please cite:

@misc{mursalov2026geoseg3d,
  author = {Mursalov, Nadir},
  title  = {GeoSeg3D: Promptable 3D Segmentation of Geological Bodies from Seismic Probability Cubes},
  year   = {2026},
  url    = {https://github.com/Cooler5D/geoseg3d},
}

License

MIT — see LICENSE in the source repository.

Downloads last month: -; Downloads are not tracked for this model. How to track