# piggy2009/dm_underwater architecture

> Architecture diagram of the GitHub repository piggy2009/dm_underwater, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-11-10). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/piggy2009/dm_underwater
- Repository: https://github.com/piggy2009/dm_underwater

## Overview

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`](https://github.com/piggy2009/dm_underwater/blob/main/train.py) (Training runner), [`infer.py`](https://github.com/piggy2009/dm_underwater/blob/main/infer.py) (Inference runner), [`model/networks.py`](https://github.com/piggy2009/dm_underwater/blob/main/model/networks.py) (Network builder), [`search_diffusion.py`](https://github.com/piggy2009/dm_underwater/blob/main/search_diffusion.py) (Step search)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_workflow["Training and inference"]
  node_train["Training runner<br/>[train.py]"]
  node_infer["Inference runner<br/>[infer.py]"]
  node_search["Step search"]
  node_config["Underwater config<br/>[underwater.json]"]
  node_results["Enhanced images"]
end

subgraph group_data["Image data"]
  node_prep["Dataset preparation<br/>[prepare_data.py]"]
  node_dataset["LR/HR dataset<br/>[LRHR_dataset.py]"]
  node_datautil["Image transforms<br/>[util.py]"]
  node_datafactory["Dataset factory<br/>[__init__.py]"]
end

subgraph group_model["Diffusion model"]
  node_modelrunner["Model wrapper<br/>[model.py]"]
  node_networks["Network builder<br/>[networks.py]"]
  node_diffusion["Gaussian diffusion<br/>[diffusion.py]"]
  node_unet["Transformer U-Net<br/>[unet.py]"]
  node_style["Style transfer<br/>[style_transfer.py]"]
  node_blocks["Attention blocks<br/>[trans_block_eca.py]"]
end

subgraph group_observability["Evaluation and logging"]
  node_metrics["Image metrics<br/>[metrics.py]"]
  node_logger["Run logging<br/>[logger.py]"]
  node_wandb["WandB reporting<br/>[wandb_logger.py]"]
end

node_user(("User"))

node_user -->|"starts training"| node_train
node_user -->|"runs inference"| node_infer
node_user -->|"searches steps"| node_search
node_user -->|"edits settings"| node_config
node_train -->|"reads config"| node_config
node_train -->|"creates datasets"| node_datafactory
node_train -->|"creates model"| node_modelrunner
node_train -->|"evaluates images"| node_metrics
node_train -->|"sets up logs"| node_logger
node_train -->|"reports metrics"| node_wandb
node_networks -->|"constructs diffusion"| node_diffusion
node_networks -->|"constructs network"| node_unet
node_dataset -->|"transforms images"| node_datautil
node_infer -->|"writes results"| node_results
node_search -.->|"supplies step sequence"| node_infer
node_prep -.->|"prepares image data"| node_dataset

click node_train "https://github.com/piggy2009/dm_underwater/blob/main/train.py"
click node_infer "https://github.com/piggy2009/dm_underwater/blob/main/infer.py"
click node_search "https://github.com/piggy2009/dm_underwater/blob/main/search_diffusion.py"
click node_config "https://github.com/piggy2009/dm_underwater/blob/main/config/underwater.json"
click node_prep "https://github.com/piggy2009/dm_underwater/blob/main/data/prepare_data.py"
click node_dataset "https://github.com/piggy2009/dm_underwater/blob/main/data/LRHR_dataset.py"
click node_datautil "https://github.com/piggy2009/dm_underwater/blob/main/data/util.py"
click node_datafactory "https://github.com/piggy2009/dm_underwater/blob/main/data/__init__.py"
click node_modelrunner "https://github.com/piggy2009/dm_underwater/blob/main/model/model.py"
click node_networks "https://github.com/piggy2009/dm_underwater/blob/main/model/networks.py"
click node_diffusion "https://github.com/piggy2009/dm_underwater/blob/main/model/ddpm_trans_modules/diffusion.py"
click node_unet "https://github.com/piggy2009/dm_underwater/blob/main/model/ddpm_trans_modules/unet.py"
click node_style "https://github.com/piggy2009/dm_underwater/blob/main/model/ddpm_trans_modules/style_transfer.py"
click node_blocks "https://github.com/piggy2009/dm_underwater/blob/main/model/ddpm_trans_modules/trans_block_eca.py"
click node_metrics "https://github.com/piggy2009/dm_underwater/blob/main/core/metrics.py"
click node_logger "https://github.com/piggy2009/dm_underwater/blob/main/core/logger.py"
click node_wandb "https://github.com/piggy2009/dm_underwater/blob/main/core/wandb_logger.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_train,node_infer,node_search,node_config,node_results,node_user toneBlue
class node_prep,node_dataset,node_datautil,node_datafactory toneAmber
class node_modelrunner,node_networks,node_diffusion,node_unet,node_style,node_blocks toneMint
class node_metrics,node_logger,node_wandb toneRose
```

## Components

### Training and inference

- [Training runner](https://github.com/piggy2009/dm_underwater/blob/main/train.py): `train.py`
- [Inference runner](https://github.com/piggy2009/dm_underwater/blob/main/infer.py): `infer.py`
- [Step search](https://github.com/piggy2009/dm_underwater/blob/main/search_diffusion.py): `search_diffusion.py`
- [Underwater config](https://github.com/piggy2009/dm_underwater/blob/main/config/underwater.json): `config/underwater.json`
- **Enhanced images**

### Image data

- [Dataset preparation](https://github.com/piggy2009/dm_underwater/blob/main/data/prepare_data.py): `data/prepare_data.py`
- [LR/HR dataset](https://github.com/piggy2009/dm_underwater/blob/main/data/LRHR_dataset.py): `data/LRHR_dataset.py`
- [Image transforms](https://github.com/piggy2009/dm_underwater/blob/main/data/util.py): `data/util.py`
- [Dataset factory](https://github.com/piggy2009/dm_underwater/blob/main/data/__init__.py): `data/__init__.py`

### Diffusion model

- [Model wrapper](https://github.com/piggy2009/dm_underwater/blob/main/model/model.py): `model/model.py`
- [Network builder](https://github.com/piggy2009/dm_underwater/blob/main/model/networks.py): `model/networks.py`
- [Gaussian diffusion](https://github.com/piggy2009/dm_underwater/blob/main/model/ddpm_trans_modules/diffusion.py): `model/ddpm_trans_modules/diffusion.py`
- [Transformer U-Net](https://github.com/piggy2009/dm_underwater/blob/main/model/ddpm_trans_modules/unet.py): `model/ddpm_trans_modules/unet.py`
- [Style transfer](https://github.com/piggy2009/dm_underwater/blob/main/model/ddpm_trans_modules/style_transfer.py): `model/ddpm_trans_modules/style_transfer.py`
- [Attention blocks](https://github.com/piggy2009/dm_underwater/blob/main/model/ddpm_trans_modules/trans_block_eca.py): `model/ddpm_trans_modules/trans_block_eca.py`

### Evaluation and logging

- [Image metrics](https://github.com/piggy2009/dm_underwater/blob/main/core/metrics.py): `core/metrics.py`
- [Run logging](https://github.com/piggy2009/dm_underwater/blob/main/core/logger.py): `core/logger.py`
- [WandB reporting](https://github.com/piggy2009/dm_underwater/blob/main/core/wandb_logger.py): `core/wandb_logger.py`

### Other components

- **User**

## Connections

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