diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..63b832d42cddb8434a0de01ed9a91d7f3ca859a3 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +data/test_data/LIME/1.bmp filter=lfs diff=lfs merge=lfs -text +data/test_data/LIME/10.bmp filter=lfs diff=lfs merge=lfs -text +data/test_data/LIME/5.bmp filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..dff4bd2f7ee304fda34d124cea3fa5af06f6534e --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +data/ diff --git a/Myloss.py b/Myloss.py new file mode 100644 index 0000000000000000000000000000000000000000..91376e32fa3e0b2384822217fd44e7b0ddbfee27 --- /dev/null +++ b/Myloss.py @@ -0,0 +1,157 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +from torchvision.models.vgg import vgg16 +import numpy as np + + +class L_color(nn.Module): + + def __init__(self): + super(L_color, self).__init__() + + def forward(self, x ): + + b,c,h,w = x.shape + + mean_rgb = torch.mean(x,[2,3],keepdim=True) + mr,mg, mb = torch.split(mean_rgb, 1, dim=1) + Drg = torch.pow(mr-mg,2) + Drb = torch.pow(mr-mb,2) + Dgb = torch.pow(mb-mg,2) + k = torch.pow(torch.pow(Drg,2) + torch.pow(Drb,2) + torch.pow(Dgb,2),0.5) + + + return k + + +class L_spa(nn.Module): + + def __init__(self): + super(L_spa, self).__init__() + # print(1)kernel = torch.FloatTensor(kernel).unsqueeze(0).unsqueeze(0) + kernel_left = torch.FloatTensor( [[0,0,0],[-1,1,0],[0,0,0]]).cuda().unsqueeze(0).unsqueeze(0) + kernel_right = torch.FloatTensor( [[0,0,0],[0,1,-1],[0,0,0]]).cuda().unsqueeze(0).unsqueeze(0) + kernel_up = torch.FloatTensor( [[0,-1,0],[0,1, 0 ],[0,0,0]]).cuda().unsqueeze(0).unsqueeze(0) + kernel_down = torch.FloatTensor( [[0,0,0],[0,1, 0],[0,-1,0]]).cuda().unsqueeze(0).unsqueeze(0) + self.weight_left = nn.Parameter(data=kernel_left, requires_grad=False) + self.weight_right = nn.Parameter(data=kernel_right, requires_grad=False) + self.weight_up = nn.Parameter(data=kernel_up, requires_grad=False) + self.weight_down = nn.Parameter(data=kernel_down, requires_grad=False) + self.pool = nn.AvgPool2d(4) + def forward(self, org , enhance ): + b,c,h,w = org.shape + + org_mean = torch.mean(org,1,keepdim=True) + enhance_mean = torch.mean(enhance,1,keepdim=True) + + org_pool = self.pool(org_mean) + enhance_pool = self.pool(enhance_mean) + + weight_diff =torch.max(torch.FloatTensor([1]).cuda() + 10000*torch.min(org_pool - torch.FloatTensor([0.3]).cuda(),torch.FloatTensor([0]).cuda()),torch.FloatTensor([0.5]).cuda()) + E_1 = torch.mul(torch.sign(enhance_pool - torch.FloatTensor([0.5]).cuda()) ,enhance_pool-org_pool) + + + D_org_letf = F.conv2d(org_pool , self.weight_left, padding=1) + D_org_right = F.conv2d(org_pool , self.weight_right, padding=1) + D_org_up = F.conv2d(org_pool , self.weight_up, padding=1) + D_org_down = F.conv2d(org_pool , self.weight_down, padding=1) + + D_enhance_letf = F.conv2d(enhance_pool , self.weight_left, padding=1) + D_enhance_right = F.conv2d(enhance_pool , self.weight_right, padding=1) + D_enhance_up = F.conv2d(enhance_pool , self.weight_up, padding=1) + D_enhance_down = F.conv2d(enhance_pool , self.weight_down, padding=1) + + D_left = torch.pow(D_org_letf - D_enhance_letf,2) + D_right = torch.pow(D_org_right - D_enhance_right,2) + D_up = torch.pow(D_org_up - D_enhance_up,2) + D_down = torch.pow(D_org_down - D_enhance_down,2) + E = (D_left + D_right + D_up +D_down) + # E = 25*(D_left + D_right + D_up +D_down) + + return E +class L_exp(nn.Module): + + def __init__(self,patch_size,mean_val): + super(L_exp, self).__init__() + # print(1) + self.pool = nn.AvgPool2d(patch_size) + self.mean_val = mean_val + def forward(self, x ): + + b,c,h,w = x.shape + x = torch.mean(x,1,keepdim=True) + mean = self.pool(x) + + d = torch.mean(torch.pow(mean- torch.FloatTensor([self.mean_val] ).cuda(),2)) + return d + +class L_TV(nn.Module): + def __init__(self,TVLoss_weight=1): + super(L_TV,self).__init__() + self.TVLoss_weight = TVLoss_weight + + def forward(self,x): + batch_size = x.size()[0] + h_x = x.size()[2] + w_x = x.size()[3] + count_h = (x.size()[2]-1) * x.size()[3] + count_w = x.size()[2] * (x.size()[3] - 1) + h_tv = torch.pow((x[:,:,1:,:]-x[:,:,:h_x-1,:]),2).sum() + w_tv = torch.pow((x[:,:,:,1:]-x[:,:,:,:w_x-1]),2).sum() + return self.TVLoss_weight*2*(h_tv/count_h+w_tv/count_w)/batch_size +class Sa_Loss(nn.Module): + def __init__(self): + super(Sa_Loss, self).__init__() + # print(1) + def forward(self, x ): + # self.grad = np.ones(x.shape,dtype=np.float32) + b,c,h,w = x.shape + # x_de = x.cpu().detach().numpy() + r,g,b = torch.split(x , 1, dim=1) + mean_rgb = torch.mean(x,[2,3],keepdim=True) + mr,mg, mb = torch.split(mean_rgb, 1, dim=1) + Dr = r-mr + Dg = g-mg + Db = b-mb + k =torch.pow( torch.pow(Dr,2) + torch.pow(Db,2) + torch.pow(Dg,2),0.5) + # print(k) + + + k = torch.mean(k) + return k + +class perception_loss(nn.Module): + def __init__(self): + super(perception_loss, self).__init__() + features = vgg16(pretrained=True).features + self.to_relu_1_2 = nn.Sequential() + self.to_relu_2_2 = nn.Sequential() + self.to_relu_3_3 = nn.Sequential() + self.to_relu_4_3 = nn.Sequential() + + for x in range(4): + self.to_relu_1_2.add_module(str(x), features[x]) + for x in range(4, 9): + self.to_relu_2_2.add_module(str(x), features[x]) + for x in range(9, 16): + self.to_relu_3_3.add_module(str(x), features[x]) + for x in range(16, 23): + self.to_relu_4_3.add_module(str(x), features[x]) + + # don't need the gradients, just want the features + for param in self.parameters(): + param.requires_grad = False + + def forward(self, x): + h = self.to_relu_1_2(x) + h_relu_1_2 = h + h = self.to_relu_2_2(h) + h_relu_2_2 = h + h = self.to_relu_3_3(h) + h_relu_3_3 = h + h = self.to_relu_4_3(h) + h_relu_4_3 = h + # out = (h_relu_1_2, h_relu_2_2, h_relu_3_3, h_relu_4_3) + return h_relu_4_3 diff --git a/README.md b/README.md index 7ffb2a0fe9e4090a9d432a094b7e35f465bb9793..5df9ce88c78c3728ef0ff11a3d87889794f220de 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,6 @@ --- -title: Zero-DCE Code -emoji: 📈 -colorFrom: green -colorTo: red +title: Zero-DCE_code +app_file: app.py sdk: gradio sdk_version: 3.35.2 -app_file: app.py -pinned: false --- - -Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference diff --git a/__pycache__/dataloader.cpython-311.pyc b/__pycache__/dataloader.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c57e8b3d594d79bfb6721d9ef0668776ceba419 Binary files /dev/null and b/__pycache__/dataloader.cpython-311.pyc differ diff --git a/__pycache__/model.cpython-311.pyc b/__pycache__/model.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ace88a9617cedb8b4becebefd6bdd4113a82eac3 Binary files /dev/null and b/__pycache__/model.cpython-311.pyc differ diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..850011f3275211fca9e6536ce426544af684d523 --- /dev/null +++ b/app.py @@ -0,0 +1,81 @@ +import torch +import torch.nn as nn +import torchvision +import torch.backends.cudnn as cudnn +import torch.optim +import os +import sys +import argparse +import time +import dataloader +import model +import numpy as np +from torchvision import transforms +from PIL import Image +import glob +import time +import gradio as gr + + + +def lowlight(image_path): + os.environ['CUDA_VISIBLE_DEVICES']='0' + data_lowlight = Image.open(image_path) + + + + data_lowlight = (np.asarray(data_lowlight)/255.0) + + + data_lowlight = torch.from_numpy(data_lowlight).float() + data_lowlight = data_lowlight.permute(2,0,1) + data_lowlight = data_lowlight.cuda().unsqueeze(0) + + DCE_net = model.enhance_net_nopool().cuda() + DCE_net.load_state_dict(torch.load('snapshots/Epoch99.pth')) + start = time.time() + _,enhanced_image,_ = DCE_net(data_lowlight) + + end_time = (time.time() - start) + print(end_time) + image_path = image_path.replace('test_data','result') + result_path = image_path + if not os.path.exists(image_path.replace('/'+image_path.split("/")[-1],'')): + os.makedirs(image_path.replace('/'+image_path.split("/")[-1],'')) + + torchvision.utils.save_image(enhanced_image, result_path) + +def predict(img): + data_lowlight = (np.asarray(img)/255.0) + + + data_lowlight = torch.from_numpy(data_lowlight).float() + data_lowlight = data_lowlight.permute(2,0,1) + data_lowlight = data_lowlight.cuda().unsqueeze(0) + + DCE_net = model.enhance_net_nopool().cuda() + DCE_net.load_state_dict(torch.load('snapshots/Epoch99.pth')) + _,enhanced_image,_ = DCE_net(data_lowlight) + + return enhanced_image + + +if __name__ == '__main__': +# test_images + with torch.no_grad(): + # filePath = 'data/test_data/' + + # file_list = os.listdir(filePath) + + # for file_name in file_list: + # test_list = glob.glob(filePath+file_name+"/*") + # for image in test_list: + # # image = image + # print(image) + # lowlight(image) + + interface = gr.Interface(fn=predict, inputs='image', outputs='image') + interface.launch() + + + diff --git a/data/test_data/DICM/01.jpg b/data/test_data/DICM/01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..44b05d9d3702ba6c95e9bb7321db33a4d928f7c0 Binary files /dev/null and b/data/test_data/DICM/01.jpg differ diff --git a/data/test_data/DICM/02.jpg b/data/test_data/DICM/02.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a436431dff1986deb3d5db4186a12f4306fb8a8c Binary files /dev/null and b/data/test_data/DICM/02.jpg differ diff --git a/data/test_data/DICM/03.jpg b/data/test_data/DICM/03.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e095c262d79976f993ec7b3843b63f11b4ef09ed Binary files /dev/null and b/data/test_data/DICM/03.jpg differ diff --git a/data/test_data/DICM/04.jpg b/data/test_data/DICM/04.jpg new file mode 100644 index 0000000000000000000000000000000000000000..14ceddad65334999ec4be2616665a9f457d40439 Binary files /dev/null and b/data/test_data/DICM/04.jpg differ diff --git a/data/test_data/DICM/05.jpg b/data/test_data/DICM/05.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1db5d169a7a4303d1d17a130f382ce3b007b785b Binary files /dev/null and b/data/test_data/DICM/05.jpg differ diff --git a/data/test_data/DICM/06.jpg b/data/test_data/DICM/06.jpg new file mode 100644 index 0000000000000000000000000000000000000000..294369c1481b75df4902dc6274e2d82943e62016 Binary files /dev/null and b/data/test_data/DICM/06.jpg differ diff --git a/data/test_data/DICM/07.jpg b/data/test_data/DICM/07.jpg new file mode 100644 index 0000000000000000000000000000000000000000..3bf3360dccb30d79537a6feb8596ba65ed5602d5 Binary files /dev/null and b/data/test_data/DICM/07.jpg differ diff --git a/data/test_data/DICM/08.jpg b/data/test_data/DICM/08.jpg new file mode 100644 index 0000000000000000000000000000000000000000..fb3e3c0163a64465e6e89fb535a3f4829c759eae Binary files /dev/null and b/data/test_data/DICM/08.jpg differ diff --git a/data/test_data/DICM/09.jpg b/data/test_data/DICM/09.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e650cf8134fe98481b797404508167597f2d177e Binary files /dev/null and b/data/test_data/DICM/09.jpg differ diff --git a/data/test_data/DICM/10.jpg b/data/test_data/DICM/10.jpg new file mode 100644 index 0000000000000000000000000000000000000000..429ca834eb49f4a02168d34c915f1bf64d2e3713 Binary files /dev/null and b/data/test_data/DICM/10.jpg differ diff --git a/data/test_data/DICM/11.jpg b/data/test_data/DICM/11.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d8abe3657939e8b41c0087752fdc6af044d1a4ee Binary files /dev/null and b/data/test_data/DICM/11.jpg differ diff --git a/data/test_data/DICM/12.jpg b/data/test_data/DICM/12.jpg new file mode 100644 index 0000000000000000000000000000000000000000..82625f0b41d1dc2687862511b0a566db4b31325f Binary files /dev/null and b/data/test_data/DICM/12.jpg differ diff --git a/data/test_data/DICM/13.jpg b/data/test_data/DICM/13.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6d8bf52b6b4af9dba2ebfb4a3dde55ecae7cfca4 Binary files /dev/null and b/data/test_data/DICM/13.jpg differ diff --git a/data/test_data/DICM/14.jpg b/data/test_data/DICM/14.jpg new file mode 100644 index 0000000000000000000000000000000000000000..c442e9f3be704a1b80c6a22dfe494c7f7bae4a5e Binary files /dev/null and b/data/test_data/DICM/14.jpg differ diff --git a/data/test_data/DICM/15.jpg b/data/test_data/DICM/15.jpg new file mode 100644 index 0000000000000000000000000000000000000000..eba788a2792a34334e129bfec8c5610109150519 Binary files /dev/null and b/data/test_data/DICM/15.jpg differ diff --git a/data/test_data/DICM/16.jpg b/data/test_data/DICM/16.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e41eb5c1cc8425e53c4f20588fe53044f1d3115c Binary files /dev/null and b/data/test_data/DICM/16.jpg differ diff --git a/data/test_data/DICM/17.jpg b/data/test_data/DICM/17.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2121823aa4296703d63955a075c8f93328fe2be1 Binary files /dev/null and b/data/test_data/DICM/17.jpg differ diff --git a/data/test_data/DICM/18.jpg b/data/test_data/DICM/18.jpg new file mode 100644 index 0000000000000000000000000000000000000000..dcf29f23cf438d476465a8db227df8a166f39c6e Binary files /dev/null and b/data/test_data/DICM/18.jpg differ diff --git a/data/test_data/DICM/19.jpg b/data/test_data/DICM/19.jpg new file mode 100644 index 0000000000000000000000000000000000000000..9b47b001c87ddd68b66b381febfbac73f46b8793 Binary files /dev/null and b/data/test_data/DICM/19.jpg differ diff --git a/data/test_data/DICM/20.jpg b/data/test_data/DICM/20.jpg new file mode 100644 index 0000000000000000000000000000000000000000..5c2b00778404b53e2c2cb783800f08cb3f028832 Binary files /dev/null and b/data/test_data/DICM/20.jpg differ diff --git a/data/test_data/DICM/21.jpg b/data/test_data/DICM/21.jpg new file mode 100644 index 0000000000000000000000000000000000000000..726930fee1c43fbef63e688deeb9fe8260ccd0d1 Binary files /dev/null and b/data/test_data/DICM/21.jpg differ diff --git a/data/test_data/DICM/22.jpg b/data/test_data/DICM/22.jpg new file mode 100644 index 0000000000000000000000000000000000000000..cb20572889def592048c92b0237bd18daf2e4a9b Binary files /dev/null and b/data/test_data/DICM/22.jpg differ diff --git a/data/test_data/DICM/25.jpg b/data/test_data/DICM/25.jpg new file mode 100644 index 0000000000000000000000000000000000000000..72373f9c11fc8d632472a30d13c4e181a16e365e Binary files /dev/null and b/data/test_data/DICM/25.jpg differ diff --git a/data/test_data/DICM/26.jpg b/data/test_data/DICM/26.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e6b12006cbcd38ee7ad4a96faf83869a4d7ed07d Binary files /dev/null and b/data/test_data/DICM/26.jpg differ diff --git a/data/test_data/DICM/27.jpg b/data/test_data/DICM/27.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1b0a81c74876b143fc8de60c473289f232165e18 Binary files /dev/null and b/data/test_data/DICM/27.jpg differ diff --git a/data/test_data/DICM/28.jpg b/data/test_data/DICM/28.jpg new file mode 100644 index 0000000000000000000000000000000000000000..89bfd8127d05d89d6036658be259fc0f6bba492c Binary files /dev/null and b/data/test_data/DICM/28.jpg differ diff --git a/data/test_data/DICM/29.jpg b/data/test_data/DICM/29.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d01332a04fae69c4f5fbd19b04221b78f06179e2 Binary files /dev/null and b/data/test_data/DICM/29.jpg differ diff --git a/data/test_data/DICM/30.jpg b/data/test_data/DICM/30.jpg new file mode 100644 index 0000000000000000000000000000000000000000..8f1a39acb7b8fed56902fcc36dd110ea16784111 Binary files /dev/null and b/data/test_data/DICM/30.jpg differ diff --git a/data/test_data/DICM/31.jpg b/data/test_data/DICM/31.jpg new file mode 100644 index 0000000000000000000000000000000000000000..0551760889fd67fa22ef3d0bfc08a56425de0fad Binary files /dev/null and b/data/test_data/DICM/31.jpg differ diff --git a/data/test_data/DICM/32.jpg b/data/test_data/DICM/32.jpg new file mode 100644 index 0000000000000000000000000000000000000000..53f66091a9f3dabcfbfe995b40ee12fd0faaf5b7 Binary files /dev/null and b/data/test_data/DICM/32.jpg differ diff --git a/data/test_data/DICM/33.jpg b/data/test_data/DICM/33.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e5994a600e87d6f9bb4490595a454ac5c147cab5 Binary files /dev/null and b/data/test_data/DICM/33.jpg differ diff --git a/data/test_data/DICM/34.jpg b/data/test_data/DICM/34.jpg new file mode 100644 index 0000000000000000000000000000000000000000..5ab7b26603e397f94c7f46c21a38752939c1b654 Binary files /dev/null and b/data/test_data/DICM/34.jpg differ diff --git a/data/test_data/DICM/35.jpg b/data/test_data/DICM/35.jpg new file mode 100644 index 0000000000000000000000000000000000000000..22374373145b747a722152a2e0bd781cc261d027 Binary files /dev/null and b/data/test_data/DICM/35.jpg differ diff --git a/data/test_data/DICM/36.jpg b/data/test_data/DICM/36.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7216b931c8fd061b70cf03dc53beec3ef6cf5693 Binary files /dev/null and b/data/test_data/DICM/36.jpg differ diff --git a/data/test_data/DICM/37.jpg b/data/test_data/DICM/37.jpg new file mode 100644 index 0000000000000000000000000000000000000000..60027f982ee3cfd6a9e46e5cfcd6b809b5da0e17 Binary files /dev/null and b/data/test_data/DICM/37.jpg differ diff --git a/data/test_data/DICM/38.jpg b/data/test_data/DICM/38.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7839fe8033e7db52bcdfd65b620a571fa1ba0076 Binary files /dev/null and b/data/test_data/DICM/38.jpg differ diff --git a/data/test_data/DICM/39.jpg b/data/test_data/DICM/39.jpg new file mode 100644 index 0000000000000000000000000000000000000000..c53cac84b6d0dbe5713094e507a24ed1a8e7755d Binary files /dev/null and b/data/test_data/DICM/39.jpg differ diff --git a/data/test_data/DICM/40.jpg b/data/test_data/DICM/40.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e606c874ec2bbed161161777335ffa34df7850db Binary files /dev/null and b/data/test_data/DICM/40.jpg differ diff --git a/data/test_data/DICM/41.jpg b/data/test_data/DICM/41.jpg new file mode 100644 index 0000000000000000000000000000000000000000..48cf100a14bf7e0fc18e7f910a6378d2f1c02a4b Binary files /dev/null and b/data/test_data/DICM/41.jpg differ diff --git a/data/test_data/DICM/42.jpg b/data/test_data/DICM/42.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a542db94a9ec2a30b89c427be8866c5d13a68fe0 Binary files /dev/null and b/data/test_data/DICM/42.jpg differ diff --git a/data/test_data/DICM/43.jpg b/data/test_data/DICM/43.jpg new file mode 100644 index 0000000000000000000000000000000000000000..21ee9833847c3daad2ea504c350e9e61e51d2d46 Binary files /dev/null and b/data/test_data/DICM/43.jpg differ diff --git a/data/test_data/DICM/44.jpg b/data/test_data/DICM/44.jpg new file mode 100644 index 0000000000000000000000000000000000000000..336baf6ea67fbece134499b8f7f604858f22f4ea Binary files /dev/null and b/data/test_data/DICM/44.jpg differ diff --git a/data/test_data/DICM/45.jpg b/data/test_data/DICM/45.jpg new file mode 100644 index 0000000000000000000000000000000000000000..5058876e0d9f4b7526b7689ce32f0ff90e8f55b2 Binary files /dev/null and b/data/test_data/DICM/45.jpg differ diff --git a/data/test_data/DICM/46.jpg b/data/test_data/DICM/46.jpg new file mode 100644 index 0000000000000000000000000000000000000000..a29250daad4edf972696224afcab9c5318281c34 Binary files /dev/null and b/data/test_data/DICM/46.jpg differ diff --git a/data/test_data/DICM/47.jpg b/data/test_data/DICM/47.jpg new file mode 100644 index 0000000000000000000000000000000000000000..496ee2d8e65ce660fa7438e3ff65ca18aa7009bf Binary files /dev/null and b/data/test_data/DICM/47.jpg differ diff --git a/data/test_data/DICM/48.jpg b/data/test_data/DICM/48.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2104003e9dbe8c3d938c3eaa89e99f72e0ef7bdf Binary files /dev/null and b/data/test_data/DICM/48.jpg differ diff --git a/data/test_data/DICM/49.jpg b/data/test_data/DICM/49.jpg new file mode 100644 index 0000000000000000000000000000000000000000..5811f989ab55c6e2e888ece3810994bf47a83399 Binary files /dev/null and b/data/test_data/DICM/49.jpg differ diff --git a/data/test_data/DICM/50.jpg b/data/test_data/DICM/50.jpg new file mode 100644 index 0000000000000000000000000000000000000000..dccca3382789c7ef286c6f470c5515977f479355 Binary files /dev/null and b/data/test_data/DICM/50.jpg differ diff --git a/data/test_data/DICM/52.jpg b/data/test_data/DICM/52.jpg new file mode 100644 index 0000000000000000000000000000000000000000..f1fe944ec3cb95b6fc7720c94a6b081d91240ebc Binary files /dev/null and b/data/test_data/DICM/52.jpg differ diff --git a/data/test_data/DICM/53.jpg b/data/test_data/DICM/53.jpg new file mode 100644 index 0000000000000000000000000000000000000000..5796d538bb24f2ea746c52f97b512bde355aeae2 Binary files /dev/null and b/data/test_data/DICM/53.jpg differ diff --git a/data/test_data/DICM/54.jpg b/data/test_data/DICM/54.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2a2d15c89ed8d87b9b18a8ed15c1fd12a26a022e Binary files /dev/null and b/data/test_data/DICM/54.jpg differ diff --git a/data/test_data/DICM/55.jpg b/data/test_data/DICM/55.jpg new file mode 100644 index 0000000000000000000000000000000000000000..82c3d79fb05bff575511a494baec38a001e4d750 Binary files /dev/null and b/data/test_data/DICM/55.jpg differ diff --git a/data/test_data/DICM/56.jpg b/data/test_data/DICM/56.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6a6c0212fcd4503b1e6b2110ebc43fa6607dd3cd Binary files /dev/null and b/data/test_data/DICM/56.jpg differ diff --git a/data/test_data/DICM/57.jpg b/data/test_data/DICM/57.jpg new file mode 100644 index 0000000000000000000000000000000000000000..3756407b7e4020fa0573c44faac90f2e580e8c0d Binary files /dev/null and b/data/test_data/DICM/57.jpg differ diff --git a/data/test_data/DICM/58.jpg b/data/test_data/DICM/58.jpg new file mode 100644 index 0000000000000000000000000000000000000000..c43c7bbd7a2fc814d59ffd3c5eeecc7b310cf4f2 Binary files /dev/null and b/data/test_data/DICM/58.jpg differ diff --git a/data/test_data/DICM/60.jpg b/data/test_data/DICM/60.jpg new file mode 100644 index 0000000000000000000000000000000000000000..b6caafaa0de4f53aac786188d3661c2b723fceb1 Binary files /dev/null and b/data/test_data/DICM/60.jpg differ diff --git a/data/test_data/DICM/61.jpg b/data/test_data/DICM/61.jpg new file mode 100644 index 0000000000000000000000000000000000000000..e7e34106609598615c836fa73170189fc1685395 Binary files /dev/null and b/data/test_data/DICM/61.jpg differ diff --git a/data/test_data/DICM/62.jpg b/data/test_data/DICM/62.jpg new file mode 100644 index 0000000000000000000000000000000000000000..db2fd9a8b083acbcd958228116705a94147b7a7e Binary files /dev/null and b/data/test_data/DICM/62.jpg differ diff --git a/data/test_data/DICM/63.jpg b/data/test_data/DICM/63.jpg new file mode 100644 index 0000000000000000000000000000000000000000..ab7e01cf757a57604d623931ce4f74b7cd9598c9 Binary files /dev/null and b/data/test_data/DICM/63.jpg differ diff --git a/data/test_data/DICM/64.jpg b/data/test_data/DICM/64.jpg new file mode 100644 index 0000000000000000000000000000000000000000..55a0fbbff2cae06f97b2eb3286a213f7d812cddb Binary files /dev/null and b/data/test_data/DICM/64.jpg differ diff --git a/data/test_data/DICM/65.jpg b/data/test_data/DICM/65.jpg new file mode 100644 index 0000000000000000000000000000000000000000..ba0b9b1be73d64c628892ae003f6268f0de4029d Binary files /dev/null and b/data/test_data/DICM/65.jpg differ diff --git a/data/test_data/DICM/66.jpg b/data/test_data/DICM/66.jpg new file mode 100644 index 0000000000000000000000000000000000000000..78235430ee1800daee24274da7f03df3658acd15 Binary files /dev/null and b/data/test_data/DICM/66.jpg differ diff --git a/data/test_data/DICM/67.jpg b/data/test_data/DICM/67.jpg new file mode 100644 index 0000000000000000000000000000000000000000..82a2897a18f6cb17efe7ce0421d5b2b69452ec09 Binary files /dev/null and b/data/test_data/DICM/67.jpg differ diff --git a/data/test_data/DICM/69.jpg b/data/test_data/DICM/69.jpg new file mode 100644 index 0000000000000000000000000000000000000000..3ea4977c1004b5bbc4b725468872e696e21c33d0 Binary files /dev/null and b/data/test_data/DICM/69.jpg differ diff --git a/data/test_data/LIME/1.bmp b/data/test_data/LIME/1.bmp new file mode 100644 index 0000000000000000000000000000000000000000..a05af90739fa8e18b6ade8d93a1b83f855892d33 --- /dev/null +++ b/data/test_data/LIME/1.bmp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4f78ab13cc1f4905ab4d9d63997b7a6856e3dab40ec3b3fb49dd03bbc937f295 +size 1468854 diff --git a/data/test_data/LIME/10.bmp b/data/test_data/LIME/10.bmp new file mode 100644 index 0000000000000000000000000000000000000000..08f8cface471100eae6dff59cde573b12cec5db5 --- /dev/null +++ b/data/test_data/LIME/10.bmp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:70bf6d23c50cb559f32c447128bd086d495503d704e32c9679be7f7f544bd1b8 +size 2461734 diff --git a/data/test_data/LIME/2.bmp b/data/test_data/LIME/2.bmp new file mode 100644 index 0000000000000000000000000000000000000000..8ae7c8ba07e83261cd74210766f03b3977ed334d Binary files /dev/null and b/data/test_data/LIME/2.bmp differ diff --git a/data/test_data/LIME/3.bmp b/data/test_data/LIME/3.bmp new file mode 100644 index 0000000000000000000000000000000000000000..261441155706d94d265294676efde3e78e736b61 Binary files /dev/null and b/data/test_data/LIME/3.bmp differ diff --git a/data/test_data/LIME/4.bmp b/data/test_data/LIME/4.bmp new file mode 100644 index 0000000000000000000000000000000000000000..59199af802cb5e2be4cd4e315e96b1ef7fe6b0ae Binary files /dev/null and b/data/test_data/LIME/4.bmp differ diff --git a/data/test_data/LIME/5.bmp b/data/test_data/LIME/5.bmp new file mode 100644 index 0000000000000000000000000000000000000000..0ae4a286d077f7e968c37d0ddfe3fc14d70dde10 --- /dev/null +++ b/data/test_data/LIME/5.bmp @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b4aecc1500fa59dd7a90be9468b679db9008be6b3fd0abff1818313f11346db4 +size 9000054 diff --git a/data/test_data/LIME/6.bmp b/data/test_data/LIME/6.bmp new file mode 100644 index 0000000000000000000000000000000000000000..0e1231fa6edae3d3be06c9b6740da053dbc5cb3e Binary files /dev/null and b/data/test_data/LIME/6.bmp differ diff --git a/data/test_data/LIME/7.bmp b/data/test_data/LIME/7.bmp new file mode 100644 index 0000000000000000000000000000000000000000..d251c452af9efd95e1256fd1ff8eca040c7f04a8 Binary files /dev/null and b/data/test_data/LIME/7.bmp differ diff --git a/data/test_data/LIME/8.bmp b/data/test_data/LIME/8.bmp new file mode 100644 index 0000000000000000000000000000000000000000..e948916ca8a8e678c9b2e082670312ac67ba4787 Binary files /dev/null and b/data/test_data/LIME/8.bmp differ diff --git a/data/test_data/LIME/9.bmp b/data/test_data/LIME/9.bmp new file mode 100644 index 0000000000000000000000000000000000000000..cb08de088863d22913737dcebf653931787accb9 Binary files /dev/null and b/data/test_data/LIME/9.bmp differ diff --git a/dataloader.py b/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..e01217cae97229ddce694a83286356be82ed495d --- /dev/null +++ b/dataloader.py @@ -0,0 +1,59 @@ +import os +import sys + +import torch +import torch.utils.data as data + +import numpy as np +from PIL import Image +import glob +import random +import cv2 + +random.seed(1143) + + +def populate_train_list(lowlight_images_path): + + + + + image_list_lowlight = glob.glob(lowlight_images_path + "*.jpg") + + train_list = image_list_lowlight + + random.shuffle(train_list) + + return train_list + + + +class lowlight_loader(data.Dataset): + + def __init__(self, lowlight_images_path): + + self.train_list = populate_train_list(lowlight_images_path) + self.size = 256 + + self.data_list = self.train_list + print("Total training examples:", len(self.train_list)) + + + + + def __getitem__(self, index): + + data_lowlight_path = self.data_list[index] + + data_lowlight = Image.open(data_lowlight_path) + + data_lowlight = data_lowlight.resize((self.size,self.size), Image.ANTIALIAS) + + data_lowlight = (np.asarray(data_lowlight)/255.0) + data_lowlight = torch.from_numpy(data_lowlight).float() + + return data_lowlight.permute(2,0,1) + + def __len__(self): + return len(self.data_list) + diff --git a/lowlight_test.py b/lowlight_test.py new file mode 100644 index 0000000000000000000000000000000000000000..850011f3275211fca9e6536ce426544af684d523 --- /dev/null +++ b/lowlight_test.py @@ -0,0 +1,81 @@ +import torch +import torch.nn as nn +import torchvision +import torch.backends.cudnn as cudnn +import torch.optim +import os +import sys +import argparse +import time +import dataloader +import model +import numpy as np +from torchvision import transforms +from PIL import Image +import glob +import time +import gradio as gr + + + +def lowlight(image_path): + os.environ['CUDA_VISIBLE_DEVICES']='0' + data_lowlight = Image.open(image_path) + + + + data_lowlight = (np.asarray(data_lowlight)/255.0) + + + data_lowlight = torch.from_numpy(data_lowlight).float() + data_lowlight = data_lowlight.permute(2,0,1) + data_lowlight = data_lowlight.cuda().unsqueeze(0) + + DCE_net = model.enhance_net_nopool().cuda() + DCE_net.load_state_dict(torch.load('snapshots/Epoch99.pth')) + start = time.time() + _,enhanced_image,_ = DCE_net(data_lowlight) + + end_time = (time.time() - start) + print(end_time) + image_path = image_path.replace('test_data','result') + result_path = image_path + if not os.path.exists(image_path.replace('/'+image_path.split("/")[-1],'')): + os.makedirs(image_path.replace('/'+image_path.split("/")[-1],'')) + + torchvision.utils.save_image(enhanced_image, result_path) + +def predict(img): + data_lowlight = (np.asarray(img)/255.0) + + + data_lowlight = torch.from_numpy(data_lowlight).float() + data_lowlight = data_lowlight.permute(2,0,1) + data_lowlight = data_lowlight.cuda().unsqueeze(0) + + DCE_net = model.enhance_net_nopool().cuda() + DCE_net.load_state_dict(torch.load('snapshots/Epoch99.pth')) + _,enhanced_image,_ = DCE_net(data_lowlight) + + return enhanced_image + + +if __name__ == '__main__': +# test_images + with torch.no_grad(): + # filePath = 'data/test_data/' + + # file_list = os.listdir(filePath) + + # for file_name in file_list: + # test_list = glob.glob(filePath+file_name+"/*") + # for image in test_list: + # # image = image + # print(image) + # lowlight(image) + + interface = gr.Interface(fn=predict, inputs='image', outputs='image') + interface.launch() + + + diff --git a/lowlight_train.py b/lowlight_train.py new file mode 100644 index 0000000000000000000000000000000000000000..7743e8de199df14d52e6dd8fa6ed3f3361b82a2a --- /dev/null +++ b/lowlight_train.py @@ -0,0 +1,124 @@ +import torch +import torch.nn as nn +import torchvision +import torch.backends.cudnn as cudnn +import torch.optim +import os +import sys +import argparse +import time +import dataloader +import model +import Myloss +import numpy as np +from torchvision import transforms + + +def weights_init(m): + classname = m.__class__.__name__ + if classname.find('Conv') != -1: + m.weight.data.normal_(0.0, 0.02) + elif classname.find('BatchNorm') != -1: + m.weight.data.normal_(1.0, 0.02) + m.bias.data.fill_(0) + + + + + +def train(config): + + os.environ['CUDA_VISIBLE_DEVICES']='0' + + DCE_net = model.enhance_net_nopool().cuda() + + DCE_net.apply(weights_init) + if config.load_pretrain == True: + DCE_net.load_state_dict(torch.load(config.pretrain_dir)) + train_dataset = dataloader.lowlight_loader(config.lowlight_images_path) + + train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config.train_batch_size, shuffle=True, num_workers=config.num_workers, pin_memory=True) + + + + L_color = Myloss.L_color() + L_spa = Myloss.L_spa() + + L_exp = Myloss.L_exp(16,0.6) + L_TV = Myloss.L_TV() + + + optimizer = torch.optim.Adam(DCE_net.parameters(), lr=config.lr, weight_decay=config.weight_decay) + + DCE_net.train() + + for epoch in range(config.num_epochs): + for iteration, img_lowlight in enumerate(train_loader): + + img_lowlight = img_lowlight.cuda() + + enhanced_image_1,enhanced_image,A = DCE_net(img_lowlight) + + Loss_TV = 200*L_TV(A) + + loss_spa = torch.mean(L_spa(enhanced_image, img_lowlight)) + + loss_col = 5*torch.mean(L_color(enhanced_image)) + + loss_exp = 10*torch.mean(L_exp(enhanced_image)) + + + # best_loss + loss = Loss_TV + loss_spa + loss_col + loss_exp + # + + + optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm(DCE_net.parameters(),config.grad_clip_norm) + optimizer.step() + + if ((iteration+1) % config.display_iter) == 0: + print("Loss at iteration", iteration+1, ":", loss.item()) + if ((iteration+1) % config.snapshot_iter) == 0: + + torch.save(DCE_net.state_dict(), config.snapshots_folder + "Epoch" + str(epoch) + '.pth') + + + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + + # Input Parameters + parser.add_argument('--lowlight_images_path', type=str, default="data/train_data/") + parser.add_argument('--lr', type=float, default=0.0001) + parser.add_argument('--weight_decay', type=float, default=0.0001) + parser.add_argument('--grad_clip_norm', type=float, default=0.1) + parser.add_argument('--num_epochs', type=int, default=200) + parser.add_argument('--train_batch_size', type=int, default=8) + parser.add_argument('--val_batch_size', type=int, default=4) + parser.add_argument('--num_workers', type=int, default=4) + parser.add_argument('--display_iter', type=int, default=10) + parser.add_argument('--snapshot_iter', type=int, default=10) + parser.add_argument('--snapshots_folder', type=str, default="snapshots/") + parser.add_argument('--load_pretrain', type=bool, default= False) + parser.add_argument('--pretrain_dir', type=str, default= "snapshots/Epoch99.pth") + + config = parser.parse_args() + + if not os.path.exists(config.snapshots_folder): + os.mkdir(config.snapshots_folder) + + + train(config) + + + + + + + + + diff --git a/model.py b/model.py new file mode 100644 index 0000000000000000000000000000000000000000..3b710a5a33b38ee742c200c69b060f05246c7d1c --- /dev/null +++ b/model.py @@ -0,0 +1,59 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +#import pytorch_colors as colors +import numpy as np + +class enhance_net_nopool(nn.Module): + + def __init__(self): + super(enhance_net_nopool, self).__init__() + + self.relu = nn.ReLU(inplace=True) + + number_f = 32 + self.e_conv1 = nn.Conv2d(3,number_f,3,1,1,bias=True) + self.e_conv2 = nn.Conv2d(number_f,number_f,3,1,1,bias=True) + self.e_conv3 = nn.Conv2d(number_f,number_f,3,1,1,bias=True) + self.e_conv4 = nn.Conv2d(number_f,number_f,3,1,1,bias=True) + self.e_conv5 = nn.Conv2d(number_f*2,number_f,3,1,1,bias=True) + self.e_conv6 = nn.Conv2d(number_f*2,number_f,3,1,1,bias=True) + self.e_conv7 = nn.Conv2d(number_f*2,24,3,1,1,bias=True) + + self.maxpool = nn.MaxPool2d(2, stride=2, return_indices=False, ceil_mode=False) + self.upsample = nn.UpsamplingBilinear2d(scale_factor=2) + + + + def forward(self, x): + + x1 = self.relu(self.e_conv1(x)) + # p1 = self.maxpool(x1) + x2 = self.relu(self.e_conv2(x1)) + # p2 = self.maxpool(x2) + x3 = self.relu(self.e_conv3(x2)) + # p3 = self.maxpool(x3) + x4 = self.relu(self.e_conv4(x3)) + + x5 = self.relu(self.e_conv5(torch.cat([x3,x4],1))) + # x5 = self.upsample(x5) + x6 = self.relu(self.e_conv6(torch.cat([x2,x5],1))) + + x_r = F.tanh(self.e_conv7(torch.cat([x1,x6],1))) + r1,r2,r3,r4,r5,r6,r7,r8 = torch.split(x_r, 3, dim=1) + + + x = x + r1*(torch.pow(x,2)-x) + x = x + r2*(torch.pow(x,2)-x) + x = x + r3*(torch.pow(x,2)-x) + enhance_image_1 = x + r4*(torch.pow(x,2)-x) + x = enhance_image_1 + r5*(torch.pow(enhance_image_1,2)-enhance_image_1) + x = x + r6*(torch.pow(x,2)-x) + x = x + r7*(torch.pow(x,2)-x) + enhance_image = x + r8*(torch.pow(x,2)-x) + r = torch.cat([r1,r2,r3,r4,r5,r6,r7,r8],1) + return enhance_image_1,enhance_image,r + + + diff --git a/snapshots/Epoch99.pth b/snapshots/Epoch99.pth new file mode 100644 index 0000000000000000000000000000000000000000..fe06003f32b620f8544953a9cd7ec397080c3d3e --- /dev/null +++ b/snapshots/Epoch99.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a4395acb874f320375d9704997cef874eaaaaa26a1777ceb29a92b70f74c3612 +size 320017