Upload 8 files

data/LOLdataset.py

@@ -0,0 +1,116 @@
+
+import os
+import random
+import torch
+import torch.utils.data as data
+import numpy as np
+from os import listdir
+from os.path import join
+from data.util import *
+from torchvision import transforms as t
+
+    
+class LOLDatasetFromFolder(data.Dataset):
+    def __init__(self, data_dir, transform=None):
+        super(LOLDatasetFromFolder, self).__init__()
+        self.data_dir = data_dir
+        self.transform = transform
+        self.norm = t.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    def __getitem__(self, index):
+
+        folder = self.data_dir+'/low'
+        folder2= self.data_dir+'/high'
+        data_filenames = [join(folder, x) for x in listdir(folder) if is_image_file(x)]
+        data_filenames2 = [join(folder2, x) for x in listdir(folder2) if is_image_file(x)]
+        num = len(data_filenames)
+
+        im1 = load_img(data_filenames[index])
+        im2 = load_img(data_filenames2[index])
+        _, file1 = os.path.split(data_filenames[index])
+        _, file2 = os.path.split(data_filenames2[index])
+        seed = random.randint(1, 1000000)
+        seed = np.random.randint(seed) # make a seed with numpy generator 
+        if self.transform:
+            random.seed(seed) # apply this seed to img tranfsorms
+            torch.manual_seed(seed) # needed for torchvision 0.7
+            im1 = self.transform(im1)
+            random.seed(seed)
+            torch.manual_seed(seed)         
+            im2 = self.transform(im2) 
+        return im1, im2, file1, file2
+
+    def __len__(self):
+        return 485
+
+    
+class LOLv2DatasetFromFolder(data.Dataset):
+    def __init__(self, data_dir, transform=None):
+        super(LOLv2DatasetFromFolder, self).__init__()
+        self.data_dir = data_dir
+        self.transform = transform
+
+    def __getitem__(self, index):
+
+        folder = self.data_dir+'/Low'
+        folder2= self.data_dir+'/Normal'
+        data_filenames = [join(folder, x) for x in listdir(folder) if is_image_file(x)]
+        data_filenames2 = [join(folder2, x) for x in listdir(folder2) if is_image_file(x)]
+        
+        im1 = load_img(data_filenames[index])
+        im2 = load_img(data_filenames2[index])
+        _, file1 = os.path.split(data_filenames[index])
+        _, file2 = os.path.split(data_filenames2[index])
+        seed = random.randint(1, 1000000)
+        seed = np.random.randint(seed) # make a seed with numpy generator 
+        if self.transform:
+            random.seed(seed) # apply this seed to img tranforms
+            torch.manual_seed(seed) # needed for torchvision 0.7
+            im1 = self.transform(im1)      
+            random.seed(seed) # apply this seed to img tranforms
+            torch.manual_seed(seed) # needed for torchvision 0.7 
+            im2 = self.transform(im2)
+        return im1, im2, file1, file2
+
+    def __len__(self):
+        return 685
+
+
+
+class LOLv2SynDatasetFromFolder(data.Dataset):
+    def __init__(self, data_dir, transform=None):
+        super(LOLv2SynDatasetFromFolder, self).__init__()
+        self.data_dir = data_dir
+        self.transform = transform
+        self.norm = t.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    def __getitem__(self, index):
+
+        folder = self.data_dir+'/Low'
+        folder2= self.data_dir+'/Normal'
+        data_filenames = [join(folder, x) for x in listdir(folder) if is_image_file(x)]
+        data_filenames2 = [join(folder2, x) for x in listdir(folder2) if is_image_file(x)]
+
+
+        im1 = load_img(data_filenames[index])
+        im2 = load_img(data_filenames2[index])
+        _, file1 = os.path.split(data_filenames[index])
+        _, file2 = os.path.split(data_filenames2[index])
+        seed = random.randint(1, 1000000)
+        seed = np.random.randint(seed) # make a seed with numpy generator 
+        if self.transform:
+            random.seed(seed) # apply this seed to img tranfsorms
+            torch.manual_seed(seed) # needed for torchvision 0.7
+            im1 = self.transform(im1)
+            random.seed(seed)
+            torch.manual_seed(seed)         
+            im2 = self.transform(im2)
+        return im1, im2, file1, file2
+
+    def __len__(self):
+        return 900
+
+
+
+    
+

