From 7ce089bdc354d45185a12ddfc5c70909359c2706 Mon Sep 17 00:00:00 2001 From: Striker Eurika Date: Wed, 7 Jan 2026 11:00:04 +0700 Subject: [PATCH 1/5] update for custom data --- .gitignore | 6 +- code/confs/LOL_smallNet_custom.yml | 124 +++++++++++++++++++++++++++++ code/data/LoL_dataset.py | 82 +++++++++++++++++++ code/data/__init__.py | 2 + code/requirements.txt | 4 +- code/train.py | 3 + 6 files changed, 217 insertions(+), 4 deletions(-) create mode 100644 code/confs/LOL_smallNet_custom.yml diff --git a/.gitignore b/.gitignore index 153c9df..8bbf4d1 100644 --- a/.gitignore +++ b/.gitignore @@ -131,4 +131,8 @@ dmypy.json # Pyre type checker .pyre/ -*.DS_Store \ No newline at end of file +*.DS_Store + +*.pth + +*_DONE \ No newline at end of file diff --git a/code/confs/LOL_smallNet_custom.yml b/code/confs/LOL_smallNet_custom.yml new file mode 100644 index 0000000..4784f3e --- /dev/null +++ b/code/confs/LOL_smallNet_custom.yml @@ -0,0 +1,124 @@ +#### general settings +name: train_custom_dataset_rebuttal_smallNet_ch32_blocks1 +use_tb_logger: true +model: LLFlow +distortion: sr +scale: 1 +gpu_ids: [0] +dataset: Custom +optimize_all_z: false +cond_encoder: ConEncoder1 +train_gt_ratio: 0.2 +avg_color_map: false + +concat_histeq: true +histeq_as_input: false +concat_color_map: false +gray_map: false # concat 1-input.mean(dim=1) to the input + +align_condition_feature: false +align_weight: 0.001 +align_maxpool: true + +to_yuv: false + +encode_color_map: false +le_curve: false +# sigmoid_output: true + +#### datasets +datasets: + train: + root: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_processed\train + quant: 32 + use_shuffle: true + n_workers: 0 # per GPU + batch_size: 16 + use_flip: true + color: RGB + use_crop: true + GT_size: 160 # 192 + noise_prob: 0 + noise_level: 5 + log_low: true + gamma_aug: false + + val: + root: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_processed\val + n_workers: 0 + quant: 32 + n_max: 20 + batch_size: 1 # must be 1 + log_low: true + +#### Test Settings +# dataroot_GT: D:\LOLdataset\eval15\high +# dataroot_LR: D:\LOLdataset\eval15\low +dataroot_unpaired: /home/data/Dataset/LOL_test/Fusion +dataroot_GT: D:\codes\projects\low-light-image-enhancement\data\lol_dataset\eval15\high +dataroot_LR: D:\codes\projects\low-light-image-enhancement\data\lol_dataset\eval15\low +model_path: D:\codes\projects\rm\LLFlow\code\models\LOL_smallNet.pth +heat: 0 # This is the standard deviation of the latent vectors + +#### network structures +network_G: + which_model_G: LLFlow + in_nc: 3 + out_nc: 3 + nf: 32 + nb: 4 # 12 for our low light encoder, 23 for LLFlow + train_RRDB: false + train_RRDB_delay: 0.5 + + flow: + K: 4 # 24.49 psnr用的12 # 16 + L: 3 # 4 + noInitialInj: true + coupling: CondAffineSeparatedAndCond + additionalFlowNoAffine: 2 + conditionInFeaDim: 64 + split: + enable: false + fea_up0: true + stackRRDB: + blocks: [1] + concat: true + +#### path +path: + # pretrain_model_G: ../pretrained_models/RRDB_DF2K_8X.pth + strict_load: true + resume_state: auto + +#### training settings: learning rate scheme, loss +train: + manual_seed: 10 + lr_G: !!float 5e-4 # normalizing flow 5e-4; l1 loss train 5e-5 + weight_decay_G: 0 # 1e-5 # 5e-5 # 1e-5 + beta1: 0.9 + beta2: 0.99 + lr_scheme: MultiStepLR + warmup_iter: -1 # no warm up + lr_steps_rel: [ 0.5, 0.75, 0.9, 0.95 ] # [0.2, 0.35, 0.5, 0.65, 0.8, 0.95] # [ 0.5, 0.75, 0.9, 0.95 ] + lr_gamma: 0.5 + + weight_l1: 0 + # flow_warm_up_iter: -1 + weight_fl: 1 + + niter: 200 #200000 + val_freq: 200 # 200 + +#### validation settings +val: + # heats: [ 0.0, 0.5, 0.75, 1.0 ] + n_sample: 4 + +test: + heats: [ 0.0, 0.7, 0.8, 0.9 ] + +#### logger +logger: + # Debug print_freq: 100 + print_freq: 100 + save_checkpoint_freq: !!float 1e3 diff --git a/code/data/LoL_dataset.py b/code/data/LoL_dataset.py index 164390c..e185682 100644 --- a/code/data/LoL_dataset.py +++ b/code/data/LoL_dataset.py @@ -284,3 +284,85 @@ def center_crop_tensor(img, size): assert border_double % 2 == 0, (img.shape, size) border = border_double // 2 return img[:, :, border:-border, border:-border] + + +# cutsom dataset for custom use +class Custom_Dataset(data.Dataset): + def __init__(self, opt, train, all_opt): + self.root = opt["root"] + self.opt = opt + self.concat_histeq = all_opt["concat_histeq"] if "concat_histeq" in all_opt.keys() else False + self.histeq_as_input = all_opt["histeq_as_input"] if "histeq_as_input" in all_opt.keys() else False + self.log_low = opt["log_low"] if "log_low" in opt.keys() else False + self.use_flip = opt["use_flip"] if "use_flip" in opt.keys() else False + self.use_rot = opt["use_rot"] if "use_rot" in opt.keys() else False + self.use_crop = opt["use_crop"] if "use_crop" in opt.keys() else False + self.use_noise = opt['noise_prob'] if "noise_prob" in opt.keys() else False + self.noise_prob = opt['noise_prob'] if self.use_noise else None + self.noise_level = opt['noise_level'] if "noise_level" in opt.keys() else 0 + self.center_crop_hr_size = opt.get("center_crop_hr_size", None) + self.crop_size = opt.get("GT_size", None) + + # Direct low/high structure + self.pairs = self.load_pairs(self.root) + self.to_tensor = ToTensor() + self.gamma_aug = opt['gamma_aug'] if 'gamma_aug' in opt.keys() else False + + def load_pairs(self, folder_path): + low_list = os.listdir(os.path.join(folder_path, 'low')) + low_list = sorted(list(filter(lambda x: 'png' in x or 'jpg' in x, low_list))) + pairs = [] + for f_name in low_list: + pairs.append([ + cv2.cvtColor(cv2.imread(os.path.join(folder_path, 'low', f_name)), cv2.COLOR_BGR2RGB), + cv2.cvtColor(cv2.imread(os.path.join(folder_path, 'high', f_name)), cv2.COLOR_BGR2RGB), + f_name.split('.')[0] + ]) + pairs[-1].append(self.hiseq_color_cv2_img(pairs[-1][0])) + return pairs + + def __len__(self): + return len(self.pairs) + + def hiseq_color_cv2_img(self, img): + (b, g, r) = cv2.split(img) + bH = cv2.equalizeHist(b) + gH = cv2.equalizeHist(g) + rH = cv2.equalizeHist(r) + result = cv2.merge((bH, gH, rH)) + return result + + def __getitem__(self, item): + lr, hr, f_name, his = self.pairs[item] + if self.histeq_as_input: + lr = his + + if self.use_crop: + hr, lr, his = random_crop(hr, lr, his, self.crop_size) + + if self.center_crop_hr_size: + hr, lr, his = center_crop(hr, self.center_crop_hr_size), center_crop(lr, self.center_crop_hr_size), center_crop(his, self.center_crop_hr_size) + + if self.use_flip: + hr, lr, his = random_flip(hr, lr, his) + + if self.use_rot: + hr, lr, his = random_rotation(hr, lr, his) + + if self.gamma_aug: + gamma = random.uniform(0.4, 2.8) + lr = gamma_aug(lr, gamma=gamma) + + hr = self.to_tensor(hr) + lr = self.to_tensor(lr) + + if self.use_noise and random.random() < self.noise_prob: + lr = torch.randn(lr.shape) * (self.noise_level / 255) + lr + if self.log_low: + lr = torch.log(torch.clamp(lr + 1e-3, min=1e-3)) + + if self.concat_histeq: + his = self.to_tensor(his) + lr = torch.cat([lr, his], dim=0) + + return {'LQ': lr, 'GT': hr, 'LQ_path': f_name, 'GT_path': f_name} \ No newline at end of file diff --git a/code/data/__init__.py b/code/data/__init__.py index f9562a8..515edc6 100644 --- a/code/data/__init__.py +++ b/code/data/__init__.py @@ -28,6 +28,8 @@ def create_dataset(dataset_opt): mode = dataset_opt['mode'] if mode == 'LoL': from data.LoL_dataset import LoL_Dataset as D + elif mode == 'Custom': + from data.LoL_dataset import Custom_Dataset as D else: raise NotImplementedError('Dataset [{:s}] is not recognized.'.format(mode)) dataset = D(dataset_opt) diff --git a/code/requirements.txt b/code/requirements.txt index 97089ad..2b4a1d8 100644 --- a/code/requirements.txt +++ b/code/requirements.txt @@ -7,9 +7,7 @@ pandas==1.2.4 PyYAML==6.0 scikit_image==0.18.1 seaborn==0.11.1 -skimage==0.0 +# skimage==0.0 tensorboard==2.6.0 tensorboardX==2.4.1 -torch==1.9.0 -torchvision==0.10.0 tqdm==4.59.0 diff --git a/code/train.py b/code/train.py index f204add..80f9bd1 100644 --- a/code/train.py +++ b/code/train.py @@ -139,6 +139,9 @@ def main(): dataset_cls = LoL_Dataset elif opt['dataset'] == 'LoL_v2': dataset_cls = LoL_Dataset_v2 + elif opt['dataset'] == 'Custom': + from data.LoL_dataset import Custom_Dataset + dataset_cls = Custom_Dataset else: raise NotImplementedError() From f263e37e725a6991a7697e936feec4d4adcfb53b Mon Sep 17 00:00:00 2001 From: Striker Eurika Date: Wed, 7 Jan 2026 14:54:07 +0700 Subject: [PATCH 2/5] Update configuration files and model paths for training and testing; remove unused custom config --- ..._custom.yml => Custom_smallNet_custom.yml} | 14 ++-- code/confs/LOL_smallNet.yml | 6 +- code/models/modules/LLFlow_arch.py | 51 +++++++----- code/test.py | 78 ++++++++++++++----- 4 files changed, 100 insertions(+), 49 deletions(-) rename code/confs/{LOL_smallNet_custom.yml => Custom_smallNet_custom.yml} (86%) diff --git a/code/confs/LOL_smallNet_custom.yml b/code/confs/Custom_smallNet_custom.yml similarity index 86% rename from code/confs/LOL_smallNet_custom.yml rename to code/confs/Custom_smallNet_custom.yml index 4784f3e..73a347f 100644 --- a/code/confs/LOL_smallNet_custom.yml +++ b/code/confs/Custom_smallNet_custom.yml @@ -29,7 +29,7 @@ le_curve: false #### datasets datasets: train: - root: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_processed\train + root: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_crops quant: 32 use_shuffle: true n_workers: 0 # per GPU @@ -44,7 +44,7 @@ datasets: gamma_aug: false val: - root: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_processed\val + root: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_crops n_workers: 0 quant: 32 n_max: 20 @@ -55,9 +55,9 @@ datasets: # dataroot_GT: D:\LOLdataset\eval15\high # dataroot_LR: D:\LOLdataset\eval15\low dataroot_unpaired: /home/data/Dataset/LOL_test/Fusion -dataroot_GT: D:\codes\projects\low-light-image-enhancement\data\lol_dataset\eval15\high -dataroot_LR: D:\codes\projects\low-light-image-enhancement\data\lol_dataset\eval15\low -model_path: D:\codes\projects\rm\LLFlow\code\models\LOL_smallNet.pth +dataroot_GT: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_crops\val\high +dataroot_LR: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_crops\val\low +model_path: D:\codes\projects\rm\LLFlow\experiments\train_custom_dataset_rebuttal_smallNet_ch32_blocks1\models\latest_G.pth heat: 0 # This is the standard deviation of the latent vectors #### network structures @@ -106,8 +106,8 @@ train: # flow_warm_up_iter: -1 weight_fl: 1 - niter: 200 #200000 - val_freq: 200 # 200 + niter: 20000 #200000 + val_freq: 1000 # 200 #### validation settings val: diff --git a/code/confs/LOL_smallNet.yml b/code/confs/LOL_smallNet.yml index 7bf4e92..941d0ca 100755 --- a/code/confs/LOL_smallNet.yml +++ b/code/confs/LOL_smallNet.yml @@ -55,9 +55,9 @@ datasets: # dataroot_GT: D:\LOLdataset\eval15\high # dataroot_LR: D:\LOLdataset\eval15\low dataroot_unpaired: /home/data/Dataset/LOL_test/Fusion -dataroot_GT: D:\Dataset\LOL-v2\LOL-v2\IntegratedTest\Test\high -dataroot_LR: D:\Dataset\LOL-v2\LOL-v2\IntegratedTest\Test\low -model_path: C:\Users\Yufei\OneDrive - Nanyang Technological University (1)\Project\AAAI2022-code-release\pretrained_model\LOL_smallNet.pth +dataroot_GT: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_crops\val\high +dataroot_LR: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_crops\val\low +model_path: D:\codes\projects\rm\LLFlow\code\models\LOL_smallNet.pth heat: 0 # This is the standard deviation of the latent vectors #### network structures diff --git a/code/models/modules/LLFlow_arch.py b/code/models/modules/LLFlow_arch.py index 3057eb0..a7db2de 100644 --- a/code/models/modules/LLFlow_arch.py +++ b/code/models/modules/LLFlow_arch.py @@ -1,6 +1,5 @@ - import math import random @@ -18,22 +17,25 @@ from models.modules.flow import unsqueeze2d, squeeze2d from torch.cuda.amp import autocast + class LLFlow(nn.Module): def __init__(self, in_nc, out_nc, nf, nb, gc=32, scale=4, K=None, opt=None, step=None): super(LLFlow, self).__init__() self.crop_size = opt['datasets']['train']['GT_size'] self.opt = opt self.quant = 255 if opt_get(opt, ['datasets', 'train', 'quant']) is \ - None else opt_get(opt, ['datasets', 'train', 'quant']) + None else opt_get(opt, ['datasets', 'train', 'quant']) if opt['cond_encoder'] == 'ConEncoder1': self.RRDB = ConEncoder1(in_nc, out_nc, nf, nb, gc, scale, opt) - elif opt['cond_encoder'] == 'NoEncoder': - self.RRDB = None # NoEncoder(in_nc, out_nc, nf, nb, gc, scale, opt) + elif opt['cond_encoder'] == 'NoEncoder': + # NoEncoder(in_nc, out_nc, nf, nb, gc, scale, opt) + self.RRDB = None elif opt['cond_encoder'] == 'RRDBNet': # if self.opt['encode_color_map']: print('Warning: ''encode_color_map'' is not implemented in RRDBNet') self.RRDB = RRDBNet(in_nc, out_nc, nf, nb, gc, scale, opt) else: - print('WARNING: Cannot find the conditional encoder %s, select RRDBNet by default.' % opt['cond_encoder']) + print('WARNING: Cannot find the conditional encoder %s, select RRDBNet by default.' % + opt['cond_encoder']) # if self.opt['encode_color_map']: print('Warning: ''encode_color_map'' is not implemented in RRDBNet') opt['cond_encoder'] = 'RRDBNet' self.RRDB = RRDBNet(in_nc, out_nc, nf, nb, gc, scale, opt) @@ -41,7 +43,8 @@ def __init__(self, in_nc, out_nc, nf, nb, gc=32, scale=4, K=None, opt=None, step if self.opt['encode_color_map']: self.color_map_encoder = ColorEncoder(nf=nf, opt=opt) - hidden_channels = opt_get(opt, ['network_G', 'flow', 'hidden_channels']) + hidden_channels = opt_get( + opt, ['network_G', 'flow', 'hidden_channels']) hidden_channels = hidden_channels or 64 self.RRDB_training = True # Default is true @@ -56,10 +59,12 @@ def __init__(self, in_nc, out_nc, nf, nb, gc=32, scale=4, K=None, opt=None, step self.i = 0 if self.opt['to_yuv']: self.A_rgb2yuv = torch.nn.Parameter(torch.tensor([[0.299, -0.14714119, 0.61497538], - [0.587, -0.28886916, -0.51496512], + [0.587, -0.28886916, - + 0.51496512], [0.114, 0.43601035, -0.10001026]]), requires_grad=False) self.A_yuv2rgb = torch.nn.Parameter(torch.tensor([[1., 1., 1.], - [0., -0.39465, 2.03211], + [0., -0.39465, + 2.03211], [1.13983, -0.58060, 0]]), requires_grad=False) if self.opt['align_maxpool']: self.max_pool = torch.nn.MaxPool2d(3) @@ -82,7 +87,7 @@ def yuv2rgb(self, yuv): rgb = torch.tensordot(yuv_, self.A_yuv2rgb, 1).transpose(1, 3) return rgb - @autocast() + # @autocast() # Disabled - causes NaN in inference def forward(self, gt=None, lr=None, z=None, eps_std=None, reverse=False, epses=None, reverse_with_grad=False, lr_enc=None, add_gt_noise=False, step=None, y_label=None, align_condition_feature=False, get_color_map=False): @@ -127,7 +132,8 @@ def normal_flow(self, gt, lr, y_onehot=None, epses=None, lr_enc=None, add_gt_noi if add_gt_noise: # Setup - noiseQuant = opt_get(self.opt, ['network_G', 'flow', 'augmentation', 'noiseQuant'], True) + noiseQuant = opt_get( + self.opt, ['network_G', 'flow', 'augmentation', 'noiseQuant'], True) if noiseQuant: z = z + ((torch.rand(z.shape, device=z.device) - 0.5) / self.quant) logdet = logdet + float(-np.log(self.quant) * pixels) @@ -147,14 +153,15 @@ def normal_flow(self, gt, lr, y_onehot=None, epses=None, lr_enc=None, add_gt_noi if 'avg_pool_color_map' in self.opt.keys() and self.opt['avg_pool_color_map']: mean = squeeze2d(F.avg_pool2d(lr_enc['color_map'], 7, 1, 3), 8) if random.random() > self.opt[ 'train_gt_ratio'] else squeeze2d(F.avg_pool2d( - gt / (gt.sum(dim=1, keepdims=True) + 1e-4), 7, 1, 3), 8) + gt / (gt.sum(dim=1, keepdims=True) + 1e-4), 7, 1, 3), 8) else: if self.RRDB is not None: mean = squeeze2d(lr_enc['color_map'], 8) if random.random() > self.opt['train_gt_ratio'] else squeeze2d( - gt/(gt.sum(dim=1, keepdims=True) + 1e-4), 8) + gt/(gt.sum(dim=1, keepdims=True) + 1e-4), 8) else: - mean = squeeze2d(lr[:,:3],8) - objective = objective + flow.GaussianDiag.logp(mean, torch.tensor(0.).to(z.device), z) + mean = squeeze2d(lr[:, :3], 8) + objective = objective + \ + flow.GaussianDiag.logp(mean, torch.tensor(0.).to(z.device), z) nll = (-objective) / float(np.log(2.) * pixels) if self.opt['encode_color_map']: @@ -167,7 +174,8 @@ def normal_flow(self, gt, lr, y_onehot=None, epses=None, lr_enc=None, add_gt_noi with torch.no_grad(): gt_enc = self.rrdbPreprocessing(gt) for k, v in gt_enc.items(): - if k in ['fea_up-1']: # ['fea_up2','fea_up1','fea_up0','fea_up-1']: + # ['fea_up2','fea_up1','fea_up0','fea_up-1']: + if k in ['fea_up-1']: if self.opt['align_maxpool']: nll = nll + (self.max_pool(gt_enc[k]) - self.max_pool(lr_enc[k])).abs().mean() * ( self.opt['align_weight'] if self.opt['align_weight'] is not None else 1) @@ -180,9 +188,11 @@ def normal_flow(self, gt, lr, y_onehot=None, epses=None, lr_enc=None, add_gt_noi def rrdbPreprocessing(self, lr): rrdbResults = self.RRDB(lr, get_steps=True) - block_idxs = opt_get(self.opt, ['network_G', 'flow', 'stackRRDB', 'blocks']) or [] + block_idxs = opt_get( + self.opt, ['network_G', 'flow', 'stackRRDB', 'blocks']) or [] if len(block_idxs) > 0: - low_level_features = [rrdbResults["block_{}".format(idx)] for idx in block_idxs] + low_level_features = [ + rrdbResults["block_{}".format(idx)] for idx in block_idxs] # low_level_features.append(rrdbResults['color_map']) concat = torch.cat(low_level_features, dim=1) @@ -195,12 +205,13 @@ def rrdbPreprocessing(self, lr): for k in keys: h = rrdbResults[k].shape[2] w = rrdbResults[k].shape[3] - rrdbResults[k] = torch.cat([rrdbResults[k], F.interpolate(concat, (h, w))], dim=1) + rrdbResults[k] = torch.cat( + [rrdbResults[k], F.interpolate(concat, (h, w))], dim=1) return rrdbResults def get_score(self, disc_loss_sigma, z): score_real = 0.5 * (1 - 1 / (disc_loss_sigma ** 2)) * thops.sum(z ** 2, dim=[1, 2, 3]) - \ - z.shape[1] * z.shape[2] * z.shape[3] * math.log(disc_loss_sigma) + z.shape[1] * z.shape[2] * z.shape[3] * math.log(disc_loss_sigma) return -score_real def reverse_flow(self, lr, z, y_onehot, eps_std, epses=None, lr_enc=None, add_gt_noise=True): @@ -214,7 +225,7 @@ def reverse_flow(self, lr, z, y_onehot, eps_std, epses=None, lr_enc=None, add_gt if lr_enc is None and self.RRDB: lr_enc = self.rrdbPreprocessing(lr) if self.opt['cond_encoder'] == "NoEncoder": - z = squeeze2d(lr[:,:3],8) + z = squeeze2d(lr[:, :3], 8) else: if 'avg_color_map' in self.opt.keys() and self.opt['avg_color_map']: z = squeeze2d(F.avg_pool2d(lr_enc['color_map'], 7, 1, 3), 8) diff --git a/code/test.py b/code/test.py index 338929d..a7e9926 100644 --- a/code/test.py +++ b/code/test.py @@ -1,6 +1,7 @@ import glob import sys from collections import OrderedDict +from datetime import datetime from natsort import natsort import argparse @@ -15,6 +16,7 @@ import os import cv2 + def fiFindByWildcard(wildcard): return natsort.natsorted(glob.glob(wildcard, recursive=True)) @@ -37,11 +39,12 @@ def predict(model, lr): return visuals.get('rlt', visuals.get('NORMAL')) -def t(array): return torch.Tensor(np.expand_dims(array.transpose([2, 0, 1]), axis=0).astype(np.float32)) / 255 +def t(array): return torch.Tensor(np.expand_dims( + array.transpose([2, 0, 1]), axis=0).astype(np.float32)) / 255 def rgb(t): return ( - np.clip((t[0] if len(t.shape) == 4 else t).detach().cpu().numpy().transpose([1, 2, 0]), 0, 1) * 255).astype( + np.clip((t[0] if len(t.shape) == 4 else t).detach().cpu().numpy().transpose([1, 2, 0]), 0, 1) * 255).astype( np.uint8) @@ -69,6 +72,7 @@ def imCropCenter(img, size): def impad(img, top=0, bottom=0, left=0, right=0, color=255): return np.pad(img, [(top, bottom), (left, right), (0, 0)], 'reflect') + def hiseq_color_cv2_img(img): (b, g, r) = cv2.split(img) bH = cv2.equalizeHist(b) @@ -77,6 +81,7 @@ def hiseq_color_cv2_img(img): result = cv2.merge((bH, gH, rH)) return result + def main(): parser = argparse.ArgumentParser() parser.add_argument("--opt", default="confs/LOL_smallNet.yml") @@ -86,19 +91,33 @@ def main(): model, opt = load_model(conf_path) model.netG = model.netG.cuda() + # Debug: Check if model loaded correctly + print(f"Model loaded from: {opt.get('model_path', 'N/A')}") + print( + f"Model has parameters: {sum(p.numel() for p in model.netG.parameters())/1e6:.2f}M") + print( + f"concat_histeq: {opt.get('concat_histeq', False)}, log_low: {opt['datasets']['train'].get('log_low', False)}") + lr_dir = opt['dataroot_LR'] hr_dir = opt['dataroot_GT'] - lr_paths = fiFindByWildcard(os.path.join(lr_dir, '*.png')) - hr_paths = fiFindByWildcard(os.path.join(hr_dir, '*.png')) + lr_paths = fiFindByWildcard(os.path.join( + lr_dir, '*.png')) + fiFindByWildcard(os.path.join(lr_dir, '*.jpg')) + hr_paths = fiFindByWildcard(os.path.join( + hr_dir, '*.png')) + fiFindByWildcard(os.path.join(hr_dir, '*.jpg')) this_dir = os.path.dirname(os.path.realpath(__file__)) test_dir = os.path.join(this_dir, '..', 'results', conf) print(f"Out dir: {test_dir}") + # ensure the output directory exists + os.makedirs(test_dir, exist_ok=True) + measure = Measure(use_gpu=False) - fname = f'measure_full.csv' + # Generate timestamp for unique filename + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + fname = f'measure_full_{timestamp}.csv' fname_tmp = fname + "_" path_out_measures = os.path.join(test_dir, fname_tmp) path_out_measures_final = os.path.join(test_dir, fname) @@ -118,37 +137,54 @@ def main(): lr = imread(lr_path) hr = imread(hr_path) - his = hiseq_color_cv2_img(lr) if opt.get("histeq_as_input", False): - lr = his - + lr = hiseq_color_cv2_img(lr) + # Pad image to be % 2 h, w, c = lr.shape lq_orig = lr.copy() lr = impad(lr, bottom=int(np.ceil(h / pad_factor) * pad_factor - h), right=int(np.ceil(w / pad_factor) * pad_factor - w)) - + lr_t = t(lr) - if opt["datasets"]["train"].get("log_low", False): + log_low_enabled = opt["datasets"]["train"].get("log_low", False) + if log_low_enabled: lr_t = torch.log(torch.clamp(lr_t + 1e-3, min=1e-3)) if opt.get("concat_histeq", False): - his = t(his) - lr_t = torch.cat([lr_t, his], dim=1) + # Compute and pad histogram equalized image to match lr dimensions + his = hiseq_color_cv2_img(imread(lr_path)) + his = impad(his, bottom=int(np.ceil(h / pad_factor) * pad_factor - h), + right=int(np.ceil(w / pad_factor) * pad_factor - w)) + his_t = t(his) + lr_t = torch.cat([lr_t, his_t], dim=1) heat = 0 - + if df is not None and len(df[(df['heat'] == heat) & (df['name'] == idx_test)]) == 1: continue - with torch.cuda.amp.autocast(): - sr_t = model.get_sr(lq=lr_t.cuda(), heat=None) + + # Debug: check input before model + print( + f"DEBUG INPUT - shape: {lr_t.shape}, min: {lr_t.min().item():.6f}, max: {lr_t.max().item():.6f}, mean: {lr_t.mean().item():.6f}, has_nan: {torch.isnan(lr_t).any().item()}") + + # Disable autocast - mixed precision can cause NaN with this model + sr_t = model.get_sr(lq=lr_t.cuda(), heat=None) + + # Debug: print output statistics + print( + f"DEBUG OUTPUT - min: {sr_t.min().item():.6f}, max: {sr_t.max().item():.6f}, mean: {sr_t.mean().item():.6f}") + + # Model outputs in [0,1] range directly (same as training GT), no need to reverse log transform # We follow a similar way of 'Kind' to finetune the overall brightness as illustrated in Line 73 (https://github.com/zhangyhuaee/KinD/blob/master/evaluate_LOLdataset.py). # A normally-exposed image can also be obtained without finetuning the global brightness and we can achvieve compatible performance in terms of SSIM and LPIPS. - mean_out = sr_t.view(sr_t.shape[0],-1).mean(dim=1) - mean_gt = cv2.cvtColor(hr.astype(np.float32), cv2.COLOR_BGR2GRAY).mean()/255 + mean_out = sr_t.view(sr_t.shape[0], -1).mean(dim=1) + mean_gt = cv2.cvtColor(hr.astype(np.float32), + cv2.COLOR_BGR2GRAY).mean()/255 sr = rgb(torch.clamp(sr_t*(mean_gt/mean_out), 0, 1)) sr = sr[:h * scale, :w * scale] - path_out_sr = os.path.join(test_dir, "{:0.2f}".format(heat).replace('.', ''), os.path.basename(hr_path)) + path_out_sr = os.path.join(test_dir, "{:0.2f}".format( + heat).replace('.', ''), os.path.basename(hr_path)) imwrite(path_out_sr, sr) @@ -161,11 +197,15 @@ def main(): str_out = format_measurements(meas) print(str_out) - df = pd.DataFrame([meas]) if df is None else pd.concat([pd.DataFrame([meas]), df]) + df = pd.DataFrame([meas]) if df is None else pd.concat( + [pd.DataFrame([meas]), df]) # df.to_csv(path_out_measures + "_", index=False) # os.rename(path_out_measures + "_", path_out_measures) + if df is None: + df = pd.DataFrame() + df.to_csv(path_out_measures, index=False) os.rename(path_out_measures, path_out_measures_final) From 3b2ca2815c963e53c50faf49b00ae2960d7adc5f Mon Sep 17 00:00:00 2001 From: Striker Eurika Date: Thu, 8 Jan 2026 12:02:56 +0700 Subject: [PATCH 3/5] Add image processing utilities and update training configuration - Introduced `crop_random.py` for generating high-quality image crops. - Added `make_data.py` for creating low-light images from ground truth. - Updated training parameters in `Custom_smallNet_custom.yml`. - Modified `.gitignore` to include `data_custom`. --- .gitignore | 4 +- code/confs/Custom_smallNet_custom.yml | 4 +- code/utils/crop_random.py | 43 ++++++++++++++++++ code/utils/make_data.py | 64 +++++++++++++++++++++++++++ 4 files changed, 112 insertions(+), 3 deletions(-) create mode 100644 code/utils/crop_random.py create mode 100644 code/utils/make_data.py diff --git a/.gitignore b/.gitignore index 8bbf4d1..bae0437 100644 --- a/.gitignore +++ b/.gitignore @@ -135,4 +135,6 @@ dmypy.json *.pth -*_DONE \ No newline at end of file +*_DONE + +data_custom \ No newline at end of file diff --git a/code/confs/Custom_smallNet_custom.yml b/code/confs/Custom_smallNet_custom.yml index 73a347f..3fb25a0 100644 --- a/code/confs/Custom_smallNet_custom.yml +++ b/code/confs/Custom_smallNet_custom.yml @@ -106,8 +106,8 @@ train: # flow_warm_up_iter: -1 weight_fl: 1 - niter: 20000 #200000 - val_freq: 1000 # 200 + niter: 1000 #200000 + val_freq: 100 # 200 #### validation settings val: diff --git a/code/utils/crop_random.py b/code/utils/crop_random.py new file mode 100644 index 0000000..596e848 --- /dev/null +++ b/code/utils/crop_random.py @@ -0,0 +1,43 @@ +import os +import random +from PIL import Image, ImageOps +from tqdm import tqdm + +def generate_high_quality_crops(input_folder, output_folder, target_w=600, target_h=400, crops_per_image=3): + if not os.path.exists(output_folder): + os.makedirs(output_folder) + + files = [f for f in os.listdir(input_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.tif'))] + + for filename in tqdm(files): + img_path = os.path.join(input_folder, filename) + try: + with Image.open(img_path) as img: + img = ImageOps.exif_transpose(img) + + # --- QUALITY STEP: DOWNSAMPLE FIRST --- + # Reducing by half (0.5) usually removes raw sensor noise + scale_factor = 0.5 + new_size = (int(img.size[0] * scale_factor), int(img.size[1] * scale_factor)) + img = img.resize(new_size, Image.Resampling.LANCZOS) + + img_w, img_h = img.size + if img_w < target_w or img_h < target_h: + continue + + for i in range(crops_per_image): + x1 = random.randint(0, img_w - target_w) + y1 = random.randint(0, img_h - target_h) + + box = (x1, y1, x1 + target_w, y1 + target_h) + cropped_img = img.crop(box) + + save_name = f"{os.path.splitext(filename)[0]}_hq_crop_{i}.png" + # Saving as PNG is vital for LLFlow to avoid JPEG artifacts + cropped_img.save(os.path.join(output_folder, save_name)) + + except Exception as e: + print(f"Error: {e}") + +if __name__ == "__main__": + generate_high_quality_crops("D:\\codes\\projects\\rm\\low-light-image-enhancement\\data\\rm_dataset", "D:\\codes\\projects\\rm\\low-light-image-enhancement\\data\\rm_dataset_crops") \ No newline at end of file diff --git a/code/utils/make_data.py b/code/utils/make_data.py new file mode 100644 index 0000000..038d2ed --- /dev/null +++ b/code/utils/make_data.py @@ -0,0 +1,64 @@ +from PIL import Image, ImageEnhance +import os +from pathlib import Path +from PIL.Image import Exif + + +# Define paths +# Replace with your source images directory +SOURCE_DIR = "D:/codes/projects/rm/low-light-image-enhancement/data/rm_dataset_crops/val" +# Replace with your ground truth images directory +GROUND_TRUTH_DIR = "D:/codes/projects/rm/low-light-image-enhancement/data/rm_dataset_crops/val/high" +# Replace with your desired low-light images directory +LOW_LIGHT_DIR = "D:/codes/projects/rm/low-light-image-enhancement/data/rm_dataset_crops/val/low" + +# Create the low-light directory if it doesn't exist +Path(LOW_LIGHT_DIR).mkdir(parents=True, exist_ok=True) +Path(GROUND_TRUTH_DIR).mkdir(parents=True, exist_ok=True) + +# Brightness reduction factor +brightness_factor = 0.15 + + +def apply_exif_rotation(img): + """Apply EXIF rotation to image if present""" + try: + exif = img._getexif() + if exif is not None: + for tag, value in exif.items(): + if tag == 274: # Orientation tag + if value == 3: + img = img.rotate(180, expand=True) + elif value == 6: + img = img.rotate(270, expand=True) + elif value == 8: + img = img.rotate(90, expand=True) + except: + pass + return img + + +# Process each image in the ground truth directory +for filename in os.listdir(SOURCE_DIR): + if filename.lower().endswith(('.png', '.jpg', '.jpeg')): # Check for image files + # Open the ground truth image + img_path = os.path.join(SOURCE_DIR, filename) + img = Image.open(img_path) + + # apply EXIF rotation if needed + img = apply_exif_rotation(img) + + # Reduce brightness to create a low-light image + enhancer = ImageEnhance.Brightness(img) + low_light_img = enhancer.enhance(brightness_factor) + + # save the ground truth image to the high directory + ground_truth_img_path = os.path.join(GROUND_TRUTH_DIR, filename) + img.save(ground_truth_img_path) + + # Save the low-light image + low_light_img_path = os.path.join(LOW_LIGHT_DIR, filename) + low_light_img.save(low_light_img_path) + + print( + f"Processed {filename}: Low-light image saved to {low_light_img_path}") From 2e8bdcf93d5b078d7087a956a1db233ff3263460 Mon Sep 17 00:00:00 2001 From: Striker Eurika Date: Thu, 8 Jan 2026 21:49:49 +0700 Subject: [PATCH 4/5] Update README and configuration files; add image enhancement usage examples and adjust dataset paths --- README.md | 17 ++++ code/confs/Custom_smallNet_custom.yml | 7 +- code/models/modules/Permutations.py | 26 +++--- code/using.py | 125 ++++++++++++++++++++++++++ code/utils/crop_random.py | 2 +- code/utils/make_data.py | 7 +- 6 files changed, 163 insertions(+), 21 deletions(-) create mode 100644 code/using.py diff --git a/README.md b/README.md index 4c6ca57..174d8e8 100755 --- a/README.md +++ b/README.md @@ -104,7 +104,24 @@ gpu_ids: [0] # Our model can be trained using a single GPU with memory>20GB. You 2. Train the network. ```bash python train.py --opt your_config_path + +``` + +### Using +You can directly use our pre-trained model for low-light image enhancement in your own project. Here is an example code snippet to load the pre-trained model and enhance a low-light image. + +**For a Single Image** + +```python +python using.py --input path/to/image.jpg --output path/to/output.jpg --model path/to/config.yml ``` + +**For a Folder of Images** + +```python +python using.py --input path/to/input_folder --output path/to/output_folder --model path/to/config.yml +``` + ## Citation If you find our work useful for your research, please cite our paper ``` diff --git a/code/confs/Custom_smallNet_custom.yml b/code/confs/Custom_smallNet_custom.yml index 3fb25a0..61c042e 100644 --- a/code/confs/Custom_smallNet_custom.yml +++ b/code/confs/Custom_smallNet_custom.yml @@ -54,9 +54,10 @@ datasets: #### Test Settings # dataroot_GT: D:\LOLdataset\eval15\high # dataroot_LR: D:\LOLdataset\eval15\low -dataroot_unpaired: /home/data/Dataset/LOL_test/Fusion -dataroot_GT: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_crops\val\high -dataroot_LR: D:\codes\projects\rm\low-light-image-enhancement\data\rm_dataset_crops\val\low +dataroot_unpaired: D:\codes\projects\rm\LLFlow\code\using\using_crops\predick + +dataroot_GT: D:\codes\projects\rm\LLFlow\code\using\using_crops\high +dataroot_LR: D:\codes\projects\rm\LLFlow\code\using\using_crops\low model_path: D:\codes\projects\rm\LLFlow\experiments\train_custom_dataset_rebuttal_smallNet_ch32_blocks1\models\latest_G.pth heat: 0 # This is the standard deviation of the latent vectors diff --git a/code/models/modules/Permutations.py b/code/models/modules/Permutations.py index 259eb3c..f6e44d3 100644 --- a/code/models/modules/Permutations.py +++ b/code/models/modules/Permutations.py @@ -1,10 +1,8 @@ - - - import numpy as np import torch from torch import nn as nn from torch.nn import functional as F +from torch.nn.parameter import Parameter from models.modules import thops @@ -14,7 +12,8 @@ def __init__(self, num_channels, LU_decomposed=False): super().__init__() w_shape = [num_channels, num_channels] w_init = np.linalg.qr(np.random.randn(*w_shape))[0].astype(np.float32) - self.register_parameter("weight", nn.Parameter(torch.Tensor(w_init))) + # Ensure weight is a Tensor + self.weight = Parameter(torch.Tensor(w_init)) self.w_shape = w_shape self.LU = LU_decomposed @@ -27,19 +26,20 @@ def get_weight(self, input, reverse): dlogdet = torch.slogdet(self.weight)[1] * pixels except Exception as e: print(e) - dlogdet = \ - torch.slogdet( - self.weight + (self.weight.mean() * torch.randn(*self.w_shape).to(input.device) * 0.001))[ - 1] * pixels + dlogdet = torch.slogdet( + self.weight + + (self.weight.mean() * torch.randn(* + self.w_shape).to(input.device) * 0.001) + )[1] * pixels if not reverse: weight = self.weight.view(w_shape[0], w_shape[1], 1, 1) else: try: - weight = torch.inverse(self.weight.double()).float() \ - .view(w_shape[0], w_shape[1], 1, 1) - except: - weight = torch.inverse(self.weight.double()+ (self.weight.mean() * torch.randn(*self.w_shape).to(input.device) * 0.001).float() \ - .view(w_shape[0], w_shape[1], 1, 1)) + weight = torch.inverse(self.weight.double()).float().view( + w_shape[0], w_shape[1], 1, 1) + except RuntimeError: + weight = torch.inverse(self.weight.double().cpu()).float().to( + input.device).view(w_shape[0], w_shape[1], 1, 1) return weight, dlogdet def forward(self, input, logdet=None, reverse=False): diff --git a/code/using.py b/code/using.py new file mode 100644 index 0000000..d5f7e18 --- /dev/null +++ b/code/using.py @@ -0,0 +1,125 @@ +import os +import argparse +import glob +import cv2 +import numpy as np +import torch +from natsort import natsorted +from utils.util import opt_get +from models import create_model +from options import options as option + + +def fiFindByWildcard(wildcard): + return natsorted(glob.glob(wildcard, recursive=True)) + + +def load_model(conf_path, model_path): + opt = option.parse(conf_path, is_train=False) + opt['gpu_ids'] = None + opt = option.dict_to_nonedict(opt) + model = create_model(opt) + + model.load_network(load_path=model_path, network=model.netG) + model.netG = model.netG.cuda() # Move model to GPU + return model, opt + + +def t(array): + return torch.Tensor(np.expand_dims(array.transpose([2, 0, 1]), axis=0).astype(np.float32)) / 255 + + +def rgb(t): + return ( + np.clip((t[0] if len(t.shape) == 4 else t).detach( + ).cpu().numpy().transpose([1, 2, 0]), 0, 1) * 255 + ).astype(np.uint8) + + +def imread(path): + return cv2.imread(path)[:, :, [2, 1, 0]] + + +def imwrite(path, img): + os.makedirs(os.path.dirname(path), exist_ok=True) + cv2.imwrite(path, img[:, :, [2, 1, 0]]) + + +def impad(img, top=0, bottom=0, left=0, right=0, color=255): + return np.pad(img, [(top, bottom), (left, right), (0, 0)], 'reflect') + + +def hiseq_color_cv2_img(img): + (b, g, r) = cv2.split(img) + bH = cv2.equalizeHist(b) + gH = cv2.equalizeHist(g) + rH = cv2.equalizeHist(r) + result = cv2.merge((bH, gH, rH)) + return result + + +def predict(model, opt, lr): + # Pad image to ensure dimensions are compatible with the model + h, w, c = lr.shape + lr = impad(lr, bottom=int(np.ceil(h / 2) * 2 - h), + right=int(np.ceil(w / 2) * 2 - w)) + + lr_t = t(lr) + lr_t = lr_t.cuda() # Move to GPU + + # If log_low is enabled, apply log transform + if opt["datasets"]["train"].get("log_low", False): + lr_t = torch.log(torch.clamp(lr_t + 1e-3, min=1e-3)) + + # If concat_histeq is enabled, concatenate histogram equalized image + if opt.get("concat_histeq", False): + # Apply histogram equalization to original image before padding + his = hiseq_color_cv2_img(lr[:h, :w]) # Use original dimensions + his = impad(his, bottom=int(np.ceil(h / 2) * 2 - h), + right=int(np.ceil(w / 2) * 2 - w)) + his_t = t(his).cuda() # Move to GPU + lr_t = torch.cat([lr_t, his_t], dim=1) + + # Use the same approach as test.py - get_sr method + sr_t = model.get_sr(lq=lr_t, heat=None) + return sr_t + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--input", required=True, + help="Path to the folder containing low-light images.") + parser.add_argument("--output", required=True, + help="Path to the folder to save enhanced images.") + parser.add_argument("--model_path", required=True, + help="Path to the trained model.") + parser.add_argument("--conf", default="confs/Custom_smallNet_custom.yml", + help="Path to the configuration file.") + args = parser.parse_args() + + model, opt = load_model(args.conf, args.model_path) + + input_dir = args.input + output_dir = args.output + + lr_paths = fiFindByWildcard(os.path.join(input_dir, '*.*')) + os.makedirs(output_dir, exist_ok=True) + + for lr_path in lr_paths: + lr = imread(lr_path) + h, w, c = lr.shape + prediction = predict(model, opt, lr) + prediction = rgb(prediction) + + # Crop to original size (remove padding) + prediction = prediction[:h, :w] + + output_path = os.path.join(output_dir, os.path.basename(lr_path)) + imwrite(output_path, prediction) + print(f"Processed: {os.path.basename(lr_path)}") + + print(f"Enhanced images saved to: {output_dir}") + + +if __name__ == "__main__": + main() diff --git a/code/utils/crop_random.py b/code/utils/crop_random.py index 596e848..e5f8b76 100644 --- a/code/utils/crop_random.py +++ b/code/utils/crop_random.py @@ -40,4 +40,4 @@ def generate_high_quality_crops(input_folder, output_folder, target_w=600, targe print(f"Error: {e}") if __name__ == "__main__": - generate_high_quality_crops("D:\\codes\\projects\\rm\\low-light-image-enhancement\\data\\rm_dataset", "D:\\codes\\projects\\rm\\low-light-image-enhancement\\data\\rm_dataset_crops") \ No newline at end of file + generate_high_quality_crops("D:\\codes\\projects\\rm\\LLFlow\\code\\using", "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\rm_dataset_crops") \ No newline at end of file diff --git a/code/utils/make_data.py b/code/utils/make_data.py index 038d2ed..86d8245 100644 --- a/code/utils/make_data.py +++ b/code/utils/make_data.py @@ -6,12 +6,11 @@ # Define paths # Replace with your source images directory -SOURCE_DIR = "D:/codes/projects/rm/low-light-image-enhancement/data/rm_dataset_crops/val" +SOURCE_DIR = "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops" # Replace with your ground truth images directory -GROUND_TRUTH_DIR = "D:/codes/projects/rm/low-light-image-enhancement/data/rm_dataset_crops/val/high" +GROUND_TRUTH_DIR = "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops\\high" # Replace with your desired low-light images directory -LOW_LIGHT_DIR = "D:/codes/projects/rm/low-light-image-enhancement/data/rm_dataset_crops/val/low" - +LOW_LIGHT_DIR = "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops\\low" # Create the low-light directory if it doesn't exist Path(LOW_LIGHT_DIR).mkdir(parents=True, exist_ok=True) Path(GROUND_TRUTH_DIR).mkdir(parents=True, exist_ok=True) From 56afa56d4292b4a2c38d0530275bf2631eddc7a7 Mon Sep 17 00:00:00 2001 From: Striker Eurika Date: Fri, 9 Jan 2026 21:32:41 +0700 Subject: [PATCH 5/5] Update dataset paths in configuration and utility files for improved data handling --- code/confs/Custom_smallNet_custom.yml | 2 +- code/test_unpaired.py | 32 ++++++++++++++++++++------- code/utils/crop_random.py | 2 +- code/utils/make_data.py | 6 ++--- 4 files changed, 29 insertions(+), 13 deletions(-) diff --git a/code/confs/Custom_smallNet_custom.yml b/code/confs/Custom_smallNet_custom.yml index 61c042e..c6dd59c 100644 --- a/code/confs/Custom_smallNet_custom.yml +++ b/code/confs/Custom_smallNet_custom.yml @@ -54,7 +54,7 @@ datasets: #### Test Settings # dataroot_GT: D:\LOLdataset\eval15\high # dataroot_LR: D:\LOLdataset\eval15\low -dataroot_unpaired: D:\codes\projects\rm\LLFlow\code\using\using_crops\predick +dataroot_unpaired: D:\codes\projects\rm\LLFlow\code\using\using_crops\only_low_not_love dataroot_GT: D:\codes\projects\rm\LLFlow\code\using\using_crops\high dataroot_LR: D:\codes\projects\rm\LLFlow\code\using\using_crops\low diff --git a/code/test_unpaired.py b/code/test_unpaired.py index 6bd97ef..b44007b 100755 --- a/code/test_unpaired.py +++ b/code/test_unpaired.py @@ -38,11 +38,12 @@ def predict(model, lr): return visuals.get('rlt', visuals.get('NORMAL')) -def t(array): return torch.Tensor(np.expand_dims(array.transpose([2, 0, 1]), axis=0).astype(np.float32)) / 255 +def t(array): return torch.Tensor(np.expand_dims( + array.transpose([2, 0, 1]), axis=0).astype(np.float32)) / 255 def rgb(t): return ( - np.clip((t[0] if len(t.shape) == 4 else t).detach().cpu().numpy().transpose([1, 2, 0]), 0, 1) * 255).astype( + np.clip((t[0] if len(t.shape) == 4 else t).detach().cpu().numpy().transpose([1, 2, 0]), 0, 1) * 255).astype( np.uint8) @@ -99,7 +100,7 @@ def main(): conf = conf_path.split('/')[-1].replace('.yml', '') model, opt = load_model(conf_path) model.netG = model.netG.cuda() - + lr_dir = opt['dataroot_unpaired'] lr_paths = fiFindByWildcard(os.path.join(lr_dir, '*.*')) @@ -117,15 +118,30 @@ def main(): lr = his lr_t = t(lr) + # Adjust lower bound in torch.clamp to avoid extreme values if opt["datasets"]["train"].get("log_low", False): - lr_t = torch.log(torch.clamp(lr_t + 1e-3, min=1e-3)) + # Changed min from 1e-3 to 1e-2 + lr_t = torch.log(torch.clamp(lr_t + 1e-3, min=1e-2)) + + # Debug: Check input tensor statistics after clamping + print( + f"DEBUG INPUT AFTER CLAMP - shape: {lr_t.shape}, min: {lr_t.min().item():.6f}, max: {lr_t.max().item():.6f}, mean: {lr_t.mean().item():.6f}, has_nan: {torch.isnan(lr_t).any().item()}") + + # Ensure the input tensor has the correct number of channels if opt.get("concat_histeq", False): - his = t(his) + his = t(hiseq_color_cv2_img(lr)) + # Concatenate histogram equalized image lr_t = torch.cat([lr_t, his], dim=1) - heat = opt['heat'] - with torch.cuda.amp.autocast(): - sr_t = model.get_sr(lq=lr_t.cuda(), heat=None) + # Debug: Check input tensor shape after concatenation + print(f"DEBUG INPUT AFTER CONCAT - shape: {lr_t.shape}") + + # Disable AMP to avoid potential instability + sr_t = model.get_sr(lq=lr_t.cuda(), heat=None) + + # Debug: Check output tensor statistics after model prediction + print( + f"DEBUG OUTPUT - shape: {sr_t.shape}, min: {sr_t.min().item():.6f}, max: {sr_t.max().item():.6f}, mean: {sr_t.mean().item():.6f}, has_nan: {torch.isnan(sr_t).any().item()}") sr = rgb(torch.clamp(sr_t, 0, 1)[:, :, padding_params[0]:sr_t.shape[2] - padding_params[1], padding_params[2]:sr_t.shape[3] - padding_params[3]]) assert raw_shape == sr.shape diff --git a/code/utils/crop_random.py b/code/utils/crop_random.py index e5f8b76..43a63fc 100644 --- a/code/utils/crop_random.py +++ b/code/utils/crop_random.py @@ -40,4 +40,4 @@ def generate_high_quality_crops(input_folder, output_folder, target_w=600, targe print(f"Error: {e}") if __name__ == "__main__": - generate_high_quality_crops("D:\\codes\\projects\\rm\\LLFlow\\code\\using", "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\rm_dataset_crops") \ No newline at end of file + generate_high_quality_crops("D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops\\only_low_not_love", "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops\\only_low_not_love_crops") \ No newline at end of file diff --git a/code/utils/make_data.py b/code/utils/make_data.py index 86d8245..c737155 100644 --- a/code/utils/make_data.py +++ b/code/utils/make_data.py @@ -6,11 +6,11 @@ # Define paths # Replace with your source images directory -SOURCE_DIR = "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops" +SOURCE_DIR = "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops\\only_low_not_love_crops" # Replace with your ground truth images directory -GROUND_TRUTH_DIR = "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops\\high" +GROUND_TRUTH_DIR = "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops\\only_low_not_love_crops_high" # Replace with your desired low-light images directory -LOW_LIGHT_DIR = "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops\\low" +LOW_LIGHT_DIR = "D:\\codes\\projects\\rm\\LLFlow\\code\\using\\using_crops\\only_low_not_love_crops_low" # Create the low-light directory if it doesn't exist Path(LOW_LIGHT_DIR).mkdir(parents=True, exist_ok=True) Path(GROUND_TRUTH_DIR).mkdir(parents=True, exist_ok=True)