Add Gradio deraining demo

.gitattributes

@@ -0,0 +1 @@
+*.pth filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,30 @@
+---
+title: Deraining Demo
+emoji: 🚀
+colorFrom: blue
+colorTo: purple
+sdk: gradio
+sdk_version: 4.0.0
+app_file: app.py
+pinned: false
+---
+
+## Deraining 데모 (Gradio)
+
+### 사용 방법
+- 이미지를 업로드하면 deraining 결과를 반환합니다.
+
+### 가중치(Weights) 준비
+이 Space는 기본적으로 `checkpoints/model_best.pth`를 찾습니다.
+
+- **옵션 A (권장)**: Space repo에 `checkpoints/model_best.pth`를 **Git LFS**로 업로드
+- **옵션 B**: Hub에서 다운로드
+  - Space 설정(Environment variables)에 아래를 추가
+    - `WEIGHTS_REPO`: 예) `your-username/your-deraining-weights`
+    - `WEIGHTS_FILENAME`: 예) `model_best.pth`
+
+### 모델 선택
+기본 모델은 `CGFWMSRNet` 입니다. 바꾸려면 환경변수로 설정하세요:
+- `MODEL_NAME`: 예) `CGFWMSRNet`
+
+

app.py

@@ -0,0 +1,144 @@
+import os
+import gradio as gr
+import numpy as np
+import torch
+import torch.nn.functional as F
+from skimage import img_as_ubyte
+
+from models import find_models_def
+
+try:
+    # Optional: download weights from the Hub when not bundled in the Space repo.
+    from huggingface_hub import hf_hub_download  # type: ignore
+except Exception:  # pragma: no cover
+    hf_hub_download = None  # type: ignore
+
+# -------------------------
+# Global (load once)
+# -------------------------
+MODEL = None
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+def _resolve_weights_path() -> str:
+    """
+    Priority:
+      1) Local bundled weights: ./checkpoints/model_best.pth
+      2) Download from Hub if WEIGHTS_REPO + WEIGHTS_FILENAME are provided
+    """
+    local_path = os.path.join(BASE_DIR, "checkpoints", "model_best.pth")
+    if os.path.exists(local_path):
+        return local_path
+
+    repo_id = os.getenv("WEIGHTS_REPO", "").strip()
+    filename = os.getenv("WEIGHTS_FILENAME", "model_best.pth").strip()
+    if repo_id and hf_hub_download is not None:
+        return hf_hub_download(repo_id=repo_id, filename=filename)
+
+    raise FileNotFoundError(
+        "Weights not found. Put weights at 'checkpoints/model_best.pth' "
+        "or set env WEIGHTS_REPO and WEIGHTS_FILENAME to download from the Hub."
+    )
+
+
+def _load_state_dict(model: torch.nn.Module, weights_path: str) -> None:
+    ckpt = torch.load(weights_path, map_location="cpu")
+    state_dict = ckpt.get("state_dict", ckpt) if isinstance(ckpt, dict) else ckpt
+
+    # Handle DataParallel 'module.' prefix
+    if isinstance(state_dict, dict) and any(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    try:
+        model.load_state_dict(state_dict, strict=True)
+    except Exception as e:
+        raise RuntimeError(
+            f"Failed to load weights from '{weights_path}'. "
+            f"Check that MODEL_NAME matches the checkpoint architecture."
+        ) from e
+
+
+def load_model():
+    global MODEL
+
+    if MODEL is not None:
+        return
+
+    # Model selection (default: CGFWMSRNet)
+    model_name = os.getenv("MODEL_NAME", "CGFWMSRNet").strip() or "CGFWMSRNet"
+    MODEL = find_models_def(model_name)()
+
+    # Load checkpoint (local first, Hub fallback)
+    weights_path = _resolve_weights_path()
+    _load_state_dict(MODEL, weights_path)
+
+    # Move to device
+    MODEL.to(DEVICE)
+    MODEL.eval()
+
+def infer_image(inp_img: np.ndarray):
+    """
+    inp_img: HxWx3 uint8 (Gradio Image numpy)
+    return: restored HxWx3 uint8
+    """
+    load_model()
+
+    if inp_img is None:
+        raise gr.Error("이미지를 업로드해주세요.")
+
+    # 입력이 RGBA로 들어올 수도 있어 RGB로 맞춤
+    if inp_img.ndim == 3 and inp_img.shape[2] == 4:
+        inp_img = inp_img[:, :, :3]
+
+    # uint8 -> float32 [0,1]
+    img = inp_img.astype(np.float32) / 255.0
+
+    # HWC -> CHW, add batch
+    img_t = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)  # (1,3,H,W)
+    img_t = img_t.to(DEVICE)
+
+    height, width = img_t.shape[2], img_t.shape[3]
+    img_multiple_of = 16
+    H = ((height + img_multiple_of) // img_multiple_of) * img_multiple_of
+    W = ((width  + img_multiple_of) // img_multiple_of) * img_multiple_of
+    padh = H - height if height % img_multiple_of != 0 else 0
+    padw = W - width  if width  % img_multiple_of != 0 else 0
+
+    img_t = F.pad(img_t, (0, padw, 0, padh), mode="reflect")
+
+    with torch.no_grad():
+        # 기존 코드: restored = model_restoration(input_)
+        restored = MODEL(img_t)
+
+        # 기존 코드: restored[0] 클램프
+        # 모델 출력 형태가 (list/tuple) 인지, tensor인지에 따라 안전하게 처리
+        if isinstance(restored, (list, tuple)):
+            restored = restored[0]
+
+        restored = torch.clamp(restored, 0, 1)
+        restored = restored[:, :, :height, :width]
+
+    # BCHW -> HWC uint8
+    restored_np = restored.squeeze(0).permute(1, 2, 0).detach().cpu().numpy()
+    restored_img = img_as_ubyte(restored_np)  # uint8
+
+    # GPU 메모리 정리(필수는 아니지만 데모에서 유용)
+    if DEVICE.startswith("cuda"):
+        torch.cuda.empty_cache()
+
+    return restored_img
+
+demo = gr.Interface(
+    fn=infer_image,
+    inputs=gr.Image(type="numpy", label="Input (Rainy)"),
+    outputs=gr.Image(type="numpy", label="Restored"),
+    title="Image Deraining Demo (Gradio + Hugging Face Spaces)",
+    description=(
+        "이미지를 업로드하면 비 제거 결과를 반환합니다. "
+        "Space에 가중치(`checkpoints/model_best.pth`)가 포함되어 있거나, "
+        "환경변수 WEIGHTS_REPO/WEIGHTS_FILENAME로 Hub에서 내려받을 수 있어야 합니다."
+    ),
+)
+
+if __name__ == "__main__":
+    demo.queue().launch()

checkpoints/model_best.pth

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

models/CGFWMSRNet.py

@@ -0,0 +1,94 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from models.modules import *
+
+
+class Model(nn.Module):
+    def __init__(
+        self,
+        in_c: int = 3,
+        out_c: int = 3,
+        n_feat: int = 40,
+        scale_unetfeats: int = 20,
+        num_cab: int = 8,
+        kernel_size: int = 3,
+        reduction: int = 4,
+        bias: bool = False,
+    ):
+        super(Model, self).__init__()
+
+        act = nn.PReLU()
+        self.shallow_feat1 = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias), CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
+        self.shallow_feat2 = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias), CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
+        self.shallow_feat3 = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias), CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
+
+        # Cross Stage Feature Fusion (CSFF)
+        self.stage1_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=False)
+        self.stage1_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
+
+        self.stage2_encoder_1 = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=True)
+        self.stage2_decoder_1 = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
+        self.stage2_encoder_2 = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=True)
+        self.stage2_decoder_2 = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
+
+        self.stage3_encoder_1 = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=True)
+        self.stage3_decoder_1 = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
+        self.stage3_encoder_2 = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=True)
+        self.stage3_decoder_2 = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
+        self.stage3_encoder_3 = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=True)
+        self.stage3_decoder_3 = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
+
+        self.CGAM12 = CGAM(n_feat, kernel_size=1, bias=bias)
+        self.CGAM23 = CGAM(n_feat, kernel_size=1, bias=bias)
+
+        self.concat12 = conv(n_feat * 2, n_feat, kernel_size, bias=bias)
+        self.concat23 = conv(n_feat * 2, n_feat, kernel_size, bias=bias)
+        self.tail = conv(n_feat, out_c, kernel_size, bias=bias)
+
+    def forward(self, x3_img: torch.Tensor):
+        # Stage 1: 1/4 resolution
+        x1_img = F.interpolate(x3_img, scale_factor=0.25, mode="bilinear", align_corners=False)
+        x1 = self.shallow_feat1(x1_img)
+        feat1 = self.stage1_encoder(x1)
+        res1 = self.stage1_decoder(feat1)
+
+        x2_CGAMfeats, stage1_img = self.CGAM12(res1[0], x1_img)
+
+        feat1 = [F.interpolate(f, scale_factor=2, mode="bilinear", align_corners=False) for f in feat1]
+        res1 = [F.interpolate(f, scale_factor=2, mode="bilinear", align_corners=False) for f in res1]
+        x2_CGAMfeats = F.interpolate(x2_CGAMfeats, scale_factor=2, mode="bilinear", align_corners=False)
+
+        # Stage 2: 1/2 resolution
+        x2_img = F.interpolate(x3_img, scale_factor=0.5, mode="bilinear", align_corners=False)
+        x2 = self.shallow_feat2(x2_img)
+        x2_cat = self.concat12(torch.cat([x2, x2_CGAMfeats], 1))
+
+        feat2_1 = self.stage2_encoder_1(x2_cat, feat1, res1)
+        res2 = self.stage2_decoder_1(feat2_1)
+        feat2_2 = self.stage2_encoder_2(res2[0], feat1, res1)
+        res2 = self.stage2_decoder_2(feat2_2)
+
+        x3_CGAMfeats, stage2_img = self.CGAM23(res2[0], x2_img)
+
+        feat2 = [F.interpolate(f, scale_factor=2, mode="bilinear", align_corners=False) for f in feat2_2]
+        res2 = [F.interpolate(f, scale_factor=2, mode="bilinear", align_corners=False) for f in res2]
+        x3_CGAMfeats = F.interpolate(x3_CGAMfeats, scale_factor=2, mode="bilinear", align_corners=False)
+
+        # Stage 3: full resolution
+        x3 = self.shallow_feat3(x3_img)
+        x3_cat = self.concat23(torch.cat([x3, x3_CGAMfeats], 1))
+
+        feat3_1 = self.stage3_encoder_1(x3_cat, feat2, res2)
+        res3 = self.stage3_decoder_1(feat3_1)
+        feat3_2 = self.stage3_encoder_2(res3[0], feat2, res2)
+        res3 = self.stage3_decoder_2(feat3_2)
+        feat3_3 = self.stage3_encoder_3(res3[0], feat2, res2)
+        res3 = self.stage3_decoder_3(feat3_3)
+
+        stage3_img = self.tail(res3[0])
+
+        return [stage3_img + x3_img, stage2_img, stage1_img]
+
+

