colie import

app.py

@@ -1,7 +1,140 @@
+import os
+import tempfile
 import gradio as gr
+from PIL import Image
+import torch
+import numpy as np
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Assuming the other files (utils.py, loss.py, siren.py, color.py, filter.py) are in the same directory,
+# and have been modified to remove all .cuda() calls and any GPU-specific code.
+# For example, in utils.py: remove .cuda() from tensors, models, etc.
+# Similarly for others. Ensure everything runs on CPU with torch.device('cpu') if needed.
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+from utils import get_image, get_v_component, replace_v_component, interpolate_image, get_coords, get_patches, filter_up
+from loss import L_exp, L_TV
+from siren import INF
+from color import rgb2hsv_torch, hsv2rgb_torch
+
+# Note: get_image is modified to take a PIL Image instead of path.
+# Add this function or modify utils.py accordingly.
+def get_image_from_pil(pil_image):
+    image = torch.from_numpy(np.array(pil_image)).float()
+    image = image / torch.max(image)
+    image = torch.movedim(image, -1, 0).unsqueeze(0)  # No .cuda()
+    return image
+
+# The enhancement function
+def enhance_image(input_image, down_size, epochs, window, L, alpha, beta, gamma, delta):
+    if input_image is None:
+        raise gr.Error("Please upload an image.")
+
+    # Process the image
+    img_rgb = get_image_from_pil(input_image)
+    img_hsv = rgb2hsv_torch(img_rgb)
+
+    img_v = get_v_component(img_hsv)
+    img_v_lr = interpolate_image(img_v, down_size, down_size)
+    coords = get_coords(down_size, down_size)
+    patches = get_patches(img_v_lr, window)
+
+    img_siren = INF(patch_dim=window**2, num_layers=4, hidden_dim=256, add_layer=2)
+    # No .cuda()
+
+    optimizer = torch.optim.Adam(img_siren.parameters(), lr=1e-5, betas=(0.9, 0.999), weight_decay=3e-4)
+
+    l_exp = L_exp(16, L)
+    l_TV = L_TV()
+
+    for epoch in range(epochs):
+        img_siren.train()
+        optimizer.zero_grad()
+
+        illu_res_lr = img_siren(patches, coords)
+        illu_res_lr = illu_res_lr.view(1, 1, down_size, down_size)
+        illu_lr = illu_res_lr + img_v_lr
+
+        img_v_fixed_lr = (img_v_lr) / (illu_lr + 1e-4)
+
+        loss_spa = torch.mean(torch.abs(torch.pow(illu_lr - img_v_lr, 2)))
+        loss_tv = l_TV(illu_lr)
+        loss_exp = torch.mean(l_exp(illu_lr))
+        loss_sparsity = torch.mean(img_v_fixed_lr)
+
+        loss = loss_spa * alpha + loss_tv * beta + loss_exp * gamma + loss_sparsity * delta
+        loss.backward()
+        optimizer.step()
+
+    img_v_fixed = filter_up(img_v_lr, img_v_fixed_lr, img_v)
+    img_hsv_fixed = replace_v_component(img_hsv, img_v_fixed)
+    img_rgb_fixed = hsv2rgb_torch(img_hsv_fixed)
+    img_rgb_fixed = img_rgb_fixed / torch.max(img_rgb_fixed)
+
+    enhanced_np = (torch.movedim(img_rgb_fixed, 1, -1)[0].detach().cpu().numpy() * 255).astype(np.uint8)
+    enhanced_pil = Image.fromarray(enhanced_np)
+
+    # Save to temp file for download
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
+        enhanced_pil.save(tmp_file.name)
+        download_path = tmp_file.name
+
+    return enhanced_pil, download_path
+
+# Description from README
+description = """
+# CoLIE: Fast Context-Based Low-Light Image Enhancement via Neural Implicit Representations
+
+**Authors:** Tomáš Chobola*, Yu Liu, Hanyi Zhang, Julia A. Schnabel, Tingying Peng*  
+*Corresponding authors*  
+**Affiliations:** Technical University of Munich, Helmholtz AI, King’s College London  
+
+Accepted to ECCV 2024.
+
+## Overview
+- **Challenges with Current Methods:** Existing deep learning methods for low-light image enhancement struggle with high-resolution images, and they often fail to meet practical visual perception needs in diverse, unseen scenarios.
+- **Introduction of CoLIE:** CoLIE (Context-Based Low-Light Image Enhancement) is a novel approach for enhancing low-light images. It works by mapping 2D coordinates of underexposed images to their illumination components, conditioned on local context.
+- **Methodology:** The method utilizes HSV color space for image reconstruction. It employs an implicit neural function along with an embedded guided filter to further reduce computational overhead.
+- **Innovations in Training:** CoLIE introduces a single image-based training loss function. This function aims to improve the model's adaptability across various scenes, enhancing its practical applicability.
+
+Upload a low-light image on the left, adjust hyperparameters, and click 'Enhance' to see the result on the right.
+"""
+
+# Gradio interface
+with gr.Blocks(title="CoLIE - Low-Light Image Enhancement") as demo:
+    gr.Markdown(description)
+
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Upload Low-Light Image")
+            down_size = gr.Slider(minimum=64, maximum=512, step=32, value=256, label="Down Size")
+            epochs = gr.Slider(minimum=10, maximum=500, step=1, value=100, label="Epochs")
+            window = gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Window Size")
+            L = gr.Slider(minimum=0.1, maximum=1.0, step=0.01, value=0.5, label="L (Optimally-Intense Threshold)")
+            alpha = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=1.0, label="Alpha (Fidelity Control)")
+            beta = gr.Slider(minimum=1.0, maximum=100.0, step=1.0, value=20.0, label="Beta (Illumination Smoothness)")
+            gamma = gr.Slider(minimum=1.0, maximum=50.0, step=1.0, value=8.0, label="Gamma (Exposure Control)")
+            delta = gr.Slider(minimum=1.0, maximum=50.0, step=1.0, value=5.0, label="Delta (Sparsity Level)")
+            enhance_btn = gr.Button("Enhance")
+
+        with gr.Column():
+            output_image = gr.Image(label="Enhanced Image")
+            output_download = gr.File(label="Download Output Image")
+            # Optional: Input download not necessary, as user uploaded it. Right-click on input image works too.
+
+    # Examples section (grid with one example)
+    # Assume you have an example image in the repo, e.g., "examples/low_light_example.png"
+    # For demo, placeholder. Replace with actual path.
+    examples = gr.Examples(
+        examples=[
+            ["examples/low_light_example.png", 256, 100, 1, 0.5, 1.0, 20.0, 8.0, 5.0]
+        ],
+        inputs=[input_image, down_size, epochs, window, L, alpha, beta, gamma, delta],
+        label="Examples (Click to load image and hyperparameters)"
+    )
+
+    enhance_btn.click(
+        enhance_image,
+        inputs=[input_image, down_size, epochs, window, L, alpha, beta, gamma, delta],
+        outputs=[output_image, output_download]
+    )
+
+demo.launch()