data/SICE_blur_SID.py

@@ -0,0 +1,134 @@
+
+import os
+import random
+import torch
+import torch.utils.data as data
+import numpy as np
+from os import listdir
+from os.path import join
+from PIL import Image
+from data.util import *
+from torchvision import transforms as t
+import torch.nn.functional as F
+
+class LOLBlurDatasetFromFolder(data.Dataset):
+    def __init__(self, data_dir, transform=None):
+        super(LOLBlurDatasetFromFolder, self).__init__()
+        self.data_dir = data_dir
+        self.transform = transform
+
+    def __getitem__(self, index):
+        while True:
+            seed = random.randint(1, 1000000)
+            random.seed(seed) 
+            index = random.randint(0, 259)
+            fill_index = str(index+1).zfill(4)
+            folder = join(self.data_dir+'/low_blur', fill_index)
+            folder2 = join(self.data_dir+'/high_sharp_scaled', fill_index)
+            if  not os.path.exists(folder):
+                continue
+            data_filenames = [join(folder, x) for x in listdir(folder) if is_image_file(x)]
+            data_filenames2 = [join(folder2, x) for x in listdir(folder2) if is_image_file(x)]
+            num = len(data_filenames)
+            if num != 0: break
+        index1 = random.randint(1,num)
+
+        im1 = load_img(data_filenames[index1-1])
+        im2 = load_img(data_filenames2[index1-1])
+        seed = random.randint(1, 1000000)
+        seed = np.random.randint(seed) # make a seed with numpy generator 
+        if self.transform:
+            random.seed(seed) # apply this seed to img tranfsorms
+            torch.manual_seed(seed) # needed for torchvision 0.7
+            im1 = self.transform(im1)
+            random.seed(seed)
+            torch.manual_seed(seed)         
+            im2 = self.transform(im2)
+        return im1, im2, data_filenames[index1-1], data_filenames2[index1-1]
+
+    def __len__(self):
+        return 10200
+    
+
+class SIDDatasetFromFolder(data.Dataset):
+    def __init__(self, data_dir, transform=None):
+        super(SIDDatasetFromFolder, self).__init__()
+        self.data_dir = data_dir
+        self.transform = transform
+
+    def __getitem__(self, index):
+        while True:
+            seed = random.randint(1, 1000000)
+            random.seed(seed) 
+            index = random.randint(0, 233)
+            fill_index = str(index+1).zfill(5)
+            folder = join(self.data_dir+'/short', fill_index)
+            folder2 = join(self.data_dir+'/long', fill_index)
+            if os.path.exists(folder): 
+                data_filenames = [join(folder, x) for x in listdir(folder) if is_image_file(x)]
+                data_filenames2 = [join(folder2, x) for x in listdir(folder2) if is_image_file(x)]
+                num = len(data_filenames)
+                break
+            else:
+                continue
+        index1 = random.randint(1,num)
+
+
+        im1 = load_img(data_filenames[index1-1])
+        im2 = load_img(data_filenames2[0])
+        _, file1 = os.path.split(data_filenames[index1-1])
+        _, file2 = os.path.split(data_filenames2[0])
+        seed = np.random.randint(random.randint(1, 1000000)) # make a seed with numpy generator 
+        if self.transform:
+            random.seed(seed) # apply this seed to img tranfsorms
+            torch.manual_seed(seed) # needed for torchvision 0.7
+            im1 = self.transform(im1)
+            random.seed(seed)
+            torch.manual_seed(seed)         
+            im2 = self.transform(im2)
+        return im1, im2, file1, file2
+
+    def __len__(self):
+        return 2099
+    
+    
+    
+class SICEDatasetFromFolder(data.Dataset):
+    def __init__(self, data_dir, transform=None):
+        super(SICEDatasetFromFolder, self).__init__()
+        self.data_dir = data_dir
+        self.transform = transform
+
+    def __getitem__(self, index):
+        while True:
+            seed = random.randint(1, 1000000)
+            random.seed(seed) 
+            index = random.randint(0, 590)
+            fill_index = str(index+1)
+            train, tail = os.path.split(self.data_dir)
+            folder = join(self.data_dir, fill_index)
+            data_gt = join(train+'/label', fill_index+'.JPG')
+            if os.path.exists(folder): 
+                data_filenames = [join(folder, x) for x in listdir(folder) if is_image_file(x)]
+                num = len(data_filenames)
+                break
+            else:
+                continue
+        index1 = random.randint(1,num)
+
+        im1 = load_img(data_filenames[index1-1])
+        im2 = load_img(data_gt)
+        _, file1 = os.path.split(data_filenames[index1-1])
+        _, file2 = os.path.split(data_gt)
+        seed = np.random.randint(random.randint(1, 1000000)) # make a seed with numpy generator 
+        if self.transform:
+            random.seed(seed) # apply this seed to img tranfsorms
+            torch.manual_seed(seed) # needed for torchvision 0.7
+            im1 = self.transform(im1)
+            random.seed(seed)
+            torch.manual_seed(seed)         
+            im2 = self.transform(im2)
+        return im1, im2, file1, file2
+
+    def __len__(self):
+        return 4803

