HyMN: Hybrid Mamba UY-Network with Mask-Encoded Prior for Laparoscopic Smoke Removal

Official implementation of

HyMN: Hybrid Mamba UY-Network with Mask-Encoded Prior for Laparoscopic Smoke Removal Mohamed Bassem and Marwan Torki. BMVC 2026.

The final model is a two-stage cascade:

1. Smoke Estimation Network (SEN) — a U-Mamba bottleneck network that predicts a dense smoke-density mask from the smoky input.
2. Image Restoration Network (IRN) — a dual-encoder Y-Net that fuses the estimated mask with the degraded image to restore the clean image.

Both stages use strided down/up-sampling. The model is trained with a combination of L1, VGG perceptual, wavelet, and mask-supervision (mask L1) losses.

The complete configuration for this architecture is in configs/cascaded.yaml.

Repository layout

configs/cascaded.yaml        # the final architecture + training recipe
models/
  cascaded_model.py          # cascaded-model (wires the two stages together)
  sen.py                     # Smoke Estimation Network  (registered as 'sen')
  irn.py                     # Image Restoration Network (registered as 'irn')
  models.py, __init__.py     # model registry
  util/                      # building blocks + lr schedulers
datasets/
  paired_data.py             # paired-dataset / mask-paired-dataset
  data_utils/                # image IO, augmentation, transforms
  list/Desmoke/*.txt         # dataset file lists (see "Data" below)
losses/                      # L1 / perceptual / wavelet / mask losses
train.py                     # training entry point
test.py                      # evaluation: reports PSNR / SSIM / LPIPS
fps_test.py                  # inference-speed benchmark
weights/wavelet_weights_c2.pkl  # filter bank required by the wavelet loss
checkpoints/                 # pretrained model (downloaded separately)

Pretrained model

The pretrained HyMN checkpoint is distributed separately (see checkpoints/README.md). Download hymn.pth into checkpoints/, then run test.py / fps_test.py against it.

Installation

python -m venv venv && source venv/bin/activate
pip install torch==2.6.0 torchvision==0.21.0   # match your CUDA build
pip install -r requirements.txt

mamba-ssm compiles CUDA kernels against the installed torch; install torch first. On some systems pip install mamba-ssm --no-build-isolation is needed.

Data

The datasets/list/Desmoke/*.txt files list one image path per line, relative to the repository root, e.g. ../datasets/synthetic/train/gt/.... Place (or symlink) your copy of the dataset so those paths resolve. The final config uses:

Split	Lists
Synthetic train	desmoke-synthetic-train-{clear,smoky,mask}.txt
Synthetic validation	desmoke-synthetic-val-{clear,smoky,mask}.txt
Real validation	desmoke-real-combined-{clear,smoky}.txt

Training

Single GPU:

python train.py --config configs/cascaded.yaml --name my-run --mixed_precision no

Multi-GPU via 🤗 Accelerate (see accelerate_{2,4,8}gpus.yaml):

accelerate launch --config_file accelerate_4gpus.yaml \
  train.py --config configs/cascaded.yaml --name my-run

Checkpoints and logs are written to save/<name>/. A SLURM example is provided in submit.sh.

Evaluation

test.py loads a trained checkpoint and reports PSNR / SSIM / LPIPS (for both the raw and the EMA weights) on a paired dataset:

# Real test set (default)
python test.py --model save/my-run/current_iter-best.pth

# Synthetic validation set
python test.py --model save/my-run/current_iter-best.pth \
  --gt ./datasets/list/Desmoke/desmoke-synthetic-val-clear.txt \
  --lq ./datasets/list/Desmoke/desmoke-synthetic-val-smoky.txt

# Also save restored images + a per-image CSV
python test.py --model save/my-run/current_iter-best.pth --save_dir results/

Inference speed

python fps_test.py --model save/my-run/current_iter-best.pth

Citation

@inproceedings{bassem2026hymn,
  title     = {HyMN: Hybrid Mamba UY-Network with Mask-Encoded Prior for Laparoscopic Smoke Removal},
  author    = {Bassem, Mohamed and Torki, Marwan},
  booktitle = {British Machine Vision Conference (BMVC)},
  year      = {2026},
}

Acknowledgments

This codebase builds on several prior works:

• EAMamba — Efficient All-Around Mamba for image restoration; our training framework and Mamba scanning components derive from it.
• MambaVision — the hybrid Mamba–Attention block used throughout the SEN and IRN backbones.
• U-Mamba — the U-Mamba bottleneck design that the Smoke Estimation Network is based on.

We thank the authors for releasing their code.