colie.py

@@ -0,0 +1,84 @@
+from utils import *
+from loss import *
+from siren import INF
+from color import rgb2hsv_torch, hsv2rgb_torch
+
+import os
+import argparse
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+
+parser = argparse.ArgumentParser(description='CoLIE')
+parser.add_argument('--input_folder', type=str, default='input/')
+parser.add_argument('--output_folder', type=str, default='output/')
+parser.add_argument('--down_size', type=int, default=256, help='downsampling size')
+parser.add_argument('--epochs', type=int, default=100)
+parser.add_argument('--window', type=int, default=1, help='context window size')
+parser.add_argument('--L', type=float, default=0.5)
+# loss fuction weigth parameters
+parser.add_argument('--alpha', type=float, required=True)
+parser.add_argument('--beta', type=float, required=True)
+parser.add_argument('--gamma', type=float, required=True)
+parser.add_argument('--delta', type=float, required=True)
+opt = parser.parse_args()
+
+
+if not os.path.exists(opt.input_folder):
+    print('input folder: {} does not exist'.format(opt.input_folder))
+    exit()
+
+if not os.path.exists(opt.output_folder):
+    os.makedirs(opt.output_folder)
+
+
+print(' > running')
+for PATH in tqdm(np.sort(os.listdir(opt.input_folder))):
+    img_rgb = get_image(os.path.join(opt.input_folder, PATH))
+    img_hsv = rgb2hsv_torch(img_rgb)
+
+    img_v = get_v_component(img_hsv)
+    img_v_lr = interpolate_image(img_v, opt.down_size, opt.down_size)
+    coords = get_coords(opt.down_size, opt.down_size)
+    patches = get_patches(img_v_lr, opt.window)
+
+
+    img_siren = INF(patch_dim=opt.window**2, num_layers=4, hidden_dim=256, add_layer=2)
+
+    optimizer = torch.optim.Adam(img_siren.parameters(), lr=1e-5, betas=(0.9, 0.999), weight_decay=3e-4)
+
+    l_exp = L_exp(16,opt.L)
+    l_TV = L_TV()
+
+    for epoch in range(opt.epochs):
+        img_siren.train()
+        optimizer.zero_grad()
+
+        illu_res_lr = img_siren(patches, coords)
+        illu_res_lr = illu_res_lr.view(1,1,opt.down_size,opt.down_size)
+        illu_lr = illu_res_lr + img_v_lr
+
+        img_v_fixed_lr = (img_v_lr) / (illu_lr + 1e-4)
+
+        loss_spa = torch.mean(torch.abs(torch.pow(illu_lr - img_v_lr, 2)))
+        loss_tv  = l_TV(illu_lr)
+        loss_exp = torch.mean(l_exp(illu_lr))
+        loss_sparsity = torch.mean(img_v_fixed_lr)
+
+
+        loss = loss_spa * opt.alpha + loss_tv * opt.beta + loss_exp * opt.gamma + loss_sparsity * opt.delta
+        loss.backward()
+        optimizer.step()
+
+
+    img_v_fixed = filter_up(img_v_lr, img_v_fixed_lr, img_v)
+    img_hsv_fixed = replace_v_component(img_hsv, img_v_fixed)
+    img_rgb_fixed = hsv2rgb_torch(img_hsv_fixed)
+    img_rgb_fixed = img_rgb_fixed / torch.max(img_rgb_fixed)
+
+    Image.fromarray(
+        (torch.movedim(img_rgb_fixed,1,-1)[0].detach().cpu().numpy() * 255).astype(np.uint8)
+    ).save(os.path.join(opt.output_folder, PATH))
+
+print(' > reconstruction done')

