feat: support SRCNN

ccrestoration/arch/__init__.py

@@ -14,3 +14,4 @@
 from ccrestoration.arch.msrswvsr_arch import MSRSWVSR  # noqa
 from ccrestoration.arch.scunet_arch import SCUNet  # noqa
 from ccrestoration.arch.dat_arch import DAT  # noqa
+from ccrestoration.arch.srcnn_arch import SRCNN  # noqa

ccrestoration/arch/srcnn_arch.py

@@ -0,0 +1,41 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from ccrestoration.arch import ARCH_REGISTRY
+from ccrestoration.type import ArchType
+from ccrestoration.util.color import rgb_to_yuv, yuv_to_rgb
+
+
+@ARCH_REGISTRY.register(name=ArchType.SRCNN)
+class SRCNN(nn.Module):
+    def __init__(self, num_channels: int = 1, scale: int = 2) -> None:
+        super(SRCNN, self).__init__()
+        self.conv1 = nn.Conv2d(num_channels, 64, kernel_size=9, padding=9 // 2)
+        self.conv2 = nn.Conv2d(64, 32, kernel_size=5, padding=5 // 2)
+        self.conv3 = nn.Conv2d(32, num_channels, kernel_size=5, padding=5 // 2)
+        self.relu = nn.ReLU(inplace=True)
+        self.num_channels = num_channels
+        self.scale = scale
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.interpolate(x, scale_factor=self.scale, mode="bilinear")
+
+        if self.num_channels == 1:
+            # RGB -> YUV
+            x = rgb_to_yuv(x)
+            y, u, v = x[:, 0:1, ...], x[:, 1:2, ...], x[:, 2:3, ...]
+
+            y = self.relu(self.conv1(y))
+            y = self.relu(self.conv2(y))
+            y = self.conv3(y)
+
+            x = torch.cat([y, u, v], dim=1)
+            # YUV -> RGB
+            x = yuv_to_rgb(x)
+        else:
+            x = self.relu(self.conv1(x))
+            x = self.relu(self.conv2(x))
+            x = self.conv3(x)
+
+        return x

ccrestoration/config/__init__.py

@@ -13,3 +13,4 @@
 from ccrestoration.config.animesr_config import AnimeSRConfig  # noqa
 from ccrestoration.config.scunet_config import SCUNetConfig  # noqa
 from ccrestoration.config.dat_config import DATConfig  # noqa
+from ccrestoration.config.srcnn_config import SRCNNConfig  # noqa

ccrestoration/config/srcnn_config.py

@@ -0,0 +1,36 @@
+from typing import Union
+
+from ccrestoration.config import CONFIG_REGISTRY
+from ccrestoration.type import ArchType, BaseConfig, ConfigType, ModelType
+
+
+class SRCNNConfig(BaseConfig):
+    arch: ArchType = ArchType.SRCNN
+    model: Union[ModelType, str] = ModelType.SRCNN
+    scale: int = 2
+    num_channels: int = 1
+
+
+SRCNNConfigs = [
+    SRCNNConfig(
+        name=ConfigType.SRCNN_2x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/SRCNN_2x.pth",
+        hash="e803ec6e0230ae12b1fa7fd1c67bd57d2e744b4f4fbbc861bf9790070fc4d19e",
+        scale=2,
+    ),
+    SRCNNConfig(
+        name=ConfigType.SRCNN_3x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/SRCNN_3x.pth",
+        hash="364ec936313d0fd1052c641b20cefd8153a2c1d89712f357f804f0119ab7ab90",
+        scale=3,
+    ),
+    SRCNNConfig(
+        name=ConfigType.SRCNN_4x,
+        url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/SRCNN_4x.pth",
+        hash="f07978e521ede367d55ef7ca83f4f4979e2339c594bead101cfbb9611023f70e",
+        scale=4,
+    ),
+]
+
+for cfg in SRCNNConfigs:
+    CONFIG_REGISTRY.register(cfg)

ccrestoration/model/__init__.py

@@ -17,3 +17,4 @@
 from ccrestoration.model.animesr_model import AnimeSRModel  # noqa
 from ccrestoration.model.scunet_model import SCUNetModel  # noqa
 from ccrestoration.model.dat_model import DATModel  # noqa
+from ccrestoration.model.srcnn_model import SRCNNModel  # noqa

ccrestoration/model/sr_base_model.py

@@ -20,19 +20,19 @@ def get_state_dict(self) -> Any:
         cfg: BaseConfig = self.config
 
         if cfg.path is not None:
-            model_path = str(cfg.path)
+            state_dict_path = str(cfg.path)
         else:
             try:
-                model_path = load_file_from_url(
+                state_dict_path = load_file_from_url(
                     config=cfg, force_download=False, model_dir=self.model_dir, gh_proxy=self.gh_proxy
                 )
             except Exception as e:
                 print(f"Error: {e}, try force download the model...")
-                model_path = load_file_from_url(
+                state_dict_path = load_file_from_url(
                     config=cfg, force_download=True, model_dir=self.model_dir, gh_proxy=self.gh_proxy
                 )
 
-        return torch.load(model_path, map_location=self.device, weights_only=True)
+        return torch.load(state_dict_path, map_location=self.device, weights_only=True)
 
     @torch.inference_mode()  # type: ignore
     def inference(self, img: torch.Tensor) -> torch.Tensor:

ccrestoration/model/srcnn_model.py

@@ -0,0 +1,28 @@
+from typing import Any
+
+from ccrestoration.arch import SRCNN
+from ccrestoration.config import SRCNNConfig
+from ccrestoration.model import MODEL_REGISTRY
+from ccrestoration.model.sr_base_model import SRBaseModel
+from ccrestoration.type import ModelType
+
+
+@MODEL_REGISTRY.register(name=ModelType.SRCNN)
+class SRCNNModel(SRBaseModel):
+    def load_model(self) -> Any:
+        cfg: SRCNNConfig = self.config
+        state_dict = self.get_state_dict()
+
+        if "params_ema" in state_dict:
+            state_dict = state_dict["params_ema"]
+        elif "params" in state_dict:
+            state_dict = state_dict["params"]
+
+        model = SRCNN(
+            num_channels=cfg.num_channels,
+            scale=cfg.scale,
+        )
+
+        model.load_state_dict(state_dict)
+        model.eval().to(self.device)
+        return model

ccrestoration/type/arch.py

@@ -12,6 +12,7 @@ class ArchType(str, Enum):
     SWINIR = "SWINIR"
     SCUNET = "SCUNET"
     DAT = "DAT"
+    SRCNN = "SRCNN"
 
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 

ccrestoration/type/config.py

@@ -72,6 +72,11 @@ class ConfigType(str, Enum):
 
     DAT_APISR_GAN_generator_4x = "DAT_APISR_GAN_generator_4x.pth"
 
+    # SRCNN
+    SRCNN_2x = "SRCNN_2x.pth"
+    SRCNN_3x = "SRCNN_3x.pth"
+    SRCNN_4x = "SRCNN_4x.pth"
+
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 
     # SpyNet

ccrestoration/type/model.py

@@ -11,6 +11,7 @@ class ModelType(str, Enum):
     SwinIR = "SwinIR"
     SCUNet = "SCUNet"
     DAT = "DAT"
+    SRCNN = "SRCNN"
 
     # ------------------------------------- Auxiliary Network ----------------------------------------------------------
 

ccrestoration/util/color.py

@@ -0,0 +1,83 @@
+import torch
+from torch import Tensor
+
+
+def rgb_to_yuv(image: Tensor) -> Tensor:
+    r"""Convert an RGB image to YUV.
+
+    .. image:: _static/img/rgb_to_yuv.png
+
+    The image data is assumed to be in the range of :math:`(0, 1)`. The range of the output is of
+    :math:`(0, 1)` to luma and the ranges of U and V are :math:`(-0.436, 0.436)` and :math:`(-0.615, 0.615)`,
+    respectively.
+
+    The YUV model adopted here follows M/PAL values (see
+    `BT.470-5 <https://www.itu.int/dms_pubrec/itu-r/rec/bt/R-REC-BT.470-5-199802-S!!PDF-E.pdf>`_, Table 2,
+    items 2.5 and 2.6).
+
+    Args:
+        image: RGB Image to be converted to YUV with shape :math:`(*, 3, H, W)`.
+
+    Returns:
+        YUV version of the image with shape :math:`(*, 3, H, W)`.
+
+    Example:
+        >>> input = torch.rand(2, 3, 4, 5)
+        >>> output = rgb_to_yuv(input)  # 2x3x4x5
+    """
+    if not isinstance(image, Tensor):
+        raise TypeError(f"Input type is not a Tensor. Got {type(image)}")
+
+    if len(image.shape) < 3 or image.shape[-3] != 3:
+        raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
+
+    r: Tensor = image[..., 0, :, :]
+    g: Tensor = image[..., 1, :, :]
+    b: Tensor = image[..., 2, :, :]
+
+    y: Tensor = 0.299 * r + 0.587 * g + 0.114 * b
+    u: Tensor = -0.147 * r - 0.289 * g + 0.436 * b
+    v: Tensor = 0.615 * r - 0.515 * g - 0.100 * b
+
+    out: Tensor = torch.stack([y, u, v], -3)
+
+    return out
+
+
+def yuv_to_rgb(image: Tensor) -> Tensor:
+    r"""Convert an YUV image to RGB.
+
+    The image data is assumed to be in the range of :math:`(0, 1)` for luma (Y). The ranges of U and V are
+    :math:`(-0.436, 0.436)` and :math:`(-0.615, 0.615)`, respectively.
+
+    YUV formula follows M/PAL values (see
+    `BT.470-5 <https://www.itu.int/dms_pubrec/itu-r/rec/bt/R-REC-BT.470-5-199802-S!!PDF-E.pdf>`_, Table 2,
+    items 2.5 and 2.6).
+
+    Args:
+        image: YUV Image to be converted to RGB with shape :math:`(*, 3, H, W)`.
+
+    Returns:
+        RGB version of the image with shape :math:`(*, 3, H, W)`.
+
+    Example:
+        >>> input = torch.rand(2, 3, 4, 5)
+        >>> output = yuv_to_rgb(input)  # 2x3x4x5
+    """
+    if not isinstance(image, Tensor):
+        raise TypeError(f"Input type is not a Tensor. Got {type(image)}")
+
+    if image.dim() < 3 or image.shape[-3] != 3:
+        raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
+
+    y: Tensor = image[..., 0, :, :]
+    u: Tensor = image[..., 1, :, :]
+    v: Tensor = image[..., 2, :, :]
+
+    r: Tensor = y + 1.14 * v  # coefficient for g is 0
+    g: Tensor = y + -0.396 * u - 0.581 * v
+    b: Tensor = y + 2.029 * u  # coefficient for b is 0
+
+    out: Tensor = torch.stack([r, g, b], -3)
+
+    return out

tests/test_srcnn.py

@@ -0,0 +1,27 @@
+import cv2
+
+from ccrestoration import AutoConfig, AutoModel, BaseConfig, ConfigType
+from ccrestoration.model import SRBaseModel
+
+from .util import ASSETS_PATH, calculate_image_similarity, compare_image_size, get_device, load_image
+
+
+class Test_SRCNN:
+    def test_official(self) -> None:
+        img1 = load_image()
+
+        for k in [
+            ConfigType.SRCNN_2x,
+            ConfigType.SRCNN_3x,
+            ConfigType.SRCNN_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)