HoqueMahmudul/convective-cloud-3d-instance-segmentation

3D Convective-Cloud Instance Segmentation — 8 Architectures

Weights for 16 volumetric segmentation networks that separate individual convective clouds in 3D atmospheric-simulation volumes. Eight architectures, each trained under two data conditions: real volumes only, and real volumes plus synthetic volumes from the companion 3D convective-cloud DDPM.

Project: 3D Convective-Cloud Generation collection

Status: research artifact from work with an associated manuscript in preparation. This repository contains weights only. Training code, methodology and evaluation are not included here and will appear in the publication.

Contents

<architecture>/<condition>/model.safetensors   # weights
<architecture>/<condition>/config.json         # constructor args, normalization, inference + post-processing settings

condition is real_only or real_plus_synth. All runs use seed 42; the saved weights are the checkpoint with the best validation panoptic quality.

Architecture	Folder	Implementation	Parameters	Loadable from this repo alone
3D U-Net	unet	monai.networks.nets.UNet	4.81M	Yes
Attention U-Net 3D	attentionunet	monai.networks.nets.AttentionUnet	5.91M	Yes
SegResNet	segresnet	monai.networks.nets.SegResNet	4.70M	Yes
SwinUNETR	swinunetr	monai.networks.nets.SwinUNETR	15.70M	Yes
ResUNet3D	resunet3d	custom	5.91M	No: definition not released
Residual Attention U-Net 3D	resattnunet3d	custom	5.93M	No: definition not released
SegFormer3D	segformer3d	custom (not Perera et al. 2024)	2.64M	No: definition not released
U-Mamba	umamba	custom, built on mamba_ssm kernels (not Ma et al. 2024)	6.50M	No: definition not released

The four MONAI models can be rebuilt from config.json and used directly. The four custom models are published for archival and reproducibility; their architecture definitions will be released with the publication.

Model contract

Input	(B, 1, D, H, W), single-channel volume, 0.5 km isotropic voxels
Normalization	x = clip((v - 302) / (1921 - 302), 0, 1) (background maps to 0)
Output	(B, 3, D, H, W) logits: channel 0 = foreground, 1 = seeds, 2 = boundary
Training patch	(32, 160, 160) (D, H, W)
Full-volume inference	MONAI sliding_window_inference, roi_size=(32,160,160), overlap=0.5, mode="gaussian" blending on logits, then sigmoid

Instances are obtained by thresholding foreground and seeds, then running a marker-controlled watershed. config.json lists two threshold sets: global_operating_point (shared by all models) and calibrated_operating_point (tuned per model on the validation volume).

Usage (MONAI architectures)

pip install monai==1.5.2 torch safetensors huggingface_hub

import json, torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from monai.networks.nets import SegResNet          # or UNet, AttentionUnet, SwinUNETR
from monai.inferers import sliding_window_inference

repo, sub = "HoqueMahmudul/convective-cloud-3d-instance-segmentation", "segresnet/real_only"
cfg = json.load(open(hf_hub_download(repo, f"{sub}/config.json")))
weights = load_file(hf_hub_download(repo, f"{sub}/model.safetensors"))

# Weights were saved from a thin wrapper, so keys carry a "network." prefix.
weights = {k.removeprefix("network."): v for k, v in weights.items()}

model = SegResNet(spatial_dims=3, in_channels=1, out_channels=3, **cfg["model_kwargs"])
model.load_state_dict(weights, strict=True)
model.eval()

lo, hi = cfg["input_normalization"]["lo"], cfg["input_normalization"]["hi"]
volume = torch.from_numpy(raw).float()             # raw: (D, H, W) numpy array
x = ((volume - lo) / (hi - lo)).clamp(0, 1)[None, None]

with torch.no_grad():
    logits = sliding_window_inference(x, roi_size=(32, 160, 160), sw_batch_size=1,
                                      predictor=model, overlap=0.5, mode="gaussian")
probs = torch.sigmoid(logits)[0]                   # (3, D, H, W): foreground, seeds, boundary

This loading recipe was verified with MONAI 1.5.2 and PyTorch 2.5.1 for all four MONAI architectures (strict=True).

Limitations

• Trained on a small number of real simulation volumes from a single simulation setup; not validated on other models, resolutions or observational data.
• One training seed per configuration.
• Not a weather-forecasting or nowcasting tool.

License

CC BY-NC 4.0, same as the companion DDPM.