color.py

@@ -0,0 +1,35 @@
+import torch
+
+def rgb2hsv_torch(rgb: torch.Tensor) -> torch.Tensor:
+    cmax, cmax_idx = torch.max(rgb, dim=1, keepdim=True)
+    cmin = torch.min(rgb, dim=1, keepdim=True)[0]
+    delta = cmax - cmin
+    hsv_h = torch.empty_like(rgb[:, 0:1, :, :])
+    cmax_idx[delta == 0] = 3
+    hsv_h[cmax_idx == 0] = (((rgb[:, 1:2] - rgb[:, 2:3]) / delta) % 6)[cmax_idx == 0]
+    hsv_h[cmax_idx == 1] = (((rgb[:, 2:3] - rgb[:, 0:1]) / delta) + 2)[cmax_idx == 1]
+    hsv_h[cmax_idx == 2] = (((rgb[:, 0:1] - rgb[:, 1:2]) / delta) + 4)[cmax_idx == 2]
+    hsv_h[cmax_idx == 3] = 0.
+    hsv_h /= 6.
+    hsv_s = torch.where(cmax == 0, torch.tensor(0.).type_as(rgb), delta / cmax)
+    hsv_v = cmax
+    return torch.cat([hsv_h, hsv_s, hsv_v], dim=1)
+
+
+def hsv2rgb_torch(hsv: torch.Tensor) -> torch.Tensor:
+    hsv_h, hsv_s, hsv_l = hsv[:, 0:1], hsv[:, 1:2], hsv[:, 2:3]
+    _c = hsv_l * hsv_s
+    _x = _c * (- torch.abs(hsv_h * 6. % 2. - 1) + 1.)
+    _m = hsv_l - _c
+    _o = torch.zeros_like(_c)
+    idx = (hsv_h * 6.).type(torch.uint8)
+    idx = (idx % 6).expand(-1, 3, -1, -1)
+    rgb = torch.empty_like(hsv)
+    rgb[idx == 0] = torch.cat([_c, _x, _o], dim=1)[idx == 0]
+    rgb[idx == 1] = torch.cat([_x, _c, _o], dim=1)[idx == 1]
+    rgb[idx == 2] = torch.cat([_o, _c, _x], dim=1)[idx == 2]
+    rgb[idx == 3] = torch.cat([_o, _x, _c], dim=1)[idx == 3]
+    rgb[idx == 4] = torch.cat([_x, _o, _c], dim=1)[idx == 4]
+    rgb[idx == 5] = torch.cat([_c, _o, _x], dim=1)[idx == 5]
+    rgb += _m
+    return rgb

