feat: huggingface_hub

app.py

@@ -9,6 +9,7 @@ import os
 import imquality.brisque as brisque
 from loss.niqe_utils import *
 import spaces
+import huggingface_hub
 
 eval_net = CIDNet().cuda()
 eval_net.trans.gated = True
@@ -19,7 +20,8 @@ def process_image(input_img,score,model_path,gamma=1.0,alpha_s=1.0,alpha_i=1.0):
     if model_path is None:
         return input_img,"Please choose a model weights."
     torch.set_grad_enabled(False)
-    eval_net.load_state_dict(torch.load(os.path.join(directory,model_path), map_location=lambda storage, loc: storage))
+    eval_net.from_pretrained(model_path)
+    # eval_net.load_state_dict(torch.load(os.path.join(directory,model_path), map_location=lambda storage, loc: storage))
     eval_net.eval()
     
     pil2tensor = transforms.Compose([transforms.ToTensor()])
@@ -44,23 +46,21 @@ def process_image(input_img,score,model_path,gamma=1.0,alpha_s=1.0,alpha_i=1.0):
     else:
         return enhanced_img,0
 
-def find_pth_files(directory):
-    pth_files = []
-    for root, dirs, files in os.walk(directory):
-        if 'train' in root.split(os.sep):
-            continue
-        for file in files:
-            if file.endswith('.pth'):
-                pth_files.append(os.path.join(root, file))
-    return pth_files
-
-def remove_weights_prefix(paths):
-    cleaned_paths = [path.replace('weights/', '') for path in paths]
-    return cleaned_paths
 
 directory = "weights"
-pth_files = find_pth_files(directory)
-pth_files2 = remove_weights_prefix(pth_files)
+pth_files = [
+    'HVI-CIDNet-Generalization',
+    'HVI-CIDNet-Sony-Total-Dark',
+    'HVI-CIDNet-LOL-Blur',
+    'HVI-CIDNet-SICE',
+    'HVI-CIDNet-LOLv2-real-bestSSIM',
+    'HVI-CIDNet-LOLv2-real-bestPSNR',
+    'HVI-CIDNet-LOLv2-syn-wperc',
+    'HVI-CIDNet-LOLv2-syn-woperc',
+    'HVI-CIDNet-LOLv1-wperc',
+    'HVI-CIDNet-LOLv1-woperc'
+]
+
 
 interface = gr.Interface(
     fn=process_image,

net/CIDNet.py

@@ -5,8 +5,9 @@ from einops import rearrange
 from net.HVI_transform import RGB_HVI
 from net.transformer_utils import *
 from net.LCA import *
+from huggingface_hub import PyTorchModelHubMixin
 
-class CIDNet(nn.Module):
+class CIDNet(nn.Module, PyTorchModelHubMixin):
     def __init__(self, 
                  channels=[36, 36, 72, 144],
                  heads=[1, 2, 4, 8],

requirements.txt

@@ -8,4 +8,5 @@ torchvision
 einops
 opencv-python
 gradio
-spaces
+spaces
+huggingface_hub

weights/LOL-Blur.pth

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

weights/LOLv1/BestLPIPS_0.0868.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:533bec252359efc8f05720ce5a73c3a8c891384e62ba62b47b457b4a4c87247e
-size 7971706

weights/LOLv1/BestSSIM_0.8631.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:267ed28a7f65af87dd7bd6f486d88cfa25e654234badb974df7bc85dbf267ee3
-size 7971706

weights/LOLv1/PSNR_24.74.pth

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

weights/LOLv1/w_perc.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:9206385b514014b60bf9396151a011cabf68abb380c2ad87ffde3d53e8227926
-size 7968002

weights/LOLv1/wo_perc.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2c5afb426fee910dd9806b4e6214f05914fb5988b7982a4624560815aefce2b5
-size 7970627

weights/LOLv2_real/best_PSNR.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:44a7f92cf02a1da322c1a34b3aa940341dc93fd53840b7f19b49f461200b00a6
-size 7971269

weights/LOLv2_real/best_SSIM.pth

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

weights/LOLv2_real/w_perc.pth

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

weights/LOLv2_syn/w_perc.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:460a9956018935cab61ddef3cea14f1dbbc02a65e12ee4d5b5ab45d704d2575b
-size 7967682

weights/LOLv2_syn/wo_perc.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:519cc5fc55491cf36b47b8295c0204dde29ba887feeed06a77e716072469e55e
-size 7974353

weights/SICE.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7b0824026124427fe0419465fc06a6be0129356c949794d1d1f52c4173091607
-size 7964800

weights/SID.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:49678f69495610a2691bb5b2665d2766b57a6a79201a236f7afa84f67c5fec5f
-size 7964543

weights/generalization.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:6395370762ec1fa0f6fe92a38c62382b1ebe21edc012f597023bfcec8b95cd27
-size 7971462