EMINIME/URetinex-Net

1import argparse2from fileinput import filename3from locale import locale_encoding_alias4import torch5import torch.nn as nn6from network.Math_Module import P, Q7from network.decom import Decom8import os9import torchvision10import torchvision.transforms as transforms11from PIL import Image12import time13from utils import *14import glob15 16"""17  As different illumination adjustment ratio will cause18different enhanced results. Certainly you can tune the ratio youself19to get the best results.20  To get better result, we use the illumination of normal light image 21to adaptively generate ratio.22  Noted that KinD and KinD++ also use ratio to guide the illumination adjustment,23for fair comparison, the ratio of their methods also generate by the illumination24of normal light image.25"""26 27def one2three(x):28    return torch.cat([x, x, x], dim=1).to(x)29 30class Inference(nn.Module):31    def __init__(self, opts):32        super().__init__()33        self.opts = opts34        # loading decomposition model 35        self.model_Decom_low = Decom()36        self.model_Decom_high = Decom()37        self.model_Decom_low = load_initialize(self.model_Decom_low, self.opts.Decom_model_low_path)38        self.model_Decom_high = load_initialize(self.model_Decom_high, self.opts.Decom_model_high_path)39        # loading R; old_model_opts; and L model40        self.unfolding_opts, self.model_R, self.model_L= load_unfolding(self.opts.unfolding_model_path)41        # loading adjustment model42        self.adjust_model = load_adjustment(self.opts.adjust_model_path)43        self.P = P()44        self.Q = Q()45        transform = [46            transforms.ToTensor(),47        ]48        self.transform = transforms.Compose(transform)49        print(self.model_Decom_low)50        print(self.model_R)51        print(self.model_L)52        print(self.adjust_model)53        #time.sleep(8)54        55    def get_ratio(self, high_l, low_l):56        ratio = (low_l / (high_l + 0.0001)).mean()57        low_ratio = torch.ones(high_l.shape).cuda() * (1/(ratio+0.0001))58        return low_ratio59 60    def unfolding(self, input_low_img):61        for t in range(self.unfolding_opts.round):      62            if t == 0: # initialize R0, L063                P, Q = self.model_Decom_low(input_low_img)64            else: # update P and Q65                w_p = (self.unfolding_opts.gamma + self.unfolding_opts.Roffset * t)66                w_q = (self.unfolding_opts.lamda + self.unfolding_opts.Loffset * t)67                P = self.P(I=input_low_img, Q=Q, R=R, gamma=w_p)68                Q = self.Q(I=input_low_img, P=P, L=L, lamda=w_q) 69            R = self.model_R(r=P, l=Q)70            L = self.model_L(l=Q)71        return R, L72    73    def lllumination_adjust(self, L, ratio):74        ratio = torch.ones(L.shape).cuda() * ratio75        return self.adjust_model(l=L, alpha=ratio)76    77    def forward(self, input_low_img, input_high_img):78        if torch.cuda.is_available():79            input_low_img = input_low_img.cuda()80            input_high_img = input_high_img.cuda()81        with torch.no_grad():82            start = time.time()  83            R, L = self.unfolding(input_low_img)84            # the ratio is calculated using the decomposed normal illumination85            _, high_L = self.model_Decom_high(input_high_img)86            ratio = self.get_ratio(high_L, L)87            High_L = self.lllumination_adjust(L, ratio)88            I_enhance = High_L * R89            p_time = (time.time() - start)90        return I_enhance, p_time91 92    def evaluate(self):93        low_files = glob.glob(self.opts.low_dir+"/*.png")94        for file in low_files:95            file_name = os.path.basename(file)96            name = file_name.split('.')[0]97            high_file = os.path.join(self.opts.high_dir, file_name)98            low_img = self.transform(Image.open(file)).unsqueeze(0)99            high_img = self.transform(Image.open(high_file)).unsqueeze(0)100            enhance, p_time = self.forward(low_img, high_img)101            if not os.path.exists(self.opts.output):102                os.makedirs(self.opts.output)103            save_path = os.path.join(self.opts.output, file_name.replace(name, "%s_URetinexNet"%(name)))104            np_save_TensorImg(enhance, save_path)  105            print("================================= time for %s: %f============================"%(file_name, p_time))106 107 108 109    110if __name__ == "__main__":111    parser = argparse.ArgumentParser(description='Configure')112    # specify your data path here!113    parser.add_argument('--low_dir', type=str, default="./test_daat/LOLdataset/eval15/low")114    parser.add_argument('--high_dir', type=str, default="./test_data/LOLdataset/eval15/high")115    parser.add_argument('--output', type=str, default="./demo/output/LOL")116    # ratio are recommended to be 3-5, bigger ratio will lead to over-exposure 117    # model path118    parser.add_argument('--Decom_model_low_path', type=str, default="./ckpt/init_low.pth")119    parser.add_argument('--Decom_model_high_path', type=str, default="./ckpt/init_high.pth")120    parser.add_argument('--unfolding_model_path', type=str, default="./ckpt/unfolding.pth")121    parser.add_argument('--adjust_model_path', type=str, default="./ckpt/L_adjust.pth")122    parser.add_argument('--gpu_id', type=int, default=0)123    124    opts = parser.parse_args()125    for k, v in vars(opts).items():126        print(k, v)127    128    os.environ['CUDA_VISIBLE_DEVICES'] = str(opts.gpu_id)129    model = Inference(opts).cuda()130    model.evaluate()131