Upload 13 files

checkpoints/FuseNet_CA_MEF_251.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c44373240c845c4a79ad49a7aa3436a53917d012165fc5a94843fc70174b25aa
+size 895263

checkpoints/FuseNet_FD_297.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2b7d29e8b90ac3f6d386d1d6c927a01ff24a95e14f79fa900d8f8fa45647237
+size 892959

checkpoints/IANet_335.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a551ba226137d681e851b91c6717ce51e731b7050fd68143f05a036e5e8e70a
+size 3409487

checkpoints/NSNet_422.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ce8998e5c9b5b3eb5d0d582cfb5874213073b247c2a577ab3416ed42d9d2a72
+size 3410639

datasets/__init__.py

@@ -0,0 +1,4 @@
+from .low_light import *
+from .low_light_test import *
+from .mef import *
+

datasets/low_light.py

@@ -0,0 +1,171 @@
+import os
+import random
+
+import torch
+import torch.utils.data as data
+import torchvision.transforms as T
+from PIL import Image
+
+
+class LowLightFDataset(data.Dataset):
+    def __init__(self, root, image_split='images_aug', targets_split='targets', training=True):
+        self.root = root
+        self.num_instances = 8
+        self.img_root = os.path.join(root, image_split)
+        self.target_root = os.path.join(root, targets_split)
+        self.training = training
+        print('----', image_split, targets_split, '----')
+        self.imgs = list(sorted(os.listdir(self.img_root)))
+        self.gts = list(sorted(os.listdir(self.target_root)))
+
+        names = [img_name.split('_')[0] + '.' + img_name.split('.')[-1] for img_name in self.imgs]
+        self.imgs = list(
+            filter(lambda img_name: img_name.split('_')[0] + '.' + img_name.split('.')[-1] in self.gts, self.imgs))
+
+        self.gts = list(filter(lambda gt: gt in names, self.gts))
+
+        print(len(self.imgs), len(self.gts))
+        self.preproc = T.Compose(
+            [T.ToTensor()]
+        )
+        self.preproc_gt = T.Compose(
+            [T.ToTensor()]
+        )
+
+    def __getitem__(self, idx):
+        fn, ext = self.gts[idx].split('.')
+        imgs = []
+        for i in range(self.num_instances):
+            img_path = os.path.join(self.img_root, f"{fn}_{i}.{ext}")
+            imgs += [self.preproc(Image.open(img_path).convert("RGB"))]
+
+        if self.training:
+            random.shuffle(imgs)
+        gt_path = os.path.join(self.target_root, self.gts[idx])
+        gt = Image.open(gt_path).convert("RGB")
+        gt = self.preproc_gt(gt)
+
+        # print(img_path, gt_path)
+        return torch.stack(imgs, dim=0), gt, fn
+
+    def __len__(self):
+        return len(self.gts)
+
+
+class LowLightFDatasetEval(data.Dataset):
+    def __init__(self, root, targets_split='targets', training=True):
+        self.root = root
+        self.num_instances = 1
+        self.img_root = os.path.join(root, 'images')
+        self.target_root = os.path.join(root, targets_split)
+        self.training = training
+
+        self.imgs = list(sorted(os.listdir(self.img_root)))
+        self.gts = list(sorted(os.listdir(self.target_root)))
+
+        self.imgs = list(filter(lambda img_name: img_name in self.gts, self.imgs))
+        self.gts = list(filter(lambda gt: gt in self.imgs, self.gts))
+
+        print(len(self.imgs), len(self.gts))
+        self.preproc = T.Compose(
+            [T.ToTensor()]
+        )
+        self.preproc_gt = T.Compose(
+            [T.ToTensor()]
+        )
+
+    def __getitem__(self, idx):
+        fn, ext = self.gts[idx].split('.')
+        imgs = []
+        for i in range(self.num_instances):
+            img_path = os.path.join(self.img_root, f"{fn}.{ext}")
+            imgs += [self.preproc(Image.open(img_path).convert("RGB"))]
+
+        gt_path = os.path.join(self.target_root, self.gts[idx])
+        gt = Image.open(gt_path).convert("RGB")
+        gt = self.preproc_gt(gt)
+
+        # print(img_path, gt_path)
+        return torch.stack(imgs, dim=0), gt, fn
+
+    def __len__(self):
+        return len(self.gts)
+
+
+class LowLightDataset(data.Dataset):
+    def __init__(self, root, targets_split='targets', color_tuning=False):
+        self.root = root
+        self.img_root = os.path.join(root, 'images')
+        self.target_root = os.path.join(root, targets_split)
+        self.color_tuning = color_tuning
+        self.imgs = list(sorted(os.listdir(self.img_root)))
+        self.gts = list(sorted(os.listdir(self.target_root)))
+
+        self.imgs = list(filter(lambda img_name: img_name in self.gts, self.imgs))
+        self.gts = list(filter(lambda gt: gt in self.imgs, self.gts))
+
+        print(len(self.imgs), len(self.gts))
+        self.preproc = T.Compose(
+            [T.ToTensor()]
+        )
+        self.preproc_gt = T.Compose(
+            [T.ToTensor()]
+        )
+
+    def __getitem__(self, idx):
+        fn, ext = self.gts[idx].split('.')
+
+        img_path = os.path.join(self.img_root, self.imgs[idx])
+        img = Image.open(img_path).convert("RGB")
+        img = self.preproc(img)
+
+        gt_path = os.path.join(self.target_root, self.gts[idx])
+        gt = Image.open(gt_path).convert("RGB")
+        gt = self.preproc_gt(gt)
+
+        if self.color_tuning:
+            return img, gt, 'a' + self.imgs[idx], 'a' + self.imgs[idx]
+        else:
+            return img, gt, fn
+
+    def __len__(self):
+        return len(self.imgs)
+
+
+class LowLightDatasetReverse(data.Dataset):
+    def __init__(self, root, targets_split='targets', color_tuning=False):
+        self.root = root
+        self.img_root = os.path.join(root, 'images')
+        self.target_root = os.path.join(root, targets_split)
+        self.color_tuning = color_tuning
+        self.imgs = list(sorted(os.listdir(self.img_root)))
+        self.gts = list(sorted(os.listdir(self.target_root)))
+
+        self.imgs = list(filter(lambda img_name: img_name in self.gts, self.imgs))
+        self.gts = list(filter(lambda gt: gt in self.imgs, self.gts))
+
+        print(len(self.imgs), len(self.gts))
+        self.preproc = T.Compose(
+            [T.ToTensor()]
+        )
+        self.preproc_gt = T.Compose(
+            [T.ToTensor()]
+        )
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.img_root, self.imgs[idx])
+        img = Image.open(img_path).convert("RGB")
+        img = self.preproc(img)
+
+        gt_path = os.path.join(self.target_root, self.gts[idx])
+        gt = Image.open(gt_path).convert("RGB")
+        gt = self.preproc_gt(gt)
+
+        if self.color_tuning:
+            return gt, img, 'a' + self.imgs[idx], 'a' + self.imgs[idx]
+        else:
+            fn, ext = os.path.splitext(self.imgs[idx])
+            return gt, img, '%03d' % int(fn) + ext
+
+    def __len__(self):
+        return len(self.imgs)

