How wangxin81/ounet_jl works
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 Training entry · pre_handle Image corrections · train_and_val.py Train and validate · my_dataset.py Dataset loader
Components 19
Paired image data
- Paired images
- Dataset loader
core/pipelines/my_dataset.py - Paired transforms
core/pipelines/transforms.py
Enhancement runtime
- Prediction entry
predict.py - Gamma UNet
core/models/unet_gamma.py - Model variants
core/models - Enhanced image
Training and evaluation
- Training entry
train.py - Distributed training
train_multi_GPU.py - Train and validate
utils/train_and_val.py - Joint losses
core/losses - LR scheduler
core/optimizer/scheduler.py - Image metrics
utils/evaluate.py
Image preparation
- Image corrections
pre_handle - CLAHE
pre_handle/clahe.py - White balance
pre_handle/wb.py - Image statistics
pre_handle/compute_mean_std.py - Gamma correction
pre_handle/gc.py
Other components
- Enhancement user
Connections 20
- 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