filter.py

@@ -0,0 +1,80 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.autograd import Variable
+
+def diff_x(input, r):
+    assert input.dim() == 4
+
+    left   = input[:, :,         r:2 * r + 1]
+    middle = input[:, :, 2 * r + 1:         ] - input[:, :,           :-2 * r - 1]
+    right  = input[:, :,        -1:         ] - input[:, :, -2 * r - 1:    -r - 1]
+
+    output = torch.cat([left, middle, right], dim=2)
+
+    return output
+
+def diff_y(input, r):
+    assert input.dim() == 4
+
+    left   = input[:, :, :,         r:2 * r + 1]
+    middle = input[:, :, :, 2 * r + 1:         ] - input[:, :, :,           :-2 * r - 1]
+    right  = input[:, :, :,        -1:         ] - input[:, :, :, -2 * r - 1:    -r - 1]
+
+    output = torch.cat([left, middle, right], dim=3)
+
+    return output
+
+class BoxFilter(nn.Module):
+    def __init__(self, r):
+        super(BoxFilter, self).__init__()
+
+        self.r = r
+
+    def forward(self, x):
+        assert x.dim() == 4
+
+        return diff_y(diff_x(x.cumsum(dim=2), self.r).cumsum(dim=3), self.r)
+
+
+class FastGuidedFilter(nn.Module):
+    def __init__(self, r, eps=1e-8):
+        super(FastGuidedFilter, self).__init__()
+
+        self.r = r
+        self.eps = eps
+        self.boxfilter = BoxFilter(r)
+
+
+    def forward(self, lr_x, lr_y, hr_x):
+        n_lrx, c_lrx, h_lrx, w_lrx = lr_x.size()
+        n_lry, c_lry, h_lry, w_lry = lr_y.size()
+        n_hrx, c_hrx, h_hrx, w_hrx = hr_x.size()
+
+        assert n_lrx == n_lry and n_lry == n_hrx
+        assert c_lrx == c_hrx and (c_lrx == 1 or c_lrx == c_lry)
+        assert h_lrx == h_lry and w_lrx == w_lry
+        assert h_lrx > 2*self.r+1 and w_lrx > 2*self.r+1
+
+        ## N
+        N = self.boxfilter(Variable(lr_x.data.new().resize_((1, 1, h_lrx, w_lrx)).fill_(1.0)))
+
+        ## mean_x
+        mean_x = self.boxfilter(lr_x) / N
+        ## mean_y
+        mean_y = self.boxfilter(lr_y) / N
+        ## cov_xy
+        cov_xy = self.boxfilter(lr_x * lr_y) / N - mean_x * mean_y
+        ## var_x
+        var_x = self.boxfilter(lr_x * lr_x) / N - mean_x * mean_x
+
+        ## A
+        A = cov_xy / (var_x + self.eps)
+        ## b
+        b = mean_y - A * mean_x
+
+        ## mean_A; mean_b
+        mean_A = F.interpolate(A, (h_hrx, w_hrx), mode='bilinear', align_corners=True)
+        mean_b = F.interpolate(b, (h_hrx, w_hrx), mode='bilinear', align_corners=True)
+
+        return mean_A*hr_x+mean_b

loss.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class L_exp(nn.Module):
+    def __init__(self, patch_size, mean_val):
+        super(L_exp, self).__init__()
+        self.pool = nn.AvgPool2d(patch_size)
+        self.mean_val = mean_val
+
+    def forward(self, x):
+        mean = self.pool(x) ** 0.5
+        d = torch.abs(torch.mean(torch.pow(mean - torch.FloatTensor([self.mean_val] ),2)))
+        return d
+
+
+class L_TV(nn.Module):
+    def __init__(self):
+        super(L_TV,self).__init__()
+
+    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 2*(h_tv/count_h+w_tv/count_w)/batch_size
+

