The code is based on MedSegDiff
We aim to utilize the capability of diffusion models for the smoke segmentation task to segment and capture the elusive boundaries of smoke. The method using denoising networks with U-Net as the main architecture, processing the segmentation map as input and the original image as a condition. Our experimental results on SMOKE5K show that our approach produces competitive visual results.
Requirement installation: pip install -r requirement.txt
Training python scripts/segmentation_training.py --data_dir dataset/train --out_dir output/ --image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --lr 5e-5 --batch_size 8
Sampling python scripts/segmentation_sample.py --data_dir dataset/test --out_dir output/ --model_path output/model.pth --image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False
Evaluation python scripts/segmentation_eval.py --inp_pth folder you save prediction images --out_pth folder you save ground truth images
The checkpoint of the model use in this work can be downloaded here: Google drive
tomeramit/SegDiff, WuJunde/MedSegDiff, hojonathanho/diffusion, openai/guided-diffusion, MIC-DKFZ/nnUNet, openai/improved-diffusion