q

.gitattributes

@@ -19,6 +19,7 @@
 *.pb filter=lfs diff=lfs merge=lfs -text
 *.pickle filter=lfs diff=lfs merge=lfs -text
 *.pkl filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

app.py

@@ -1,7 +1,45 @@
 import gradio as gr
+from inference import infer
 
-def greet(name):
-    return "Hello " + name + "!!"
+def greet(image, prompt):
+    restore_img = infer(img=image, text_prompt=prompt)
+    return restore_img
 
-demo = gr.Interface(fn=greet, inputs=[gr.components.Image(), "Text Instruction"], outputs=gr.components.Image())
-demo.launch()
+
+title = "🖼️ ICDR 🖼️"
+description = ''' ## ICDR: Image Restoration Framework for Composite Degradation following Human Instructions
+Our Github : https://github.com/
+
+Siwon Kim, Donghyeon Yoon
+
+Ajou Univ
+'''
+
+
+article = "<p style='text-align: center'><a href='https://github.com/' target='_blank'>ICDR</a></p>"
+
+#### Image,Prompts examples
+examples = [['input/00010.png', "I love this photo, could you remove the haze and more brighter?"],
+            ['input/00058.png', "I have to post an emotional shot on Instagram, but it was shot too foggy and too dark. Change it like a sunny day and brighten it up!"]]
+
+css = """
+    .image-frame img, .image-container img {
+        width: auto;
+        height: auto;
+        max-width: none;
+    }
+"""
+
+
+demo = gr.Interface(
+    fn=greet, 
+    inputs=[gr.Image(type="pil", label="Input"),
+            gr.Text(label="Prompt") ],
+    outputs=[gr.Image(type="pil", label="Ouput")],
+    title=title,
+    description=description,
+    article=article,
+    examples=examples,
+    css=css,
+    )
+demo.launch(share=True)

ckpt/epoch_287.pth

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

inference.py

@@ -0,0 +1,59 @@
+import argparse
+import subprocess
+from tqdm import tqdm
+import numpy as np
+
+import torch
+from torch.utils.data import DataLoader
+
+from utils.dataset_utils_CDD import DerainDehazeDataset
+from utils.val_utils import AverageMeter, compute_psnr_ssim
+from utils.image_io import save_image_tensor
+
+from text_net.model import AirNet
+
+def test_Derain_Dehaze(opt, net, dataset, task="derain"):
+    output_path = opt.output_path + task + '/'
+    subprocess.check_output(['mkdir', '-p', output_path])
+
+    # dataset.set_dataset(task)
+    testloader = DataLoader(dataset, batch_size=1, pin_memory=True, shuffle=False, num_workers=0)
+    print(len(testloader))
+
+    with torch.no_grad():
+        for ([degraded_name], degradation, degrad_patch, clean_patch, text_prompt) in tqdm(testloader):
+            degrad_patch, clean_patch = degrad_patch.cuda(), clean_patch.cuda()
+            restored = net(x_query=degrad_patch, x_key=degrad_patch, text_prompt = text_prompt)
+
+        return save_image_tensor(restored)
+
+
+def infer(text_prompt = "", img=None):
+    parser = argparse.ArgumentParser()
+    # Input Parameters
+    parser.add_argument('--cuda', type=int, default=0)
+    parser.add_argument('--derain_path', type=str, default="data/Test_prompting/", help='save path of test raining images')
+    parser.add_argument('--output_path', type=str, default="output/demo11", help='output save path')
+    parser.add_argument('--ckpt_path', type=str, default="ckpt/epoch_287.pth", help='checkpoint save path')
+    # parser.add_argument('--text_prompt', type=str, default="derain")
+
+    opt = parser.parse_args()
+    # opt.text_prompt = text_prompt
+
+    np.random.seed(0)
+    torch.manual_seed(0)
+    torch.cuda.set_device(opt.cuda)
+
+    opt.batch_size = 7
+    ckpt_path = opt.ckpt_path
+
+    derain_set = DerainDehazeDataset(opt, img=img, text_prompt = text_prompt)
+
+    # Make network
+    net = AirNet(opt).cuda()
+    net.eval()
+    net.load_state_dict(torch.load(ckpt_path, map_location=torch.device(opt.cuda)))
+
+    restored = test_Derain_Dehaze(opt, net, derain_set, task="derain")
+
+    return restored

text_net/DGRN.py

