Upload folder using huggingface_hub

RetinexFormer_FiveK.yml

@@ -0,0 +1,129 @@
+# general settings
+name: Enhancement_RetinexFormer_FiveK
+model_type: ImageCleanModel
+scale: 1
+num_gpu: 1  # set num_gpu: 0 for cpu mode
+manual_seed: 100
+
+# dataset and data loader settings
+datasets:
+  train:
+    name: TrainSet
+    type: Dataset_PairedImage
+    dataroot_gt: data/FiveK/train/target
+    dataroot_lq: data/FiveK/train/target
+    geometric_augs: true
+    cache_data: true
+    filename_tmpl: '{}'
+    io_backend:
+      type: disk
+
+    # data loader
+    use_shuffle: true
+    num_worker_per_gpu: 8
+    batch_size_per_gpu: 8
+
+    ### -------------Progressive training--------------------------
+    # mini_batch_sizes: [8,5,4,2,1,1]             # Batch size per gpu   
+    # iters: [92000,64000,48000,36000,36000,24000]
+    # gt_size: 384   # Max patch size for progressive training
+    # gt_sizes: [128,160,192,256,320,384]  # Patch sizes for progressive training.
+    ### ------------------------------------------------------------
+
+    ### ------- Training on single fixed-patch size 128x128---------
+    mini_batch_sizes: [8]   
+    iters: [300000]
+    gt_size: 280   
+    gt_sizes: [280]
+    ### ------------------------------------------------------------
+
+    dataset_enlarge_ratio: 1
+    prefetch_mode: ~
+
+  val:
+    name: ValSet
+    type: Dataset_PairedImage
+    # dataroot_gt: data/FiveK/test/target
+    # dataroot_lq: data/FiveK/test/input
+    dataroot_gt: data/Faded/Low
+    dataroot_lq: data/Faded/Normal
+    io_backend:
+      type: disk
+
+# network structures
+network_g:
+  type: RetinexFormer
+  in_channels: 3
+  out_channels: 3
+  n_feat: 40
+  stage: 1
+  num_blocks: [1,2,2]
+
+
+
+# path
+path:
+  pretrain_network_g: ~
+  strict_load_g: true
+  resume_state: ~
+
+# training settings
+train:
+  total_iter: 300000
+  warmup_iter: -1 # no warm up
+  use_grad_clip: true
+
+  # Split 300k iterations into two cycles. 
+  # 1st cycle: fixed 3e-4 LR for 92k iters. 
+  # 2nd cycle: cosine annealing (3e-4 to 1e-6) for 208k iters.
+  scheduler:
+    type: CosineAnnealingRestartCyclicLR
+    periods: [92000, 208000]       
+    restart_weights: [1,1]
+    eta_mins: [0.0003,0.000001]   
+  
+  mixing_augs:
+    mixup: true
+    mixup_beta: 1.2
+    use_identity: true
+
+  optim_g:
+    type: Adam
+    lr: !!float 2e-4
+    # weight_decay: !!float 1e-4
+    betas: [0.9, 0.999]
+  
+  # losses
+  pixel_opt:
+    type: L1Loss
+    loss_weight: 1
+    reduction: mean
+
+# validation settings
+val:
+  window_size: 4
+  val_freq: !!float 4e3
+  save_img: false
+  rgb2bgr: true
+  use_image: false
+  max_minibatch: 8
+
+  metrics:
+    psnr: # metric name, can be arbitrary
+      type: calculate_psnr
+      crop_border: 0
+      test_y_channel: false
+
+# logging settings
+logger:
+  print_freq: 1000
+  save_checkpoint_freq: !!float 4e3
+  use_tb_logger: true
+  wandb:
+    project: low_light
+    resume_id: ~
+
+# dist training settings
+dist_params:
+  backend: nccl
+  port: 29500

inference.py