datasets/low_light_test.py

@@ -0,0 +1,41 @@
+import os
+
+import torch.utils.data as data
+import torchvision.transforms as T
+from PIL import Image
+
+
+class LowLightDatasetTest(data.Dataset):
+    def __init__(self, root, reside=False):
+        self.root = root
+        self.items = []
+
+        subsets = os.listdir(root)
+        for subset in subsets:
+            img_root = os.path.join(root, subset)
+            img_names = list(sorted(os.listdir(img_root)))
+
+            for img_name in img_names:
+                self.items.append((
+                    os.path.join(img_root, img_name),
+                    subset,
+                    img_name
+                ))
+
+        self.preproc = T.Compose(
+            [T.ToTensor()]
+        )
+        self.preproc_raw = T.Compose(
+            [T.ToTensor()]
+        )
+
+    def __getitem__(self, idx):
+        img_path, subset, img_name = self.items[idx]
+        img = Image.open(img_path).convert("RGB")
+        img = img.resize((img.width // 8 * 8, img.height // 8 * 8), Image.ANTIALIAS)
+        img_raw = self.preproc_raw(img)
+
+        return img_raw, subset, img_name
+
+    def __len__(self):
+        return len(self.items)

datasets/mef.py

@@ -0,0 +1,36 @@
+import os
+import random
+
+import torch.utils.data as data
+import torchvision.transforms as T
+from PIL import Image
+
+
+class MEFDataset(data.Dataset):
+    def __init__(self, root):
+        self.img_root = root
+
+        self.numbers = list(sorted(os.listdir(self.img_root)))
+        print(len(self.numbers))
+
+        self.preproc = T.Compose(
+            [T.ToTensor()]
+        )
+
+    def __getitem__(self, idx):
+        number = self.numbers[idx]
+        im_dir = os.path.join(self.img_root, number)
+        fn1, fn2 = tuple(random.sample(os.listdir(im_dir), k=2))
+        fp1 = os.path.join(im_dir, fn1)
+        fp2 = os.path.join(im_dir, fn2)
+        img1 = Image.open(fp1).convert("RGB")
+        img2 = Image.open(fp2).convert("RGB")
+        img1 = self.preproc(img1)
+        img2 = self.preproc(img2)
+
+        fn1 = f'{number}_{fn1}'
+        fn2 = f'{number}_{fn2}'
+        return img1, img2, fn1, fn2
+
+    def __len__(self):
+        return len(self.numbers)

tools/TensorboardWriter.py

@@ -0,0 +1,16 @@
+import threading
+from tensorboardX import SummaryWriter
+
+
+class SingleSummaryWriter(SummaryWriter):
+    _instance_lock = threading.Lock()
+
+    def __init__(self, logdir=None, **kwargs):
+        super().__init__(logdir, **kwargs)
+
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(SingleSummaryWriter, "_instance"):
+            with SingleSummaryWriter._instance_lock:
+                if not hasattr(SingleSummaryWriter, "_instance"):
+                    SingleSummaryWriter._instance = object.__new__(cls)
+        return SingleSummaryWriter._instance

tools/__init__.py

@@ -0,0 +1,4 @@
+from .TensorboardWriter import *
+from .model_utils import *
+from .saver import *
+from .mutils import *

tools/model_utils.py

@@ -0,0 +1,11 @@
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def count_conv_layers(model):
+    cnt = 0
+    for mo in model().modules():
+        if type(mo).__name__ == 'Conv2d':
+            cnt += 1
+
+    print(model.__name__, cnt, count_parameters(model()))

tools/mutils.py

@@ -0,0 +1,30 @@
+import datetime
+import os
+import shutil
+import time
+
+
+def contains(key, lst):
+    flag = False
+    for item in lst:
+        if key == item:
+            flag = True
+    return flag
+
+
+def make_empty_dir(new_dir):
+    if os.path.exists(new_dir):
+        shutil.rmtree(new_dir)
+    os.makedirs(new_dir, exist_ok=True)
+
+
+def get_timestamp():
+    return str(time.time()).replace('.', '')
+
+
+def get_formatted_time():
+    return datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+
+
+if __name__ == '__main__':
+    pass

tools/saver.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn as nn
+import os
+import time
+from tools import mutils
+
+saved_grad = None
+saved_name = None
+
+base_url = './results'
+os.makedirs(base_url, exist_ok=True)
+
+
+def normalize_tensor_mm(tensor):
+    return (tensor - tensor.min()) / (tensor.max() - tensor.min())
+
+
+def normalize_tensor_sigmoid(tensor):
+    return nn.functional.sigmoid(tensor)
+
+
+def save_image(tensor, name=None, save_path=None, exit_flag=False, timestamp=False, norm=False):
+    import torchvision.utils as vutils
+    os.makedirs(base_url, exist_ok=True)
+    if norm:
+        tensor = normalize_tensor_mm(tensor)
+    grid = vutils.make_grid(tensor.detach().cpu(), nrow=4)
+
+    if save_path:
+        vutils.save_image(grid, save_path)
+    else:
+        if timestamp:
+            vutils.save_image(grid, f'{base_url}/{name}_{mutils.get_timestamp()}.png')
+        else:
+            vutils.save_image(grid, f'{base_url}/{name}.png')
+    if exit_flag:
+        exit(0)
+
+
+def save_feature(tensor, name, exit_flag=False, timestamp=False):
+    import torchvision.utils as vutils
+    # tensors = [tensor, normalize_tensor_mm(tensor), normalize_tensor_sigmoid(tensor)]
+    tensors = [tensor]
+    titles = ['original', 'min-max', 'sigmoid']
+    os.makedirs(base_url, exist_ok=True)
+    if timestamp:
+        name += '_' + str(time.time()).replace('.', '')
+
+    for index, tensor in enumerate(tensors):
+        _data = tensor.detach().cpu().squeeze(0).unsqueeze(1)
+        num_per_row = 8
+        grid = vutils.make_grid(_data, nrow=num_per_row)
+        vutils.save_image(grid, f'{base_url}/{name}_{titles[index]}.png')
+    if exit_flag:
+        exit(0)