models/__init__.py

@@ -0,0 +1,9 @@
+import importlib
+
+
+def find_models_def(model_name: str):
+    module_name = f"models.{model_name}"
+    module = importlib.import_module(module_name)
+    return getattr(module, "Model")
+
+

models/modules.py

@@ -0,0 +1,177 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def conv(in_channels, out_channels, kernel_size, bias=False, stride=1):
+    return nn.Conv2d(
+        in_channels,
+        out_channels,
+        kernel_size,
+        padding=(kernel_size // 2),
+        bias=bias,
+        stride=stride,
+    )
+
+
+class CALayer(nn.Module):
+    def __init__(self, channel, reduction=16, bias=False):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv_du = nn.Sequential(
+            nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=bias),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=bias),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        y = self.avg_pool(x)
+        y = self.conv_du(y)
+        return x * y
+
+
+class CAB(nn.Module):
+    def __init__(self, n_feat, kernel_size, reduction, bias, act):
+        super().__init__()
+        self.body = nn.Sequential(
+            conv(n_feat, n_feat, kernel_size, bias=bias),
+            act,
+            conv(n_feat, n_feat, kernel_size, bias=bias),
+        )
+        self.CA = CALayer(n_feat, reduction, bias=bias)
+
+    def forward(self, x):
+        res = self.body(x)
+        res = self.CA(res)
+        res += x
+        return res
+
+
+class CG(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.sobel_x = torch.tensor([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]]).unsqueeze(0).unsqueeze(0)
+        self.sobel_y = torch.tensor([[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]]).unsqueeze(0).unsqueeze(0)
+
+    def forward(self, image):
+        device = image.device
+        sobel_x = self.sobel_x.to(device)
+        sobel_y = self.sobel_y.to(device)
+
+        gradients = []
+        for c in range(image.shape[1]):
+            grad_x = F.conv2d(image[:, c : c + 1, :, :], sobel_x, padding=1)
+            grad_y = F.conv2d(image[:, c : c + 1, :, :], sobel_y, padding=1)
+            gradient = torch.sqrt(torch.clamp(grad_x**2 + grad_y**2, min=1e-6))
+            gradients.append(gradient)
+        return torch.cat(gradients, dim=1)
+
+
+class CGAM(nn.Module):
+    def __init__(self, n_feat, kernel_size, bias):
+        super().__init__()
+        self.conv1 = conv(n_feat, n_feat, kernel_size, bias=bias)
+        self.conv2 = conv(n_feat, 3, kernel_size, bias=bias)
+        self.conv3 = conv(3, n_feat, kernel_size, bias=bias)
+        self.gradfilter = CG()
+
+    def forward(self, x, x_img):
+        x1 = self.conv1(x)
+        img = self.conv2(x) + x_img
+        rain_grad = self.gradfilter(x_img)
+        clean_grad = self.gradfilter(img)
+        x2 = self.conv3(torch.abs(rain_grad - clean_grad))
+        grad_att = torch.sigmoid(x1 * x2)
+        x1 = x1 + x1 * grad_att
+        return x1, img
+
+
+class DownSample(nn.Module):
+    def __init__(self, in_channels, s_factor):
+        super().__init__()
+        self.down = nn.Sequential(
+            nn.Upsample(scale_factor=0.5, mode="bilinear", align_corners=False),
+            nn.Conv2d(in_channels, in_channels + s_factor, 1, stride=1, padding=0, bias=False),
+        )
+
+    def forward(self, x):
+        return self.down(x)
+
+
+class SkipUpSample(nn.Module):
+    def __init__(self, in_channels, s_factor):
+        super().__init__()
+        self.up = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
+            nn.Conv2d(in_channels + s_factor, in_channels, 1, stride=1, padding=0, bias=False),
+        )
+
+    def forward(self, x, y):
+        x = self.up(x)
+        return x + y
+
+
+class Encoder(nn.Module):
+    def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff):
+        super().__init__()
+        self.encoder_level1 = nn.Sequential(*[CAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(2)])
+        self.encoder_level2 = nn.Sequential(*[CAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act) for _ in range(2)])
+        self.encoder_level3 = nn.Sequential(*[CAB(n_feat + (scale_unetfeats * 2), kernel_size, reduction, bias=bias, act=act) for _ in range(2)])
+
+        self.down12 = DownSample(n_feat, scale_unetfeats)
+        self.down23 = DownSample(n_feat + scale_unetfeats, scale_unetfeats)
+
+        if csff:
+            self.csff_enc1 = nn.Conv2d(n_feat, n_feat, kernel_size=1, bias=bias)
+            self.csff_enc2 = nn.Conv2d(n_feat + scale_unetfeats, n_feat + scale_unetfeats, kernel_size=1, bias=bias)
+            self.csff_enc3 = nn.Conv2d(n_feat + (scale_unetfeats * 2), n_feat + (scale_unetfeats * 2), kernel_size=1, bias=bias)
+
+            self.csff_dec1 = nn.Conv2d(n_feat, n_feat, kernel_size=1, bias=bias)
+            self.csff_dec2 = nn.Conv2d(n_feat + scale_unetfeats, n_feat + scale_unetfeats, kernel_size=1, bias=bias)
+            self.csff_dec3 = nn.Conv2d(n_feat + (scale_unetfeats * 2), n_feat + (scale_unetfeats * 2), kernel_size=1, bias=bias)
+
+    def forward(self, x, encoder_outs=None, decoder_outs=None):
+        enc1 = self.encoder_level1(x)
+        if (encoder_outs is not None) and (decoder_outs is not None):
+            enc1 = enc1 + self.csff_enc1(encoder_outs[0]) + self.csff_dec1(decoder_outs[0])
+
+        x = self.down12(enc1)
+        enc2 = self.encoder_level2(x)
+        if (encoder_outs is not None) and (decoder_outs is not None):
+            enc2 = enc2 + self.csff_enc2(encoder_outs[1]) + self.csff_dec2(decoder_outs[1])
+
+        x = self.down23(enc2)
+        enc3 = self.encoder_level3(x)
+        if (encoder_outs is not None) and (decoder_outs is not None):
+            enc3 = enc3 + self.csff_enc3(encoder_outs[2]) + self.csff_dec3(decoder_outs[2])
+
+        return [enc1, enc2, enc3]
+
+
+class Decoder(nn.Module):
+    def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats):
+        super().__init__()
+        self.decoder_level1 = nn.Sequential(*[CAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(2)])
+        self.decoder_level2 = nn.Sequential(*[CAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act) for _ in range(2)])
+        self.decoder_level3 = nn.Sequential(*[CAB(n_feat + (scale_unetfeats * 2), kernel_size, reduction, bias=bias, act=act) for _ in range(2)])
+
+        self.skip_attn1 = CAB(n_feat, kernel_size, reduction, bias=bias, act=act)
+        self.skip_attn2 = CAB(n_feat + scale_unetfeats, kernel_size, reduction, bias=bias, act=act)
+
+        self.up21 = SkipUpSample(n_feat, scale_unetfeats)
+        self.up32 = SkipUpSample(n_feat + scale_unetfeats, scale_unetfeats)
+
+    def forward(self, outs):
+        enc1, enc2, enc3 = outs
+        dec3 = self.decoder_level3(enc3)
+
+        x = self.up32(dec3, self.skip_attn2(enc2))
+        dec2 = self.decoder_level2(x)
+
+        x = self.up21(dec2, self.skip_attn1(enc1))
+        dec1 = self.decoder_level1(x)
+
+        return [dec1, dec2, dec3]
+
+

requirements.txt

@@ -0,0 +1,6 @@
+gradio>=4.0,<6.0
+numpy
+torch
+scikit-image
+huggingface_hub
+