Upload 14 files

.gitignore

@@ -0,0 +1,129 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,3 +1,60 @@
----
-license: gpl-3.0
----
+## [Low-light Image Enhancement via Breaking Down the Darkness](https://arxiv.org/abs/2111.15557)
+by Xiaojie Guo, Qiming Hu.
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mingcv/Bread/blob/main/bread_demo_uploader.ipynb) (Online Demo)
+
+<!-- ![figure_tease](https://github.com/mingcv/Bread/blob/main/figures/figure_tease.png) --> 
+
+### 1. Dependencies
+* Python3
+* PyTorch>=1.0
+* OpenCV-Python, TensorboardX
+* NVIDIA GPU+CUDA
+
+### 2. Network Architecture
+![figure_arch](https://github.com/mingcv/Bread/blob/main/figures/Bread_architecture_full.png)
+
+### 3. Data Preparation
+
+#### 3.1. Training dataset
+* 485 low/high-light image pairs from our485 of [LOL dataset](https://daooshee.github.io/BMVC2018website/), each low image of which is augmented by our [exposure_augment.py](https://github.com/mingcv/Bread/blob/main/exposure_augment.py) to generate 8 images under different exposures. ([Download Link for Augmented LOL](https://drive.google.com/file/d/1gyX2kYJWuj3C00eobd49MjRuNbZ29dqN/view?usp=sharing))
+* To train the MECAN (if it is desired), 559 randomly-selected multi-exposure sequences from [SICE](https://github.com/csjcai/SICE) are adopted ([Download Link for a resized version](https://drive.google.com/file/d/1OTNP-QJ3Nade5my04A2iYVTY77IQBEMf/view?usp=sharing)).
+
+#### 3.2. Tesing dataset
+The images for testing can be downloaded in [this link](https://github.com/mingcv/Bread/releases/download/checkpoints/data.zip).
+
+<!-- * 15 low/high-light image pairs from eval15 of [LOL dataset](https://daooshee.github.io/BMVC2018website/).
+* 44 low-light images from DICM.
+* 8 low-light images from NPE.
+* 24 low-light images from VV. -->
+
+### 4. Usage
+
+#### 4.1. Training 
+* Multi-exposure data synthesis: ```python exposure_augment.py```
+* Train IAN: ```python train_IAN.py -m IAN --comment IAN_train --batch_size 1 --val_interval 1 --num_epochs 500 --lr 0.001 --no_sche```
+* Train ANSN: ```python train_ANSN.py -m1 IAN -m2 ANSN --comment ANSN_train --batch_size 1 --val_interval 1 --num_epochs 500 --lr 0.001 --no_sche -m1w ./checkpoints/IAN_335.pth```
+* Train CAN: ```python train_CAN.py -m1 IAN -m3 FuseNet --comment CAN_train --batch_size 1 --val_interval 1 --num_epochs 500 --lr 0.001 --no_sche -m1w ./checkpoints/IAN_335.pth```
+* Train MECAN on SICE: ```python train_MECAN.py -m FuseNet --comment MECAN_train --batch_size 1 --val_interval 1 --num_epochs 500 --lr 0.001 --no_sche```
+* Finetune MECAN on SICE and LOL datasets:  ```python train_MECAN_finetune.py -m FuseNet --comment MECAN_finetune --batch_size 1 --val_interval 1 --num_epochs 500 --lr 1e-4 --no_sche -mw ./checkpoints/FuseNet_MECAN_for_Finetuning_404.pth```
+
+#### 4.2. Testing 
+* *\[Tips\]: Using gamma correction for evaluation with parameter --gc; Show extra intermediate outputs with parameter --save_extra*
+* Evaluation: ```python eval_Bread.py -m1 IAN -m2 ANSN -m3 FuseNet -m4 FuseNet --mef --comment Bread+NFM+ME[eval] --batch_size 1 -m1w ./checkpoints/IAN_335.pth -m2w ./checkpoints/ANSN_422.pth -m3w ./checkpoints/FuseNet_MECAN_251.pth -m4w ./checkpoints/FuseNet_NFM_297.pth```
+* Testing: ```python test_Bread.py -m1 IAN -m2 ANSN -m3 FuseNet -m4 FuseNet --mef --comment Bread+NFM+ME[test] --batch_size 1 -m1w ./checkpoints/IAN_335.pth -m2w ./checkpoints/ANSN_422.pth -m3w ./checkpoints/FuseNet_MECAN_251.pth -m4w ./checkpoints/FuseNet_NFM_297.pth```
+* Remove NFM: ```python test_Bread_NoNFM.py -m1 IAN -m2 ANSN -m3 FuseNet --mef -a 0.10 --comment Bread+ME[test] --batch_size 1 -m1w ./checkpoints/IAN_335.pth -m2w ./checkpoints/ANSN_422.pth -m3w ./checkpoints/FuseNet_MECAN_251.pth```
+
+#### 4.3. Trained weights
+Please refer to [our release](https://github.com/mingcv/Bread/releases/tag/checkpoints). 
+
+### 5. Quantitative comparison on eval15
+![table_eval](https://github.com/mingcv/Bread/blob/main/figures/table_eval.png)
+
+### 6. Visual comparison on eval15
+![figure_eval](https://github.com/mingcv/Bread/blob/main/figures/figure_eval.png)
+
+### 7. Visual comparison on DICM
+![figure_test_dicm](https://github.com/mingcv/Bread/blob/main/figures/figure_test_dicm.png)
+
+### 8. Visual comparison on VV and MEF-DS
+![figure_test_vv_mefds](https://github.com/mingcv/Bread/blob/main/figures/figure_test_vv_mefds.png)

eval_Bread.py

@@ -0,0 +1,218 @@
+import argparse
+import os
+
+import kornia
+import torch
+import torch.nn.functional as F
+import tqdm
+from torch import nn
+from torch.utils.data import DataLoader
+
+import models
+from datasets import LowLightDataset
+from tools import saver, mutils
+from models import PSNR, SSIM
+import numpy as np
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Breaking Downing the Darkness')
+    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus being used')
+    parser.add_argument('--num_workers', type=int, default=12, help='num_workers of dataloader')
+    parser.add_argument('--batch_size', type=int, default=4, help='The number of images per batch among all devices')
+    parser.add_argument('-m1', '--model1', type=str, default='IANet', help='Model1 Name')
+    parser.add_argument('-m2', '--model2', type=str, default='NSNet', help='Model2 Name')
+    parser.add_argument('-m3', '--model3', type=str, default='FuseNet', help='Model3 Name')
+    parser.add_argument('-m4', '--model4', type=str, default=None, help='Model4 Name')
+
+    parser.add_argument('-m1w', '--model1_weight', type=str, default=None, help='Model weight of IAN')
+    parser.add_argument('-m2w', '--model2_weight', type=str, default=None, help='Model weight of ANSN')
+    parser.add_argument('-m3w', '--model3_weight', type=str, default=None, help='Model weight of CAN')
+    parser.add_argument('-m4w', '--model4_weight', type=str, default=None, help='Model weight of NFM')
+
+    parser.add_argument('--mef', action='store_true', help='using color adation based MEF data or not')
+    parser.add_argument('--gc', action='store_true', help='using gamma correction or not')
+    parser.add_argument('--save_extra', action='store_true', help='save intermediate outputs or not')
+
+    parser.add_argument('--comment', type=str, default='default',
+                        help='Project comment')
+
+    parser.add_argument('--alpha', '-a', type=float, default=0.10)
+    parser.add_argument('--lr', type=float, default=0.01)
+    parser.add_argument('--optim', type=str, default='adamw', help='select optimizer for training, '
+                                                                   'suggest using \'admaw\' until the'
+                                                                   ' very final stage then switch to \'sgd\'')
+    parser.add_argument('--data_path', type=str, default='./data/LOL/eval',
+                        help='the root folder of dataset')
+    parser.add_argument('--log_path', type=str, default='logs/')
+    parser.add_argument('--saved_path', type=str, default='logs/')
+    args = parser.parse_args()
+    return args
+
+
+class ModelBreadNet(nn.Module):
+    def __init__(self, model1, model2, model3, model4):
+        super().__init__()
+        self.eps = 1e-6
+        self.model_ianet = model1(in_channels=1, out_channels=1)
+        self.model_nsnet = model2(in_channels=2, out_channels=1)
+        self.model_canet = model3(in_channels=4, out_channels=2) if opt.mef else model3(in_channels=6, out_channels=2)
+        self.model_fdnet = model4(in_channels=3, out_channels=1) if opt.model4 else None
+        self.load_weight(self.model_ianet, opt.model1_weight)
+        self.load_weight(self.model_nsnet, opt.model2_weight)
+        self.load_weight(self.model_canet, opt.model3_weight)
+        self.load_weight(self.model_fdnet, opt.model4_weight)
+
+    def load_weight(self, model, weight_pth):
+        if model is not None:
+            state_dict = torch.load(weight_pth)
+            ret = model.load_state_dict(state_dict, strict=True)
+            print(ret)
+
+    def noise_syn_exp(self, illumi, strength):
+        return torch.exp(-illumi) * strength
+
+    def forward(self, image, image_gt):
+        # Color space mapping
+        texture_in, cb_in, cr_in = torch.split(kornia.color.rgb_to_ycbcr(image), 1, dim=1)
+        texture_gt, _, _ = torch.split(kornia.color.rgb_to_ycbcr(image_gt), 1, dim=1)
+
+        # Illumination prediction
+        texture_in_down = F.interpolate(texture_in, scale_factor=0.5, mode='bicubic', align_corners=True)
+        texture_illumi = self.model_ianet(texture_in_down)
+        texture_illumi = F.interpolate(texture_illumi, scale_factor=2, mode='bicubic', align_corners=True)
+
+        # Illumination adjustment
+        texture_illumi = torch.clamp(texture_illumi, 0., 1.)
+        texture_ia = texture_in / torch.clamp_min(texture_illumi, self.eps)
+        texture_ia = torch.clamp(texture_ia, 0., 1.)
+
+        # Noise suppression and fusion
+        texture_nss = []
+        for strength in [0., 0.05, 0.1]:
+            attention = self.noise_syn_exp(texture_illumi, strength=strength)
+            texture_res = self.model_nsnet(torch.cat([texture_ia, attention], dim=1))
+            texture_ns = texture_ia + texture_res
+            texture_nss.append(texture_ns)
+        texture_nss = torch.cat(texture_nss, dim=1).detach()
+        texture_fd = self.model_fdnet(texture_nss)
+
+        # Gamma correction to align the brightness with ground truth;
+        # other methods involved in our main paper are also conducted the same correction for evaluation.
+        if opt.gc:
+            max_psnr = 0
+            best = None
+            for ga in np.arange(0.1, 2.0, 0.01):
+                tx_en = texture_fd ** ga
+                psnr = PSNR(tx_en, texture_gt)
+                if psnr > max_psnr:
+                    max_psnr = psnr
+                    best = tx_en
+
+            texture_fd = torch.clamp(best, 0, 1)
+
+        # Color adaption
+        if not opt.mef:
+            image_ia_ycbcr = kornia.color.rgb_to_ycbcr(torch.clamp(image / (texture_illumi + self.eps), 0, 1))
+            _, cb_ia, cr_ia = torch.split(image_ia_ycbcr, 1, dim=1)
+            colors = self.model_canet(torch.cat([texture_in, cb_in, cr_in, texture_fd, cb_ia, cr_ia], dim=1))
+        else:
+            colors = self.model_canet(
+                torch.cat([texture_in, cb_in, cr_in, texture_fd], dim=1))
+
+        cb_out, cr_out = torch.split(colors, 1, dim=1)
+        cb_out = torch.clamp(cb_out, 0, 1)
+        cr_out = torch.clamp(cr_out, 0, 1)
+
+        # Color space mapping
+        image_out = kornia.color.ycbcr_to_rgb(
+            torch.cat([texture_fd, cb_out, cr_out], dim=1))
+        image_out = torch.clamp(image_out, 0, 1)
+
+        # Calculating image quality metrics
+        psnr = PSNR(image_out, image_gt)
+        ssim = SSIM(image_out, image_gt).item()
+
+        return texture_ia, texture_nss, texture_fd, image_out, texture_illumi, texture_res, psnr, ssim
+
+
+def evaluation(opt):
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    else:
+        torch.manual_seed(42)
+
+    timestamp = mutils.get_formatted_time()
+    opt.saved_path = opt.saved_path + f'/{opt.comment}/{timestamp}'
+    os.makedirs(opt.saved_path, exist_ok=True)
+
+    val_params = {'batch_size': 1,
+                  'shuffle': False,
+                  'drop_last': False,
+                  'num_workers': opt.num_workers}
+
+    val_set = LowLightDataset(opt.data_path)
+
+    val_generator = DataLoader(val_set, **val_params)
+    val_generator = tqdm.tqdm(val_generator)
+
+    model1 = getattr(models, opt.model1)
+    model2 = getattr(models, opt.model2)
+    model3 = getattr(models, opt.model3)
+    model4 = getattr(models, opt.model4) if opt.model4 else None
+
+    model = ModelBreadNet(model1, model2, model3, model4)
+    print(model)
+
+    if opt.num_gpus > 0:
+        model = model.cuda()
+        if opt.num_gpus > 1:
+            model = nn.DataParallel(model)
+
+    model.eval()
+    psnrs, ssims, fns = [], [], []
+    for iter, (data, target, name) in enumerate(val_generator):
+        saver.base_url = os.path.join(opt.saved_path, 'results')
+        with torch.no_grad():
+            if opt.num_gpus == 1:
+                data = data.cuda()
+                target = target.cuda()
+            texture_in, _, _ = torch.split(kornia.color.rgb_to_ycbcr(data), 1, dim=1)
+            texture_gt, _, _ = torch.split(kornia.color.rgb_to_ycbcr(target), 1, dim=1)
+            texture_ia, texture_nss, texture_fd, image_out, \
+            texture_illumi, texture_res, psnr, ssim = model(data, target)
+            if opt.save_extra:
+                saver.save_image(data, name=os.path.splitext(name[0])[0] + '_im_in')
+                saver.save_image(target, name=os.path.splitext(name[0])[0] + '_im_gt')
+
+                saver.save_image(texture_in, name=os.path.splitext(name[0])[0] + '_y_in')
+                saver.save_image(texture_gt, name=os.path.splitext(name[0])[0] + '_y_gt')
+
+                saver.save_image(texture_ia, name=os.path.splitext(name[0])[0] + '_ia')
+                for i in range(texture_nss.shape[1]):
+                    saver.save_image(texture_nss[:, i, ...], name=os.path.splitext(name[0])[0] + f'_ns_{i}')
+                saver.save_image(texture_fd, name=os.path.splitext(name[0])[0] + '_fd')
+
+                saver.save_image(texture_illumi, name=os.path.splitext(name[0])[0] + '_illumi')
+                saver.save_image(texture_res, name=os.path.splitext(name[0])[0] + '_res')
+
+                saver.save_image(image_out, name=os.path.splitext(name[0])[0] + '_out')
+            else:
+                saver.save_image(image_out, name=os.path.splitext(name[0])[0] + '_Bread')
+
+            psnrs.append(psnr)
+            ssims.append(ssim)
+            fns.append(name[0])
+
+    results = list(zip(psnrs, ssims, fns))
+    results.sort(key=lambda item: item[0])
+    for r in results:
+        print(*r)
+    psnr = np.mean(np.array(psnrs))
+    ssim = np.mean(np.array(ssims))
+    print('psnr: ', psnr, ', ssim: ', ssim)
+
+
+if __name__ == '__main__':
+    opt = get_args()
+    evaluation(opt)

exposure_augment.py

@@ -0,0 +1,60 @@
+import math
+import os
+
+import PIL.Image as Image
+import numpy as np
+import torch
+import torchvision.transforms as vtrans
+import tqdm
+
+
+def main(fip, fod):
+    max_overex_rate = 0.25
+    steps = 20
+    num_gen = 4
+
+    im = Image.open(fip)
+    im = vtrans.ToTensor()(im)
+    im_max = torch.flatten(torch.max(im, dim=0, keepdim=True).values)
+    mag = 1. / torch.topk(im_max, math.floor(len(im_max) * max_overex_rate + 1)).values
+    mag = mag[range(0, len(mag), int(len(mag) * (1. / steps)))]
+    mag_diff = torch.diff(mag, 1)
+    mag = mag[:-1]
+
+    top_mag_diff = torch.topk(mag_diff, num_gen).values
+    min_gain = top_mag_diff[top_mag_diff > 0][-1]
+    min_mag = mag[0]
+    max_mag = mag[mag_diff > min_gain][-1]
+    fn, ext = os.path.basename(fip).split('.')
+    bar.set_description(f'{fn}: {min_gain}')
+    ma = np.arange(1, min_mag - min_gain, min_gain * 2)
+    if len(ma) > num_gen:
+        mags = np.append(np.linspace(1, min_mag - min_gain, num_gen),
+                         np.linspace(min_mag, max_mag, num_gen))
+    elif len(ma) == num_gen:
+        mags = np.append(ma, np.linspace(min_mag, max_mag, num_gen))
+    else:
+        mags = np.linspace(1, max_mag, num_gen * 2)
+
+    im = Image.open(fip)
+    im_raw = vtrans.ToTensor()(im)
+
+    for i, mag in enumerate(mags):
+        im = im_raw * mag
+        im.clamp_max_(1.)
+        fop = os.path.join(fod, f'{fn}_{i}.{ext}')
+
+        if not os.path.exists(fop):
+            vtrans.ToPILImage()(im).save(fop)
+
+
+if __name__ == '__main__':
+    # one needs to download it online
+    fid = './data/LOL/train/images'
+    fod = './data/LOL/train/images_aug'
+    os.makedirs(fod, exist_ok=True)
+
+    bar = tqdm.tqdm(os.listdir(fid))
+    for fn in bar:
+        fip = os.path.join(fid, fn)
+        main(fip, fod)

scripts.sh

@@ -0,0 +1,42 @@
+#######################################
+# Breaking Down the Darkness : Testing
+#######################################
+
+# Using gamma correction for evaluation with parameter --gc
+# Show extra intermediate outputs with parameter --save_extra
+CUDA_VISIBLE_DEVICES=0 python eval_Bread.py -m1 IAN -m2 ANSN -m3 FuseNet -m4 FuseNet --mef --comment Bread+NFM+ME[eval] --batch_size 1 -m1w ./checkpoints/IAN_335.pth -m2w ./checkpoints/ANSN_422.pth -m3w ./checkpoints/FuseNet_MECAN_251.pth -m4w ./checkpoints/FuseNet_NFM_297.pth
+CUDA_VISIBLE_DEVICES=0 python test_Bread.py -m1 IAN -m2 ANSN -m3 FuseNet -m4 FuseNet --mef --comment Bread+NFM+ME[test] --batch_size 1 -m1w ./checkpoints/IAN_335.pth -m2w ./checkpoints/ANSN_422.pth -m3w ./checkpoints/FuseNet_MECAN_251.pth -m4w ./checkpoints/FuseNet_NFM_297.pth
+
+# Using CAN w/o MEF data
+CUDA_VISIBLE_DEVICES=0 python eval_Bread.py -m1 IAN -m2 ANSN -m3 IAN -m4 FuseNet --comment Bread+NFM[eval] --batch_size 1 -m1w ./checkpoints/IAN_335.pth -m2w ./checkpoints/ANSN_422.pth -m3w ./checkpoints/IANet_CAN_51.pth -m4w ./checkpoints/FuseNet_NFM_297.pth
+CUDA_VISIBLE_DEVICES=0 python test_Bread.py -m1 IAN -m2 ANSN -m3 IAN -m4 FuseNet --comment Bread+NFM[test] --batch_size 1 -m1w ./checkpoints/IAN_335.pth -m2w ./checkpoints/ANSN_422.pth -m3w ./checkpoints/IANet_CAN_51.pth -m4w ./checkpoints/FuseNet_NFM_297.pth
+
+# Remove NFM for much better generalization performance
+# changing parameter a for an ideal denoising strength
+CUDA_VISIBLE_DEVICES=0 python test_Bread_NoNFM.py -m1 IAN -m2 ANSN -m3 FuseNet --mef -a 0.10 --comment Bread+ME[test] --batch_size 1 -m1w ./checkpoints/IAN_335.pth -m2w ./checkpoints/ANSN_422.pth -m3w ./checkpoints/FuseNet_MECAN_251.pth
+
+##############################################################
+# Breaking Down the Darkness : Training
+# SICE dataset and LOL dataset are need to be download online
+##############################################################
+
+# Multi-exposure data synthesis
+python exposure_augment.py
+
+# Train IAN
+CUDA_VISIBLE_DEVICES=0 python train_IAN.py -m IAN --comment IAN_train --batch_size 1 --val_interval 1 --num_epochs 500 --lr 0.001 --no_sche
+
+# Train ANSN
+CUDA_VISIBLE_DEVICES=0 python train_ANSN.py -m1 IAN -m2 ANSN --comment ANSN_train --batch_size 1 --val_interval 1 --num_epochs 500 --lr 0.001 --no_sche -m1w ./checkpoints/IAN_335.pth
+
+# Train CAN
+CUDA_VISIBLE_DEVICES=0 python train_CAN.py -m1 IAN -m3 FuseNet --comment CAN_train --batch_size 1 --val_interval 1 --num_epochs 500 --lr 0.001 --no_sche -m1w ./checkpoints/IAN_335.pth
+
+# Train MECAN on SICE
+CUDA_VISIBLE_DEVICES=0 python train_MECAN.py -m FuseNet --comment MECAN_train --batch_size 1 --val_interval 1 --num_epochs 500 --lr 0.001 --no_sche
+
+# Finetune MECAN on SICE and LOL datasets
+CUDA_VISIBLE_DEVICES=0 python train_MECAN_finetune.py -m FuseNet --comment MECAN_finetune --batch_size 1 --val_interval 1 --num_epochs 500 --lr 1e-4 --no_sche -mw ./checkpoints/FuseNet_MECAN_for_Finetuning_404.pth
+
+
+

test_Bread.py

@@ -0,0 +1,195 @@
+import argparse
+import os
+
+import kornia
+import torch
+import torch.nn.functional as F
+import tqdm
+from torch import nn
+from torch.utils.data import DataLoader
+
+import models
+from datasets import LowLightDatasetTest
+from tools import saver, mutils
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Breaking Downing the Darkness')
+    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus being used')
+    parser.add_argument('--num_workers', type=int, default=12, help='num_workers of dataloader')
+    parser.add_argument('--batch_size', type=int, default=4, help='The number of images per batch among all devices')
+    parser.add_argument('-m1', '--model1', type=str, default='IANet', help='Model1 Name')
+    parser.add_argument('-m2', '--model2', type=str, default='NSNet', help='Model2 Name')
+    parser.add_argument('-m3', '--model3', type=str, default='FuseNet', help='Model3 Name')
+    parser.add_argument('-m4', '--model4', type=str, default=None, help='Model4 Name')
+
+    parser.add_argument('-m1w', '--model1_weight', type=str, default=None, help='Model weight of IAN')
+    parser.add_argument('-m2w', '--model2_weight', type=str, default=None, help='Model weight of ANSN')
+    parser.add_argument('-m3w', '--model3_weight', type=str, default=None, help='Model weight of CAN')
+    parser.add_argument('-m4w', '--model4_weight', type=str, default=None, help='Model weight of NFM')
+
+    parser.add_argument('--mef', action='store_true')
+    parser.add_argument('--save_extra', action='store_true', help='save intermediate outputs or not')
+
+    parser.add_argument('--comment', type=str, default='default',
+                        help='Project comment')
+
+    parser.add_argument('--alpha', '-a', type=float, default=0.10)
+    parser.add_argument('--lr', type=float, default=0.01)
+    parser.add_argument('--optim', type=str, default='adamw', help='select optimizer for training, '
+                                                                   'suggest using \'admaw\' until the'
+                                                                   ' very final stage then switch to \'sgd\'')
+    parser.add_argument('--data_path', type=str, default='./data/test',
+                        help='the root folder of dataset')
+    parser.add_argument('--log_path', type=str, default='logs/')
+    parser.add_argument('--saved_path', type=str, default='logs/')
+    args = parser.parse_args()
+    return args
+
+
+class ModelBreadNet(nn.Module):
+    def __init__(self, model1, model2, model3, model4):
+        super().__init__()
+        self.eps = 1e-6
+        self.model_ianet = model1(in_channels=1, out_channels=1)
+        self.model_nsnet = model2(in_channels=2, out_channels=1)
+        self.model_canet = model3(in_channels=4, out_channels=2) if opt.mef else model3(in_channels=6, out_channels=2)
+        self.model_fdnet = model4(in_channels=3, out_channels=1) if opt.model4 else None
+
+        self.load_weight(self.model_ianet, opt.model1_weight)
+        self.load_weight(self.model_nsnet, opt.model2_weight)
+        self.load_weight(self.model_canet, opt.model3_weight)
+        self.load_weight(self.model_fdnet, opt.model4_weight)
+
+    def load_weight(self, model, weight_pth):
+        if model is not None:
+            state_dict = torch.load(weight_pth)
+            ret = model.load_state_dict(state_dict, strict=True)
+            print(ret)
+
+    def noise_syn_exp(self, illumi, strength):
+        return torch.exp(-illumi) * strength
+
+    def forward(self, image):
+        # Color space mapping
+        texture_in, cb_in, cr_in = torch.split(kornia.color.rgb_to_ycbcr(image), 1, dim=1)
+
+        # Illumination prediction
+        texture_in_down = F.interpolate(texture_in, scale_factor=0.5, mode='bicubic', align_corners=True)
+        texture_illumi = self.model_ianet(texture_in_down)
+        texture_illumi = F.interpolate(texture_illumi, scale_factor=2, mode='bicubic', align_corners=True)
+
+        # Illumination adjustment
+        texture_illumi = torch.clamp(texture_illumi, 0., 1.)
+        texture_ia = texture_in / torch.clamp_min(texture_illumi, self.eps)
+        texture_ia = torch.clamp(texture_ia, 0., 1.)
+
+        # Noise suppression and fusion
+        texture_nss = []
+        for strength in [0., 0.05, 0.1]:
+            attention = self.noise_syn_exp(texture_illumi, strength=strength)
+            texture_res = self.model_nsnet(torch.cat([texture_ia, attention], dim=1))
+            texture_ns = texture_ia + texture_res
+            texture_nss.append(texture_ns)
+        texture_nss = torch.cat(texture_nss, dim=1).detach()
+        texture_fd = self.model_fdnet(texture_nss)
+
+        # Further preserve the texture under brighter illumination
+        texture_fd = texture_illumi * texture_in + (1 - texture_illumi) * texture_fd
+        texture_fd = torch.clamp(texture_fd, 0, 1)
+
+        # Color adaption
+        if not opt.mef:
+            image_ia_ycbcr = kornia.color.rgb_to_ycbcr(torch.clamp(image / (texture_illumi + self.eps), 0, 1))
+            _, cb_ia, cr_ia = torch.split(image_ia_ycbcr, 1, dim=1)
+            colors = self.model_canet(torch.cat([texture_in, cb_in, cr_in, texture_fd, cb_ia, cr_ia], dim=1))
+        else:
+            colors = self.model_canet(
+                torch.cat([texture_in, cb_in, cr_in, texture_fd], dim=1))
+        cb_out, cr_out = torch.split(colors, 1, dim=1)
+        cb_out = torch.clamp(cb_out, 0, 1)
+        cr_out = torch.clamp(cr_out, 0, 1)
+
+        # Color space mapping
+        image_out = kornia.color.ycbcr_to_rgb(
+            torch.cat([texture_fd, cb_out, cr_out], dim=1))
+
+        # Further preserve the color under brighter illumination
+        img_fusion = texture_illumi * image + (1 - texture_illumi) * image_out
+        _, cb_fuse, cr_fuse = torch.split(kornia.color.rgb_to_ycbcr(img_fusion), 1, dim=1)
+        image_out = kornia.color.ycbcr_to_rgb(
+            torch.cat([texture_fd, cb_fuse, cr_fuse], dim=1))
+        image_out = torch.clamp(image_out, 0, 1)
+
+        return texture_ia, texture_nss, texture_fd, image_out, texture_illumi, texture_res
+
+
+def test(opt):
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    else:
+        torch.manual_seed(42)
+
+    timestamp = mutils.get_formatted_time()
+    opt.saved_path = opt.saved_path + f'/{opt.comment}/{timestamp}'
+    os.makedirs(opt.saved_path, exist_ok=True)
+
+    test_params = {'batch_size': 1,
+                   'shuffle': False,
+                   'drop_last': False,
+                   'num_workers': opt.num_workers}
+
+    test_set = LowLightDatasetTest(opt.data_path)
+
+    test_generator = DataLoader(test_set, **test_params)
+    test_generator = tqdm.tqdm(test_generator)
+
+    model1 = getattr(models, opt.model1)
+    model2 = getattr(models, opt.model2)
+    model3 = getattr(models, opt.model3)
+    model4 = getattr(models, opt.model4)
+
+    model = ModelBreadNet(model1, model2, model3, model4)
+    print(model)
+
+    if opt.num_gpus > 0:
+        model = model.cuda()
+        if opt.num_gpus > 1:
+            model = nn.DataParallel(model)
+
+    model.eval()
+
+    for iter, (data, subset, name) in enumerate(test_generator):
+        saver.base_url = os.path.join(opt.saved_path, 'results', subset[0])
+        with torch.no_grad():
+            if opt.num_gpus == 1:
+                data = data.cuda()
+            texture_in, _, _ = torch.split(kornia.color.rgb_to_ycbcr(data), 1, dim=1)
+
+            texture_ia, texture_nss, texture_fd, image_out, texture_illumi, texture_res = model(data)
+
+            if opt.save_extra:
+                saver.save_image(data, name=os.path.splitext(name[0])[0] + '_im_in')
+                saver.save_image(texture_in, name=os.path.splitext(name[0])[0] + '_y_in')
+                saver.save_image(texture_ia, name=os.path.splitext(name[0])[0] + '_ia')
+                for i in range(texture_nss.shape[1]):
+                    saver.save_image(texture_nss[:, i, ...], name=os.path.splitext(name[0])[0] + f'_ns_{i}')
+                saver.save_image(texture_fd, name=os.path.splitext(name[0])[0] + '_fd')
+
+                saver.save_image(texture_illumi, name=os.path.splitext(name[0])[0] + '_illumi')
+                saver.save_image(texture_res, name=os.path.splitext(name[0])[0] + '_res')
+                saver.save_image(image_out, name=os.path.splitext(name[0])[0] + '_out')
+            else:
+                saver.save_image(image_out, name=os.path.splitext(name[0])[0] + '_Bread')
+
+
+def save_checkpoint(model, name):
+    if isinstance(model, nn.DataParallel):
+        torch.save(model.module3.model_nsnet.state_dict(), os.path.join(opt.saved_path, name))
+    else:
+        torch.save(model.model_nsnet.state_dict(), os.path.join(opt.saved_path, name))
+
+
+if __name__ == '__main__':
+    opt = get_args()
+    test(opt)

test_Bread_NoNFM.py

@@ -0,0 +1,167 @@
+import argparse
+import os
+
+import kornia
+import torch
+import torch.nn.functional as F
+import tqdm
+from torch import nn
+from torch.utils.data import DataLoader
+
+import models
+from datasets import LowLightDatasetTest
+from tools import saver, mutils
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Breaking Downing the Darkness')
+    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus being used')
+    parser.add_argument('--num_workers', type=int, default=12, help='num_workers of dataloader')
+    parser.add_argument('--batch_size', type=int, default=4, help='The number of images per batch among all devices')
+    parser.add_argument('-m1', '--model1', type=str, default='IAN', help='Model1 Name')
+    parser.add_argument('-m2', '--model2', type=str, default='ANSN', help='Model2 Name')
+    parser.add_argument('-m3', '--model3', type=str, default='FuseNet', help='Model3 Name')
+
+    parser.add_argument('-m1w', '--model1_weight', type=str, default=None, help='Model weight of IAN')
+    parser.add_argument('-m2w', '--model2_weight', type=str, default=None, help='Model weight of ANSN')
+    parser.add_argument('-m3w', '--model3_weight', type=str, default=None, help='Model weight of CAN')
+
+    parser.add_argument('--mef', action='store_true')
+    parser.add_argument('--save_extra', action='store_true', help='save intermediate outputs or not')
+
+    parser.add_argument('--comment', type=str, default='default',
+                        help='Project comment')
+
+    parser.add_argument('--alpha', '-a', type=float, default=0.10)
+
+    parser.add_argument('--data_path', type=str, default='./data/test',
+                        help='the root folder of dataset')
+    parser.add_argument('--log_path', type=str, default='logs/')
+    parser.add_argument('--saved_path', type=str, default='logs/')
+    args = parser.parse_args()
+    return args
+
+
+class ModelBreadNet(nn.Module):
+    def __init__(self, model1, model2, model3):
+        super().__init__()
+        self.eps = 1e-6
+        self.model_ianet = model1(in_channels=1, out_channels=1)
+        self.model_nsnet = model2(in_channels=2, out_channels=1)
+        self.model_canet = model3(in_channels=4, out_channels=2) if opt.mef else model3(in_channels=6, out_channels=2)
+
+        self.load_weight(self.model_ianet, opt.model1_weight)
+        self.load_weight(self.model_nsnet, opt.model2_weight)
+        self.load_weight(self.model_canet, opt.model3_weight)
+
+    def load_weight(self, model, weight_pth):
+        if model is not None:
+            state_dict = torch.load(weight_pth)
+            ret = model.load_state_dict(state_dict, strict=True)
+            print(ret)
+
+    def noise_syn_exp(self, illumi, strength):
+        return torch.exp(-illumi) * strength
+
+    def forward(self, image):
+        # Color space mapping
+        texture_in, cb_in, cr_in = torch.split(kornia.color.rgb_to_ycbcr(image), 1, dim=1)
+
+        # Illumination prediction
+        texture_in_down = F.interpolate(texture_in, scale_factor=0.5, mode='bicubic', align_corners=True)
+        texture_illumi = self.model_ianet(texture_in_down)
+        texture_illumi = F.interpolate(texture_illumi, scale_factor=2, mode='bicubic', align_corners=True)
+
+        # Illumination adjustment
+        texture_illumi = torch.clamp(texture_illumi, 0., 1.)
+        texture_ia = texture_in / torch.clamp_min(texture_illumi, self.eps)
+        texture_ia = torch.clamp(texture_ia, 0., 1.)
+
+        # Noise suppression and fusion
+        attention = self.noise_syn_exp(texture_illumi, strength=opt.alpha)
+        texture_res = self.model_nsnet(torch.cat([texture_ia, attention], dim=1))
+        texture_ns = texture_ia + texture_res
+
+        # Further preserve the texture under brighter illumination
+        texture_ns = texture_illumi * texture_in + (1 - texture_illumi) * texture_ns
+        texture_ns = torch.clamp(texture_ns, 0, 1)
+
+        # Color adaption
+        colors = self.model_canet(
+            torch.cat([texture_in, cb_in, cr_in, texture_ns], dim=1))
+        cb_out, cr_out = torch.split(colors, 1, dim=1)
+        cb_out = torch.clamp(cb_out, 0, 1)
+        cr_out = torch.clamp(cr_out, 0, 1)
+
+        # Color space mapping
+        image_out = kornia.color.ycbcr_to_rgb(
+            torch.cat([texture_ns, cb_out, cr_out], dim=1))
+
+        # Further preserve the color under brighter illumination
+        img_fusion = texture_illumi * image + (1 - texture_illumi) * image_out
+        _, cb_fuse, cr_fuse = torch.split(kornia.color.rgb_to_ycbcr(img_fusion), 1, dim=1)
+        image_out = kornia.color.ycbcr_to_rgb(
+            torch.cat([texture_ns, cb_fuse, cr_fuse], dim=1))
+        image_out = torch.clamp(image_out, 0, 1)
+
+        return texture_ia, texture_ns, image_out, texture_illumi, texture_res
+
+
+def test(opt):
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    else:
+        torch.manual_seed(42)
+
+    timestamp = mutils.get_formatted_time()
+    opt.saved_path = opt.saved_path + f'/{opt.comment}/{timestamp}'
+    os.makedirs(opt.saved_path, exist_ok=True)
+
+    test_params = {'batch_size': 1,
+                   'shuffle': False,
+                   'drop_last': False,
+                   'num_workers': opt.num_workers}
+
+    test_set = LowLightDatasetTest(opt.data_path)
+
+    test_generator = DataLoader(test_set, **test_params)
+    test_generator = tqdm.tqdm(test_generator)
+
+    model1 = getattr(models, opt.model1)
+    model2 = getattr(models, opt.model2)
+    model3 = getattr(models, opt.model3)
+
+    model = ModelBreadNet(model1, model2, model3)
+    print(model)
+
+    if opt.num_gpus > 0:
+        model = model.cuda()
+        if opt.num_gpus > 1:
+            model = nn.DataParallel(model)
+
+    model.eval()
+
+    for iter, (data, subset, name) in enumerate(test_generator):
+        saver.base_url = os.path.join(opt.saved_path, 'results', subset[0])
+        with torch.no_grad():
+            if opt.num_gpus == 1:
+                data = data.cuda()
+            texture_in, _, _ = torch.split(kornia.color.rgb_to_ycbcr(data), 1, dim=1)
+
+            texture_ia, texture_ns, image_out, texture_illumi, texture_res = model(data)
+
+            if opt.save_extra:
+                saver.save_image(data, name=os.path.splitext(name[0])[0] + '_im_in')
+                saver.save_image(texture_in, name=os.path.splitext(name[0])[0] + '_y_in')
+                saver.save_image(texture_ia, name=os.path.splitext(name[0])[0] + '_ia')
+                saver.save_image(texture_ns, name=os.path.splitext(name[0])[0] + '_ns')
+
+                saver.save_image(texture_illumi, name=os.path.splitext(name[0])[0] + '_illumi')
+                saver.save_image(texture_res, name=os.path.splitext(name[0])[0] + '_res')
+                saver.save_image(image_out, name=os.path.splitext(name[0])[0] + '_out')
+            else:
+                saver.save_image(image_out, name=os.path.splitext(name[0])[0] + '_Bread')
+
+if __name__ == '__main__':
+    opt = get_args()
+    test(opt)

train_ANSN.py

@@ -0,0 +1,264 @@
+import argparse
+import datetime
+import os
+import traceback
+
+import kornia
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+
+import models
+from datasets import LowLightFDataset, LowLightFDatasetEval
+from models import PSNR, SSIM, CosineLR
+from tools import SingleSummaryWriter
+from tools import saver, mutils
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Breaking Downing the Darkness')
+    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus being used')
+    parser.add_argument('--num_workers', type=int, default=12, help='num_workers of dataloader')
+    parser.add_argument('--batch_size', type=int, default=1, help='The number of images per batch among all devices')
+    parser.add_argument('-m1', '--model1', type=str, default='INet',
+                        help='Model1 Name')
+    parser.add_argument('-m2', '--model2', type=str, default='NSNet',
+                        help='Model1 Name')
+    parser.add_argument('-m1w', '--model1_weight', type=str, default=None,
+                        help='Model Name')
+
+    parser.add_argument('--comment', type=str, default='default',
+                        help='Project comment')
+    parser.add_argument('--graph', action='store_true')
+    parser.add_argument('--no_sche', action='store_true')
+    parser.add_argument('--sampling', action='store_true')
+
+    parser.add_argument('--slope', type=float, default=2.)
+    parser.add_argument('--lr', type=float, default=0.001)
+    parser.add_argument('--optim', type=str, default='adam', help='select optimizer for training, '
+                                                                  'suggest using \'admaw\' until the'
+                                                                  ' very final stage then switch to \'sgd\'')
+    parser.add_argument('--num_epochs', type=int, default=500)
+    parser.add_argument('--val_interval', type=int, default=1, help='Number of epoches between valing phases')
+    parser.add_argument('--save_interval', type=int, default=500, help='Number of steps between saving')
+    parser.add_argument('--data_path', type=str, default='./data/LOL',
+                        help='the root folder of dataset')
+    parser.add_argument('--log_path', type=str, default='logs/')
+    parser.add_argument('--saved_path', type=str, default='logs/')
+    args = parser.parse_args()
+    return args
+
+
+class ModelNSNet(nn.Module):
+    def __init__(self, model1, model2):
+        super().__init__()
+        self.texture_loss = models.MSELoss()
+        self.model_ianet = model1(in_channels=1, out_channels=1)
+        self.model_nsnet = model2(in_channels=2, out_channels=1)
+
+        assert opt.model1_weight is not None
+        self.load_weight(self.model_ianet, opt.model1_weight)
+        self.model_ianet.eval()
+        self.eps = 1e-2
+
+    def load_weight(self, model, weight_pth):
+        state_dict = torch.load(weight_pth)
+        ret = model.load_state_dict(state_dict, strict=True)
+        print(ret)
+
+    def noise_syn(self, illumi, strength):
+        return torch.exp(-illumi) * strength
+
+    def forward(self, image, image_gt, training=True):
+        with torch.no_grad():
+            image = image.squeeze(0)
+            texture_in, _, _ = torch.split(kornia.color.rgb_to_ycbcr(image), 1, dim=1)
+            texture_gt, _, _ = torch.split(kornia.color.rgb_to_ycbcr(image_gt), 1, dim=1)
+
+            texture_in_down = F.interpolate(texture_in, scale_factor=0.5, mode='bicubic', align_corners=True)
+            illumi = self.model_ianet(texture_in_down)
+            illumi = F.interpolate(illumi, scale_factor=2, mode='bicubic', align_corners=True)
+
+            attention = self.noise_syn(illumi, strength=0.1)
+
+            noise = torch.normal(mean=0., std=attention)
+            noisy_gt = torch.clamp(texture_gt + noise, 0., 1.)
+
+        texture_res = self.model_nsnet(torch.cat([noisy_gt, attention], dim=1))
+        restor_loss = self.texture_loss(texture_res, texture_gt - noisy_gt)
+
+        texture_ns = noisy_gt + texture_res
+
+        psnr = PSNR(texture_ns, texture_gt)
+        ssim = SSIM(texture_ns, texture_gt).item()
+        return noisy_gt, texture_ns, texture_res, illumi, restor_loss, psnr, ssim
+
+
+def train(opt):
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    else:
+        torch.manual_seed(42)
+
+    # params.project_name = params.project_name + str(time.time()).replace('.', '')
+    timestamp = mutils.get_formatted_time()
+    opt.saved_path = opt.saved_path + f'/{opt.comment}/{timestamp}'
+    opt.log_path = opt.log_path + f'/{opt.comment}/{timestamp}/tensorboard/'
+    os.makedirs(opt.log_path, exist_ok=True)
+    os.makedirs(opt.saved_path, exist_ok=True)
+
+    training_params = {'batch_size': opt.batch_size,
+                       'shuffle': True,
+                       'drop_last': True,
+                       'num_workers': opt.num_workers}
+
+    val_params = {'batch_size': 1,
+                  'shuffle': False,
+                  'drop_last': True,
+                  'num_workers': opt.num_workers}
+
+    training_set = LowLightFDataset(os.path.join(opt.data_path, 'train'))
+    training_generator = DataLoader(training_set, **training_params)
+
+    val_set = LowLightFDatasetEval(os.path.join(opt.data_path, 'eval'))
+    val_generator = DataLoader(val_set, **val_params)
+
+    model1 = getattr(models, opt.model1)
+    model2 = getattr(models, opt.model2)
+    writer = SingleSummaryWriter(opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/')
+
+    model = ModelNSNet(model1, model2)
+    print(model)
+
+    if opt.num_gpus > 0:
+        model = model.cuda()
+        if opt.num_gpus > 1:
+            model = nn.DataParallel(model)
+
+    if opt.optim == 'adam':
+        optimizer = torch.optim.Adam(model.model_nsnet.parameters(), opt.lr)
+    else:
+        optimizer = torch.optim.SGD(model.model_nsnet.parameters(), opt.lr, momentum=0.9, nesterov=True)
+
+    scheduler = CosineLR(optimizer, opt.lr, opt.num_epochs)
+    epoch = 0
+    step = 0
+    model.model_nsnet.train()
+
+    num_iter_per_epoch = len(training_generator)
+
+    try:
+        for epoch in range(opt.num_epochs):
+            last_epoch = step // num_iter_per_epoch
+            if epoch < last_epoch:
+                continue
+
+            epoch_loss = []
+            progress_bar = tqdm(training_generator)
+
+            saver.base_url = os.path.join(opt.saved_path, 'results', '%03d' % epoch)
+            if not opt.sampling:
+                for iter, (data, target, name) in enumerate(progress_bar):
+                    if iter < step - last_epoch * num_iter_per_epoch:
+                        progress_bar.update()
+                        continue
+                    try:
+                        if opt.num_gpus == 1:
+                            data = data.cuda()
+                            target = target.cuda()
+
+                        optimizer.zero_grad()
+
+                        noisy_gt, texture_ns, texture_res, illumi, \
+                        restor_loss, psnr, ssim = model(data, target, training=True)
+
+                        loss = restor_loss
+
+                        loss.backward()
+                        optimizer.step()
+
+                        epoch_loss.append(float(loss))
+
+                        progress_bar.set_description(
+                            'Step: {}. Epoch: {}/{}. Iteration: {}/{}. restor_loss: {:.5f},  psnr: {:.5f}, ssim: {:.5f}'.format(
+                                step, epoch, opt.num_epochs, iter + 1, num_iter_per_epoch, restor_loss.item(), psnr,
+                                ssim))
+                        writer.add_scalar('Loss/train', loss, step)
+                        writer.add_scalar('PSNR/train', psnr, step)
+                        writer.add_scalar('SSIM/train', ssim, step)
+
+                        # log learning_rate
+                        current_lr = optimizer.param_groups[0]['lr']
+                        writer.add_scalar('learning_rate', current_lr, step)
+
+                        step += 1
+
+                    except Exception as e:
+                        print('[Error]', traceback.format_exc())
+                        print(e)
+                        continue
+
+            if not opt.no_sche:
+                scheduler.step()
+
+            if epoch % opt.val_interval == 0:
+                model.model_nsnet.eval()
+                loss_ls = []
+                psnrs = []
+                ssims = []
+
+                for iter, (data, target, name) in enumerate(val_generator):
+                    with torch.no_grad():
+                        if opt.num_gpus == 1:
+                            data = data.cuda()
+                            target = target.cuda()
+
+                        noisy_gt, texture_ns, texture_res, \
+                        illumi, restor_loss, psnr, ssim = model(data, target, training=False)
+                        texture_gt, _, _ = torch.split(kornia.color.rgb_to_ycbcr(target), 1, dim=1)
+
+                        saver.save_image(noisy_gt, name=os.path.splitext(name[0])[0] + '_in')
+                        saver.save_image(texture_ns, name=os.path.splitext(name[0])[0] + '_ns')
+                        saver.save_image(texture_res, name=os.path.splitext(name[0])[0] + '_res')
+                        saver.save_image(illumi, name=os.path.splitext(name[0])[0] + '_ill')
+                        saver.save_image(target, name=os.path.splitext(name[0])[0] + '_gt')
+
+                        loss = restor_loss
+                        loss_ls.append(loss.item())
+                        psnrs.append(psnr)
+                        ssims.append(ssim)
+
+                loss = np.mean(np.array(loss_ls))
+                psnr = np.mean(np.array(psnrs))
+                ssim = np.mean(np.array(ssims))
+
+                print(
+                    'Val. Epoch: {}/{}. Loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
+                        epoch, opt.num_epochs, loss, psnr, ssim))
+                writer.add_scalar('Loss/val', loss, step)
+                writer.add_scalar('PSNR/val', psnr, step)
+                writer.add_scalar('SSIM/val', ssim, step)
+
+                save_checkpoint(model, f'{opt.model2}_{"%03d" % epoch}_{psnr}_{ssim}_{step}.pth')
+
+                model.model_nsnet.train()
+
+    except KeyboardInterrupt:
+        save_checkpoint(model, f'{opt.model2}_{epoch}_{step}_keyboardInterrupt.pth')
+        writer.close()
+    writer.close()
+
+
+def save_checkpoint(model, name):
+    if isinstance(model, nn.DataParallel):
+        torch.save(model.module.model_nsnet.state_dict(), os.path.join(opt.saved_path, name))
+    else:
+        torch.save(model.model_nsnet.state_dict(), os.path.join(opt.saved_path, name))
+
+
+if __name__ == '__main__':
+    opt = get_args()
+    train(opt)

train_CAN.py

@@ -0,0 +1,276 @@
+import argparse
+import datetime
+import os
+import traceback
+
+import kornia
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+
+import models
+from datasets import LowLightFDataset, LowLightFDatasetEval
+from models import PSNR, SSIM, CosineLR
+from tools import SingleSummaryWriter
+from tools import saver, mutils
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Breaking Downing the Darkness')
+    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus being used')
+    parser.add_argument('--num_workers', type=int, default=12, help='num_workers of dataloader')
+    parser.add_argument('--batch_size', type=int, default=1, help='The number of images per batch among all devices')
+    parser.add_argument('-m1', '--model1', type=str, default='INet',
+                        help='Model Name')
+    parser.add_argument('-m3', '--model3', type=str, default='INet',
+                        help='Model Name')
+    parser.add_argument('-m1w', '--model1_weight', type=str, default=None,
+                        help='Model Name')
+    parser.add_argument('-m3w', '--model3_weight', type=str, default=None,
+                        help='Model Name')
+    parser.add_argument('-ts', '--targets_split', type=str, default='targets',
+                        help='dir of targets')
+    parser.add_argument('--comment', type=str, default='default',
+                        help='Project comment')
+    parser.add_argument('--graph', action='store_true')
+    parser.add_argument('--scratch', action='store_true')
+    parser.add_argument('--sampling', action='store_true')
+    parser.add_argument('--test_on_start', action='store_true')
+
+    parser.add_argument('--lr', type=float, default=0.01)
+    parser.add_argument('--no_sche', action='store_true')
+
+    parser.add_argument('--optim', type=str, default='adam', help='select optimizer for training, '
+                                                                  'suggest using \'admaw\' until the'
+                                                                  ' very final stage then switch to \'sgd\'')
+    parser.add_argument('--num_epochs', type=int, default=500)
+    parser.add_argument('--val_interval', type=int, default=1, help='Number of epoches between valing phases')
+    parser.add_argument('--save_interval', type=int, default=500, help='Number of steps between saving')
+    parser.add_argument('--data_path', type=str, default='./data/LOL',
+                        help='the root folder of dataset')
+    parser.add_argument('--log_path', type=str, default='logs/')
+    parser.add_argument('--saved_path', type=str, default='logs/')
+    args = parser.parse_args()
+    return args
+
+
+def compute_gradient(img):
+    gradx = img[..., 1:, :] - img[..., :-1, :]
+    grady = img[..., 1:] - img[..., :-1]
+    return gradx, grady
+
+
+class ModelCANet(nn.Module):
+    def __init__(self, model1, model3):
+        super().__init__()
+        self.color_loss = models.L1Loss()
+        self.restor_loss = models.MSSSIML1Loss(channels=3)
+        self.model_ianet = model1(in_channels=1, out_channels=1)
+        self.model_canet = model3(in_channels=6, out_channels=2)
+        self.eps = 1e-2
+        self.load_weight(self.model_ianet, opt.model1_weight)
+        if opt.model3_weight is not None:
+            self.load_weight(self.model_canet, opt.model3_weight)
+        self.model_ianet.eval()
+
+    def load_weight(self, model, weight_pth):
+        state_dict = torch.load(weight_pth)
+        ret = model.load_state_dict(state_dict, strict=True)
+        print(ret)
+
+    def forward(self, image, image_gt, training=True):
+        if training:
+            image = image.squeeze(0)
+            image_gt = image_gt.repeat(8, 1, 1, 1)
+
+        texture_in, cb_in, cr_in = torch.split(kornia.color.rgb_to_ycbcr(image), 1, dim=1)
+
+        texture_in_down = F.interpolate(texture_in, scale_factor=0.5, mode='bicubic', align_corners=True)
+        texture_illumi = self.model_ianet(texture_in_down)
+        texture_illumi = F.interpolate(texture_illumi, scale_factor=2, mode='bicubic', align_corners=True)
+
+        texture_en, cb_en, cr_en = torch.split(kornia.color.rgb_to_ycbcr(image / torch.clamp_min(texture_illumi, self.eps)),
+                                               1, dim=1)
+        texture_gt, cb_gt, cr_gt = torch.split(kornia.color.rgb_to_ycbcr(image_gt), 1, dim=1)
+
+        colors = self.model_canet(torch.cat([texture_in, cb_in, cr_in, texture_gt, cb_en, cr_en], dim=1))
+
+        cb, cr = torch.split(colors, 1, dim=1)
+
+        color_loss1 = self.color_loss(cb, cb_gt)
+        color_loss2 = self.color_loss(cr, cr_gt)
+
+        image_out = kornia.color.ycbcr_to_rgb(torch.cat([texture_gt, cb, cr], dim=1))
+        restor_loss = self.restor_loss(image_out, image_gt) * 1.0
+
+        psnr = PSNR(image_out, image_gt)
+        ssim = SSIM(image_out, image_gt).item()
+        return image_out, color_loss1, color_loss2, restor_loss, psnr, ssim
+
+
+def train(opt):
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    else:
+        torch.manual_seed(42)
+
+    timestamp = mutils.get_formatted_time()
+    opt.saved_path = opt.saved_path + f'/{opt.comment}/{timestamp}'
+    opt.log_path = opt.log_path + f'/{opt.comment}/{timestamp}/tensorboard/'
+    os.makedirs(opt.log_path, exist_ok=True)
+    os.makedirs(opt.saved_path, exist_ok=True)
+
+    training_params = {'batch_size': opt.batch_size,
+                       'shuffle': True,
+                       'drop_last': True,
+                       'num_workers': opt.num_workers}
+
+    val_params = {'batch_size': 1,
+                  'shuffle': False,
+                  'drop_last': False,
+                  'num_workers': opt.num_workers}
+
+    training_set = LowLightFDataset(os.path.join(opt.data_path, 'train'), targets_split=opt.targets_split,
+                                    training=True)
+    training_generator = DataLoader(training_set, **training_params)
+
+    val_set = LowLightFDatasetEval(os.path.join(opt.data_path, 'eval'), training=False)
+    val_generator = DataLoader(val_set, **val_params)
+
+    model1 = getattr(models, opt.model1)
+    model3 = getattr(models, opt.model3)
+
+    model = ModelCANet(model1, model3)
+    print(model)
+
+    writer = SingleSummaryWriter(opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/')
+
+    if opt.num_gpus > 0:
+        model = model.cuda()
+        if opt.num_gpus > 1:
+            model = nn.DataParallel(model)
+
+    if opt.optim == 'adam':
+        optimizer = torch.optim.Adam(model.model_canet.parameters(), opt.lr)
+    else:
+        optimizer = torch.optim.SGD(model.model_canet.parameters(), opt.lr, momentum=0.9, nesterov=True)
+
+    scheduler = CosineLR(optimizer, opt.lr, opt.num_epochs)
+    epoch = 0
+    step = 0
+    model.model_canet.train()
+
+    num_iter_per_epoch = len(training_generator)
+
+    try:
+        for epoch in range(opt.num_epochs):
+            last_epoch = step // num_iter_per_epoch
+            if epoch < last_epoch:
+                continue
+
+            epoch_loss = []
+            progress_bar = tqdm(training_generator)
+            if not opt.sampling and not opt.test_on_start:
+                for iter, (data, target, name) in enumerate(progress_bar):
+                    if iter < step - last_epoch * num_iter_per_epoch:
+                        progress_bar.update()
+                        continue
+                    try:
+                        if opt.num_gpus == 1:
+                            data, target = data.cuda(), target.cuda()
+                        optimizer.zero_grad()
+
+                        image_out, color_loss1, color_loss2, \
+                        restor_loss, psnr, ssim = model(data, target, training=True)
+                        loss = color_loss1 + color_loss2 + restor_loss
+                        loss.backward()
+                        optimizer.step()
+
+                        epoch_loss.append(float(loss))
+
+                        progress_bar.set_description(
+                            'Step: {}. Epoch: {}/{}. Iteration: {}/{}. color_loss1: {:1.5f}, color_loss2: {:1.5f}, restor_loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
+                                step, epoch, opt.num_epochs, iter + 1, num_iter_per_epoch,
+                                color_loss1.item(), color_loss2.item(),
+                                restor_loss.item(), psnr, ssim))
+                        writer.add_scalar('Loss/train', loss, step)
+                        writer.add_scalar('PSNR/train', psnr, step)
+                        writer.add_scalar('SSIM/train', ssim, step)
+
+                        # log learning_rate
+                        current_lr = optimizer.param_groups[0]['lr']
+                        writer.add_scalar('learning_rate', current_lr, step)
+
+                        step += 1
+
+                    except Exception as e:
+                        print('[Error]', traceback.format_exc())
+                        print(e)
+                        continue
+                    # scheduler.step(np.mean(epoch_loss))
+
+            if opt.no_sche:
+                scheduler.step()
+
+            saver.base_url = os.path.join(opt.saved_path, 'results', '%03d' % epoch)
+
+            if epoch % opt.val_interval == 0:
+                model.model_canet.eval()
+                loss_ls = []
+                psnrs = []
+                ssims = []
+
+                for iter, (data, target, name) in enumerate(val_generator):
+                    with torch.no_grad():
+                        if opt.num_gpus == 1:
+                            data = data.squeeze(0).cuda()
+                            target = target.cuda()
+
+                        image_out, color_loss1, color_loss2, restor_loss, \
+                        psnr, ssim = model(data, target, training=False)
+                        saver.save_image(image_out, name=os.path.splitext(name[0])[0] + '_out')
+                        saver.save_image(data, name=os.path.splitext(name[0])[0] + '_in')
+                        saver.save_image(target, name=os.path.splitext(name[0])[0] + '_gt')
+
+                        loss = restor_loss + color_loss1 + color_loss2
+                        loss_ls.append(loss.item())
+                        psnrs.append(psnr)
+                        ssims.append(ssim)
+
+                loss = np.mean(np.array(loss_ls))
+                psnr = np.mean(np.array(psnrs))
+                ssim = np.mean(np.array(ssims))
+
+                print(
+                    'Val. Epoch: {}/{}. Loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
+                        epoch, opt.num_epochs, loss, psnr, ssim))
+                writer.add_scalar('Loss/val', loss, step)
+                writer.add_scalar('PSNR/val', psnr, step)
+                writer.add_scalar('SSIM/val', ssim, step)
+
+                save_checkpoint(model, f'{opt.model3}_{"%03d" % epoch}_{psnr}_{ssim}_{step}.pth')
+
+                model.model_canet.train()
+
+            opt.test_on_start = False
+            if opt.sampling:
+                exit(0)
+    except KeyboardInterrupt:
+        save_checkpoint(model, f'{opt.model3}_{epoch}_{step}_keyboardInterrupt.pth')
+        writer.close()
+    writer.close()
+
+
+def save_checkpoint(model, name):
+    if isinstance(model, nn.DataParallel):
+        torch.save(model.module.model_canet.state_dict(), os.path.join(opt.saved_path, name))
+    else:
+        torch.save(model.model_canet.state_dict(), os.path.join(opt.saved_path, name))
+
+
+if __name__ == '__main__':
+    opt = get_args()
+    train(opt)

train_IAN.py

@@ -0,0 +1,255 @@
+import argparse
+import datetime
+import os
+import traceback
+
+import kornia
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+
+import models
+from datasets import LowLightDataset, LowLightFDataset
+from models import PSNR, SSIM, CosineLR
+from tools import SingleSummaryWriter
+from tools import saver, mutils
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Breaking Downing the Darkness')
+    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus being used')
+    parser.add_argument('--num_workers', type=int, default=12, help='num_workers of dataloader')
+    parser.add_argument('--batch_size', type=int, default=1, help='The number of images per batch among all devices')
+    parser.add_argument('-m', '--model', type=str, default='INet',
+                        help='Model Name')
+    parser.add_argument('--comment', type=str, default='default',
+                        help='Project comment')
+    parser.add_argument('--graph', action='store_true')
+    parser.add_argument('--scratch', action='store_true')
+
+    parser.add_argument('--lr', type=float, default=0.01)
+    parser.add_argument('--no_sche', action='store_true')
+
+    parser.add_argument('--optim', type=str, default='adam', help='select optimizer for training, '
+                                                                  'suggest using \'admaw\' until the'
+                                                                  ' very final stage then switch to \'sgd\'')
+    parser.add_argument('--num_epochs', type=int, default=500)
+    parser.add_argument('--val_interval', type=int, default=1, help='Number of epoches between valing phases')
+    parser.add_argument('--save_interval', type=int, default=500, help='Number of steps between saving')
+    parser.add_argument('--data_path', type=str, default='./data/LOL',
+                        help='the root folder of dataset')
+    parser.add_argument('--log_path', type=str, default='logs/')
+    parser.add_argument('--saved_path', type=str, default='logs/')
+    args = parser.parse_args()
+    return args
+
+
+def compute_gradient(img):
+    gradx = img[..., 1:, :] - img[..., :-1, :]
+    grady = img[..., 1:] - img[..., :-1]
+    return gradx, grady
+
+
+class ModelINet(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.restor_loss = models.MSELoss()
+        self.wtv_loss = models.WTVLoss2()
+        self.model = model(in_channels=1, out_channels=1)
+        self.eps = 1e-2
+
+    def forward(self, image, image_gt, training=True):
+        if training:
+            image = image.squeeze(0)
+            image_gt = image_gt.repeat(8, 1, 1, 1)
+
+        texture_in, _, _ = torch.split(kornia.color.rgb_to_ycbcr(image), 1, dim=1)
+        texture_gt, _, _ = torch.split(kornia.color.rgb_to_ycbcr(image_gt), 1, dim=1)
+
+        texture_in_down = F.interpolate(texture_in, scale_factor=0.5, mode='bicubic', align_corners=True)
+        texture_gt_down = F.interpolate(texture_gt, scale_factor=0.5, mode='bicubic', align_corners=True)
+
+        illumi = self.model(texture_in_down)
+
+        texture_out = texture_in_down / torch.clamp_min(illumi, self.eps)
+        restor_loss = self.restor_loss(texture_out, texture_gt_down)
+        restor_loss += self.restor_loss(texture_in_down, texture_gt_down * illumi)
+
+        tv_loss = self.wtv_loss(illumi, texture_in_down)
+        if training:
+            psnr = 0.0
+            ssim = 0.0
+        else:
+            illumi = F.interpolate(illumi, scale_factor=2, mode='bicubic', align_corners=True)
+            texture_out = texture_in / torch.clamp_min(illumi, self.eps)
+
+            psnr = PSNR(texture_out, texture_gt)
+            ssim = SSIM(texture_out, texture_gt).item()
+        return texture_out, illumi, restor_loss, tv_loss, psnr, ssim
+
+
+def train(opt):
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    else:
+        torch.manual_seed(42)
+
+    timestamp = mutils.get_formatted_time()
+    opt.saved_path = opt.saved_path + f'/{opt.comment}/{timestamp}'
+    opt.log_path = opt.log_path + f'/{opt.comment}/{timestamp}/tensorboard/'
+    os.makedirs(opt.log_path, exist_ok=True)
+    os.makedirs(opt.saved_path, exist_ok=True)
+
+    training_params = {'batch_size': opt.batch_size,
+                       'shuffle': True,
+                       'drop_last': True,
+                       'num_workers': opt.num_workers}
+
+    val_params = {'batch_size': 1,
+                  'shuffle': False,
+                  'drop_last': True,
+                  'num_workers': opt.num_workers}
+
+    training_set = LowLightFDataset(os.path.join(opt.data_path, 'train'), image_split='images_aug',
+                                    targets_split='targets')
+    training_generator = DataLoader(training_set, **training_params)
+
+    val_set = LowLightDataset(os.path.join(opt.data_path, 'eval'), targets_split='targets')
+    val_generator = DataLoader(val_set, **val_params)
+
+    model = getattr(models, opt.model)
+
+    model = ModelINet(model)
+    print(model)
+    # load last weights
+
+    writer = SingleSummaryWriter(opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/')
+
+    if opt.num_gpus > 0:
+        model = model.cuda()
+        if opt.num_gpus > 1:
+            model = nn.DataParallel(model)
+
+    if opt.optim == 'adam':
+        optimizer = torch.optim.Adam(model.parameters(), opt.lr)
+    else:
+        optimizer = torch.optim.SGD(model.parameters(), opt.lr, momentum=0.9, nesterov=True)
+
+    scheduler = CosineLR(optimizer, opt.lr, opt.num_epochs)
+    epoch = 0
+    step = 0
+    model.train()
+
+    num_iter_per_epoch = len(training_generator)
+
+    try:
+        for epoch in range(opt.num_epochs):
+            last_epoch = step // num_iter_per_epoch
+            if epoch < last_epoch:
+                continue
+
+            epoch_loss = []
+            progress_bar = tqdm(training_generator)
+            for iter, (data, target, name) in enumerate(progress_bar):
+                if iter < step - last_epoch * num_iter_per_epoch:
+                    progress_bar.update()
+                    continue
+                try:
+                    if opt.num_gpus == 1:
+                        data, target = data.cuda(), target.cuda()
+
+                    optimizer.zero_grad()
+
+                    texture_out, texture_attention, restor_loss, \
+                    tv_loss, psnr, ssim = model(data, target, training=True)
+                    loss = restor_loss + tv_loss
+                    loss.backward()
+                    optimizer.step()
+
+                    epoch_loss.append(float(loss))
+
+                    progress_bar.set_description(
+                        'Step: {}. Epoch: {}/{}. Iteration: {}/{}. var: {:.5f}, res_loss: {:.5f}, tv_loss: {:.5f}, psnr: {:.3f}, ssim: {:.3f}'.format(
+                            step, epoch, opt.num_epochs, iter + 1, num_iter_per_epoch, torch.var(texture_attention),
+                            restor_loss.item(),
+                            tv_loss.item(), psnr, ssim))
+                    writer.add_scalar('Loss/train', loss, step)
+                    writer.add_scalar('PSNR/train', psnr, step)
+                    writer.add_scalar('SSIM/train', ssim, step)
+
+                    # log learning_rate
+                    current_lr = optimizer.param_groups[0]['lr']
+                    writer.add_scalar('learning_rate', current_lr, step)
+
+                    step += 1
+
+                except Exception as e:
+                    print('[Error]', traceback.format_exc())
+                    print(e)
+                    continue
+
+            if opt.no_sche:
+                scheduler.step()
+
+            saver.base_url = os.path.join(opt.saved_path, 'results', '%03d' % epoch)
+
+            if epoch % opt.val_interval == 0:
+                model.eval()
+                loss_ls = []
+                psnrs = []
+                ssims = []
+
+                for iter, (data, target, name) in enumerate(val_generator):
+                    with torch.no_grad():
+                        if opt.num_gpus == 1:
+                            data = data.cuda()
+                            target = target.cuda()
+                        texture_in, _, _ = torch.split(kornia.color.rgb_to_ycbcr(data), 1, dim=1)
+                        texture_gt, _, _ = torch.split(kornia.color.rgb_to_ycbcr(target), 1, dim=1)
+
+                        texture_out, texture_attention, restor_loss, \
+                        tv_loss, psnr, ssim = model(data, target, training=False)
+                        saver.save_image(texture_out, name=os.path.splitext(name[0])[0] + '_out')
+                        saver.save_image(texture_in, name=os.path.splitext(name[0])[0] + '_in')
+                        saver.save_image(texture_gt, name=os.path.splitext(name[0])[0] + '_gt')
+                        saver.save_image(texture_attention, name=os.path.splitext(name[0])[0] + '_att')
+
+                        loss = restor_loss + tv_loss
+                        loss_ls.append(loss.item())
+                        psnrs.append(psnr)
+                        ssims.append(ssim)
+
+                loss = np.mean(np.array(loss_ls))
+                psnr = np.mean(np.array(psnrs))
+                ssim = np.mean(np.array(ssims))
+
+                print(
+                    'Val. Epoch: {}/{}. Loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
+                        epoch, opt.num_epochs, loss, psnr, ssim))
+                writer.add_scalar('Loss/val', loss, step)
+                writer.add_scalar('PSNR/val', psnr, step)
+                writer.add_scalar('SSIM/val', ssim, step)
+
+                save_checkpoint(model, f'{opt.model}_{"%03d" % epoch}_{psnr}_{ssim}_{step}.pth')
+
+                model.train()
+
+    except KeyboardInterrupt:
+        save_checkpoint(model, f'{opt.model}_{epoch}_{step}_keyboardInterrupt.pth')
+        writer.close()
+    writer.close()
+
+
+def save_checkpoint(model, name):
+    if isinstance(model, nn.DataParallel):
+        torch.save(model.module.model.state_dict(), os.path.join(opt.saved_path, name))
+    else:
+        torch.save(model.model.state_dict(), os.path.join(opt.saved_path, name))
+
+
+if __name__ == '__main__':
+    opt = get_args()
+    train(opt)

train_MECAN.py

@@ -0,0 +1,261 @@
+import argparse
+import datetime
+import os
+import traceback
+
+import kornia
+import numpy as np
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+
+import models
+from datasets import MEFDataset
+from models import PSNR, SSIM, CosineLR
+from tools import SingleSummaryWriter
+from tools import saver, mutils
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Breaking Downing the Darkness')
+    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus being used')
+    parser.add_argument('--num_workers', type=int, default=12, help='num_workers of dataloader')
+    parser.add_argument('--batch_size', type=int, default=1, help='The number of images per batch among all devices')
+    parser.add_argument('-m', '--model', type=str, default='INet',
+                        help='Model Name')
+    parser.add_argument('-ts', '--targets_split', type=str, default='targets',
+                        help='dir of targets')
+    parser.add_argument('--comment', type=str, default='default',
+                        help='Project comment')
+    parser.add_argument('--graph', action='store_true')
+    parser.add_argument('--scratch', action='store_true')
+    parser.add_argument('--sampling', action='store_true')
+    parser.add_argument('--test_on_start', action='store_true')
+
+    parser.add_argument('--lr', type=float, default=0.01)
+    parser.add_argument('--no_sche', action='store_true')
+
+    parser.add_argument('--optim', type=str, default='adam', help='select optimizer for training, '
+                                                                  'suggest using \'admaw\' until the'
+                                                                  ' very final stage then switch to \'sgd\'')
+    parser.add_argument('--num_epochs', type=int, default=500)
+    parser.add_argument('--val_interval', type=int, default=1, help='Number of epoches between valing phases')
+    parser.add_argument('--save_interval', type=int, default=500, help='Number of steps between saving')
+    parser.add_argument('--data_path', type=str, default='./data/SICE',
+                        help='the root folder of dataset')
+    parser.add_argument('--log_path', type=str, default='logs/')
+    parser.add_argument('--saved_path', type=str, default='logs/')
+    parser.add_argument('-r', '--resize', type=int, default=-1, help='resize of training images')
+    args = parser.parse_args()
+    return args
+
+
+def compute_gradient(img):
+    gradx = img[..., 1:, :] - img[..., :-1, :]
+    grady = img[..., 1:] - img[..., :-1]
+    return gradx, grady
+
+
+class ModelINet(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.color_loss = models.L1Loss()
+        self.restor_loss = models.MSSSIML1Loss(channels=3)
+        self.model_canet = model(in_channels=4, out_channels=2)
+        self.eps = 1e-2
+
+    def load_weight(self, model, weight_pth):
+        state_dict = torch.load(weight_pth)
+        ret = model.load_state_dict(state_dict, strict=True)
+        print(ret)
+
+    def forward(self, image, image_gt, training=True):
+        texture_in, cb_in, cr_in = torch.split(kornia.color.rgb_to_ycbcr(image), 1, dim=1)
+        texture_gt, cb_gt, cr_gt = torch.split(kornia.color.rgb_to_ycbcr(image_gt), 1, dim=1)
+
+        colors = self.model_canet(torch.cat([texture_in, cb_in, cr_in, texture_gt], dim=1))
+        cb, cr = torch.split(colors, 1, dim=1)
+
+        color_loss1 = self.color_loss(cb, cb_gt)
+        color_loss2 = self.color_loss(cr, cr_gt)
+
+        image_out = kornia.color.ycbcr_to_rgb(torch.cat([texture_gt, cb, cr], dim=1))
+        restor_loss = self.restor_loss(image_out, image_gt)
+
+        psnr = PSNR(image_out, image_gt)
+        ssim = SSIM(image_out, image_gt).item()
+        return image_out, color_loss1, color_loss2, restor_loss, psnr, ssim
+
+
+def train(opt):
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    else:
+        torch.manual_seed(42)
+
+    # params.project_name = params.project_name + str(time.time()).replace('.', '')
+    timestamp = mutils.get_formatted_time()
+    opt.saved_path = opt.saved_path + f'/{opt.comment}/{timestamp}'
+    opt.log_path = opt.log_path + f'/{opt.comment}/{timestamp}/tensorboard/'
+    os.makedirs(opt.log_path, exist_ok=True)
+    os.makedirs(opt.saved_path, exist_ok=True)
+
+    training_params = {'batch_size': opt.batch_size,
+                       'shuffle': True,
+                       'drop_last': True,
+                       'num_workers': opt.num_workers}
+
+    val_params = {'batch_size': 1,
+                  'shuffle': False,
+                  'drop_last': True,
+                  'num_workers': opt.num_workers}
+
+    training_set = MEFDataset(os.path.join(opt.data_path, 'train'))
+    training_generator = DataLoader(training_set, **training_params)
+
+    val_set = MEFDataset(os.path.join(opt.data_path, 'eval'))
+    val_generator = DataLoader(val_set, **val_params)
+
+    model = getattr(models, opt.model)
+
+    model = ModelINet(model)
+    print(model)
+
+    writer = SingleSummaryWriter(opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/')
+
+    if opt.num_gpus > 0:
+        model = model.cuda()
+        if opt.num_gpus > 1:
+            model = nn.DataParallel(model)
+
+    if opt.optim == 'adam':
+        optimizer = torch.optim.Adam(model.model_canet.parameters(), opt.lr)
+    else:
+        optimizer = torch.optim.SGD(model.model_canet.parameters(), opt.lr, momentum=0.9, nesterov=True)
+
+    scheduler = CosineLR(optimizer, opt.lr, opt.num_epochs)
+    epoch = 0
+    step = 0
+    model.model_canet.train()
+
+    num_iter_per_epoch = len(training_generator)
+
+    try:
+        for epoch in range(opt.num_epochs):
+            last_epoch = step // num_iter_per_epoch
+            if epoch < last_epoch:
+                continue
+
+            epoch_loss = []
+            progress_bar = tqdm(training_generator)
+            if not opt.sampling and not opt.test_on_start:
+                for iter, (data, target, name1, name2) in enumerate(progress_bar):
+                    if iter < step - last_epoch * num_iter_per_epoch:
+                        progress_bar.update()
+                        continue
+                    try:
+                        if opt.num_gpus == 1:
+                            data, target = data.cuda(), target.cuda()
+                        optimizer.zero_grad()
+
+                        image_out, color_loss1, color_loss2, \
+                        restor_loss, psnr, ssim = model(data, target, training=True)
+                        loss = color_loss1 + color_loss2 + restor_loss
+
+                        loss.backward()
+                        optimizer.step()
+
+                        epoch_loss.append(float(loss))
+
+                        progress_bar.set_description(
+                            'Step: {}. Epoch: {}/{}. Iteration: {}/{}. color_loss1: {:1.5f}, color_loss2: {:1.5f}, restor_loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
+                                step, epoch, opt.num_epochs, iter + 1, num_iter_per_epoch,
+                                color_loss1.item(), color_loss2.item(),
+                                restor_loss.item(), psnr, ssim))
+                        writer.add_scalar('Loss/train', loss, step)
+                        writer.add_scalar('PSNR/train', psnr, step)
+                        writer.add_scalar('SSIM/train', ssim, step)
+
+                        # log learning_rate
+                        current_lr = optimizer.param_groups[0]['lr']
+                        writer.add_scalar('learning_rate', current_lr, step)
+
+                        step += 1
+
+                    except Exception as e:
+                        print('[Error]', traceback.format_exc())
+                        print(e)
+                        continue
+
+            if opt.no_sche:
+                scheduler.step()
+
+            saver.base_url = os.path.join(opt.saved_path, 'results', '%03d' % epoch)
+
+            if epoch % opt.val_interval == 0:
+                model.model_canet.eval()
+                loss_ls = []
+                psnrs = []
+                ssims = []
+
+                for iter, (data, target, name1, name2) in enumerate(val_generator):
+                    with torch.no_grad():
+                        if opt.num_gpus == 1:
+                            data, target = data.cuda(), target.cuda()
+
+                        image_out, color_loss1, color_loss2, restor_loss, \
+                        psnr, ssim = model(data, target, training=False)
+                        saver.save_image(data, name=os.path.splitext(name1[0])[0] + '_im1')
+                        saver.save_image(target, name=os.path.splitext(name2[0])[0] + '_im2')
+                        saver.save_image(image_out, name=os.path.splitext(name2[0])[0] + '_im2_pred')
+
+                        loss = restor_loss + color_loss1 + color_loss2
+                        loss_ls.append(loss.item())
+                        psnrs.append(psnr)
+                        ssims.append(ssim)
+
+                        # reverse
+                        image_out, color_loss1, color_loss2, restor_loss, \
+                        psnr, ssim = model(target, data, training=False)
+                        saver.save_image(image_out, name=os.path.splitext(name1[0])[0] + '_im1_pred')
+
+                        loss = restor_loss + color_loss1 + color_loss2
+                        loss_ls.append(loss.item())
+                        psnrs.append(psnr)
+                        ssims.append(ssim)
+
+                loss = np.mean(np.array(loss_ls))
+                psnr = np.mean(np.array(psnrs))
+                ssim = np.mean(np.array(ssims))
+
+                print(
+                    'Val. Epoch: {}/{}. Loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
+                        epoch, opt.num_epochs, loss, psnr, ssim))
+                writer.add_scalar('Loss/val', loss, step)
+                writer.add_scalar('PSNR/val', psnr, step)
+                writer.add_scalar('SSIM/val', ssim, step)
+
+                save_checkpoint(model, f'{opt.model}_{"%03d" % epoch}_{psnr}_{ssim}_{step}.pth')
+
+                model.model_canet.train()
+
+            opt.test_on_start = False
+            if opt.sampling:
+                exit(0)
+    except KeyboardInterrupt:
+        save_checkpoint(model, f'{opt.model}_{epoch}_{step}_keyboardInterrupt.pth')
+        writer.close()
+    writer.close()
+
+
+def save_checkpoint(model, name):
+    if isinstance(model, nn.DataParallel):
+        torch.save(model.module.model_canet.state_dict(), os.path.join(opt.saved_path, name))
+    else:
+        torch.save(model.model_canet.state_dict(), os.path.join(opt.saved_path, name))
+
+
+if __name__ == '__main__':
+    opt = get_args()
+    train(opt)

train_MECAN_finetune.py

@@ -0,0 +1,268 @@
+# original author: signatrix
+# adapted from https://github.com/signatrix/efficientdet/blob/master/train.py
+# modified by Zylo117
+
+import argparse
+import datetime
+import os
+import traceback
+
+import kornia
+import numpy as np
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+
+import models
+from datasets import MEFDataset, LowLightDataset, LowLightDatasetReverse
+from models import PSNR, SSIM, CosineLR
+from tools import SingleSummaryWriter
+from tools import saver, mutils
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Breaking Downing the Darkness')
+    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus being used')
+    parser.add_argument('--num_workers', type=int, default=12, help='num_workers of dataloader')
+    parser.add_argument('--batch_size', type=int, default=1, help='The number of images per batch among all devices')
+    parser.add_argument('-m', '--model', type=str, default='INet',
+                        help='Model Name')
+    parser.add_argument('-mw', '--model_weight', type=str, default=None,
+                        help='Model weight')
+    parser.add_argument('-ts', '--targets_split', type=str, default='targets',
+                        help='dir of targets')
+    parser.add_argument('--comment', type=str, default='default',
+                        help='Project comment')
+    parser.add_argument('--graph', action='store_true')
+    parser.add_argument('--scratch', action='store_true')
+    parser.add_argument('--sampling', action='store_true')
+    parser.add_argument('--test_on_start', action='store_true')
+
+    parser.add_argument('--lr', type=float, default=0.01)
+    parser.add_argument('--no_sche', action='store_true')
+
+    parser.add_argument('--optim', type=str, default='adam', help='select optimizer for training, '
+                                                                  'suggest using \'admaw\' until the'
+                                                                  ' very final stage then switch to \'sgd\'')
+    parser.add_argument('--num_epochs', type=int, default=500)
+    parser.add_argument('--val_interval', type=int, default=1, help='Number of epoches between valing phases')
+    parser.add_argument('--save_interval', type=int, default=500, help='Number of steps between saving')
+    parser.add_argument('--data_path1', type=str, default='./data/SICE',
+                        help='the root folder of dataset')
+    parser.add_argument('--data_path2', type=str, default='./data/LOL',
+                        help='the root folder of dataset')
+    parser.add_argument('--log_path', type=str, default='logs/')
+
+    parser.add_argument('--saved_path', type=str, default='logs/')
+    args = parser.parse_args()
+    return args
+
+
+def compute_gradient(img):
+    gradx = img[..., 1:, :] - img[..., :-1, :]
+    grady = img[..., 1:] - img[..., :-1]
+    return gradx, grady
+
+
+class ModelINet(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.color_loss = models.SSIML1Loss(channels=1)
+        self.restor_loss = models.SSIML1Loss(channels=3)
+        self.model_canet = model(in_channels=4, out_channels=2)
+        self.eps = 1e-2
+        self.load_weight(self.model_canet, opt.model_weight)
+
+    def load_weight(self, model, weight_pth):
+        state_dict = torch.load(weight_pth)
+        ret = model.load_state_dict(state_dict, strict=True)
+        print(ret)
+
+    def forward(self, image, image_gt, training=True):
+        texture_in, cb_in, cr_in = torch.split(kornia.color.rgb_to_ycbcr(image), 1, dim=1)
+        texture_gt, cb_gt, cr_gt = torch.split(kornia.color.rgb_to_ycbcr(image_gt), 1, dim=1)
+
+        colors = self.model_canet(torch.cat([texture_in, cb_in, cr_in, texture_gt], dim=1))
+        cb, cr = torch.split(colors, 1, dim=1)
+
+        color_loss1 = self.color_loss(cb, cb_gt)
+        color_loss2 = self.color_loss(cr, cr_gt)
+
+        image_out = kornia.color.ycbcr_to_rgb(torch.cat([texture_gt, cb, cr], dim=1))
+        restor_loss = self.restor_loss(image_out, image_gt)
+
+        psnr = PSNR(image_out, image_gt)
+        ssim = SSIM(image_out, image_gt).item()
+        return image_out, color_loss1, color_loss2, restor_loss, psnr, ssim
+
+
+def train(opt):
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    else:
+        torch.manual_seed(42)
+
+    # params.project_name = params.project_name + str(time.time()).replace('.', '')
+    timestamp = mutils.get_formatted_time()
+    opt.saved_path = opt.saved_path + f'/{opt.comment}/{timestamp}'
+    opt.log_path = opt.log_path + f'/{opt.comment}/{timestamp}/tensorboard/'
+    os.makedirs(opt.log_path, exist_ok=True)
+    os.makedirs(opt.saved_path, exist_ok=True)
+
+    training_params = {'batch_size': opt.batch_size,
+                       'shuffle': True,
+                       'drop_last': True,
+                       'num_workers': opt.num_workers}
+
+    val_params = {'batch_size': 1,
+                  'shuffle': False,
+                  'drop_last': True,
+                  'num_workers': opt.num_workers}
+
+    training_set1 = MEFDataset(os.path.join(opt.data_path1, 'train'))
+    training_set2 = LowLightDataset(os.path.join(opt.data_path2, 'train'), color_tuning=True,
+                                    targets_split=opt.targets_split)
+    training_set3 = LowLightDatasetReverse(os.path.join(opt.data_path2, 'train'), color_tuning=True,
+                                           targets_split=opt.targets_split)
+    training_set = torch.utils.data.ConcatDataset([training_set1, training_set2, training_set3])
+    training_generator = DataLoader(training_set, **training_params)
+
+    # val_set = MEFDataset(os.path.join(opt.data_path1, 'eval'))
+    val_set = LowLightDataset(os.path.join(opt.data_path2, 'eval'), color_tuning=True)
+    # val_set = torch.utils.data.ConcatDataset([val_set1, val_set2])
+    val_generator = DataLoader(val_set, **val_params)
+
+    model = getattr(models, opt.model)
+
+    model = ModelINet(model)
+    print(model)
+
+    writer = SingleSummaryWriter(opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/')
+
+
+    if opt.num_gpus > 0:
+        model = model.cuda()
+        if opt.num_gpus > 1:
+            model = nn.DataParallel(model)
+
+    if opt.optim == 'adam':
+        optimizer = torch.optim.Adam(model.model_canet.parameters(), opt.lr)
+    else:
+        optimizer = torch.optim.SGD(model.model_canet.parameters(), opt.lr, momentum=0.9, nesterov=True)
+
+    scheduler = CosineLR(optimizer, opt.lr, opt.num_epochs)
+    epoch = 0
+    step = 0
+    model.model_canet.train()
+
+    num_iter_per_epoch = len(training_generator)
+
+    try:
+        for epoch in range(opt.num_epochs):
+            last_epoch = step // num_iter_per_epoch
+            if epoch < last_epoch:
+                continue
+
+            epoch_loss = []
+            progress_bar = tqdm(training_generator)
+            if not opt.sampling and not opt.test_on_start:
+                for iter, (data, target, name1, name2) in enumerate(progress_bar):
+                    if iter < step - last_epoch * num_iter_per_epoch:
+                        progress_bar.update()
+                        continue
+                    try:
+                        if opt.num_gpus == 1:
+                            data, target = data.cuda(), target.cuda()
+                        optimizer.zero_grad()
+
+                        image_out, color_loss1, color_loss2, \
+                        restor_loss, psnr, ssim = model(data, target, training=True)
+                        loss = color_loss1 + color_loss2 + restor_loss
+                        loss.backward()
+                        optimizer.step()
+
+                        epoch_loss.append(float(loss))
+
+                        progress_bar.set_description(
+                            'Step: {}. Epoch: {}/{}. Iteration: {}/{}. color_loss1: {:1.5f}, color_loss2: {:1.5f}, restor_loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
+                                step, epoch, opt.num_epochs, iter + 1, num_iter_per_epoch,
+                                color_loss1.item(), color_loss2.item(),
+                                restor_loss.item(), psnr, ssim))
+                        writer.add_scalar('Loss/train', loss, step)
+                        writer.add_scalar('PSNR/train', psnr, step)
+                        writer.add_scalar('SSIM/train', ssim, step)
+
+                        # log learning_rate
+                        current_lr = optimizer.param_groups[0]['lr']
+                        writer.add_scalar('learning_rate', current_lr, step)
+
+                        step += 1
+
+                    except Exception as e:
+                        print('[Error]', traceback.format_exc())
+                        print(e)
+                        continue
+                    # scheduler.step(np.mean(epoch_loss))
+
+            if opt.no_sche:
+                scheduler.step()
+
+            saver.base_url = os.path.join(opt.saved_path, 'results', '%03d' % epoch)
+
+            if epoch % opt.val_interval == 0:
+                model.model_canet.eval()
+                loss_ls = []
+                psnrs = []
+                ssims = []
+
+                for iter, (data, target, name1, name2) in enumerate(val_generator):
+                    with torch.no_grad():
+                        if opt.num_gpus == 1:
+                            data, target = data.cuda(), target.cuda()
+
+                        image_out, color_loss1, color_loss2, restor_loss, \
+                        psnr, ssim = model(data, target, training=False)
+                        saver.save_image(data, name=os.path.splitext(name1[0])[0] + '_im1')
+                        saver.save_image(target, name=os.path.splitext(name2[0])[0] + '_im2')
+                        saver.save_image(image_out, name=os.path.splitext(name2[0])[0] + '_im2_pred')
+
+                        loss = restor_loss + color_loss1 + color_loss2
+                        loss_ls.append(loss.item())
+                        psnrs.append(psnr)
+                        ssims.append(ssim)
+
+                loss = np.mean(np.array(loss_ls))
+                psnr = np.mean(np.array(psnrs))
+                ssim = np.mean(np.array(ssims))
+
+                print(
+                    'Val. Epoch: {}/{}. Loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
+                        epoch, opt.num_epochs, loss, psnr, ssim))
+                writer.add_scalar('Loss/val', loss, step)
+                writer.add_scalar('PSNR/val', psnr, step)
+                writer.add_scalar('SSIM/val', ssim, step)
+
+                save_checkpoint(model, f'{opt.model}_{"%03d" % epoch}_{psnr}_{ssim}_{step}.pth')
+
+                model.model_canet.train()
+
+            opt.test_on_start = False
+            if opt.sampling:
+                exit(0)
+    except KeyboardInterrupt:
+        save_checkpoint(model, f'{opt.model}_{epoch}_{step}_keyboardInterrupt.pth')
+        writer.close()
+    writer.close()
+
+
+def save_checkpoint(model, name):
+    if isinstance(model, nn.DataParallel):
+        torch.save(model.module.model_canet.state_dict(), os.path.join(opt.saved_path, name))
+    else:
+        torch.save(model.model_canet.state_dict(), os.path.join(opt.saved_path, name))
+
+
+if __name__ == '__main__':
+    opt = get_args()
+    train(opt)

train_NFM.py

@@ -0,0 +1,279 @@
+import argparse
+import datetime
+import os
+import traceback
+
+import kornia
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+
+import models
+from datasets import LowLightDataset, LowLightFDataset
+from models import PSNR, SSIM, CosineLR
+from tools import SingleSummaryWriter
+from tools import saver, mutils
+
+
+def get_args():
+    parser = argparse.ArgumentParser('Breaking Downing the Darkness')
+    parser.add_argument('--num_gpus', type=int, default=1, help='number of gpus being used')
+    parser.add_argument('--num_workers', type=int, default=12, help='num_workers of dataloader')
+    parser.add_argument('--batch_size', type=int, default=1, help='The number of images per batch among all devices')
+    parser.add_argument('-m1', '--model1', type=str, default='INet',
+                        help='Model1 Name')
+    parser.add_argument('-m2', '--model2', type=str, default='NSNet',
+                        help='Model2 Name')
+    parser.add_argument('-m3', '--model3', type=str, default='NSNet',
+                        help='Model3 Name')
+
+    parser.add_argument('-m1w', '--model1_weight', type=str, default=None,
+                        help='Model Name')
+    parser.add_argument('-m2w', '--model2_weight', type=str, default=None,
+                        help='Model Name')
+
+    parser.add_argument('--comment', type=str, default='default',
+                        help='Project comment')
+    parser.add_argument('--graph', action='store_true')
+    parser.add_argument('--no_sche', action='store_true')
+    parser.add_argument('--sampling', action='store_true')
+
+    parser.add_argument('--slope', type=float, default=2.)
+    parser.add_argument('--lr', type=float, default=1e-4)
+    parser.add_argument('--optim', type=str, default='adam', help='select optimizer for training, '
+                                                                  'suggest using \'admaw\' until the'
+                                                                  ' very final stage then switch to \'sgd\'')
+    parser.add_argument('--num_epochs', type=int, default=500)
+    parser.add_argument('--val_interval', type=int, default=1, help='Number of epoches between valing phases')
+    parser.add_argument('--save_interval', type=int, default=500, help='Number of steps between saving')
+    parser.add_argument('--data_path', type=str, default='./data/LOL',
+                        help='the root folder of dataset')
+    parser.add_argument('--log_path', type=str, default='logs/')
+    parser.add_argument('--saved_path', type=str, default='logs/')
+    args = parser.parse_args()
+    return args
+
+
+class ModelNSNet(nn.Module):
+    def __init__(self, model1, model2, model3):
+        super().__init__()
+        self.texture_loss = models.SSIML1Loss(channels=1)
+        self.model_ianet = model1(in_channels=1, out_channels=1)
+        self.model_nsnet = model2(in_channels=2, out_channels=1)
+        self.model_fusenet = model3(in_channels=3, out_channels=1)
+
+        assert opt.model1_weight is not None
+        self.load_weight(self.model_ianet, opt.model1_weight)
+        self.load_weight(self.model_nsnet, opt.model2_weight)
+        self.model_ianet.eval()
+        self.model_nsnet.eval()
+        self.eps = 1e-2
+
+    def load_weight(self, model, weight_pth):
+        state_dict = torch.load(weight_pth)
+        ret = model.load_state_dict(state_dict, strict=True)
+        print(ret)
+
+    def noise_syn(self, illumi, strength):
+        return torch.exp(-illumi) * strength
+
+    def forward(self, image, image_gt, training=True):
+        texture_nss = []
+        with torch.no_grad():
+            if training:
+                image = image.squeeze(0)
+                image_gt = image_gt.repeat(8, 1, 1, 1)
+
+            texture_in, _, _ = torch.split(kornia.color.rgb_to_ycbcr(image), 1, dim=1)
+            texture_gt, _, _ = torch.split(kornia.color.rgb_to_ycbcr(image_gt), 1, dim=1)
+
+            texture_in_down = F.interpolate(texture_in, scale_factor=0.5, mode='bicubic', align_corners=True)
+            illumi = self.model_ianet(texture_in_down)
+            illumi = F.interpolate(illumi, scale_factor=2, mode='bicubic', align_corners=True)
+            noisy_gt = texture_in / torch.clamp_min(illumi, self.eps)
+
+            for strength in [0, 0.05, 0.1]:
+                illumi = torch.clamp(illumi, 0., 1.)
+                attention = self.noise_syn(illumi, strength=strength)
+                texture_res = self.model_nsnet(torch.cat([noisy_gt, attention], dim=1))
+                texture_ns = noisy_gt + texture_res
+                texture_nss.append(texture_ns)
+
+            texture_nss = torch.cat(texture_nss, dim=1).detach()
+
+        texture_fuse = self.model_fusenet(texture_nss)
+        restor_loss = self.texture_loss(texture_fuse, texture_gt)
+        psnr = PSNR(texture_fuse, texture_gt)
+        ssim = SSIM(texture_fuse, texture_gt).item()
+        return noisy_gt, texture_nss, texture_fuse, texture_res, illumi, restor_loss, psnr, ssim
+
+
+def train(opt):
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+    else:
+        torch.manual_seed(42)
+
+    timestamp = mutils.get_formatted_time()
+    opt.saved_path = opt.saved_path + f'/{opt.comment}/{timestamp}'
+    opt.log_path = opt.log_path + f'/{opt.comment}/{timestamp}/tensorboard/'
+    os.makedirs(opt.log_path, exist_ok=True)
+    os.makedirs(opt.saved_path, exist_ok=True)
+
+    training_params = {'batch_size': opt.batch_size,
+                       'shuffle': True,
+                       'drop_last': True,
+                       'num_workers': opt.num_workers}
+
+    val_params = {'batch_size': 1,
+                  'shuffle': False,
+                  'drop_last': True,
+                  'num_workers': opt.num_workers}
+
+    training_set = LowLightFDataset(os.path.join(opt.data_path, 'train'), image_split='images_aug')
+    training_generator = DataLoader(training_set, **training_params)
+
+    val_set = LowLightDataset(os.path.join(opt.data_path, 'eval'))
+    val_generator = DataLoader(val_set, **val_params)
+
+    model1 = getattr(models, opt.model1)
+    model2 = getattr(models, opt.model2)
+    model3 = getattr(models, opt.model3)
+    writer = SingleSummaryWriter(opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/')
+
+    model = ModelNSNet(model1, model2, model3)
+    print(model)
+
+    if opt.num_gpus > 0:
+        model = model.cuda()
+        if opt.num_gpus > 1:
+            model = nn.DataParallel(model)
+
+    if opt.optim == 'adam':
+        optimizer = torch.optim.Adam(model.model_fusenet.parameters(), opt.lr)
+    else:
+        optimizer = torch.optim.SGD(model.model_fusenet.parameters(), opt.lr, momentum=0.9, nesterov=True)
+
+    scheduler = CosineLR(optimizer, opt.lr, opt.num_epochs)
+    epoch = 0
+    step = 0
+    model.model_fusenet.train()
+
+    num_iter_per_epoch = len(training_generator)
+
+    try:
+        for epoch in range(opt.num_epochs):
+            last_epoch = step // num_iter_per_epoch
+            if epoch < last_epoch:
+                continue
+
+            epoch_loss = []
+            progress_bar = tqdm(training_generator)
+
+            saver.base_url = os.path.join(opt.saved_path, 'results', '%03d' % epoch)
+            if not opt.sampling:
+                for iter, (data, target, name) in enumerate(progress_bar):
+                    if iter < step - last_epoch * num_iter_per_epoch:
+                        progress_bar.update()
+                        continue
+                    try:
+                        if opt.num_gpus == 1:
+                            data = data.cuda()
+                            target = target.cuda()
+
+                        optimizer.zero_grad()
+
+                        noisy_gt, texture_nss, texture_fuse, texture_res, \
+                        illumi, restor_loss, psnr, ssim = model(data, target, training=True)
+
+                        loss = restor_loss
+                        loss.backward()
+                        optimizer.step()
+
+                        epoch_loss.append(float(loss))
+
+                        progress_bar.set_description(
+                            'Step: {}. Epoch: {}/{}. Iteration: {}/{}. restor_loss: {:.5f},  psnr: {:.5f}, ssim: {:.5f}'.format(
+                                step, epoch, opt.num_epochs, iter + 1, num_iter_per_epoch, restor_loss.item(), psnr,
+                                ssim))
+                        writer.add_scalar('Loss/train', loss, step)
+                        writer.add_scalar('PSNR/train', psnr, step)
+                        writer.add_scalar('SSIM/train', ssim, step)
+
+                        # log learning_rate
+                        current_lr = optimizer.param_groups[0]['lr']
+                        writer.add_scalar('learning_rate', current_lr, step)
+
+                        step += 1
+
+                    except Exception as e:
+                        print('[Error]', traceback.format_exc())
+                        print(e)
+                        continue
+
+            if not opt.no_sche:
+                scheduler.step()
+
+            if epoch % opt.val_interval == 0:
+                model.model_fusenet.eval()
+                loss_ls = []
+                psnrs = []
+                ssims = []
+
+                for iter, (data, target, name) in enumerate(val_generator):
+                    with torch.no_grad():
+                        if opt.num_gpus == 1:
+                            data = data.cuda()
+                            target = target.cuda()
+
+                        noisy_gt, texture_nss, texture_fuse, texture_res, \
+                        illumi, restor_loss, psnr, ssim = model(data, target, training=False)
+                        texture_gt, _, _ = torch.split(kornia.color.rgb_to_ycbcr(target), 1, dim=1)
+
+                        saver.save_image(noisy_gt, name=os.path.splitext(name[0])[0] + '_in')
+                        saver.save_image(texture_nss.transpose(0, 1), name=os.path.splitext(name[0])[0] + '_ns')
+                        saver.save_image(texture_fuse, name=os.path.splitext(name[0])[0] + '_fuse')
+                        saver.save_image(texture_res, name=os.path.splitext(name[0])[0] + '_res')
+                        saver.save_image(illumi, name=os.path.splitext(name[0])[0] + '_ill')
+                        saver.save_image(target, name=os.path.splitext(name[0])[0] + '_gt')
+
+                        loss = restor_loss
+                        loss_ls.append(loss.item())
+                        psnrs.append(psnr)
+                        ssims.append(ssim)
+
+                loss = np.mean(np.array(loss_ls))
+                psnr = np.mean(np.array(psnrs))
+                ssim = np.mean(np.array(ssims))
+
+                print(
+                    'Val. Epoch: {}/{}. Loss: {:1.5f}, psnr: {:.5f}, ssim: {:.5f}'.format(
+                        epoch, opt.num_epochs, loss, psnr, ssim))
+                writer.add_scalar('Loss/val', loss, step)
+                writer.add_scalar('PSNR/val', psnr, step)
+                writer.add_scalar('SSIM/val', ssim, step)
+
+                save_checkpoint(model, f'{opt.model3}_{"%03d" % epoch}_{psnr}_{ssim}_{step}.pth')
+
+                model.model_fusenet.train()
+            if opt.sampling:
+                exit(0)
+    except KeyboardInterrupt:
+        save_checkpoint(model, f'{opt.model3}_{epoch}_{step}_keyboardInterrupt.pth')
+        writer.close()
+    writer.close()
+
+
+def save_checkpoint(model, name):
+    if isinstance(model, nn.DataParallel):
+        torch.save(model.module.model_fusenet.state_dict(), os.path.join(opt.saved_path, name))
+    else:
+        torch.save(model.model_fdnet.state_dict(), os.path.join(opt.saved_path, name))
+
+
+if __name__ == '__main__':
+    opt = get_args()
+    train(opt)