EMINIME/URetinex-Net

1import torchvision2from torch.nn import init3import numpy as np4import os5import time6import torch7from PIL import Image8import glob9 10def save_TensorImg(img_tensor, path, nrow=1):11    torchvision.utils.save_image(img_tensor, path, nrow=nrow)12 13def np_save_TensorImg(img_tensor, path):14    img = np.squeeze(img_tensor.cpu().permute(0, 2, 3, 1).numpy())15    im = Image.fromarray(np.clip(img*255, 0, 255.0).astype('uint8'))16#    print('&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&')17#    print(type(im))18#    print(im)19    im.save(path, 'png')20 21 22# 这个函数复制自上面这一个函数，目的是要返回一个图片 numpy 数组 23def result_for_gradio(img_tensor):24    img = np.squeeze(img_tensor.cpu().permute(0, 2, 3, 1).numpy())25    # print('&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&&')26    27    im=np.clip(img*255, 0, 255.0).astype('uint8')28    # print(im)29    # im = Image.fromarray(np.clip(img*255, 0, 255.0).astype('uint8'))30    return im31 32def define_modelR(opts):33    if opts.R_model == "HalfDnCNNSE":34        from network.restoration import HalfDnCNNSE35        model_R = HalfDnCNNSE(opts)36    return model_R37 38def define_modelL(opts):39    if opts.L_model == "Illumination_Alone":40        from network.illumination_enhance import Illumination_Alone41        model_L = Illumination_Alone(opts)42    return model_L43 44def define_modelA(opts):45    if opts.A_model == "naive":46        from network.illumination_adjustment import Adjust_naive47        model_A = Adjust_naive(opts)48    return model_A49 50 51 52def load_initialize(model, decom_model_path):53    if os.path.exists(decom_model_path):54        # torch.load with map_location=torch.device('cpu')55        checkpoint_Decom_low = torch.load(decom_model_path,map_location ='cpu')56        model.load_state_dict(checkpoint_Decom_low['state_dict']['model_R'])57        # to freeze the params of Decomposition Model58        for param in model.parameters():59            param.requires_grad = False   60        return model61    else:62        print("pretrained Initialize Model does not exist, check ---> %s " % decom_model_path)63        exit()64 65def load_unfolding(unfolding_model_path):66    if os.path.exists(unfolding_model_path):67        checkpoint = torch.load(unfolding_model_path,map_location ='cpu')68        old_opts = checkpoint["opts"]69        model_R = define_modelR(old_opts)70        model_L = define_modelL(old_opts)71        model_R.load_state_dict(checkpoint['state_dict']['model_R'])72        model_L.load_state_dict(checkpoint['state_dict']['model_L'])73        for param_R in model_R.parameters():74            param_R.requires_grad = False75        for param_L in model_L.parameters():76            param_L.requires_grad = False77        return old_opts, model_R, model_L78    else:79        print("pretrained Unfolding Model does not exist, check ---> %s"%unfolding_model_path)80        exit()81 82def load_adjustment(adjust_model_path):83    if os.path.exists(adjust_model_path):84        checkpoint_Adjust = torch.load(adjust_model_path,map_location ='cpu')85        model_A = define_modelA(checkpoint_Adjust['opts'])86        model_A.load_state_dict(checkpoint_Adjust['state_dict']['model_A'])87        print(" ===========>  loading pretrained Illumination Adjustment Model from: %s " % adjust_model_path)88        # to freeze the params of Decomposition Model89        for param in model_A.parameters():90            param.requires_grad = False       91        return model_A92    else:93        print("pretrained Adjustment Model does not exist, check ---> %s"%adjust_model_path)94        exit()95    96 97        98 99def param_all(model,  net_input):100    import torchsummary101    shape = net_input.shape102    torchsummary.summary(model, (shape[1], shape[2], shape[3]))103 104def param_self_compute(model):105    parmas = 0106    for p in model.parameters():107        #print(p)108        parmas += p.numel()109    return parmas110 111 112 113 114