Add weights for 16 3D cloud instance-segmentation models (8 architectures x 2 conditions)
67305fd verified | license: cc-by-nc-4.0 | |
| tags: | |
| - segmentation | |
| - instance-segmentation | |
| - 3d | |
| - volumetric | |
| - monai | |
| - atmospheric-science | |
| # 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](https://huggingface.co/HoqueMahmudul/convective-cloud-denoising-diffusion-probabilistic-model-3d). | |
| **Project:** [3D Convective-Cloud Generation collection](https://huggingface.co/collections/HoqueMahmudul/3d-convective-cloud-generation-ddpm-190m-params-6a6ba230c89356ff617d481a) | |
| > **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) | |
| ```bash | |
| pip install monai==1.5.2 torch safetensors huggingface_hub | |
| ``` | |
| ```python | |
| 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. | |