@@ -0,0 +1,232 @@
+import torch.nn as nn
+import torch
+from .deform_conv import DCN_layer
+import clip
+
+clip_model, preprocess = clip.load("ViT-B/32", device='cuda')
+
+# 동적으로 텍스트 임베딩 차원 가져오기
+text_embed_dim = clip_model.text_projection.shape[1]
+
+
+def default_conv(in_channels, out_channels, kernel_size, bias=True):
+    return nn.Conv2d(in_channels, out_channels, kernel_size, padding=(kernel_size // 2), bias=bias)
+
+
+class DGM(nn.Module):
+    def __init__(self, channels_in, channels_out, kernel_size):
+        super(DGM, self).__init__()
+        self.channels_out = channels_out
+        self.channels_in = channels_in
+        self.kernel_size = kernel_size
+
+        self.dcn = DCN_layer(self.channels_in, self.channels_out, kernel_size,
+                             padding=(kernel_size - 1) // 2, bias=False)
+        self.sft = SFT_layer(self.channels_in, self.channels_out)
+
+        self.relu = nn.LeakyReLU(0.1, True)
+
+    def forward(self, x, inter, text_prompt):
+        '''
+        :param x: feature map: B * C * H * W
+        :inter: degradation map: B * C * H * W
+        '''
+        dcn_out = self.dcn(x, inter)
+        sft_out = self.sft(x, inter, text_prompt)
+        out = dcn_out + sft_out
+        out = x + out
+
+        return out
+
+# Projection Head 정의
+class TextProjectionHead(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super(TextProjectionHead, self).__init__()
+        self.proj = nn.Sequential(
+            nn.Linear(input_dim, output_dim),
+            nn.ReLU(),
+            nn.Linear(output_dim, output_dim)
+        ).float()
+    
+    def forward(self, x):
+        return self.proj(x.float())
+
+
+
+class SFT_layer(nn.Module):
+    def __init__(self, channels_in, channels_out):
+        super(SFT_layer, self).__init__()
+        self.conv_gamma = nn.Sequential(
+            nn.Conv2d(channels_in, channels_out, 1, 1, 0, bias=False),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(channels_out, channels_out, 1, 1, 0, bias=False),
+        )
+        self.conv_beta = nn.Sequential(
+            nn.Conv2d(channels_in, channels_out, 1, 1, 0, bias=False),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(channels_out, channels_out, 1, 1, 0, bias=False),
+        )
+
+        self.text_proj_head = TextProjectionHead(text_embed_dim, channels_out)
+
+        '''
+        self.text_gamma = nn.Sequential(
+            nn.Conv2d(channels_out, channels_out, 1, 1, 0, bias=False),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(channels_out, channels_out, 1, 1, 0, bias=False),
+        ).float()
+        self.text_beta = nn.Sequential(
+            nn.Conv2d(channels_out, channels_out, 1, 1, 0, bias=False),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(channels_out, channels_out, 1, 1, 0, bias=False),
+        ).float()
+        '''
+
+        self.cross_attention = nn.MultiheadAttention(embed_dim=channels_out, num_heads=2)
+
+
+    def forward(self, x, inter, text_prompt):
+        '''
+        :param x: degradation representation: B * C
+        :param inter: degradation intermediate representation map: B * C * H * W
+        '''
+        # img_gamma = self.conv_gamma(inter)
+        # img_beta = self.conv_beta(inter)
+
+        B, C, H, W = inter.shape #cross attention
+        
+
+        text_tokens = clip.tokenize(text_prompt).to(x.device)  # Tokenize the text prompts (Batch size)
+        with torch.no_grad():
+            text_embed = clip_model.encode_text(text_tokens)
+
+        text_proj = self.text_proj_head(text_embed).float()
+
+        # 텍스트 임베딩 차원 확장: (B, C, H, W)로 변경 #concat 
+        # text_proj_expanded = text_proj.unsqueeze(-1).unsqueeze(-1).expand(B, self.conv_gamma[0].out_channels, H, W)
+        text_proj_expanded = text_proj.unsqueeze(-1).unsqueeze(-1).expand(B, C, H, W)
+        
+        # 이미지 중간 표현과 텍스트 임베딩 결합 (concat)
+        combined = inter * text_proj_expanded
+        # combined = torch.cat([inter, text_proj_expanded], dim=1)
+
+        # 이미지와 텍스트 기반 gamma와 beta 계산
+        img_gamma = self.conv_gamma(combined)
+        img_beta = self.conv_beta(combined)
+        
+        ''' simple concat
+        text_gamma = self.text_gamma(text_proj.unsqueeze(-1).unsqueeze(-1))  # Reshape to match (B, C, H, W)
+        text_beta = self.text_beta(text_proj.unsqueeze(-1).unsqueeze(-1))  # Reshape to match (B, C, H, W)
+        '''
+
+        '''
+        text_proj = text_proj.unsqueeze(1).expand(-1, H*W, -1)  # B * (H*W) * C
+
+        # 이미지 중간 표현 변환: B * (H*W) * C로 변경
+        inter_flat = inter.view(B, C, -1).permute(2, 0, 1)  # (H*W) * B * C
+        
+        # Cross-attention 적용
+        attn_output, _ = self.cross_attention(text_proj.permute(1, 0, 2), inter_flat, inter_flat)
+        attn_output = attn_output.permute(1, 2, 0).view(B, C, H, W)  # B * C * H * W
+
+        # Gamma와 Beta 계산
+        img_gamma = self.conv_gamma(attn_output)
+        img_beta = self.conv_beta(attn_output)
+'''
+        # concat으로 text 결합 실험
+        return x * img_gamma + img_beta
+
+
+class DGB(nn.Module):
+    def __init__(self, conv, n_feat, kernel_size):
+        super(DGB, self).__init__()
+
+        # self.da_conv1 = DGM(n_feat, n_feat, kernel_size)
+        # self.da_conv2 = DGM(n_feat, n_feat, kernel_size)
+        self.dgm1 = DGM(n_feat, n_feat, kernel_size)
+        self.dgm2 = DGM(n_feat, n_feat, kernel_size)
+        self.conv1 = conv(n_feat, n_feat, kernel_size)
+        self.conv2 = conv(n_feat, n_feat, kernel_size)
+
+        self.relu = nn.LeakyReLU(0.1, True)
+
+    def forward(self, x, inter, text_prompt):
+        '''
+        :param x: feature map: B * C * H * W
+        :param inter: degradation representation: B * C * H * W
+        '''
+
+        out = self.relu(self.dgm1(x, inter, text_prompt))
+        out = self.relu(self.conv1(out))
+        out = self.relu(self.dgm2(out, inter, text_prompt))
+        out = self.conv2(out) + x
+
+        return out
+
+
+class DGG(nn.Module):
+    def __init__(self, conv, n_feat, kernel_size, n_blocks):
+        super(DGG, self).__init__()
+        self.n_blocks = n_blocks
+        modules_body = [
+            DGB(conv, n_feat, kernel_size) \
+            for _ in range(n_blocks)
+        ]
+        modules_body.append(conv(n_feat, n_feat, kernel_size))
+
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x, inter, text_prompt):
+        '''
+        :param x: feature map: B * C * H * W
+        :param inter: degradation representation: B * C * H * W
+        '''
+        res = x
+        for i in range(self.n_blocks):
+            res = self.body[i](res, inter, text_prompt)
+        res = self.body[-1](res)
+        res = res + x
+
+        return res
+
+
+class DGRN(nn.Module):
+    def __init__(self, opt, conv=default_conv):
+        super(DGRN, self).__init__()
+
+        self.n_groups = 5
+        n_blocks = 5
+        n_feats = 64
+        kernel_size = 3
+
+        # head module
+        modules_head = [conv(3, n_feats, kernel_size)]
+        self.head = nn.Sequential(*modules_head)
+
+        # body
+        modules_body = [
+            DGG(default_conv, n_feats, kernel_size, n_blocks) \
+            for _ in range(self.n_groups)
+        ]
+        modules_body.append(conv(n_feats, n_feats, kernel_size))
+        self.body = nn.Sequential(*modules_body)
+
+        # tail
+        modules_tail = [conv(n_feats, 3, kernel_size)]
+        self.tail = nn.Sequential(*modules_tail)
+
+    def forward(self, x, inter, text_prompt):
+        # head
+        x = self.head(x)
+
+        # body
+        res = x
+        for i in range(self.n_groups):
+            res = self.body[i](res, inter, text_prompt)
+        res = self.body[-1](res)
+        res = res + x
+
+        # tail
+        x = self.tail(res)
+
+        return x

text_net/deform_conv.py

@@ -0,0 +1,65 @@
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn.modules.utils import _pair
+
+from mmcv.ops import modulated_deform_conv2d
+
+
+class DCN_layer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1,
+                 groups=1, deformable_groups=1, bias=True, extra_offset_mask=True):
+        super(DCN_layer, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.groups = groups
+        self.deformable_groups = deformable_groups
+        self.with_bias = bias
+
+        self.weight = nn.Parameter(
+            torch.Tensor(out_channels, in_channels // groups, *self.kernel_size))
+
+        self.extra_offset_mask = extra_offset_mask
+        self.conv_offset_mask = nn.Conv2d(
+            self.in_channels * 2,
+            self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1],
+            kernel_size=self.kernel_size, stride=_pair(self.stride), padding=_pair(self.padding),
+            bias=True
+        )
+
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self.init_offset()
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        n = self.in_channels
+        for k in self.kernel_size:
+            n *= k
+        stdv = 1. / math.sqrt(n)
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.zero_()
+
+    def init_offset(self):
+        self.conv_offset_mask.weight.data.zero_()
+        self.conv_offset_mask.bias.data.zero_()
+
+    def forward(self, input_feat, inter):
+        feat_degradation = torch.cat([input_feat, inter], dim=1)
+
+        out = self.conv_offset_mask(feat_degradation)
+        o1, o2, mask = torch.chunk(out, 3, dim=1)
+        offset = torch.cat((o1, o2), dim=1)
+        mask = torch.sigmoid(mask)
+
+        return modulated_deform_conv2d(input_feat.contiguous(), offset, mask, self.weight, self.bias, self.stride,
+                                       self.padding, self.dilation, self.groups, self.deformable_groups)

text_net/encoder.py

@@ -0,0 +1,67 @@
+from torch import nn
+from text_net.moco import MoCo
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_feat, out_feat, stride=1):
+        super(ResBlock, self).__init__()
+        self.backbone = nn.Sequential(
+            nn.Conv2d(in_feat, out_feat, kernel_size=3, stride=stride, padding=1, bias=False),
+            nn.BatchNorm2d(out_feat),
+            nn.LeakyReLU(0.1, True),
+            nn.Conv2d(out_feat, out_feat, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_feat),
+        )
+        self.shortcut = nn.Sequential(
+            nn.Conv2d(in_feat, out_feat, kernel_size=1, stride=stride, bias=False),
+            nn.BatchNorm2d(out_feat)
+        )
+
+    def forward(self, x):
+        return nn.LeakyReLU(0.1, True)(self.backbone(x) + self.shortcut(x))
+
+
+class ResEncoder(nn.Module):
+    def __init__(self):
+        super(ResEncoder, self).__init__()
+
+        self.E_pre = ResBlock(in_feat=3, out_feat=64, stride=1)
+        self.E = nn.Sequential(
+            ResBlock(in_feat=64, out_feat=128, stride=2),
+            ResBlock(in_feat=128, out_feat=256, stride=2),
+            nn.AdaptiveAvgPool2d(1)
+        )
+
+        self.mlp = nn.Sequential(
+            nn.Linear(256, 256),
+            nn.LeakyReLU(0.1, True),
+            nn.Linear(256, 256),
+        )
+
+    def forward(self, x):
+        inter = self.E_pre(x)
+        fea = self.E(inter).squeeze(-1).squeeze(-1)
+        out = self.mlp(fea)
+
+        return fea, out, inter
+
+
+class CBDE(nn.Module):
+    def __init__(self, opt):
+        super(CBDE, self).__init__()
+
+        dim = 256
+
+        # Encoder
+        self.E = MoCo(base_encoder=ResEncoder, dim=dim, K=opt.batch_size * dim)
+
+    def forward(self, x_query, x_key):
+        if self.training:
+            # degradation-aware represenetion learning
+            fea, logits, labels, inter = self.E(x_query, x_key)
+
+            return fea, logits, labels, inter
+        else:
+            # degradation-aware represenetion learning
+            fea, inter = self.E(x_query, x_query)
+            return fea, inter

text_net/moco.py

@@ -0,0 +1,166 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import torch
+import torch.nn as nn
+
+
+class MoCo(nn.Module):
+    """
+    Build a MoCo model with: a query encoder, a key encoder, and a queue
+    https://arxiv.org/abs/1911.05722
+    """
+    def __init__(self, base_encoder, dim=256, K=3*256, m=0.999, T=0.07, mlp=False):
+        """
+        dim: feature dimension (default: 128)
+        K: queue size; number of negative keys (default: 65536)
+        m: moco momentum of updating key encoder (default: 0.999)
+        T: softmax temperature (default: 0.07)
+        """
+        super(MoCo, self).__init__()
+
+        self.K = K
+        self.m = m
+        self.T = T
+
+        # create the encoders
+        # num_classes is the output fc dimension
+        self.encoder_q = base_encoder()
+        self.encoder_k = base_encoder()
+
+        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
+            param_k.data.copy_(param_q.data)  # initialize
+            param_k.requires_grad = False  # not update by gradient
+
+        # create the queue
+        self.register_buffer("queue", torch.randn(dim, K))
+        self.queue = nn.functional.normalize(self.queue, dim=0)
+
+        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
+
+    @torch.no_grad()
+    def _momentum_update_key_encoder(self):
+        """
+        Momentum update of the key encoder
+        """
+        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
+            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)
+
+    @torch.no_grad()
+    def _dequeue_and_enqueue(self, keys):
+        # gather keys before updating queue
+        # keys = concat_all_gather(keys)
+        batch_size = keys.shape[0]
+
+        ptr = int(self.queue_ptr)
+        assert self.K % batch_size == 0  # for simplicity
+
+        # replace the keys at ptr (dequeue and enqueue)
+        self.queue[:, ptr:ptr + batch_size] = keys.transpose(0, 1)
+        ptr = (ptr + batch_size) % self.K  # move pointer
+
+        self.queue_ptr[0] = ptr
+
+    @torch.no_grad()
+    def _batch_shuffle_ddp(self, x):
+        """
+        Batch shuffle, for making use of BatchNorm.
+        *** Only support DistributedDataParallel (DDP) model. ***
+        """
+        # gather from all gpus
+        batch_size_this = x.shape[0]
+        x_gather = concat_all_gather(x)
+        batch_size_all = x_gather.shape[0]
+
+        num_gpus = batch_size_all // batch_size_this
+
+        # random shuffle index
+        idx_shuffle = torch.randperm(batch_size_all).cuda()
+
+        # broadcast to all gpus
+        torch.distributed.broadcast(idx_shuffle, src=0)
+
+        # index for restoring
+        idx_unshuffle = torch.argsort(idx_shuffle)
+
+        # shuffled index for this gpu
+        gpu_idx = torch.distributed.get_rank()
+        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]
+
+        return x_gather[idx_this], idx_unshuffle
+
+    @torch.no_grad()
+    def _batch_unshuffle_ddp(self, x, idx_unshuffle):
+        """
+        Undo batch shuffle.
+        *** Only support DistributedDataParallel (DDP) model. ***
+        """
+        # gather from all gpus
+        batch_size_this = x.shape[0]
+        x_gather = concat_all_gather(x)
+        batch_size_all = x_gather.shape[0]
+
+        num_gpus = batch_size_all // batch_size_this
+
+        # restored index for this gpu
+        gpu_idx = torch.distributed.get_rank()
+        idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]
+
+        return x_gather[idx_this]
+
+    def forward(self, im_q, im_k):
+        """
+        Input:
+            im_q: a batch of query images
+            im_k: a batch of key images
+        Output:
+            logits, targets
+        """
+        if self.training:
+            # compute query features
+            embedding, q, inter = self.encoder_q(im_q)  # queries: NxC
+            q = nn.functional.normalize(q, dim=1)
+
+            # compute key features
+            with torch.no_grad():  # no gradient to keys
+                self._momentum_update_key_encoder()  # update the key encoder
+
+                _, k, _ = self.encoder_k(im_k)  # keys: NxC
+                k = nn.functional.normalize(k, dim=1)
+
+            # compute logits
+            # Einstein sum is more intuitive
+            # positive logits: Nx1
+            l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
+            # negative logits: NxK
+            l_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])
+
+            # logits: Nx(1+K)
+            logits = torch.cat([l_pos, l_neg], dim=1)
+
+            # apply temperature
+            logits /= self.T
+
+            # labels: positive key indicators
+            labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()
+            # dequeue and enqueue
+            self._dequeue_and_enqueue(k)
+
+            return embedding, logits, labels, inter
+        else:
+            embedding, _, inter = self.encoder_q(im_q)
+
+            return embedding, inter
+
+
+# utils
+@torch.no_grad()
+def concat_all_gather(tensor):
+    """
+    Performs all_gather operation on the provided tensors.
+    *** Warning ***: torch.distributed.all_gather has no gradient.
+    """
+    tensors_gather = [torch.ones_like(tensor)
+        for _ in range(torch.distributed.get_world_size())]
+    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
+
+    output = torch.cat(tensors_gather, dim=0)
+    return output

