{"id":23955559,"url":"https://github.com/tnwei/waternet","last_synced_at":"2025-07-22T23:04:11.758Z","repository":{"id":39671111,"uuid":"492883376","full_name":"tnwei/waternet","owner":"tnwei","description":"Modern replication of WaterNet from \"An Underwater Image Enhancement Benchmark Dataset and Beyond\", IEEE TIP 2019","archived":false,"fork":false,"pushed_at":"2024-06-16T06:11:09.000Z","size":9117,"stargazers_count":54,"open_issues_count":2,"forks_count":10,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-04-13T03:43:36.898Z","etag":null,"topics":["deep-learning","pytorch","underwater-image-enhancement","underwater-image-restoration"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/tnwei.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2022-05-16T14:59:28.000Z","updated_at":"2025-04-09T01:32:57.000Z","dependencies_parsed_at":"2025-01-06T15:39:51.003Z","dependency_job_id":null,"html_url":"https://github.com/tnwei/waternet","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/tnwei/waternet","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tnwei%2Fwaternet","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tnwei%2Fwaternet/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tnwei%2Fwaternet/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tnwei%2Fwaternet/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/tnwei","download_url":"https://codeload.github.com/tnwei/waternet/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tnwei%2Fwaternet/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":266586905,"owners_count":23952205,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-07-22T02:00:09.085Z","response_time":66,"last_error":null,"robots_txt_status":null,"robots_txt_updated_at":null,"robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","pytorch","underwater-image-enhancement","underwater-image-restoration"],"created_at":"2025-01-06T15:35:37.009Z","updated_at":"2025-07-22T23:04:11.737Z","avatar_url":"https://github.com/tnwei.png","language":"Jupyter Notebook","funding_links":[],"categories":["Jupyter Notebook"],"sub_categories":[],"readme":"# Replication of WaterNet, IEEE TIP 2019\n\nModern replication of WaterNet in Pytorch from \"An Underwater Image Enhancement Benchmark Dataset and Beyond\", IEEE TIP 2019. \n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tnwei/waternet/blob/main/colab-example-waternet.ipynb)\n\n![](docs/frames-0968-small.jpeg)\n\n![](docs/173-small.jpeg)\n\n_Source data: [A Surprise Diving Encounter with a Giant Humpback Whale, YouTube](https://www.youtube.com/watch?v=vAZkzV41sCc)_\n\nUpdate 2 years later: For the messy experimental code used to perform this replication, see the [dev branch](https://github.com/tnwei/waternet/tree/dev)\n\n## Usage in commandline\n\nGit clone this repo, then do either `pip install -r requirements.txt` or `conda env create -f env.yaml`. Download weights from the training replication results section if needed.\n\nUse `inference.py`: \n\n```bash\n$ python inference.py --help\nusage: inference.py [-h] [--source SOURCE] [--weights WEIGHTS] [--name NAME] [--show-split]\n\noptions:\n-h, --help         show this help message and exit\n--source SOURCE    Path to input image/video/directory, supports image formats: bmp, jpg, jpeg, png, gif, and video formats: mp4, mpeg, avi\n--weights WEIGHTS  (Optional) Path to model weights, defaults to waternet_exported_state_dict-daa0ee.pt. Auto-downloads pretrained weights if not available.\n--name NAME        (Optional) Subfolder name to save under `./output`.\n--show-split       (Optional) Left/right of output is original/processed. Adds before/after watermark. \n```\n\nExample:\n\n```bash\n$ python inference.py --source data/example/173.jpeg\nUsing device: cuda\nNo weights specified in --weights, using default: waternet_exported_state_dict-daa0ee.pt\nDownloading: \"https://www.dropbox.com/s/j8ida1d86hy5tm4/waternet_exported_state_dict-daa0ee.pt?dl=1\" to /home/tnwei/projects/waternet/waternet_exported_state_dict-daa0ee.pt\nTotal images/videos: 1\nSaved output to /home/tnwei/projects/waternet/output/40!  \n```\n\n## Usage in Python via torchhub\n\nLoad pretrained weights for your own usage via torchhub:\n\n```python\nimport torch\nimport cv2\n\n# Load from torchhub\npreprocess, postprocess, model = torch.hub.load('tnwei/waternet', 'waternet')\nmodel.eval();\n\n# Load one image using OpenCV\nim = cv2.imread(\"example.png\")\nrgb_im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)\n\n# Inference -\u003e return numpy array (1, 3, H, W)\nrgb_ten, wb_ten, he_ten, gc_ten = preprocess(rgb_im)\nout_ten = model(rgb_ten, wb_ten, he_ten, gc_ten)\nout_im = postprocess(out_ten)\n```\n\n## Training\n\nGit clone this repo, then do either `pip install -r requirements.txt` or `conda env create -f env.yaml`. \n\nDownload the UIEB data raws and references from the [project website](https://li-chongyi.github.io/proj_benchmark.html). Unzip the RAR files, and place `raw-890` and `reference-890` in `data/`.\n\nUse `train.py`:\n\n```bash\n$ python train.py --help\nusage: train.py [-h] [--epochs EPOCHS] [--batch-size BATCH_SIZE] [--height HEIGHT] [--width WIDTH] [--weights WEIGHTS]\n\noptional arguments:\n  -h, --help            show this help message and exit\n  --epochs EPOCHS       (Optional) Num epochs, defaults to 400\n  --batch-size BATCH_SIZE\n                        (Optional) Batch size, defaults to 16\n  --height HEIGHT       (Optional) Image height, defaults to 112\n  --width WIDTH         (Optional) Image width, defaults to 112\n  --weights WEIGHTS     (Optional) Starting weights for training\n  --seed SEED           (Optional) Seed to pass to `torch.random_seed()` for reproducibility, defaults to None i.e.\n                        random\n```\n\nExample:\n\n```bash\n$ python train.py\nUsing device: cuda\nEpoch 1/400: 100%|####################| 50/50 [01:11\u003c00:00,  1.43s/it, loss=3.39e+3]\nValidation: 100%|####################| 6/6 [00:05\u003c00:00,  1.01it/s]\n    Train || mse: 2.4e+03   ssim: 0.458   psnr: 15.2   perceptual_loss: 4.93e+04   loss: 4.86e+03\n    Val   || mse: 1.24e+03   ssim: 0.592   psnr: 17.2   perceptual_loss: 6.98e+03\n\n...\n...\n\nEpoch 400/400: 100%|####################| 50/50 [01:02\u003c00:00,  1.25s/it, loss=873]\nValidation: 100%|####################| 6/6 [00:05\u003c00:00,  1.08it/s]\n    Train || mse: 435   ssim: 0.919   psnr: 21.9   perceptual_loss: 8.54e+03   loss: 862\n    Val   || mse: 446   ssim: 0.915   psnr: 21.7   perceptual_loss: 1.94e+03\n\nMetrics and weights saved to /home/tnwei/projects/waternet/training/0\nTotal time: 28182.797650701832s\n```\n\n## Scoring\n\n`score.py` was derived from `train.py` to evaluate trained network weights on the UIEB dataset.\n\n## Paper summary\n\nIn [this paper](https://arxiv.org/abs/1901.05495), the authors constructed the Underwater Image Enhancement Benchmark dataset, which provides one-to-one pairings of original and reference underwater images. This was done by restoring underwater images via a multitude of SOTA methods, and have volunteers thoroughly vet and select the best ones as reference. he UIEB dataset can be downloaded from the [author's website](https://li-chongyi.github.io/proj_benchmark.html).\n\nThe UIEB dataset is subsequently used to train WaterNet, a fully convolutional network to perform image-to-image translation for reversing the effects of underwater light scattering. WaterNet is a gated fusion network that consists of two branches: input refinement and confidence map generation. First, transformed versions of the original input image are derived (namely, white balance correction, gamma correction and histogram equalization). \n\n- Input refinement: For each transformed image, a stack of conv layers generate refined inputs from the original image and the transformed image. \n- Confidence map generation: The original input image and the transformed images are forward propagated through a stack of conv layers to generate three confidence maps, one for each refined input from the previous branch.\n\nWaterNet's final output image is a weighted sum of the refined inputs with the associated confidence maps. \n\n\u003c!-- TODO: arch diagram below: --\u003e\n\nSource code for WaterNet is [shared on github](https://github.com/Li-Chongyi/Water-Net_Code), which is what this repo seeks to replicate. \n\n## Replication notes\n\nThe original repo uses MATLAB for generating transformed images, and Tensorflow 1 for the neural network itself.  \n\nHere's what I've done in this repo:\n\n- Reimplemented the network in Pytorch, converted the original pretrained weights.\n- Replaced MATLAB image transforms in Python. The original repo later added image transforms in Python, but they are not equivalent to the original MATLAB implementation. Gamma correction and white balance correction logic are reimplemented as-is. Histogram equalization however relied on MATLAB's built-in `adapthisteq` function, and thus can't be reimplemented faithfully. Ended up using OpenCV's implementation of contrast-limited adaptive histogram equalization. Testing revealed that both implementations are close but not equivalent. \n- Training and inferencing CLI.\n- Training replication (see below)\n\n## Training replication results\n\nUsing dataset split with random seed 0 (replicated weights shared via Dropbox):\n\n| Method | MSE  (x1e3) ↓ | PSNR (dB) ↑ | SSIM  ↑ |\n| ------ | -------------- | ----------- | ------- |\n| Original WaterNet weights, from paper | 0.7976 | 19.1130 | 0.7971\n| Original WaterNet weights, scored in this repo | 0.6260 | 20.3766 | 0.8415 |\n| Replicated WaterNet, 400 epochs / [link to weights](https://www.dropbox.com/s/0nzt1jowxavbkwa/replicated-waternet-20220528.pt?dl=0) | 0.4674 | 21.7300 | 0.9183 |\n| **Replicated WaterNet, 400 epochs @ 256x256** / [link to weights](https://www.dropbox.com/s/g3ttnbjgi1pjp0v/replicated-waternet-256x256-20220529.pt?dl=0) | **0.4574** | **21.7964** | **0.9195** |\n\n## Acknowledgements\n\nPlumbing for structuring model training and inferencing inspired by [ultralytics/yolov5](https://github.com/ultralytics/yolov5/)\n\n------------------\n\n*Project based on the [cookiecutter-datascience-lite](https://github.com/tnwei/cookiecutter-datascience-lite/) template.*\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftnwei%2Fwaternet","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ftnwei%2Fwaternet","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftnwei%2Fwaternet/lists"}