data/data.py

@@ -0,0 +1,53 @@
+from torchvision.transforms import Compose, ToTensor, RandomCrop, RandomHorizontalFlip, RandomVerticalFlip
+from data.LOLdataset import *
+from data.eval_sets import *
+from data.SICE_blur_SID import *
+from data.fivek import *
+
+def transform1(size=256):
+    return Compose([
+        RandomCrop((size, size)),
+        RandomHorizontalFlip(),
+        RandomVerticalFlip(),
+        ToTensor(),
+    ])
+
+def transform2():
+    return Compose([ToTensor()])
+
+
+
+def get_lol_training_set(data_dir,size):
+    return LOLDatasetFromFolder(data_dir, transform=transform1(size))
+
+
+def get_lol_v2_training_set(data_dir,size):
+    return LOLv2DatasetFromFolder(data_dir, transform=transform1(size))
+
+
+def get_training_set_blur(data_dir,size):
+    return LOLBlurDatasetFromFolder(data_dir, transform=transform1(size))
+
+
+def get_lol_v2_syn_training_set(data_dir,size):
+    return LOLv2SynDatasetFromFolder(data_dir, transform=transform1(size))
+
+
+def get_SID_training_set(data_dir,size):
+    return SIDDatasetFromFolder(data_dir, transform=transform1(size))
+
+
+def get_SICE_training_set(data_dir,size):
+    return SICEDatasetFromFolder(data_dir, transform=transform1(size))
+
+def get_SICE_eval_set(data_dir):
+    return SICEDatasetFromFolderEval(data_dir, transform=transform2())
+
+def get_eval_set(data_dir):
+    return DatasetFromFolderEval(data_dir, transform=transform2())
+
+def get_fivek_training_set(data_dir,size):
+    return FiveKDatasetFromFolder(data_dir, transform=transform1(size))
+
+def get_fivek_eval_set(data_dir):
+    return SICEDatasetFromFolderEval(data_dir, transform=transform2())

data/eval_sets.py