siren.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class SirenLayer(nn.Module):
+    def __init__(self, in_f, out_f, w0=30, is_first=False, is_last=False):
+        super().__init__()
+        self.in_f = in_f
+        self.w0 = w0
+        self.linear = nn.Linear(in_f, out_f)
+        self.is_first = is_first
+        self.is_last = is_last
+        if not self.is_last: 
+            self.init_weights()
+    
+    def init_weights(self):
+        b = 1 / self.in_f if self.is_first else np.sqrt(6 / self.in_f) / self.w0
+        with torch.no_grad():
+            self.linear.weight.uniform_(-b, b)
+
+    def forward(self, x):
+        x = self.linear(x)
+        return nn.Sigmoid()(x) if self.is_last else torch.sin(self.w0 * x)
+
+
+class INF(nn.Module):
+    def __init__(self, patch_dim, num_layers, hidden_dim, add_layer, weight_decay=None):
+        super().__init__()
+        '''
+        `add_layer` should be in range of  [1, num_layers-2]
+        '''
+
+        patch_layers = [SirenLayer(patch_dim, hidden_dim, is_first=True)]
+        spatial_layers = [SirenLayer(2, hidden_dim, is_first=True)]
+        output_layers = []
+        
+        for _ in range(1, add_layer - 2):
+            patch_layers.append(SirenLayer(hidden_dim, hidden_dim))
+            spatial_layers.append(SirenLayer(hidden_dim, hidden_dim))
+        patch_layers.append(SirenLayer(hidden_dim, hidden_dim//2))
+        spatial_layers.append(SirenLayer(hidden_dim, hidden_dim//2))
+        
+        for _ in range(add_layer, num_layers - 1):
+            output_layers.append(SirenLayer(hidden_dim, hidden_dim))
+        output_layers.append(SirenLayer(hidden_dim, 1, is_last=True))
+
+        self.patch_net = nn.Sequential(*patch_layers)
+        self.spatial_net = nn.Sequential(*spatial_layers)
+        self.output_net = nn.Sequential(*output_layers)
+        
+        if not weight_decay:
+            weight_decay = [0.1, 0.0001, 0.001]
+            
+        self.params = []
+        self.params += [{'params':self.spatial_net.parameters(),'weight_decay':weight_decay[0]}]
+        self.params += [{'params':self.patch_net.parameters(),'weight_decay':weight_decay[1]}]
+        self.params += [{'params':self.output_net.parameters(),'weight_decay':weight_decay[2]}]
+
+    def forward(self, patch, spatial):
+        return self.output_net(torch.cat((self.patch_net(patch), self.spatial_net(spatial)), -1))

utils.py

@@ -0,0 +1,76 @@
+from PIL import Image
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from filter import FastGuidedFilter
+
+
+def get_image(path):
+    """
+    Reads and returns RGB image, (1,3,H,W).
+    """
+    image = torch.from_numpy(np.array(Image.open(path))).float()
+    image = image / torch.max(image)
+    image = torch.movedim(image, -1, 0).unsqueeze(0)
+    return image
+
+
+def get_v_component(img_hsv):
+    """
+    Assumes (1,3,H,W) HSV image.
+    """
+    return img_hsv[:,-1].unsqueeze(0)
+
+
+def replace_v_component(img_hsv, v_new):
+    """
+    Replaces the V component of a HSV image (1,3,H,W).
+    """
+    img_hsv[:,-1] = v_new
+    return img_hsv
+
+
+def interpolate_image(img, H, W):
+    """
+    Reshapes the image based on new resolution.
+    """
+    return F.interpolate(img, size=(H,W))
+
+
+def get_coords(H, W):
+    """
+    Creates a coordinates grid for INF.
+    """
+    coords = np.dstack(np.meshgrid(np.linspace(0, 1, H), np.linspace(0, 1, W)))
+    coords = torch.from_numpy(coords).float()
+    return coords
+
+
+def get_patches(img, KERNEL_SIZE):
+    """
+    Creates a tensor where the channel contains patch information.
+    """
+    kernel = torch.zeros((KERNEL_SIZE ** 2, 1, KERNEL_SIZE, KERNEL_SIZE))
+
+    for i in range(KERNEL_SIZE):
+        for j in range(KERNEL_SIZE):
+            kernel[int(torch.sum(kernel).item()),0,i,j] = 1
+
+    pad = nn.ReflectionPad2d(KERNEL_SIZE//2)
+    im_padded = pad(img)
+
+    extracted = torch.nn.functional.conv2d(im_padded, kernel, padding=0).squeeze(0)
+
+    return torch.movedim(extracted, 0, -1)
+
+
+def filter_up(x_lr, y_lr, x_hr, r=1):
+    """
+    Applies the guided filter to upscale the predicted image.
+    """
+    guided_filter = FastGuidedFilter(r=r)
+    y_hr = guided_filter(x_lr, y_lr, x_hr)
+    y_hr = torch.clip(y_hr, 0, 1)
+    return y_hr