Spaces:
Sleeping
Sleeping
| import argparse | |
| import torch | |
| import torch.nn as nn | |
| from network.Math_Module import P, Q | |
| from network.decom import Decom | |
| import os | |
| #import torchvision | |
| import torchvision.transforms as transforms | |
| from PIL import Image | |
| import time | |
| from utils import * | |
| import cv2 | |
| def one2three(x): | |
| return torch.cat([x, x, x], dim=1).to(x) | |
| class Inference(nn.Module): | |
| def __init__(self, opts): | |
| super().__init__() | |
| self.opts = opts | |
| # loading decomposition model | |
| self.model_Decom_low = Decom() | |
| self.model_Decom_low = load_initialize(self.model_Decom_low, | |
| self.opts.Decom_model_low_path) | |
| # loading R; old_model_opts; and L model | |
| self.unfolding_opts, self.model_R, self.model_L = load_unfolding( | |
| self.opts.unfolding_model_path) | |
| # loading adjustment model | |
| self.adjust_model = load_adjustment(self.opts.adjust_model_path) | |
| self.P = P() | |
| self.Q = Q() | |
| transform = [ | |
| transforms.ToTensor(), | |
| ] | |
| self.transform = transforms.Compose(transform) | |
| print(self.model_Decom_low) | |
| print(self.model_R) | |
| print(self.model_L) | |
| print(self.adjust_model) | |
| #time.sleep(8) | |
| def unfolding(self, input_low_img): | |
| for t in range(self.unfolding_opts.round): | |
| if t == 0: # initialize R0, L0 | |
| P, Q = self.model_Decom_low(input_low_img) | |
| else: # update P and Q | |
| w_p = (self.unfolding_opts.gamma + | |
| self.unfolding_opts.Roffset * t) | |
| w_q = (self.unfolding_opts.lamda + | |
| self.unfolding_opts.Loffset * t) | |
| P = self.P(I=input_low_img, Q=Q, R=R, gamma=w_p) | |
| Q = self.Q(I=input_low_img, P=P, L=L, lamda=w_q) | |
| R = self.model_R(r=P, l=Q) | |
| L = self.model_L(l=Q) | |
| return R, L | |
| def lllumination_adjust(self, L, ratio): | |
| ratio = torch.ones(L.shape) * self.opts.ratio | |
| return self.adjust_model(l=L, alpha=ratio) | |
| def forward(self, input_low_img): | |
| if torch.cuda.is_available(): | |
| input_low_img = input_low_img.cuda() | |
| with torch.no_grad(): | |
| start = time.time() | |
| R, L = self.unfolding(input_low_img) | |
| High_L = self.lllumination_adjust(L, self.opts.ratio) | |
| I_enhance = High_L * R | |
| p_time = (time.time() - start) | |
| return I_enhance, p_time | |
| def run(self, low_img_path): | |
| file_name = os.path.basename(self.opts.img_path) | |
| name = file_name.split('.')[0] | |
| low_img = self.transform(Image.open(low_img_path)).unsqueeze(0) | |
| # print('**************************************************************************') | |
| # print(low_img) | |
| # print(type(low_img)) | |
| # print(type(Image.open(low_img_path))) | |
| # print(Image.open(low_img_path)) | |
| enhance, p_time = self.forward(input_low_img=low_img) | |
| if not os.path.exists(self.opts.output): | |
| os.makedirs(self.opts.output) | |
| save_path = os.path.join( | |
| self.opts.output, | |
| file_name.replace(name, | |
| "%s_%d_URetinexNet" % (name, self.opts.ratio))) | |
| np_save_TensorImg(enhance, save_path) | |
| print( | |
| "================================= time for %s: %f============================" | |
| % (file_name, p_time)) | |
| # 这是我自己修改的 run 函数 | |
| # 避免了把图片储存到硬盘上面 | |
| # 后续也可以修改把图片储存到硬盘上面 | |
| def runForWeb(self, image): | |
| # 首先对输入的图片进行下采样直到符合最低运行像素限制 | |
| max_pixel_limit=600*600 | |
| pyr_down_times=0 | |
| while True: | |
| a=len(image) | |
| b=len(image[0]) | |
| c=a*b | |
| if(c<=max_pixel_limit): | |
| break | |
| pyr_down_times+=1 | |
| image=cv2.pyrDown(image) | |
| print(image.shape) | |
| # 输入 | |
| low_img = self.transform(Image.fromarray(np.uint8(image))).unsqueeze(0) | |
| # low_img=Image.fromarray(image.astype('uint8')).convert('RGB') | |
| # print('#############################################') | |
| # print(type(low_img)) | |
| # print(low_img) | |
| # 训练 | |
| enhance, p_time = self.forward(input_low_img=low_img) | |
| # print('UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU') | |
| # 输出 | |
| # 这里需要修改一下 utils.py 的结果放回函数，参考上面 run 函数 np_save_TensorImg 这里需要修改一下的位置 | |
| # 退训练结果进行上采样，还原原图大小 | |
| result_image=result_for_gradio(enhance) | |
| for i in range(pyr_down_times): | |
| result_image=cv2.pyrUp(result_image) | |
| # return result_for_gradio(enhance) | |
| print(result_image.shape) | |
| return result_image | |
| # 这是提供给 gradio 框架调用的接口 | |
| # gradio 框架负责提供后端操控和前端的页面展示 | |
| def functionForGradio(image): | |
| parser = argparse.ArgumentParser(description='Configure') | |
| # specify your data path here! | |
| parser.add_argument('--img_path', type=str, default="./demo/input/3.png") | |
| parser.add_argument('--output', type=str, default="./demo/output") | |
| # ratio are recommended to be 3-5, bigger ratio will lead to over-exposure | |
| parser.add_argument('--ratio', type=int, default=5) | |
| # model path | |
| parser.add_argument('--Decom_model_low_path', | |
| type=str, | |
| default="./ckpt/init_low.pth") | |
| parser.add_argument('--unfolding_model_path', | |
| type=str, | |
| default="./ckpt/unfolding.pth") | |
| parser.add_argument('--adjust_model_path', | |
| type=str, | |
| default="./ckpt/L_adjust.pth") | |
| parser.add_argument('--gpu_id', type=int, default=0) | |
| opts = parser.parse_args() | |
| for k, v in vars(opts).items(): | |
| print(k, v) | |
| os.environ['CUDA_VISIBLE_DEVICES'] = str(opts.gpu_id) | |
| model = Inference(opts) | |
| # 这里传入 numpy 数组然后开始训练 | |
| return model.runForWeb(image) | |
| # 这是算法本来的主函数，上面提供的 gradio 框架调用的接口就是修改自主函数 | |
| # if __name__ == "__main__": | |
| # parser = argparse.ArgumentParser(description='Configure') | |
| # # specify your data path here! | |
| # parser.add_argument('--img_path', type=str, default="./demo/input/3.png") | |
| # parser.add_argument('--output', type=str, default="./demo/output") | |
| # # ratio are recommended to be 3-5, bigger ratio will lead to over-exposure | |
| # parser.add_argument('--ratio', type=int, default=5) | |
| # # model path | |
| # parser.add_argument('--Decom_model_low_path', | |
| # type=str, | |
| # default="./ckpt/init_low.pth") | |
| # parser.add_argument('--unfolding_model_path', | |
| # type=str, | |
| # default="./ckpt/unfolding.pth") | |
| # parser.add_argument('--adjust_model_path', | |
| # type=str, | |
| # default="./ckpt/L_adjust.pth") | |
| # parser.add_argument('--gpu_id', type=int, default=0) | |
| # opts = parser.parse_args() | |
| # for k, v in vars(opts).items(): | |
| # print(k, v) | |
| # os.environ['CUDA_VISIBLE_DEVICES'] = str(opts.gpu_id) | |
| # model = Inference(opts) | |
| # model.run(opts.img_path) | |