Adding first version of app

app.py

@@ -0,0 +1,47 @@
+import torch
+import torch.optim
+import model
+import numpy as np
+from PIL import Image
+import streamlit as st
+from torchvision import transforms
+
+scale_factor = 1
+
+@st.cache
+def load_model() -> torch.nn.Module:
+    DCE_net = model.enhance_net_nopool(scale_factor)
+    DCE_net.load_state_dict(torch.load("lowlight-dce-snapshot.pth", map_location=torch.device('cpu')))
+    
+    return DCE_net
+
+def fix_lowlight(image: Image.Image) -> Image.Image:
+    DCE_net = load_model()
+    data_lowlight = np.asarray(image) / 255.0
+
+    data_lowlight = torch.from_numpy(data_lowlight).float()
+
+    h = (data_lowlight.shape[0] // scale_factor) * scale_factor
+    w = (data_lowlight.shape[1] // scale_factor) * scale_factor
+    data_lowlight = data_lowlight[0:h, 0:w, :]
+    data_lowlight = data_lowlight.permute(2, 0, 1)
+    data_lowlight = data_lowlight.unsqueeze(0)
+    
+    enhanced_image, _ = DCE_net(data_lowlight)
+    im = transforms.ToPILImage()(enhanced_image[0]).convert("RGB")
+
+    return im
+
+
+def main():
+    st.title("Lowlight Enhancement")
+    st.write("This is a simple lowlight enhancement app with great performance and does not require paired images to train.")
+    uploaded_file = st.file_uploader("Lowlight Image")
+    if uploaded_file:
+        data_lowlight = Image.open(uploaded_file)
+        
+        fixed_img = fix_lowlight(data_lowlight)
+        
+        st.image(fixed_img, caption="Enhanced Image", use_column_width=True)
+
+main()

loss.py

@@ -0,0 +1,139 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models.vgg import vgg16
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+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, 4)
+        Drb = torch.pow(mr-mb, 4)
+        Dgb = torch.pow(mb-mg, 4)
+
+        return torch.sqrt(Drg + Drb + Dgb)
+
+			
+class L_spa(nn.Module):
+    def __init__(self):
+        super(L_spa, self).__init__()
+        kernel_left = torch.tensor( [[0,0,0], [-1,1,0], [0,0,0]], dtype=torch.float, device=device).unsqueeze(0).unsqueeze(0)
+        kernel_right = torch.tensor( [[0,0,0], [0,1,-1], [0,0,0]], dtype=torch.float, device=device).unsqueeze(0).unsqueeze(0)
+        kernel_up = torch.tensor( [[0,-1,0], [0,1,0], [0,0,0]], dtype=torch.float, device=device).unsqueeze(0).unsqueeze(0)
+        kernel_down = torch.tensor( [[0,0,0], [0,1,0], [0,-1,0]], dtype=torch.float, device=device).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)
+
+        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):
+        super(L_exp, self).__init__()
+        self.pool = nn.AvgPool2d(patch_size)
+        # self.mean_val = mean_val
+    def forward(self, x, mean_val):
+        b,c,h,w = x.shape
+        x = torch.mean(x,1,keepdim=True)
+        mean = self.pool(x)
+
+        d = torch.mean(torch.pow(mean - torch.tensor([mean_val], dtype=torch.float, device=device).to(), 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__()
+
+    def forward(self, x ):
+        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)
+
+        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])
+        
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, x):
+        h = self.to_relu_1_2(x)
+        h = self.to_relu_2_2(h)
+        h = self.to_relu_3_3(h)
+        h = self.to_relu_4_3(h)
+        
+        return h

model.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CSDN_Tem(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(CSDN_Tem, self).__init__()
+        self.depth_conv = nn.Conv2d(
+            in_channels=in_ch,
+            out_channels=in_ch,
+            kernel_size=3,
+            padding=1,
+            groups=in_ch
+        )
+        self.point_conv = nn.Conv2d(
+            in_channels=in_ch,
+            out_channels=out_ch,
+            kernel_size=1
+        )
+
+    def forward(self, input):
+        out = self.depth_conv(input)
+        out = self.point_conv(out)
+        return out
+
+class enhance_net_nopool(nn.Module):
+	def __init__(self,scale_factor):
+		super(enhance_net_nopool, self).__init__()
+
+		self.relu = nn.ReLU(inplace=True)
+		self.scale_factor = scale_factor
+		self.upsample = nn.UpsamplingBilinear2d(scale_factor=self.scale_factor)
+		number_f = 32
+
+#   zerodce DWC + p-shared
+		self.e_conv1 = CSDN_Tem(3, number_f) 
+		self.e_conv2 = CSDN_Tem(number_f, number_f) 
+		self.e_conv3 = CSDN_Tem(number_f, number_f) 
+		self.e_conv4 = CSDN_Tem(number_f, number_f) 
+		self.e_conv5 = CSDN_Tem(number_f * 2, number_f) 
+		self.e_conv6 = CSDN_Tem(number_f * 2, number_f) 
+		self.e_conv7 = CSDN_Tem(number_f * 2, 3) 
+
+	def enhance(self, x, x_r):
+		for _ in range(8): x = x + x_r * (torch.pow(x, 2) - x)
+
+		return x
+		
+	def forward(self, x):
+		x_down = x if self.scale_factor==1 else F.interpolate(x, scale_factor = 1 / self.scale_factor, mode='bilinear')
+
+		x1 = self.relu(self.e_conv1(x_down))
+		x2 = self.relu(self.e_conv2(x1))
+		x3 = self.relu(self.e_conv3(x2))
+		x4 = self.relu(self.e_conv4(x3))
+		x5 = self.relu(self.e_conv5(torch.cat([x3, x4], 1)))
+		x6 = self.relu(self.e_conv6(torch.cat([x2, x5], 1)))
+		x_r = torch.tanh(self.e_conv7(torch.cat([x1, x6], 1)))
+
+		x_r = x_r if self.scale_factor==1 else self.upsample(x_r)
+		enhance_image = self.enhance(x, x_r)
+
+		return enhance_image, x_r

requirements.txt

@@ -0,0 +1,2 @@
+torch==1.9.0
+torchvision==0.2.2