Geoseg3d

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

Model details

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 >2563 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 803, three target classes, mean ± std):

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

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

1. Real-data validation covers Zechstein salt only; channels and reefs are validated on synthetic only.
2. The concat-at-input architecture doesn't scale beyond ~2563 cubes on 32 GB VRAM; predict_logits_tiled is the production path but is not validated end-to-end.
3. 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).
4. 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.