@@ -0,0 +1,52 @@
+
+import os
+import torch.utils.data as data
+from os import listdir
+from os.path import join
+from data.util import *
+import torch.nn.functional as F
+
+class SICEDatasetFromFolderEval(data.Dataset):
+    def __init__(self, data_dir, transform=None):
+        super(SICEDatasetFromFolderEval, self).__init__()
+        data_filenames = [join(data_dir, x) for x in listdir(data_dir) if is_image_file(x)]
+        data_filenames.sort()
+        self.data_filenames = data_filenames
+        self.transform = transform
+
+    def __getitem__(self, index):
+        input = load_img(self.data_filenames[index])
+        _, file = os.path.split(self.data_filenames[index])
+
+        if self.transform:
+            input = self.transform(input)
+            factor = 8
+            h, w = input.shape[1], input.shape[2]
+            H, W = ((h + factor) // factor) * factor, ((w + factor) // factor) * factor
+            padh = H - h if h % factor != 0 else 0
+            padw = W - w if w % factor != 0 else 0
+            input = F.pad(input.unsqueeze(0), (0,padw,0,padh), 'reflect').squeeze(0)
+        return input, file, h, w
+
+    def __len__(self):
+        return len(self.data_filenames)
+    
+    
+class DatasetFromFolderEval(data.Dataset):
+    def __init__(self, data_dir, transform=None):
+        super(DatasetFromFolderEval, self).__init__()
+        data_filenames = [join(data_dir, x) for x in listdir(data_dir) if is_image_file(x)]
+        data_filenames.sort()
+        self.data_filenames = data_filenames
+        self.transform = transform
+
+    def __getitem__(self, index):
+        input = load_img(self.data_filenames[index])
+        _, file = os.path.split(self.data_filenames[index])
+
+        if self.transform:
+            input = self.transform(input)
+        return input, file
+
+    def __len__(self):
+        return len(self.data_filenames)

data/fivek.py

@@ -0,0 +1,42 @@
+# Add new fivek dataset follow Retinexformer(https://github.com/caiyuanhao1998/Retinexformer)
+
+import os
+import random
+import torch
+import torch.utils.data as data
+import numpy as np
+from os import listdir
+from os.path import join
+from data.util import *
+
+class FiveKDatasetFromFolder(data.Dataset):
+    def __init__(self, data_dir, transform=None):
+        super(FiveKDatasetFromFolder, self).__init__()
+        self.data_dir = data_dir
+        self.transform = transform
+
+    def __getitem__(self, index):
+
+        folder = self.data_dir+'/input'
+        folder2= self.data_dir+'/target'
+        data_filenames = [join(folder, x) for x in listdir(folder) if is_image_file(x)]
+        data_filenames2 = [join(folder2, x) for x in listdir(folder2) if is_image_file(x)]
+
+
+        im1 = load_img(data_filenames[index])
+        im2 = load_img(data_filenames2[index])
+        _, file1 = os.path.split(data_filenames[index])
+        _, file2 = os.path.split(data_filenames2[index])
+        seed = random.randint(1, 1000000)
+        seed = np.random.randint(seed) # make a seed with numpy generator 
+        if self.transform:
+            random.seed(seed) # apply this seed to img tranfsorms
+            torch.manual_seed(seed) # needed for torchvision 0.7
+            im1 = self.transform(im1)
+            random.seed(seed)
+            torch.manual_seed(seed)         
+            im2 = self.transform(im2)
+        return im1, im2, file1, file2
+
+    def __len__(self):
+        return 4500

data/options.py

@@ -0,0 +1,92 @@
+import argparse
+
+def _str2bool(v):
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+def option():
+    # Training settings
+    parser = argparse.ArgumentParser(description='CIDNet')
+    parser.add_argument('--batchSize', type=int, default=8, help='training batch size')
+    parser.add_argument('--cropSize', type=int, default=256, help='image crop size (patch size)')
+    parser.add_argument('--nEpochs', type=int, default=1000, help='number of epochs to train for end')
+    parser.add_argument('--start_epoch', type=int, default=0, help='number of epochs to start, >0 is retrained a pre-trained pth')
+    parser.add_argument('--snapshots', type=int, default=10, help='Snapshots for save checkpoints pth')
+    parser.add_argument('--lr', type=float, default=1e-4, help='Learning Rate')
+    parser.add_argument('--gpu_mode', type=_str2bool, default=True)
+    parser.add_argument('--shuffle', type=_str2bool, default=True)
+    parser.add_argument('--threads', type=int, default=16, help='number of threads for dataloader to use')
+
+    # choose a scheduler
+    parser.add_argument('--cos_restart_cyclic', type=_str2bool, default=False)
+    parser.add_argument('--cos_restart', type=_str2bool, default=True)
+
+    # warmup training
+    parser.add_argument('--warmup_epochs', type=int, default=3, help='warmup_epochs')
+    parser.add_argument('--start_warmup', type=_str2bool, default=True, help='turn False to train without warmup') 
+
+    # train datasets
+    parser.add_argument('--data_train_lol_blur'     , type=str, default='./datasets/LOL_blur/train')
+    parser.add_argument('--data_train_lol_v1'       , type=str, default='./datasets/LOLdataset/our485')
+    parser.add_argument('--data_train_lolv2_real'   , type=str, default='./datasets/LOLv2/Real_captured/Train')
+    parser.add_argument('--data_train_lolv2_syn'    , type=str, default='./datasets/LOLv2/Synthetic/Train')
+    parser.add_argument('--data_train_SID'          , type=str, default='./datasets/Sony_total_dark/train')
+    parser.add_argument('--data_train_SICE'         , type=str, default='./datasets/SICE/Dataset/train')
+    parser.add_argument('--data_train_fivek'        , type=str, default='./datasets/FiveK/train')
+
+    # validation input
+    parser.add_argument('--data_val_lol_blur'       , type=str, default='./datasets/LOL_blur/eval/low_blur')
+    parser.add_argument('--data_val_lol_v1'         , type=str, default='./datasets/LOLdataset/eval15/low')
+    parser.add_argument('--data_val_lolv2_real'     , type=str, default='./datasets/LOLv2/Real_captured/Test/Low')
+    parser.add_argument('--data_val_lolv2_syn'      , type=str, default='./datasets/LOLv2/Synthetic/Test/Low')
+    parser.add_argument('--data_val_SID'            , type=str, default='./datasets/Sony_total_dark/eval/short')
+    parser.add_argument('--data_val_SICE_mix'       , type=str, default='./datasets/SICE/Dataset/eval/test')
+    parser.add_argument('--data_val_SICE_grad'      , type=str, default='./datasets/SICE/Dataset/eval/test')
+    parser.add_argument('--data_test_fivek'         , type=str, default='./datasets/FiveK/test/input')
+
+    # validation groundtruth
+    parser.add_argument('--data_valgt_lol_blur'     , type=str, default='./datasets/LOL_blur/eval/high_sharp_scaled/')
+    parser.add_argument('--data_valgt_lol_v1'       , type=str, default='./datasets/LOLdataset/eval15/high/')
+    parser.add_argument('--data_valgt_lolv2_real'   , type=str, default='./datasets/LOLv2/Real_captured/Test/Normal/')
+    parser.add_argument('--data_valgt_lolv2_syn'    , type=str, default='./datasets/LOLv2/Synthetic/Test/Normal/')
+    parser.add_argument('--data_valgt_SID'          , type=str, default='./datasets/Sony_total_dark/eval/long/')
+    parser.add_argument('--data_valgt_SICE_mix'     , type=str, default='./datasets/SICE/Dataset/eval/target/')
+    parser.add_argument('--data_valgt_SICE_grad'    , type=str, default='./datasets/SICE/Dataset/eval/target/')
+    parser.add_argument('--data_valgt_fivek'        , type=str, default='./datasets/FiveK/test/target/')
+
+    parser.add_argument('--val_folder', default='./results/', help='Location to save validation datasets')
+
+    # loss weights
+    parser.add_argument('--HVI_weight', type=float, default=1.0)
+    parser.add_argument('--L1_weight', type=float, default=1.0)
+    parser.add_argument('--D_weight',  type=float, default=0.5)
+    parser.add_argument('--E_weight',  type=float, default=50.0)
+    parser.add_argument('--P_weight',  type=float, default=1e-2)
+    
+    # use random gamma function (enhancement curve) to improve generalization
+    parser.add_argument('--gamma', type=_str2bool, default=False)
+    parser.add_argument('--start_gamma', type=int, default=60)
+    parser.add_argument('--end_gamma', type=int, default=120)
+
+    # auto grad, turn off to speed up training
+    parser.add_argument('--grad_detect', type=_str2bool, default=False, help='if gradient explosion occurs, turn-on it')
+    parser.add_argument('--grad_clip', type=_str2bool, default=True, help='if gradient fluctuates too much, turn-on it')
+    
+    
+    # choose which dataset you want to train
+    parser.add_argument('--dataset', type=str, default='lol_v1',
+    choices=['lol_v1',
+             'lolv2_real',
+             'lolv2_syn',
+             'lol_blur', 
+             'SID',
+             'SICE_mix',
+             'SICE_grad',
+             'fivek'],
+    help='Select the dataset to train on (default: %(default)s)')
+
+    return parser

data/scheduler.py

@@ -0,0 +1,172 @@
+from torch.optim.lr_scheduler import _LRScheduler
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+import math
+
+class GradualWarmupScheduler(_LRScheduler):
+    """ Gradually warm-up(increasing) learning rate in optimizer.
+    Proposed in 'Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour'.
+
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        multiplier: target learning rate = base lr * multiplier if multiplier > 1.0. if multiplier = 1.0, lr starts from 0 and ends up with the base_lr.
+        total_epoch: target learning rate is reached at total_epoch, gradually
+        after_scheduler: after target_epoch, use this scheduler(eg. ReduceLROnPlateau)
+    """
+
+    def __init__(self, optimizer, multiplier, total_epoch, after_scheduler=None):
+        self.multiplier = multiplier
+        if self.multiplier < 1.:
+            raise ValueError('multiplier should be greater thant or equal to 1.')
+        self.total_epoch = total_epoch
+        self.after_scheduler = after_scheduler
+        self.finished = False
+        super(GradualWarmupScheduler, self).__init__(optimizer)
+
+    def get_lr(self):
+        if self.last_epoch > self.total_epoch:
+            if self.after_scheduler:
+                if not self.finished:
+                    self.after_scheduler.base_lrs = [base_lr * self.multiplier for base_lr in self.base_lrs]
+                    self.finished = True
+                return self.after_scheduler.get_lr()
+            return [base_lr * self.multiplier for base_lr in self.base_lrs]
+
+        if self.multiplier == 1.0:
+            return [base_lr * (float(self.last_epoch) / self.total_epoch) for base_lr in self.base_lrs]
+        else:
+            return [base_lr * ((self.multiplier - 1.) * self.last_epoch / self.total_epoch + 1.) for base_lr in self.base_lrs]
+
+    def step_ReduceLROnPlateau(self, metrics, epoch=None):
+        if epoch is None:
+            epoch = self.last_epoch + 1
+        self.last_epoch = epoch if epoch != 0 else 1  # ReduceLROnPlateau is called at the end of epoch, whereas others are called at beginning
+        if self.last_epoch <= self.total_epoch:
+            warmup_lr = [base_lr * ((self.multiplier - 1.) * self.last_epoch / self.total_epoch + 1.) for base_lr in self.base_lrs]
+            for param_group, lr in zip(self.optimizer.param_groups, warmup_lr):
+                param_group['lr'] = lr
+        else:
+            if epoch is None:
+                self.after_scheduler.step(metrics, None)
+            else:
+                self.after_scheduler.step(metrics, epoch - self.total_epoch)
+
+    def step(self, epoch=None, metrics=None):
+        if type(self.after_scheduler) != ReduceLROnPlateau:
+            if self.finished and self.after_scheduler:
+                if epoch is None:
+                    self.after_scheduler.step(None)
+                else:
+                    self.after_scheduler.step(epoch - self.total_epoch)
+            else:
+                return super(GradualWarmupScheduler, self).step(epoch)
+        else:
+            self.step_ReduceLROnPlateau(metrics, epoch)
+            
+def get_position_from_periods(iteration, cumulative_period):
+    """Get the position from a period list.
+
+    It will return the index of the right-closest number in the period list.
+    For example, the cumulative_period = [100, 200, 300, 400],
+    if iteration == 50, return 0;
+    if iteration == 210, return 2;
+    if iteration == 300, return 2.
+
+    Args:
+        iteration (int): Current iteration.
+        cumulative_period (list[int]): Cumulative period list.
+
+    Returns:
+        int: The position of the right-closest number in the period list.
+    """
+    for i, period in enumerate(cumulative_period):
+        if iteration <= period:
+            return i
+        
+class CosineAnnealingRestartCyclicLR(_LRScheduler):
+    """ Cosine annealing with restarts learning rate scheme.
+    An example of config:
+    periods = [10, 10, 10, 10]
+    restart_weights = [1, 0.5, 0.5, 0.5]
+    eta_min=1e-7
+    It has four cycles, each has 10 iterations. At 10th, 20th, 30th, the
+    scheduler will restart with the weights in restart_weights.
+    Args:
+        optimizer (torch.nn.optimizer): Torch optimizer.
+        periods (list): Period for each cosine anneling cycle.
+        restart_weights (list): Restart weights at each restart iteration.
+            Default: [1].
+        eta_min (float): The mimimum lr. Default: 0.
+        last_epoch (int): Used in _LRScheduler. Default: -1.
+    """
+
+    def __init__(self,
+                 optimizer,
+                 periods,
+                 restart_weights=(1, ),
+                 eta_mins=(0, ),
+                 last_epoch=-1):
+        self.periods = periods
+        self.restart_weights = restart_weights
+        self.eta_mins = eta_mins
+        assert (len(self.periods) == len(self.restart_weights)
+                ), 'periods and restart_weights should have the same length.'
+        self.cumulative_period = [
+            sum(self.periods[0:i + 1]) for i in range(0, len(self.periods))
+        ]
+        super(CosineAnnealingRestartCyclicLR, self).__init__(optimizer, last_epoch)
+        
+    def get_lr(self):
+        idx = get_position_from_periods(self.last_epoch,
+                                        self.cumulative_period)
+        current_weight = self.restart_weights[idx]
+        nearest_restart = 0 if idx == 0 else self.cumulative_period[idx - 1]
+        current_period = self.periods[idx]
+        eta_min = self.eta_mins[idx]
+
+        return [
+            eta_min + current_weight * 0.5 * (base_lr - eta_min) *
+            (1 + math.cos(math.pi * (
+                (self.last_epoch - nearest_restart) / current_period)))
+            for base_lr in self.base_lrs
+        ]
+
+class CosineAnnealingRestartLR(_LRScheduler):
+    """ Cosine annealing with restarts learning rate scheme.
+
+    An example of config:
+    periods = [10, 10, 10, 10]
+    restart_weights = [1, 0.5, 0.5, 0.5]
+    eta_min=1e-7
+
+    It has four cycles, each has 10 iterations. At 10th, 20th, 30th, the
+    scheduler will restart with the weights in restart_weights.
+
+    Args:
+        optimizer (torch.nn.optimizer): Torch optimizer.
+        periods (list): Period for each cosine anneling cycle.
+        restart_weights (list): Restart weights at each restart iteration.
+            Default: [1].
+        eta_min (float): The mimimum lr. Default: 0.
+        last_epoch (int): Used in _LRScheduler. Default: -1.
+    """
+
+    def __init__(self, optimizer, periods, restart_weights=(1, ), eta_min=0, last_epoch=-1):
+        self.periods = periods
+        self.restart_weights = restart_weights
+        self.eta_min = eta_min
+        assert (len(self.periods) == len(
+            self.restart_weights)), 'periods and restart_weights should have the same length.'
+        self.cumulative_period = [sum(self.periods[0:i + 1]) for i in range(0, len(self.periods))]
+        super(CosineAnnealingRestartLR, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        idx = get_position_from_periods(self.last_epoch, self.cumulative_period)
+        current_weight = self.restart_weights[idx]
+        nearest_restart = 0 if idx == 0 else self.cumulative_period[idx - 1]
+        current_period = self.periods[idx]
+
+        return [
+            self.eta_min + current_weight * 0.5 * (base_lr - self.eta_min) *
+            (1 + math.cos(math.pi * ((self.last_epoch - nearest_restart) / current_period)))
+            for base_lr in self.base_lrs
+        ]

data/util.py

@@ -0,0 +1,9 @@
+
+from PIL import Image
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in [".png", ".jpg", ".bmp", ".JPG", ".jpeg"])
+
+def load_img(filepath):
+    img = Image.open(filepath).convert('RGB')
+    return img