@@ -0,0 +1,118 @@
+import numpy as np
+import os
+import argparse
+import cv2
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from basicsr.models import create_model
+from basicsr.utils.options import parse
+from skimage import img_as_ubyte
+from torch.cuda.amp import autocast, GradScaler
+def load_img(filepath):
+    return cv2.cvtColor(cv2.imread(filepath), cv2.COLOR_BGR2RGB)
+def save_img(filepath, img):
+    cv2.imwrite(filepath, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+def self_ensemble(x, model):
+    def forward_transformed(x, hflip, vflip, rotate, model):
+        if hflip:
+            x = torch.flip(x, (-2,))
+        if vflip:
+            x = torch.flip(x, (-1,))
+        if rotate:
+            x = torch.rot90(x, dims=(-2, -1))
+        x = model(x)
+        if rotate:
+            x = torch.rot90(x, dims=(-2, -1), k=3)
+        if vflip:
+            x = torch.flip(x, (-1,))
+        if hflip:
+            x = torch.flip(x, (-2,))
+        return x
+    t = []
+    for hflip in [False, True]:
+        for vflip in [False, True]:
+            for rot in [False, True]:
+                t.append(forward_transformed(x, hflip, vflip, rot, model))
+    t = torch.stack(t)
+    return torch.mean(t, dim=0)
+
+
+# Set GPU
+gpu_list = ','.join(str(x) for x in '0')
+os.environ['CUDA_VISIBLE_DEVICES'] = gpu_list
+print('export CUDA_VISIBLE_DEVICES=' + gpu_list)
+
+# Load YAML configuration
+opt = parse("RetinexFormer_FiveK.yml", is_train=False)
+opt['dist'] = False
+print(opt)
+
+# Load model
+model_restoration = create_model(opt).net_g
+checkpoint = torch.load("pretrained_weights/FiveK.pth")
+model_restoration.load_state_dict(checkpoint['params'])
+print("===>Testing using weights: ", "pretrained_weights/FiveK.pth")
+model_restoration.cuda()
+model_restoration = nn.DataParallel(model_restoration)
+model_restoration.eval()
+
+
+def process_image(inp_path, model_restoration, out_dir, factor=4):
+    torch.cuda.ipc_collect()
+    torch.cuda.empty_cache()
+
+    img = np.float32(load_img(inp_path)) / 255.
+
+
+    # Resize image to have height 1024px while maintaining aspect ratio
+    max_height = 1024
+    aspect_ratio = img.shape[1] / img.shape[0]
+    new_width = int(max_height * aspect_ratio)
+    img = cv2.resize(img, (new_width, max_height), interpolation=cv2.INTER_AREA)
+
+
+    img = torch.from_numpy(img).permute(2, 0, 1)
+    input_ = img.unsqueeze(0).cuda()
+
+    # Padding in case images are not multiples of 4
+    b, c, h, w = input_.shape
+    H, W = ((h + factor) // factor) * factor, ((w + factor) // factor) * factor
+    padh = H - h if h % factor != 0 else 0
+    padw = W - w if w % factor != 0 else 0
+    input_ = F.pad(input_, (0, padw, 0, padh), 'reflect')
+
+    scaler = GradScaler()  # for mixed precision
+
+    with autocast():  # enable mixed precision
+        if h < 3000 and w < 3000:
+            if 1 == 0:
+                restored = self_ensemble(input_, model_restoration)
+            else:
+                restored = model_restoration(input_)
+        else:
+            # split and test
+            input_1 = input_[:, :, :, 1::2]
+            input_2 = input_[:, :, :, 0::2]
+            if 1 == 0:
+                restored_1 = self_ensemble(input_1, model_restoration)
+                restored_2 = self_ensemble(input_2, model_restoration)
+            else:
+                restored_1 = model_restoration(input_1)
+                restored_2 = model_restoration(input_2)
+            restored = torch.zeros_like(input_)
+            restored[:, :, :, 1::2] = restored_1
+            restored[:, :, :, 0::2] = restored_2
+
+        # Unpad images to original dimensions
+        restored = restored[:, :, :h, :w]
+
+    restored = torch.clamp(restored, 0, 1).cpu().detach().permute(0, 2, 3, 1).squeeze(0).numpy()
+
+
+    if True:
+       save_img(os.path.join(out_dir, os.path.splitext(os.path.split(inp_path)[-1])[0] + '.png'), img_as_ubyte(restored))
+
+

pretrained_weights/pytorch_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:800f6a9281fe8d95daca3108f2b826d5a2adead09031e0e998d30a615286d9c1
+size 6478393

push.py

@@ -0,0 +1,19 @@
+from huggingface_hub import Repository, HfApi
+
+# # Initialize repository
+# repo = Repository(
+#     local_dir="/home/duyht/ShotXRetouch",
+#     clone_from="DuySota/retouch"
+# )
+
+# # Add, commit, and push files
+# repo.git_add(auto_lfs_track=True)
+# repo.git_commit("Initial commit")
+# repo.git_push()
+
+# Alternatively, you can use the HfApi for more granular control
+api = HfApi()
+api.upload_folder(
+    repo_id="DuySota/retouch",
+    folder_path="/home/duyht/ShotXRetouch"
+)