text_net/model.py

@@ -0,0 +1,29 @@
+from torch import nn
+
+from text_net.encoder import CBDE
+from text_net.DGRN import DGRN
+
+
+class AirNet(nn.Module):
+    def __init__(self, opt):
+        super(AirNet, self).__init__()
+
+        # Restorer
+        self.R = DGRN(opt)
+
+        # Encoder
+        self.E = CBDE(opt)
+
+    def forward(self, x_query, x_key, text_prompt):
+        if self.training:
+            fea, logits, labels, inter = self.E(x_query, x_key)
+
+            restored = self.R(x_query, inter, text_prompt)
+
+            return restored, logits, labels
+        else:
+            fea, inter = self.E(x_query, x_query)
+
+            restored = self.R(x_query, inter, text_prompt)
+
+            return restored

utils/dataset_utils.py

@@ -0,0 +1,309 @@
+import os
+import random
+import copy
+from PIL import Image
+import numpy as np
+
+from torch.utils.data import Dataset
+from torchvision.transforms import ToPILImage, Compose, RandomCrop, ToTensor
+
+from utils.image_utils import random_augmentation, crop_img
+from utils.degradation_utils import Degradation
+
+
+class TrainDataset(Dataset):
+    def __init__(self, args):
+        super(TrainDataset, self).__init__()
+        self.args = args
+        self.rs_ids = []
+        self.hazy_ids = []
+        self.D = Degradation(args)
+        self.de_temp = 0
+        self.de_type = self.args.de_type
+        self.image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.gif']
+
+        self.de_dict = {'denoise_15': 0, 'denoise_25': 1, 'denoise_50': 2, 'derain': 3, 'dehaze': 4}
+
+        self._init_ids()
+
+        self.crop_transform = Compose([
+            ToPILImage(),
+            RandomCrop(args.patch_size),
+        ])
+
+        self.toTensor = ToTensor()
+
+    def _init_ids(self):
+        if 'denoise_15' in self.de_type or 'denoise_25' in self.de_type or 'denoise_50' in self.de_type:
+            self._init_clean_ids()
+        if 'derain' in self.de_type:
+            self._init_rs_ids()
+        if 'dehaze' in self.de_type:
+            self._init_hazy_ids()
+
+        random.shuffle(self.de_type)
+
+    def _init_clean_ids(self):
+        clean_ids = []
+        # 파일 목록 중 이미지 파일만 필터링
+        name_list = os.listdir(self.args.denoise_dir)
+        name_list = [file for file in name_list if os.path.splitext(file)[1].lower() in self.image_extensions]
+
+        clean_ids += [self.args.denoise_dir + id_ for id_ in name_list]
+
+        if 'denoise_15' in self.de_type:
+            self.s15_ids = copy.deepcopy(clean_ids)
+            random.shuffle(self.s15_ids)
+            self.s15_counter = 0
+        if 'denoise_25' in self.de_type:
+            self.s25_ids = copy.deepcopy(clean_ids)
+            random.shuffle(self.s25_ids)
+            self.s25_counter = 0
+        if 'denoise_50' in self.de_type:
+            self.s50_ids = copy.deepcopy(clean_ids)
+            random.shuffle(self.s50_ids)
+            self.s50_counter = 0
+
+        # print(clean_ids)
+
+        self.num_clean = len(clean_ids)
+
+    def _init_hazy_ids(self):
+        # 파일 목록 중 이미지 파일만 필터링
+        dehaze_ids = []
+        name_list = os.listdir(self.args.dehaze_dir)
+        name_list = [file for file in name_list if os.path.splitext(file)[1].lower() in self.image_extensions]
+        dehaze_ids += [self.args.dehaze_dir + id_ for id_ in name_list]
+        self.hazy_ids = dehaze_ids
+
+        self.hazy_counter = 0
+        self.num_hazy = len(self.hazy_ids)
+
+    def _init_rs_ids(self):
+        # 파일 목록 중 이미지 파일만 필터링
+        derain_ids = []
+        name_list = os.listdir(self.args.derain_dir)
+        name_list = [file for file in name_list if os.path.splitext(file)[1].lower() in self.image_extensions]
+        derain_ids += [self.args.derain_dir + id_ for id_ in name_list]
+        self.rs_ids = derain_ids
+
+        self.rl_counter = 0
+        # print(derain_ids)
+
+        self.num_rl = len(self.rs_ids)
+
+    def _crop_patch(self, img_1, img_2):
+        H = img_1.shape[0]
+        W = img_1.shape[1]
+        ind_H = random.randint(0, H - self.args.patch_size)
+        ind_W = random.randint(0, W - self.args.patch_size)
+
+        patch_1 = img_1[ind_H:ind_H + self.args.patch_size, ind_W:ind_W + self.args.patch_size]
+        patch_2 = img_2[ind_H:ind_H + self.args.patch_size, ind_W:ind_W + self.args.patch_size]
+
+        return patch_1, patch_2
+
+    def _get_gt_name(self, rainy_name):
+        gt_name = 'data/' + 'Target/Derain/norain-' + rainy_name.split('rain-')[-1]
+        return gt_name
+
+    def _get_nonhazy_name(self, hazy_name):
+        gt_name = 'data/' + 'Target/Dehaze/nohaze-' + rainy_name.split('haze-')[-1]
+        return gt_name
+
+    def __getitem__(self, _):
+        de_id = self.de_dict[self.de_type[self.de_temp]]
+
+        if de_id < 3:
+            if de_id == 0:
+                clean_id = self.s15_ids[self.s15_counter]
+                self.s15_counter = (self.s15_counter + 1) % self.num_clean
+                if self.s15_counter == 0:
+                    random.shuffle(self.s15_ids)
+            elif de_id == 1:
+                clean_id = self.s25_ids[self.s25_counter]
+                self.s25_counter = (self.s25_counter + 1) % self.num_clean
+                if self.s25_counter == 0:
+                    random.shuffle(self.s25_ids)
+            elif de_id == 2:
+                clean_id = self.s50_ids[self.s50_counter]
+                self.s50_counter = (self.s50_counter + 1) % self.num_clean
+                if self.s50_counter == 0:
+                    random.shuffle(self.s50_ids)
+
+            # clean_id = random.randint(0, len(self.clean_ids) - 1)
+            clean_img = crop_img(np.array(Image.open(clean_id).convert('RGB')), base=16)
+            clean_patch_1, clean_patch_2 = self.crop_transform(clean_img), self.crop_transform(clean_img)
+            clean_patch_1, clean_patch_2 = np.array(clean_patch_1), np.array(clean_patch_2)
+
+            # clean_name = self.clean_ids[clean_id].split("/")[-1].split('.')[0]
+            clean_name = clean_id.split("/")[-1].split('.')[0]
+
+            clean_patch_1, clean_patch_2 = random_augmentation(clean_patch_1, clean_patch_2)
+            degrad_patch_1, degrad_patch_2 = self.D.degrade(clean_patch_1, clean_patch_2, de_id)
+        else:
+            if de_id == 3:
+                # Rain Streak Removal
+                # rl_id = random.randint(0, len(self.rl_ids) - 1)
+                degrad_img = crop_img(np.array(Image.open(self.rs_ids[self.rl_counter]).convert('RGB')), base=16)
+                clean_name = self._get_gt_name(self.rs_ids[self.rl_counter])
+                clean_img = crop_img(np.array(Image.open(clean_name).convert('RGB')), base=16)
+
+                self.rl_counter = (self.rl_counter + 1) % self.num_rl
+                if self.rl_counter == 0:
+                    random.shuffle(self.rs_ids)
+            elif de_id == 4:
+                # Dehazing with SOTS outdoor training set
+                # hazy_id = random.randint(0, len(self.hazy_ids) - 1)
+                degrad_img = crop_img(np.array(Image.open(self.hazy_ids[self.hazy_counter]).convert('RGB')), base=16)
+                clean_name = self._get_nonhazy_name(self.hazy_ids[self.hazy_counter])
+                clean_img = crop_img(np.array(Image.open(clean_name).convert('RGB')), base=16)
+
+                self.hazy_counter = (self.hazy_counter + 1) % self.num_hazy
+                if self.hazy_counter == 0:
+                    random.shuffle(self.hazy_ids)
+            degrad_patch_1, clean_patch_1 = random_augmentation(*self._crop_patch(degrad_img, clean_img))
+            degrad_patch_2, clean_patch_2 = random_augmentation(*self._crop_patch(degrad_img, clean_img))
+
+        clean_patch_1, clean_patch_2 = self.toTensor(clean_patch_1), self.toTensor(clean_patch_2)
+        degrad_patch_1, degrad_patch_2 = self.toTensor(degrad_patch_1), self.toTensor(degrad_patch_2)
+
+        self.de_temp = (self.de_temp + 1) % len(self.de_type)
+        if self.de_temp == 0:
+            random.shuffle(self.de_type)
+
+        return [clean_name, de_id], degrad_patch_1, degrad_patch_2, clean_patch_1, clean_patch_2
+
+    def __len__(self):
+        return 400 * len(self.args.de_type)
+
+
+class DenoiseTestDataset(Dataset):
+    def __init__(self, args):
+        super(DenoiseTestDataset, self).__init__()
+        self.args = args
+        self.clean_ids = []
+        self.sigma = 15
+        self.image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.gif']
+
+        self._init_clean_ids()
+
+        self.toTensor = ToTensor()
+
+    def _init_clean_ids(self):
+        clean_ids = []
+        # 파일 목록 중 이미지 파일만 필터링
+        name_list = os.listdir(self.args.denoise_path)
+        name_list = [file for file in name_list if os.path.splitext(file)[1].lower() in self.image_extensions]
+        self.clean_ids += [self.args.denoise_path + id_ for id_ in name_list]
+
+        self.num_clean = len(self.clean_ids)
+
+    def _add_gaussian_noise(self, clean_patch):
+        noise = np.random.randn(*clean_patch.shape)
+        noisy_patch = np.clip(clean_patch + noise * self.sigma, 0, 255).astype(np.uint8)
+        return noisy_patch, clean_patch
+
+    def set_sigma(self, sigma):
+        self.sigma = sigma
+
+    def __getitem__(self, clean_id):
+        clean_img = crop_img(np.array(Image.open(self.clean_ids[clean_id]).convert('RGB')), base=16)
+        clean_name = self.clean_ids[clean_id].split("/")[-1].split('.')[0]
+
+        noisy_img, _ = self._add_gaussian_noise(clean_img)
+        clean_img, noisy_img = self.toTensor(clean_img), self.toTensor(noisy_img)
+
+        return [clean_name], noisy_img, clean_img
+
+    def __len__(self):
+        return self.num_clean
+
+
+class DerainDehazeDataset(Dataset):
+    def __init__(self, args, task="derain"):
+        super(DerainDehazeDataset, self).__init__()
+        self.ids = []
+        self.task_idx = 0
+        self.args = args
+        self.image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.gif']
+
+        self.task_dict = {'derain': 0, 'dehaze': 1}
+        self.toTensor = ToTensor()
+
+        self.set_dataset(task)
+
+    def _init_input_ids(self):
+        if self.task_idx == 0:
+            self.ids = []
+            # 파일 목록 중 이미지 파일만 필터링
+            name_list = os.listdir(self.args.derain_path + 'input/')
+            name_list = [file for file in name_list if os.path.splitext(file)[1].lower() in self.image_extensions]
+            self.ids += [self.args.derain_path + 'input/' + id_ for id_ in name_list]
+        elif self.task_idx == 1:
+            self.ids = []
+            # 파일 목록 중 이미지 파일만 필터링
+            name_list = os.listdir(self.args.dehaze_path + 'input/')
+            name_list = [file for file in name_list if os.path.splitext(file)[1].lower() in self.image_extensions]
+            self.ids += [self.args.dehaze_path + 'input/' + id_ for id_ in name_list]
+
+        self.length = len(self.ids)
+
+    def _get_gt_path(self, degraded_name):
+        if self.task_idx == 0:
+            gt_name = '/'.join(degraded_name.replace("input", "target").split('/')[:-1] + degraded_name.replace("input", "target").replace("rain", "norain").split('/')[-1:])
+            print(gt_name)
+        elif self.task_idx == 1:
+            dir_name = degraded_name.split("input")[0] + 'target/'
+            name = degraded_name.split('/')[-1].split('_')[0] + '.png'
+            gt_name = dir_name + name
+        return gt_name
+
+    def set_dataset(self, task):
+        self.task_idx = self.task_dict[task]
+        self._init_input_ids()
+
+    def __getitem__(self, idx):
+        degraded_path = self.ids[idx]
+        clean_path = self._get_gt_path(degraded_path)
+
+        degraded_img = crop_img(np.array(Image.open(degraded_path).convert('RGB')), base=16)
+        clean_img = crop_img(np.array(Image.open(clean_path).convert('RGB')), base=16)
+
+        clean_img, degraded_img = self.toTensor(clean_img), self.toTensor(degraded_img)
+        degraded_name = degraded_path.split('/')[-1][:-4]
+
+        return [degraded_name], degraded_img, clean_img
+
+    def __len__(self):
+        return self.length
+
+
+class TestSpecificDataset(Dataset):
+    def __init__(self, args):
+        super(TestSpecificDataset, self).__init__()
+        self.args = args
+        self.degraded_ids = []
+        self._init_clean_ids(args.test_path)
+
+        self.toTensor = ToTensor()
+
+    def _init_clean_ids(self, root):
+        degraded_ids = []
+        # 파일 목록 중 이미지 파일만 필터링
+        name_list = os.listdir(root)
+        name_list = [file for file in name_list if os.path.splitext(file)[1].lower() in self.image_extensions]
+        self.degraded_ids += [root + id_ for id_ in name_list]
+
+        self.num_img = len(self.degraded_ids)
+
+    def __getitem__(self, idx):
+        degraded_img = crop_img(np.array(Image.open(self.degraded_ids[idx]).convert('RGB')), base=16)
+        name = self.degraded_ids[idx].split('/')[-1][:-4]
+
+        degraded_img = self.toTensor(degraded_img)
+
+        return [name], degraded_img
+
+    def __len__(self):
+        return self.num_img