feat: support APISR

TensoRaws · Oct 7, 2024 · 0b1b234 · 0b1b234
1 parent 2b154f2
commit 0b1b234
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 48 deletions.
diff --git a/ccrestoration/config/edsr_config.py b/ccrestoration/config/edsr_config.py
@@ -38,30 +38,30 @@ class EDSRConfig(BaseConfig):
         scale=4,
     ),
     # Official Large size models
-    EDSRConfig(
-        name=ConfigType.EDSR_Lx2_f256b32_DIV2K_official_2x,
-        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/EDSR_Lx2_f256b32_DIV2K_official_2x.pth",
-        hash="be38e77dcff9ec95225cea6326b5a616d57869824688674da317df37f3d87d1b",
-        scale=2,
-        num_feat=256,
-        num_block=32,
-    ),
-    EDSRConfig(
-        name=ConfigType.EDSR_Lx3_f256b32_DIV2K_official_3x,
-        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/EDSR_Lx3_f256b32_DIV2K_official_3x.pth",
-        hash="3660f70d306481ef4867731500c4e3d901a1b8547996cf4245a09ffbc151b70b",
-        scale=3,
-        num_feat=256,
-        num_block=32,
-    ),
-    EDSRConfig(
-        name=ConfigType.EDSR_Lx4_f256b32_DIV2K_official_4x,
-        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/EDSR_Lx4_f256b32_DIV2K_official_4x.pth",
-        hash="76ee1c8f48813f46024bee8d2700f417f6b2db070e899954ff1552fbae343e93",
-        scale=4,
-        num_feat=256,
-        num_block=32,
-    ),
+    # EDSRConfig(
+    #     name=ConfigType.EDSR_Lx2_f256b32_DIV2K_official_2x,
+    #     url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/EDSR_Lx2_f256b32_DIV2K_official_2x.pth",
+    #     hash="be38e77dcff9ec95225cea6326b5a616d57869824688674da317df37f3d87d1b",
+    #     scale=2,
+    #     num_feat=256,
+    #     num_block=32,
+    # ),
+    # EDSRConfig(
+    #     name=ConfigType.EDSR_Lx3_f256b32_DIV2K_official_3x,
+    #     url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/EDSR_Lx3_f256b32_DIV2K_official_3x.pth",
+    #     hash="3660f70d306481ef4867731500c4e3d901a1b8547996cf4245a09ffbc151b70b",
+    #     scale=3,
+    #     num_feat=256,
+    #     num_block=32,
+    # ),
+    # EDSRConfig(
+    #     name=ConfigType.EDSR_Lx4_f256b32_DIV2K_official_4x,
+    #     url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/EDSR_Lx4_f256b32_DIV2K_official_4x.pth",
+    #     hash="76ee1c8f48813f46024bee8d2700f417f6b2db070e899954ff1552fbae343e93",
+    #     scale=4,
+    #     num_feat=256,
+    #     num_block=32,
+    # ),
 ]
 
 for cfg in EDSRConfigs:

diff --git a/ccrestoration/type/config.py b/ccrestoration/type/config.py
@@ -38,9 +38,6 @@ class ConfigType(str, Enum):
     EDSR_Mx2_f64b16_DIV2K_official_2x = "EDSR_Mx2_f64b16_DIV2K_official_2x.pth"
     EDSR_Mx3_f64b16_DIV2K_official_3x = "EDSR_Mx3_f64b16_DIV2K_official_3x.pth"
     EDSR_Mx4_f64b16_DIV2K_official_4x = "EDSR_Mx4_f64b16_DIV2K_official_4x.pth"
-    EDSR_Lx2_f256b32_DIV2K_official_2x = "EDSR_Lx2_f256b32_DIV2K_official_2x.pth"
-    EDSR_Lx3_f256b32_DIV2K_official_3x = "EDSR_Lx3_f256b32_DIV2K_official_3x.pth"
-    EDSR_Lx4_f256b32_DIV2K_official_4x = "EDSR_Lx4_f256b32_DIV2K_official_4x.pth"
 
     # SwinIR
     SwinIR_classicalSR_DF2K_s64w8_SwinIR_M_2x = "SwinIR_classicalSR_DF2K_s64w8_SwinIR_M_2x.pth"

diff --git a/tests/test_edsr.py b/tests/test_edsr.py
@@ -1,5 +1,4 @@
 import cv2
-import pytest
 
 from ccrestoration import AutoConfig, AutoModel, BaseConfig, ConfigType
 from ccrestoration.model import SRBaseModel
@@ -26,23 +25,3 @@ def test_official_M(self) -> None:
 
             assert calculate_image_similarity(img1, img2)
             assert compare_image_size(img1, img2, cfg.scale)
-
-    @pytest.mark.skip("Skip because it's too large")
-    def test_official_L(self) -> None:
-        img1 = load_image()
-
-        for k in [
-            ConfigType.EDSR_Lx2_f256b32_DIV2K_official_2x,
-            ConfigType.EDSR_Lx3_f256b32_DIV2K_official_3x,
-            ConfigType.EDSR_Lx4_f256b32_DIV2K_official_4x,
-        ]:
-            print(f"Testing {k}")
-            cfg: BaseConfig = AutoConfig.from_pretrained(k)
-            model: SRBaseModel = AutoModel.from_config(config=cfg, fp16=False, device=get_device())
-            print(model.device)
-
-            img2 = model.inference_image(img1)
-            cv2.imwrite(str(ASSETS_PATH / f"test_{k}_out.jpg"), img2)
-
-            assert calculate_image_similarity(img1, img2)
-            assert compare_image_size(img1, img2, cfg.scale)