# wangxin81/ounet_jl architecture

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

- Interactive diagram: https://gitdiagram.com/wangxin81/ounet_jl
- Repository: https://github.com/wangxin81/ounet_jl

## Overview

This repository is an underwater-image enhancement framework: it trains a UNet-style model on paired images and targets, then uses the model to enhance input images. Training combines reconstruction, structural, perceptual, and total-variation losses, with PSNR/SSIM validation; the README also identifies UIEB and UFO-120 as datasets. The tree documents additional model variants, image-preparation utilities, and image-quality metrics. Training and loss wiring is supported by sampled source; prediction and several other subsystems are represented by their tree paths, but their detailed wiring could not be verified from the excerpts.

Main entry points: [`train.py`](https://github.com/wangxin81/ounet_jl/blob/main/train.py) (Training entry), [`pre_handle`](https://github.com/wangxin81/ounet_jl/tree/main/pre_handle) (Image corrections), [`utils/train_and_val.py`](https://github.com/wangxin81/ounet_jl/blob/main/utils/train_and_val.py) (Train and validate), [`core/pipelines/my_dataset.py`](https://github.com/wangxin81/ounet_jl/blob/main/core/pipelines/my_dataset.py) (Dataset loader)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Paired image data"]
  node_paired["Paired images"]
  node_dataset["Dataset loader<br/>[my_dataset.py]"]
  node_transforms["Paired transforms<br/>[transforms.py]"]
end

subgraph group_enhance["Enhancement runtime"]
  node_predict["Prediction entry<br/>[predict.py]"]
  node_gamma["Gamma UNet<br/>[unet_gamma.py]"]
  node_modelvariants["Model variants"]
  node_enhanced["Enhanced image"]
end

subgraph group_learn["Training and evaluation"]
  node_trainer["Training entry<br/>[train.py]"]
  node_multigpu["Distributed training<br/>[train_multi_GPU.py]"]
  node_trainval["Train and validate<br/>[train_and_val.py]"]
  node_losses["Joint losses"]
  node_scheduler["LR scheduler<br/>[scheduler.py]"]
  node_metrics["Image metrics<br/>[evaluate.py]"]
end

subgraph group_prepare["Image preparation"]
  node_preprocess["Image corrections"]
  node_clahe["CLAHE<br/>[clahe.py]"]
  node_whitebalance["White balance<br/>[wb.py]"]
  node_stats["Image statistics"]
  node_gc["Gamma correction<br/>[gc.py]"]
end

node_user(("Enhancement user"))

node_paired -->|"loaded by"| node_dataset
node_dataset -->|"transformed by"| node_transforms
node_trainer -->|"runs training"| node_trainval
node_trainer -->|"loads data"| node_dataset
node_trainer -->|"creates model"| node_gamma
node_trainval -->|"calls model"| node_gamma
node_trainval -->|"calculates loss"| node_losses
node_trainer -->|"creates schedule"| node_scheduler
node_trainer -->|"validates with"| node_trainval
node_multigpu -->|"loads data"| node_dataset
node_multigpu -->|"creates schedule"| node_scheduler
node_user -->|"requests enhancement"| node_predict
node_predict -->|"runs model"| node_gamma
node_predict -->|"produces"| node_enhanced
node_trainer -.->|"alternative entry"| node_multigpu
node_metrics -->|"compares images"| node_paired
node_preprocess -->|"includes"| node_clahe
node_preprocess -->|"includes"| node_whitebalance
node_preprocess -->|"includes"| node_stats
node_preprocess -->|"includes"| node_gc

click node_dataset "https://github.com/wangxin81/ounet_jl/blob/main/core/pipelines/my_dataset.py"
click node_transforms "https://github.com/wangxin81/ounet_jl/blob/main/core/pipelines/transforms.py"
click node_predict "https://github.com/wangxin81/ounet_jl/blob/main/predict.py"
click node_gamma "https://github.com/wangxin81/ounet_jl/blob/main/core/models/unet_gamma.py"
click node_modelvariants "https://github.com/wangxin81/ounet_jl/tree/main/core/models"
click node_trainer "https://github.com/wangxin81/ounet_jl/blob/main/train.py"
click node_multigpu "https://github.com/wangxin81/ounet_jl/blob/main/train_multi_GPU.py"
click node_trainval "https://github.com/wangxin81/ounet_jl/blob/main/utils/train_and_val.py"
click node_losses "https://github.com/wangxin81/ounet_jl/tree/main/core/losses"
click node_scheduler "https://github.com/wangxin81/ounet_jl/blob/main/core/optimizer/scheduler.py"
click node_metrics "https://github.com/wangxin81/ounet_jl/blob/main/utils/evaluate.py"
click node_preprocess "https://github.com/wangxin81/ounet_jl/tree/main/pre_handle"
click node_clahe "https://github.com/wangxin81/ounet_jl/blob/main/pre_handle/clahe.py"
click node_whitebalance "https://github.com/wangxin81/ounet_jl/blob/main/pre_handle/wb.py"
click node_stats "https://github.com/wangxin81/ounet_jl/blob/main/pre_handle/compute_mean_std.py"
click node_gc "https://github.com/wangxin81/ounet_jl/blob/main/pre_handle/gc.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_paired,node_dataset,node_transforms,node_user toneBlue
class node_predict,node_gamma,node_modelvariants,node_enhanced toneAmber
class node_trainer,node_multigpu,node_trainval,node_losses,node_scheduler,node_metrics toneMint
class node_preprocess,node_clahe,node_whitebalance,node_stats,node_gc toneRose
```

## Components

### Paired image data

- **Paired images**
- [Dataset loader](https://github.com/wangxin81/ounet_jl/blob/main/core/pipelines/my_dataset.py): `core/pipelines/my_dataset.py`
- [Paired transforms](https://github.com/wangxin81/ounet_jl/blob/main/core/pipelines/transforms.py): `core/pipelines/transforms.py`

### Enhancement runtime

- [Prediction entry](https://github.com/wangxin81/ounet_jl/blob/main/predict.py): `predict.py`
- [Gamma UNet](https://github.com/wangxin81/ounet_jl/blob/main/core/models/unet_gamma.py): `core/models/unet_gamma.py`
- [Model variants](https://github.com/wangxin81/ounet_jl/tree/main/core/models): `core/models`
- **Enhanced image**

### Training and evaluation

- [Training entry](https://github.com/wangxin81/ounet_jl/blob/main/train.py): `train.py`
- [Distributed training](https://github.com/wangxin81/ounet_jl/blob/main/train_multi_GPU.py): `train_multi_GPU.py`
- [Train and validate](https://github.com/wangxin81/ounet_jl/blob/main/utils/train_and_val.py): `utils/train_and_val.py`
- [Joint losses](https://github.com/wangxin81/ounet_jl/tree/main/core/losses): `core/losses`
- [LR scheduler](https://github.com/wangxin81/ounet_jl/blob/main/core/optimizer/scheduler.py): `core/optimizer/scheduler.py`
- [Image metrics](https://github.com/wangxin81/ounet_jl/blob/main/utils/evaluate.py): `utils/evaluate.py`

### Image preparation

- [Image corrections](https://github.com/wangxin81/ounet_jl/tree/main/pre_handle): `pre_handle`
- [CLAHE](https://github.com/wangxin81/ounet_jl/blob/main/pre_handle/clahe.py): `pre_handle/clahe.py`
- [White balance](https://github.com/wangxin81/ounet_jl/blob/main/pre_handle/wb.py): `pre_handle/wb.py`
- [Image statistics](https://github.com/wangxin81/ounet_jl/blob/main/pre_handle/compute_mean_std.py): `pre_handle/compute_mean_std.py`
- [Gamma correction](https://github.com/wangxin81/ounet_jl/blob/main/pre_handle/gc.py): `pre_handle/gc.py`

### Other components

- **Enhancement user**

## Connections

- Paired images → Dataset loader: loaded by
- Dataset loader → Paired transforms: transformed by
- Training entry → Train and validate: runs training
- Training entry → Dataset loader: loads data
- Training entry → Gamma UNet: creates model
- Train and validate → Gamma UNet: calls model
- Train and validate → Joint losses: calculates loss
- Training entry → LR scheduler: creates schedule
- Training entry → Train and validate: validates with
- Distributed training → Dataset loader: loads data
- Distributed training → LR scheduler: creates schedule
- Enhancement user → Prediction entry: requests enhancement
- Prediction entry → Gamma UNet: runs model
- Prediction entry → Enhanced image: produces
- Training entry → Distributed training: alternative entry
- Image metrics → Paired images: compares images
- Image corrections → CLAHE: includes
- Image corrections → White balance: includes
- Image corrections → Image statistics: includes
- Image corrections → Gamma correction: includes