CD-DCU-Net

PyTorch code and checkpoints for the CD-DCU-Net experiments. The repository contains two parts:

• CD-DCU-Net with Concrete Dropout for uncertainty estimation;
• CD-DCU-Net and four backbone comparison networks without Concrete Dropout.

Contents

The public repository is arranged as follows:

.
|-- CD-DCU-Net(Incorporate CD-module)/
|   |-- CD-DCU-Net.py
|   |-- dataloader.py
|   |-- train_with_α.py
|   |-- train_without_α.py
|   |-- predict.py
|   |-- porposed.pth
|   `-- proposed(without_α).pth
|
|-- dataset_upload/
|   |-- input/
|   |   |-- 1/
|   |   |-- 2/
|   |   `-- ...
|   `-- output/
|       |-- 1/
|       |-- 2/
|       `-- ...
|
`-- Different networks compare(without CD-model)/
    |-- dataloader.py
    |-- train.py
    |-- predict.py
    |-- CD-DCU-Net/
    |   |-- CD-DCU-Net.py
    |   |-- proposed_initial.pth
    |   `-- kaiming_initial.pth
    |-- DFPP-Net/
    |   |-- DFPP-Net.py
    |   `-- DFPP-Net.pth
    |-- DR-U-Net/
    |   |-- DR-U-Net.py
    |   `-- DR-U-Net.pth
    |-- ISP-Net/
    |   |-- ISP-Net.py
    |   `-- ISP-Net.pth
    `-- TS-U-Net/
        |-- TS-U-Net.py
        `-- TS-U-Net.pth

Concrete Dropout Experiments

CD-DCU-Net(Incorporate CD-module)/CD-DCU-Net.py adds Concrete Dropout after each encoder and decoder stage and at the bottleneck. Its output is

{
    "M_mean": M_mean,
    "M_var": M_var,
    "D_mean": D_mean,
    "D_var": D_var,
}

The two training scripts differ in whether the temperature constant is applied to the log-variance term of the NLL loss:

train_with_α.py
    10e-4 * log(var + eps) + (target - mean)^2 / (var + eps)

train_without_α.py
    log(var + eps) + (target - mean)^2 / (var + eps)

The temperature constant is 10e-4 and is applied to log(var + eps). It is not the BNNLoss(alpha=...) argument, which is present in both scripts and scales the full prediction loss.

10e-4 is the literal used by the code and evaluates to 1e-3 in Python. It is left unchanged in this repository.

The corresponding checkpoints are:

Training setting	Checkpoint
With the temperature constant in the log-variance term	porposed.pth
Without the temperature constant	proposed(without_α).pth

The spelling of porposed.pth matches the provided file.

Backbone Comparisons

The comparison directory uses one shared data loader, training script, and prediction script. The available models are:

Model	Model file	Checkpoint
CD-DCU-Net, proposed initialization	CD-DCU-Net/CD-DCU-Net.py	CD-DCU-Net/proposed_initial.pth
CD-DCU-Net, Kaiming initialization	CD-DCU-Net/CD-DCU-Net.py	CD-DCU-Net/kaiming_initial.pth
DFPP-Net	DFPP-Net/DFPP-Net.py	DFPP-Net/DFPP-Net.pth
DR-U-Net	DR-U-Net/DR-U-Net.py	DR-U-Net/DR-U-Net.pth
ISP-Net	ISP-Net/ISP-Net.py	ISP-Net/ISP-Net.pth
TS-U-Net	TS-U-Net/TS-U-Net.py	TS-U-Net/TS-U-Net.pth

CD-DCU-Net is the proposed backbone. kaiming_initial.pth is the Kaiming-initialized result. Checkpoints without an explicit Kaiming label use the proposed initialization setting.

All comparison models return the same interface:

{"M_mean": M_mean, "D_mean": D_mean}

The shared train.py and predict.py are currently configured for DFPP-Net. See Selecting a model before running another backbone.

Environment

Required packages:

torch
torchvision
numpy
scipy
pillow
matplotlib
tqdm
opencv-python
torchsummary

Install PyTorch for the required CUDA version, then install the remaining packages:

pip install numpy scipy pillow matplotlib tqdm opencv-python torchsummary

Multi-GPU training uses the NCCL backend.

Data Format

In the public repository layout, both experiment directories use one shared dataset/ folder at the repository root. The dataset is not duplicated inside either experiment directory. input/ contains the three input images and output/ contains the corresponding labels. Sample folder names are numeric.

dataset/
|-- input/
|   |-- 1/
|   |   |-- image1.png
|   |   |-- image2.png
|   |   `-- image3.png
|   |-- 2/
|   |   |-- image1.png
|   |   |-- image2.png
|   |   `-- image3.png
|   `-- ...
`-- output/
    |-- 1/
    |   `-- label.mat
    |-- 2/
    |   `-- label.mat
    `-- ...

Each label.mat must contain two arrays named M and D. For every sample:

• the three PNG files must be single-channel images of the same size;
• M and D must match the image height and width;
• the numeric folder name must be the same under input/ and output/.

The data loader divides input images and labels by 1000.0. It applies synchronized geometric augmentation only to the training split.

Train/validation split

The training and validation sets are generated from the sample folders rather than stored in separate directories. All current training scripts use an 80/20 split:

train_ratio = 0.8
seed = 42

The split is produced in the following order:

