piggy2009/dm_underwater

How piggy2009/dm_underwater works

DM_underwater is an underwater-image enhancement project built around a transformer-based diffusion model. Its documented workflows are training from configured image data and inference from a pretrained checkpoint; inference produces enhanced images, while an evolutionary search tool can find diffusion time-step sequences. The graph also includes dataset preparation and validation/logging capabilities. Training-to-data/model and reporting relationships are visible in the training excerpt; several inference and model-internal connections are documented only at a high level, so they are not expanded into speculative call chains.

Main entry points: train.py Training runner · infer.py Inference runner · networks.py Network builder · search_diffusion.py Step search

Components 19

Training and inference

Training runner train.py
Inference runner infer.py
Step search search_diffusion.py
Underwater config config/underwater.json
Enhanced images

Image data

Dataset preparation data/prepare_data.py
LR/HR dataset data/LRHR_dataset.py
Image transforms data/util.py
Dataset factory data/__init__.py

Diffusion model

Model wrapper model/model.py
Network builder model/networks.py
Gaussian diffusion model/ddpm_trans_modules/diffusion.py
Transformer U-Net model/ddpm_trans_modules/unet.py
Style transfer model/ddpm_trans_modules/style_transfer.py
Attention blocks model/ddpm_trans_modules/trans_block_eca.py

Evaluation and logging

Image metrics core/metrics.py
Run logging core/logger.py
WandB reporting core/wandb_logger.py

Other components

User

Connections 16

User → Training runner: starts training README.md
User → Inference runner: runs inference README.md
User → Step search: searches steps README.md
User → Underwater config: edits settings README.md
Training runner → Underwater config: reads config train.py
Training runner → Dataset factory: creates datasets train.py
Training runner → Model wrapper: creates model train.py
Training runner → Image metrics: evaluates images train.py
Training runner → Run logging: sets up logs train.py
Training runner → WandB reporting: reports metrics train.py
Network builder → Gaussian diffusion: constructs diffusion model/networks.py
Network builder → Transformer U-Net: constructs network model/networks.py
LR/HR dataset → Image transforms: transforms images data/LRHR_dataset.py
Inference runner → Enhanced images: writes results README.md
Step search → Inference runner: supplies step sequence README.md
Dataset preparation → LR/HR dataset: prepares image data

Drawn by GitDiagram on November 10, 2026 from the repository's file tree, README and source files. Read as Markdown ·