1. Read and sort the folder names under dataset/input/ with Python sorted().
2. Create one index for each sorted folder name.
3. Shuffle the indices with np.random.RandomState(seed).
4. Use the first int(0.8 * number_of_samples) indices for training.
5. Use all remaining indices, approximately 20% of the samples, for validation.

The following script prints the exact sample folders assigned to each set:

import os
import numpy as np

data_dir = "dataset"
train_ratio = 0.8
seed = 42

samples = sorted(os.listdir(os.path.join(data_dir, "input")))
indices = list(range(len(samples)))

rng = np.random.RandomState(seed)
rng.shuffle(indices)

train_size = int(train_ratio * len(samples))
train_samples = [samples[i] for i in indices[:train_size]]
val_samples = [samples[i] for i in indices[train_size:]]

print("Training samples:", train_samples)
print("Validation samples:", val_samples)

Run this script from the repository root. It reproduces the membership generated by create_data_loaders(). The distributed sampler may change the order in which training samples are read during each epoch, but it does not change which samples belong to the training or validation set.

Folder names are sorted as strings because the data loader uses sorted(os.listdir(...)). For example, non-padded names may be ordered as 1, 10, 11, 2. The script above deliberately uses the same rule so that its result matches the training code exactly.

Selecting a Model

The architecture filenames contain hyphens, which cannot be used in a normal Python import. Copy or rename the selected file to an import-safe name before running the scripts.

Concrete Dropout model

Inside CD-DCU-Net(Incorporate CD-module)/, copy or rename CD-DCU-Net.py to cd_dcu_net.py. Then use the same import in both training scripts and in predict.py:

from cd_dcu_net import UNet

The constructor is:

UNet(n_channels=3, bilinear=False)

This replaces from mcunet1 import UNet in the training scripts and from proposed_uncertainty import UNet in the prediction script.

Comparison models

Copy the selected model file into Different networks compare(without CD-model)/, give it the module name shown below, and update both train.py and predict.py.

Model	Module and import	Constructor
CD-DCU-Net	from cd_dcu_net import UNet	UNet(n_channels=3, bilinear=False)
DFPP-Net	from dfppnet import DFPPNet3FrameMD	DFPPNet3FrameMD(n_feat=4, stage0_align=True, stage1_align=False, stage2_align=False, return_aux=False)
DR-U-Net	from dr_unet import UNet	UNet(n_channels=3, out_channels=2, bilinear=False)
ISP-Net	from isp_net import UNet	UNet(n_channels=3, bilinear=False)
TS-U-Net	from ts_unet import UNet	UNet(n_channels=3, bilinear=False)

Change the constructor in two places:

• the main block of train.py;
• Predictor._load_model() in predict.py.

The architecture and constructor must match the checkpoint.

Training

Edit the configuration block at the bottom of the selected training script:

data_dir = "../dataset"
batch_size = 4
train_ratio = 0.8
crop_size = (1360, 672)
num_workers = 16
num_epochs = 300
learning_rate = 0.0001
pretrained_path = "path/to/initialization_checkpoint.pth"

Also set save_dir in the Trainer(...) call. crop_size is ordered as (height, width) and must not exceed the image size.

Training is initialized from a checkpoint:

checkpoint = torch.load(pretrained_path, map_location="cpu")
model.load_state_dict(checkpoint["model_state_dict"], strict=True)

The checkpoint must contain model_state_dict and must match the selected architecture.

Run one of the Concrete Dropout experiments from its directory:

python "train_with_α.py"
python "train_without_α.py"

Run a backbone comparison after updating the shared model import and constructor:

cd "Different networks compare(without CD-model)"
python train.py

For multi-GPU training:

torchrun --standalone --nproc_per_node=2 "train_with_α.py"

The scripts save epoch checkpoints, best_model.pth, training curves, and the five most recent rolling checkpoints. Concrete Dropout training also saves the learned dropout-rate curves.

Prediction

Set the checkpoint, crop size, device, and three image paths at the bottom of predict.py:

predictor = Predictor(
    model_path="path/to/model_checkpoint.pth",
    crop_size=(1360, 672),
    device="cuda",
)

img1_path = "../dataset/input/1/image1.png"
img2_path = "../dataset/input/1/image2.png"
img3_path = "../dataset/input/1/image3.png"

Use the same crop size as training. Set device="cpu" for CPU inference.

Run:

python predict.py

Single prediction

result = predictor.predict_single(img1_path, img2_path, img3_path)
predictor.save_prediction(result, "predictions/single_result.mat", mode="single")

Single prediction returns M and D.

Uncertainty prediction

Monte Carlo prediction is used for the Concrete Dropout model:

results = predictor.predict_with_uncertainty(
    img1_path,
    img2_path,
    img3_path,
    n_samples=50,
)
predictor.save_prediction(results, "predictions/mc_result.mat", mode="mc")

During Monte Carlo prediction, predict.py enables Concrete Dropout, runs n_samples stochastic forward passes, and restores evaluation mode afterward. The output contains mean predictions and model, data, and total uncertainty maps for both M and D.

The comparison backbones return only M_mean and D_mean; use predict_single() for those checkpoints.

Hugging Face Upload

The checkpoint files are large. Track them with Git LFS before pushing the repository:

git lfs install
git lfs track "*.pth"
git add .gitattributes
git add .
git commit -m "Add code and checkpoints"