birth

README.md

@@ -0,0 +1,83 @@
+# [TEAM ACVLAB][NTIRE 2025 Image Shadow Removal Challenge](https://cvlai.net/ntire/2025/) @ [CVPR 2025](https://cvpr.thecvf.com/)
+
+## Link to the codes/executables of the solution(s):
+* [Checkpoints](https://drive.google.com/file/d/1USD5sLvEcgFqIg7BDzc1OuInzSx3GnUN/view?usp=drive_link)
+* Input / Output file
+
+## Environments
+```bash
+conda create -n ntire_shadow python=3.9 -y
+
+conda activate ntire_shadow
+
+pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
+
+pip install -r requirements.txt
+
+```
+
+## Folder Structure
+```bash
+test_dir
+├── Origin          <- Put the shadow affected images in this folder
+│   ├── 0000.png
+│   ├── 0001.png
+│   ├── ...
+├── Depth
+├── Normal
+
+
+output_dir
+├── 0000.png
+├── 0001.png
+├──...
+```
+
+## How to test?
+1. Clone [Depth anything v2](https://github.com/DepthAnything/Depth-Anything-V2.git)
+
+```bash
+git clone https://github.com/DepthAnything/Depth-Anything-V2.git
+```
+2. Download the [pretrain model of depth anything v2](https://huggingface.co/depth-anything/Depth-Anything-V2-Large/resolve/main/depth_anything_v2_vitl.pth?download=true)
+
+3. Run ```python Depth-Anything-V2/get_depth_normap.py```
+
+Now folder structure will be
+```bash
+test_dir
+├── Origin
+│   ├── 0000.png
+│   ├── 0001.png
+│   ├── ...
+├── Depth
+│   ├── 0000.npy
+│   ├── 0001.npy
+│   ├── ...
+├── Normal
+│   ├── 0000.npy
+│   ├── 0001.npy
+│   ├── ...
+
+output_dir
+├── 0000.png
+├── 0001.png
+├──...
+```
+
+4. Clone [DINOv2](https://github.com/facebookresearch/dinov2.git)
+```bash
+git clone https://github.com/facebookresearch/dinov2.git
+```
+
+5. Download [shadow removal weight](https://drive.google.com/file/d/1USD5sLvEcgFqIg7BDzc1OuInzSx3GnUN/view?usp=drive_link)
+
+```bash 
+gdown 1USD5sLvEcgFqIg7BDzc1OuInzSx3GnUN
+```
+
+6. Run ```run_test.sh``` to get inference results.
+
+```bash
+bash run_test.sh
+```

dataset.py

@@ -0,0 +1,228 @@
+import numpy as np
+import os
+from torch.utils.data import Dataset
+import torch
+from utils import load_normal, load_ssao, load_img, depthToPoint, process_normal, load_depth, Augment_RGB_torch
+import torch.nn.functional as F
+import random
+
+augment = Augment_RGB_torch()
+transforms_aug = [method for method in dir(augment) if callable(getattr(augment, method)) if not method.startswith('_')] 
+
+##################################################################################################
+class DataLoaderTrain(Dataset):
+    def __init__(self, rgb_dir, img_options=None, target_transform=None, debug=False):
+        super(DataLoaderTrain, self).__init__()
+
+        self.target_transform = target_transform
+        
+        gt_dir = 'shadow_free'
+        input_dir = 'origin'
+        depth_dir = 'depth'
+        normal_dir = 'normal'
+
+        clean_files = sorted(os.listdir(os.path.join(rgb_dir, gt_dir)))       # shadow free
+        noisy_files = sorted(os.listdir(os.path.join(rgb_dir, input_dir)))    # origin
+        depth_files = sorted(os.listdir(os.path.join(rgb_dir, depth_dir)))    # depth
+        normal_files = sorted(os.listdir(os.path.join(rgb_dir, normal_dir)))  # noraml map
+
+        self.clean_filenames = [os.path.join(rgb_dir, gt_dir, x) for x in clean_files]        # shadow free
+        self.noisy_filenames = [os.path.join(rgb_dir, input_dir, x) for x in noisy_files]     # origin
+        self.depth_filenames = [os.path.join(rgb_dir, depth_dir, x) for x in depth_files]     # depth
+        self.normal_filenames = [os.path.join(rgb_dir, normal_dir, x) for x in normal_files]  # noraml map
+        self.img_options = img_options
+
+        if debug:
+            self.tar_size = 100 
+        else:
+            self.tar_size = len(self.noisy_filenames)  
+
+    def __len__(self):
+        return self.tar_size
+
+    def __getitem__(self, index):
+        tar_index   = index % self.tar_size
+        
+        clean = np.float32(load_img(self.clean_filenames[tar_index]))
+        noisy = np.float32(load_img(self.noisy_filenames[tar_index]))
+        depth = np.float32(load_depth(self.depth_filenames[tar_index]))
+        normal = np.float32(load_normal(self.normal_filenames[tar_index]))
+
+        point = depthToPoint(60, depth)  
+
+        normal = process_normal(normal)
+
+        clean = torch.from_numpy(clean)
+        noisy = torch.from_numpy(noisy)
+        depth = torch.from_numpy(depth)
+        point = torch.from_numpy(point)
+        normal =  torch.from_numpy(normal)
+
+        point = point / (2 * point[:,:,2].mean())
+
+        clean  = clean.permute(2,0,1)
+        noisy  = noisy.permute(2,0,1)
+        point  = point.permute(2,0,1)
+        normal  = normal.permute(2,0,1)
+
+        clean_filename = os.path.split(self.clean_filenames[tar_index])[-1]
+        noisy_filename = os.path.split(self.noisy_filenames[tar_index])[-1]
+        depth_filename = os.path.split(self.depth_filenames[tar_index])[-1]
+        normal_filename = os.path.split(self.normal_filenames[tar_index])[-1]
+
+
+        augment.rotate = random.randint(-20,20)
+        apply_trans = transforms_aug[random.randint(0, 2)]
+
+        # [0, 1]
+        clean = getattr(augment, apply_trans)(clean)
+        noisy = getattr(augment, apply_trans)(noisy) 
+        point = getattr(augment, apply_trans)(point)
+        normal = getattr(augment, apply_trans)(normal)
+
+
+        #Crop Input and Target
+        ps = self.img_options['patch_size']
+        scale = 1#random.uniform(1, 1.5)
+
+        H = noisy.shape[1]
+        W = noisy.shape[2]
+        scaled_ps = (int)(scale * ps)
+        if H - scaled_ps != 0 or W - scaled_ps != 0:
+            r = np.random.randint(0, H - scaled_ps + 1)
+            c = np.random.randint(0, W - scaled_ps + 1)
+            clean   = clean  [:, r:r + scaled_ps, c:c + scaled_ps]
+            noisy   = noisy  [:, r:r + scaled_ps, c:c + scaled_ps]
+            point   = point  [:, r:r + scaled_ps, c:c + scaled_ps]
+            normal  = normal [:, r:r + scaled_ps, c:c + scaled_ps]
+
+        # scale back to the patch_size
+        if scale != 1:
+            clean = F.interpolate(clean.unsqueeze(0), size=[ps, ps], mode='bilinear')
+            noisy = F.interpolate(noisy.unsqueeze(0), size=[ps, ps], mode='bilinear')
+            point = F.interpolate(point.unsqueeze(0), size=[ps, ps], mode='nearest')
+            normal = F.interpolate(normal.unsqueeze(0), size=[ps, ps], mode='nearest')
+            return clean.squeeze(0), noisy.squeeze(0), point.squeeze(0), normal.squeeze(0), noisy_filename
+
+        return clean, noisy, point, normal, clean_filename, noisy_filename
+
+
+##################################################################################################
+class DataLoaderVal(Dataset):
+    def __init__(self, rgb_dir, target_transform=None, debug=False):
+        super(DataLoaderVal, self).__init__()
+
+        self.target_transform = target_transform
+        
+        gt_dir = 'shadow_free'
+        input_dir = 'origin'
+        depth_dir = 'depth'
+        normal_dir = 'normal'
+        
+        clean_files = sorted(os.listdir(os.path.join(rgb_dir, gt_dir)))
+        noisy_files = sorted(os.listdir(os.path.join(rgb_dir, input_dir)))
+        depth_files = sorted(os.listdir(os.path.join(rgb_dir, depth_dir)))
+        normal_files = sorted(os.listdir(os.path.join(rgb_dir, normal_dir)))
+
+        self.clean_filenames = [os.path.join(rgb_dir, gt_dir, x) for x in clean_files]
+        self.noisy_filenames = [os.path.join(rgb_dir, input_dir, x) for x in noisy_files]
+        self.depth_filenames = [os.path.join(rgb_dir, depth_dir, x) for x in depth_files]
+        self.normal_filenames = [os.path.join(rgb_dir, normal_dir, x) for x in normal_files]
+
+        if debug:
+            self.tar_size = 10
+        else:
+            self.tar_size = len(self.noisy_filenames)
+
+    def __len__(self):
+        return self.tar_size 
+
+    def __getitem__(self, index):
+        tar_index   = index % self.tar_size
+        clean  = np.float32(load_img(self.clean_filenames[tar_index]))
+        noisy  = np.float32(load_img(self.noisy_filenames[tar_index]))
+        depth  = np.float32(load_depth(self.depth_filenames[tar_index]))
+        normal = np.float32(load_normal(self.normal_filenames[tar_index]))
+
+        point = depthToPoint(60, depth)   
+        normal = process_normal(normal)
+        point = point / (2 * point[:,:,2].mean())
+
+        clean_filename = os.path.split(self.clean_filenames[tar_index])[-1]
+        noisy_filename = os.path.split(self.noisy_filenames[tar_index])[-1]
+
+        clean  = torch.from_numpy(clean)
+        noisy  = torch.from_numpy(noisy)
+        point  = torch.from_numpy(point)
+        normal = torch.from_numpy(normal)
+
+        
+        clean = clean.permute(2,0,1)
+        noisy = noisy.permute(2,0,1)
+        point  = point.permute(2,0,1)
+        normal = normal.permute(2,0,1)
+
+
+        return clean, noisy, point, normal, clean_filename, noisy_filename
+    
+
+
+
+##################################################################################################
+class DataLoaderTest(Dataset):
+    def __init__(self, rgb_dir, target_transform=None, debug=False):
+        super(DataLoaderTest, self).__init__()
+
+        self.target_transform = target_transform
+        
+        # gt_dir = 'shadow_free'
+        input_dir = 'origin'
+        depth_dir = 'depth'
+        normal_dir = 'normal'
+        
+        # clean_files = sorted(os.listdir(os.path.join(rgb_dir, gt_dir)))
+        noisy_files = sorted(os.listdir(os.path.join(rgb_dir, input_dir)))
+        depth_files = sorted(os.listdir(os.path.join(rgb_dir, depth_dir)))
+        normal_files = sorted(os.listdir(os.path.join(rgb_dir, normal_dir)))
+
+        # self.clean_filenames = [os.path.join(rgb_dir, gt_dir, x) for x in clean_files]
+        self.noisy_filenames = [os.path.join(rgb_dir, input_dir, x) for x in noisy_files]
+        self.depth_filenames = [os.path.join(rgb_dir, depth_dir, x) for x in depth_files]
+        self.normal_filenames = [os.path.join(rgb_dir, normal_dir, x) for x in normal_files]
+
+        if debug:
+            self.tar_size = 10
+        else:
+            self.tar_size = len(self.noisy_filenames)
+
+    def __len__(self):
+        return self.tar_size 
+
+    def __getitem__(self, index):
+        tar_index   = index % self.tar_size
+        # clean  = np.float32(load_img(self.clean_filenames[tar_index]))
+        noisy  = np.float32(load_img(self.noisy_filenames[tar_index]))
+        depth  = np.float32(load_depth(self.depth_filenames[tar_index]))
+        normal = np.float32(load_normal(self.normal_filenames[tar_index]))
+
+        point = depthToPoint(60, depth)   
+        normal = process_normal(normal)
+        point = point / (2 * point[:,:,2].mean())
+
+        # clean_filename = os.path.split(self.clean_filenames[tar_index])[-1]
+        noisy_filename = os.path.split(self.noisy_filenames[tar_index])[-1]
+
+        # clean  = torch.from_numpy(clean)
+        noisy  = torch.from_numpy(noisy)
+        point  = torch.from_numpy(point)
+        normal = torch.from_numpy(normal)
+
+        
+        # clean = clean.permute(2,0,1)
+        noisy = noisy.permute(2,0,1)
+        point  = point.permute(2,0,1)
+        normal = normal.permute(2,0,1)
+
+
+        return noisy, noisy, point, normal, noisy_filename, noisy_filename
+

freqfusion.py

@@ -0,0 +1,441 @@
+# TPAMI 2024：Frequency-aware Feature Fusion for Dense Image Prediction
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+import warnings
+import numpy as np
+
+try:
+    from mmcv.ops.carafe import normal_init, xavier_init, carafe
+except ImportError:
+
+    def xavier_init(module: nn.Module,
+                    gain: float = 1,
+                    bias: float = 0,
+                    distribution: str = 'normal') -> None:
+        assert distribution in ['uniform', 'normal']
+        if hasattr(module, 'weight') and module.weight is not None:
+            if distribution == 'uniform':
+                nn.init.xavier_uniform_(module.weight, gain=gain)
+            else:
+                nn.init.xavier_normal_(module.weight, gain=gain)
+        if hasattr(module, 'bias') and module.bias is not None:
+            nn.init.constant_(module.bias, bias)
+
+    def carafe(x, normed_mask, kernel_size, group=1, up=1):
+            b, c, h, w = x.shape
+            _, m_c, m_h, m_w = normed_mask.shape
+            # print('x', x.shape)
+            # print('normed_mask', normed_mask.shape)
+            # assert m_c == kernel_size ** 2 * up ** 2
+            assert m_h == up * h
+            assert m_w == up * w
+            pad = kernel_size // 2
+            # print(pad)
+            pad_x = F.pad(x, pad=[pad] * 4, mode='reflect')
+            # print(pad_x.shape)
+            unfold_x = F.unfold(pad_x, kernel_size=(kernel_size, kernel_size), stride=1, padding=0)
+            # unfold_x = unfold_x.reshape(b, c, 1, kernel_size, kernel_size, h, w).repeat(1, 1, up ** 2, 1, 1, 1, 1)
+            unfold_x = unfold_x.reshape(b, c * kernel_size * kernel_size, h, w)
+            unfold_x = F.interpolate(unfold_x, scale_factor=up, mode='nearest')
+            # normed_mask = normed_mask.reshape(b, 1, up ** 2, kernel_size, kernel_size, h, w)
+            unfold_x = unfold_x.reshape(b, c, kernel_size * kernel_size, m_h, m_w)
+            normed_mask = normed_mask.reshape(b, 1, kernel_size * kernel_size, m_h, m_w)
+            res = unfold_x * normed_mask
+            # test
+            # res[:, :, 0] = 1
+            # res[:, :, 1] = 2
+            # res[:, :, 2] = 3
+            # res[:, :, 3] = 4
+            res = res.sum(dim=2).reshape(b, c, m_h, m_w)
+            # res = F.pixel_shuffle(res, up)
+            # print(res.shape)
+            # print(res)
+            return res
+
+    def normal_init(module, mean=0, std=1, bias=0):
+        if hasattr(module, 'weight') and module.weight is not None:
+            nn.init.normal_(module.weight, mean, std)
+        if hasattr(module, 'bias') and module.bias is not None:
+            nn.init.constant_(module.bias, bias)
+
+
+def constant_init(module, val, bias=0):
+    if hasattr(module, 'weight') and module.weight is not None:
+        nn.init.constant_(module.weight, val)
+    if hasattr(module, 'bias') and module.bias is not None:
+        nn.init.constant_(module.bias, bias)
+
+def resize(input,
+           size=None,
+           scale_factor=None,
+           mode='nearest',
+           align_corners=None,
+           warning=True):
+    if warning:
+        if size is not None and align_corners:
+            input_h, input_w = tuple(int(x) for x in input.shape[2:])
+            output_h, output_w = tuple(int(x) for x in size)
+            if output_h > input_h or output_w > input_w:
+                if ((output_h > 1 and output_w > 1 and input_h > 1
+                     and input_w > 1) and (output_h - 1) % (input_h - 1)
+                        and (output_w - 1) % (input_w - 1)):
+                    warnings.warn(
+                        f'When align_corners={align_corners}, '
+                        'the output would more aligned if '
+                        f'input size {(input_h, input_w)} is `x+1` and '
+                        f'out size {(output_h, output_w)} is `nx+1`')
+    return F.interpolate(input, size, scale_factor, mode, align_corners)
+
+def hamming2D(M, N):
+    """
+    生成二维Hamming窗
+
+    参数：
+    - M：窗口的行数
+    - N：窗口的列数
+
+    返回：
+    - 二维Hamming窗
+    """
+    # 生成水平和垂直方向上的Hamming窗
+    # hamming_x = np.blackman(M)
+    # hamming_x = np.kaiser(M)
+    hamming_x = np.hamming(M)
+    hamming_y = np.hamming(N)
+    # 通过外积生成二维Hamming窗
+    hamming_2d = np.outer(hamming_x, hamming_y)
+    return hamming_2d
+
+class FreqFusion(nn.Module):
+    def __init__(self,
+                hr_channels,
+                lr_channels,
+                scale_factor=1,
+                lowpass_kernel=5,
+                highpass_kernel=3,
+                up_group=1,
+                encoder_kernel=3,
+                encoder_dilation=1,
+                compressed_channels=64,        
+                align_corners=False,
+                upsample_mode='nearest',
+                feature_resample=False, # use offset generator or not
+                feature_resample_group=4,
+                comp_feat_upsample=True, # use ALPF & AHPF for init upsampling
+                use_high_pass=True,
+                use_low_pass=True,
+                hr_residual=True,
+                semi_conv=True,
+                hamming_window=True, # for regularization, do not matter really
+                feature_resample_norm=True,
+                **kwargs):
+        super().__init__()
+        self.scale_factor = scale_factor
+        self.lowpass_kernel = lowpass_kernel
+        self.highpass_kernel = highpass_kernel
+        self.up_group = up_group
+        self.encoder_kernel = encoder_kernel
+        self.encoder_dilation = encoder_dilation
+        self.compressed_channels = compressed_channels
+        self.hr_channel_compressor = nn.Conv2d(hr_channels, self.compressed_channels,1)
+        self.lr_channel_compressor = nn.Conv2d(lr_channels, self.compressed_channels,1)
+        self.content_encoder = nn.Conv2d( # ALPF generator
+            self.compressed_channels,
+            lowpass_kernel ** 2 * self.up_group * self.scale_factor * self.scale_factor,
+            self.encoder_kernel,
+            padding=int((self.encoder_kernel - 1) * self.encoder_dilation / 2),
+            dilation=self.encoder_dilation,
+            groups=1)
+        
+        self.align_corners = align_corners
+        self.upsample_mode = upsample_mode
+        self.hr_residual = hr_residual
+        self.use_high_pass = use_high_pass
+        self.use_low_pass = use_low_pass
+        self.semi_conv = semi_conv
+        self.feature_resample = feature_resample
+        self.comp_feat_upsample = comp_feat_upsample
+        if self.feature_resample:
+            self.dysampler = LocalSimGuidedSampler(in_channels=compressed_channels, scale=2, style='lp', groups=feature_resample_group, use_direct_scale=True, kernel_size=encoder_kernel, norm=feature_resample_norm)
+        if self.use_high_pass:
+            self.content_encoder2 = nn.Conv2d( # AHPF generator
+                self.compressed_channels,
+                highpass_kernel ** 2 * self.up_group * self.scale_factor * self.scale_factor,
+                self.encoder_kernel,
+                padding=int((self.encoder_kernel - 1) * self.encoder_dilation / 2),
+                dilation=self.encoder_dilation,
+                groups=1)
+        self.hamming_window = hamming_window
+        lowpass_pad=0
+        highpass_pad=0
+        if self.hamming_window:
+            self.register_buffer('hamming_lowpass', torch.FloatTensor(hamming2D(lowpass_kernel + 2 * lowpass_pad, lowpass_kernel + 2 * lowpass_pad))[None, None,])
+            self.register_buffer('hamming_highpass', torch.FloatTensor(hamming2D(highpass_kernel + 2 * highpass_pad, highpass_kernel + 2 * highpass_pad))[None, None,])
+        else:
+            self.register_buffer('hamming_lowpass', torch.FloatTensor([1.0]))
+            self.register_buffer('hamming_highpass', torch.FloatTensor([1.0]))
+        self.init_weights()
+
+    def init_weights(self):
+        for m in self.modules():
+            # print(m)
+            if isinstance(m, nn.Conv2d):
+                xavier_init(m, distribution='uniform')
+        normal_init(self.content_encoder, std=0.001)
+        if self.use_high_pass:
+            normal_init(self.content_encoder2, std=0.001)
+
+    def kernel_normalizer(self, mask, kernel, scale_factor=None, hamming=1):
+        if scale_factor is not None:
+            mask = F.pixel_shuffle(mask, self.scale_factor)
+        n, mask_c, h, w = mask.size()
+        mask_channel = int(mask_c / float(kernel**2)) # group
+        # mask = mask.view(n, mask_channel, -1, h, w)
+        # mask = F.softmax(mask, dim=2, dtype=mask.dtype)
+        # mask = mask.view(n, mask_c, h, w).contiguous()
+
+        mask = mask.view(n, mask_channel, -1, h, w)
+        mask = F.softmax(mask, dim=2, dtype=mask.dtype)
+        mask = mask.view(n, mask_channel, kernel, kernel, h, w)
+        mask = mask.permute(0, 1, 4, 5, 2, 3).view(n, -1, kernel, kernel)
+        # mask = F.pad(mask, pad=[padding] * 4, mode=self.padding_mode) # kernel + 2 * padding
+        mask = mask * hamming
+        mask /= mask.sum(dim=(-1, -2), keepdims=True)
+        # print(hamming)
+        # print(mask.shape)
+        mask = mask.view(n, mask_channel, h, w, -1)
+        mask =  mask.permute(0, 1, 4, 2, 3).view(n, -1, h, w).contiguous()
+        return mask
+
+    def forward(self, hr_feat, lr_feat, use_checkpoint=False): # use check_point to save GPU memory
+        if use_checkpoint:
+            return checkpoint(self._forward, hr_feat, lr_feat)
+        else:
+            return self._forward(hr_feat, lr_feat)
+
+    def _forward(self, hr_feat, lr_feat):
+        compressed_hr_feat = self.hr_channel_compressor(hr_feat)
+        compressed_lr_feat = self.lr_channel_compressor(lr_feat)
+        if self.semi_conv:
+            if self.comp_feat_upsample:
+                if self.use_high_pass:
+                    mask_hr_hr_feat = self.content_encoder2(compressed_hr_feat) #从hr_feat得到初始高通滤波特征
+                    mask_hr_init = self.kernel_normalizer(mask_hr_hr_feat, self.highpass_kernel, hamming=self.hamming_highpass) #kernel归一化得到初始高通滤波
+                    compressed_hr_feat = compressed_hr_feat + compressed_hr_feat - carafe(compressed_hr_feat, mask_hr_init, self.highpass_kernel, self.up_group, 1) #利用初始高通滤波对压缩hr_feat的高频增强 （x-x的低通结果=x的高通结果）
+                    
+                    mask_lr_hr_feat = self.content_encoder(compressed_hr_feat) #从hr_feat得到初始低通滤波特征
+                    mask_lr_init = self.kernel_normalizer(mask_lr_hr_feat, self.lowpass_kernel, hamming=self.hamming_lowpass) #kernel归一化得到初始低通滤波
+                    
+                    mask_lr_lr_feat_lr = self.content_encoder(compressed_lr_feat) #从hr_feat得到另一部分初始低通滤波特征
+                    mask_lr_lr_feat = F.interpolate( #利用初始低通滤波对另一部分初始低通滤波特征上采样
+                        carafe(mask_lr_lr_feat_lr, mask_lr_init, self.lowpass_kernel, self.up_group, 2), size=compressed_hr_feat.shape[-2:], mode='nearest')
+                    mask_lr = mask_lr_hr_feat + mask_lr_lr_feat #将两部分初始低通滤波特征合在一起
+
+                    mask_lr_init = self.kernel_normalizer(mask_lr, self.lowpass_kernel, hamming=self.hamming_lowpass) #得到初步融合的初始低通滤波
+                    mask_hr_lr_feat = F.interpolate( #使用初始低通滤波对lr_feat处理，分辨率得到提高
+                        carafe(self.content_encoder2(compressed_lr_feat), mask_lr_init, self.lowpass_kernel, self.up_group, 2), size=compressed_hr_feat.shape[-2:], mode='nearest')
+                    mask_hr = mask_hr_hr_feat + mask_hr_lr_feat # 最终高通滤波特征
+                else: raise NotImplementedError
+            else:
+                mask_lr = self.content_encoder(compressed_hr_feat) + F.interpolate(self.content_encoder(compressed_lr_feat), size=compressed_hr_feat.shape[-2:], mode='nearest')
+                if self.use_high_pass:
+                    mask_hr = self.content_encoder2(compressed_hr_feat) + F.interpolate(self.content_encoder2(compressed_lr_feat), size=compressed_hr_feat.shape[-2:], mode='nearest')
+        else:
+            compressed_x = F.interpolate(compressed_lr_feat, size=compressed_hr_feat.shape[-2:], mode='nearest') + compressed_hr_feat
+            mask_lr = self.content_encoder(compressed_x)
+            if self.use_high_pass: 
+                mask_hr = self.content_encoder2(compressed_x)
+        
+        mask_lr = self.kernel_normalizer(mask_lr, self.lowpass_kernel, hamming=self.hamming_lowpass)
+        if self.semi_conv:
+                lr_feat = carafe(lr_feat, mask_lr, self.lowpass_kernel, self.up_group, 2)
+        else:
+            lr_feat = resize(
+                input=lr_feat,
+                size=hr_feat.shape[2:],
+                mode=self.upsample_mode,
+                align_corners=None if self.upsample_mode == 'nearest' else self.align_corners)
+            lr_feat = carafe(lr_feat, mask_lr, self.lowpass_kernel, self.up_group, 1)
+
+        if self.use_high_pass:
+            mask_hr = self.kernel_normalizer(mask_hr, self.highpass_kernel, hamming=self.hamming_highpass)
+            hr_feat_hf = hr_feat - carafe(hr_feat, mask_hr, self.highpass_kernel, self.up_group, 1)
+            if self.hr_residual:
+                # print('using hr_residual')
+                hr_feat = hr_feat_hf + hr_feat
+            else:
+                hr_feat = hr_feat_hf
+
+        if self.feature_resample:
+            # print(lr_feat.shape)
+            lr_feat = self.dysampler(hr_x=compressed_hr_feat, 
+                                     lr_x=compressed_lr_feat, feat2sample=lr_feat)
+                
+        return  mask_lr, hr_feat, lr_feat
+
+
+
+class LocalSimGuidedSampler(nn.Module):
+    """
+    offset generator in FreqFusion
+    """
+    def __init__(self, in_channels, scale=2, style='lp', groups=4, use_direct_scale=True, kernel_size=1, local_window=3, sim_type='cos', norm=True, direction_feat='sim_concat'):
+        super().__init__()
+        assert scale==2
+        assert style=='lp'
+
+        self.scale = scale
+        self.style = style
+        self.groups = groups
+        self.local_window = local_window
+        self.sim_type = sim_type
+        self.direction_feat = direction_feat
+
+        if style == 'pl':
+            assert in_channels >= scale ** 2 and in_channels % scale ** 2 == 0
+        assert in_channels >= groups and in_channels % groups == 0
+
+        if style == 'pl':
+            in_channels = in_channels // scale ** 2
+            out_channels = 2 * groups
+        else:
+            out_channels = 2 * groups * scale ** 2
+        if self.direction_feat == 'sim':
+            self.offset = nn.Conv2d(local_window**2 - 1, out_channels, kernel_size=kernel_size, padding=kernel_size//2)
+        elif self.direction_feat == 'sim_concat':
+            self.offset = nn.Conv2d(in_channels + local_window**2 - 1, out_channels, kernel_size=kernel_size, padding=kernel_size//2)
+        else: raise NotImplementedError
+        normal_init(self.offset, std=0.001)
+        if use_direct_scale:
+            if self.direction_feat == 'sim':
+                self.direct_scale = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size//2)
+            elif self.direction_feat == 'sim_concat':
+                self.direct_scale = nn.Conv2d(in_channels + local_window**2 - 1, out_channels, kernel_size=kernel_size, padding=kernel_size//2)
+            else: raise NotImplementedError
+            constant_init(self.direct_scale, val=0.)
+
+        out_channels = 2 * groups
+        if self.direction_feat == 'sim':
+            self.hr_offset = nn.Conv2d(local_window**2 - 1, out_channels, kernel_size=kernel_size, padding=kernel_size//2)
+        elif self.direction_feat == 'sim_concat':
+            self.hr_offset = nn.Conv2d(in_channels + local_window**2 - 1, out_channels, kernel_size=kernel_size, padding=kernel_size//2)
+        else: raise NotImplementedError
+        normal_init(self.hr_offset, std=0.001)
+        
+        if use_direct_scale:
+            if self.direction_feat == 'sim':
+                self.hr_direct_scale = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size//2)
+            elif self.direction_feat == 'sim_concat':
+                self.hr_direct_scale = nn.Conv2d(in_channels + local_window**2 - 1, out_channels, kernel_size=kernel_size, padding=kernel_size//2)
+            else: raise NotImplementedError
+            constant_init(self.hr_direct_scale, val=0.)
+
+        self.norm = norm
+        if self.norm:
+            self.norm_hr = nn.GroupNorm(in_channels // 8, in_channels)
+            self.norm_lr = nn.GroupNorm(in_channels // 8, in_channels)
+        else:
+            self.norm_hr = nn.Identity()
+            self.norm_lr = nn.Identity()
+        self.register_buffer('init_pos', self._init_pos())
+
+    def _init_pos(self):
+        h = torch.arange((-self.scale + 1) / 2, (self.scale - 1) / 2 + 1) / self.scale
+        return torch.stack(torch.meshgrid([h, h])).transpose(1, 2).repeat(1, self.groups, 1).reshape(1, -1, 1, 1)
+    
+    def sample(self, x, offset, scale=None):
+        if scale is None: scale = self.scale
+        B, _, H, W = offset.shape
+        offset = offset.view(B, 2, -1, H, W)
+        coords_h = torch.arange(H) + 0.5
+        coords_w = torch.arange(W) + 0.5
+        coords = torch.stack(torch.meshgrid([coords_w, coords_h])
+                             ).transpose(1, 2).unsqueeze(1).unsqueeze(0).type(x.dtype).to(x.device)
+        normalizer = torch.tensor([W, H], dtype=x.dtype, device=x.device).view(1, 2, 1, 1, 1)
+        coords = 2 * (coords + offset) / normalizer - 1
+        coords = F.pixel_shuffle(coords.view(B, -1, H, W), scale).view(
+            B, 2, -1, scale * H, scale * W).permute(0, 2, 3, 4, 1).contiguous().flatten(0, 1)
+        return F.grid_sample(x.reshape(B * self.groups, -1, x.size(-2), x.size(-1)), coords, mode='bilinear',
+                             align_corners=False, padding_mode="border").view(B, -1, scale * H, scale * W)
+    
+    def forward(self, hr_x, lr_x, feat2sample):
+        hr_x = self.norm_hr(hr_x)
+        lr_x = self.norm_lr(lr_x)
+
+        if self.direction_feat == 'sim':
+            hr_sim = compute_similarity(hr_x, self.local_window, dilation=2, sim='cos')
+            lr_sim = compute_similarity(lr_x, self.local_window, dilation=2, sim='cos')
+        elif self.direction_feat == 'sim_concat':
+            hr_sim = torch.cat([hr_x, compute_similarity(hr_x, self.local_window, dilation=2, sim='cos')], dim=1)
+            lr_sim = torch.cat([lr_x, compute_similarity(lr_x, self.local_window, dilation=2, sim='cos')], dim=1)
+            hr_x, lr_x = hr_sim, lr_sim
+        # offset = self.get_offset(hr_x, lr_x)
+        offset = self.get_offset_lp(hr_x, lr_x, hr_sim, lr_sim)
+        return self.sample(feat2sample, offset)
+    
+    # def get_offset_lp(self, hr_x, lr_x):
+    def get_offset_lp(self, hr_x, lr_x, hr_sim, lr_sim):
+        if hasattr(self, 'direct_scale'):
+            # offset = (self.offset(lr_x) + F.pixel_unshuffle(self.hr_offset(hr_x), self.scale)) * (self.direct_scale(lr_x) + F.pixel_unshuffle(self.hr_direct_scale(hr_x), self.scale)).sigmoid() + self.init_pos
+            offset = (self.offset(lr_sim) + F.pixel_unshuffle(self.hr_offset(hr_sim), self.scale)) * (self.direct_scale(lr_x) + F.pixel_unshuffle(self.hr_direct_scale(hr_x), self.scale)).sigmoid() + self.init_pos
+            # offset = (self.offset(lr_sim) + F.pixel_unshuffle(self.hr_offset(hr_sim), self.scale)) * (self.direct_scale(lr_sim) + F.pixel_unshuffle(self.hr_direct_scale(hr_sim), self.scale)).sigmoid() + self.init_pos
+        else:
+            offset =  (self.offset(lr_x) + F.pixel_unshuffle(self.hr_offset(hr_x), self.scale)) * 0.25 + self.init_pos
+        return offset
+
+    def get_offset(self, hr_x, lr_x):
+        if self.style == 'pl':
+            raise NotImplementedError
+        return self.get_offset_lp(hr_x, lr_x)
+    
+
+def compute_similarity(input_tensor, k=3, dilation=1, sim='cos'):
+    """
+    计算输入张量中每一点与周围KxK范围内的点的余弦相似度。
+
+    参数：
+    - input_tensor: 输入张量，形状为[B, C, H, W]
+    - k: 范围大小，表示周围KxK范围内的点
+
+    返回：
+    - 输出张量，形状为[B, KxK-1, H, W]
+    """
+    B, C, H, W = input_tensor.shape
+    # 使用零填充来处理边界情况
+    # padded_input = F.pad(input_tensor, (k // 2, k // 2, k // 2, k // 2), mode='constant', value=0)
+
+    # 展平输入张量中每个点及其周围KxK范围内的点
+    unfold_tensor = F.unfold(input_tensor, k, padding=(k // 2) * dilation, dilation=dilation) # B, CxKxK, HW
+    # print(unfold_tensor.shape)
+    unfold_tensor = unfold_tensor.reshape(B, C, k**2, H, W)
+
+    # 计算余弦相似度
+    if sim == 'cos':
+        similarity = F.cosine_similarity(unfold_tensor[:, :, k * k // 2:k * k // 2 + 1], unfold_tensor[:, :, :], dim=1)
+    elif sim == 'dot':
+        similarity = unfold_tensor[:, :, k * k // 2:k * k // 2 + 1] * unfold_tensor[:, :, :]
+        similarity = similarity.sum(dim=1)
+    else:
+        raise NotImplementedError
+
+    # 移除中心点的余弦相似度，得到[KxK-1]的结果
+    similarity = torch.cat((similarity[:, :k * k // 2], similarity[:, k * k // 2 + 1:]), dim=1)
+
+    # 将结果重塑回[B, KxK-1, H, W]的形状
+    similarity = similarity.view(B, k * k - 1, H, W)
+    return similarity
+
+
+if __name__ == '__main__':
+    # x = torch.rand(4, 128, 16, 16)
+    # mask = torch.rand(4, 4 * 25, 16, 16)
+    # carafe(x, mask, kernel_size=5, group=1, up=2)
+
+    hr_feat = torch.rand(1, 128, 512, 512)
+    lr_feat = torch.rand(1, 128, 256, 256)
+    model = FreqFusion(hr_channels=128, lr_channels=128)
+    mask_lr, hr_feat, lr_feat = model(hr_feat=hr_feat, lr_feat=lr_feat)
+    print(mask_lr.shape)

model.py

@@ -0,0 +1,1418 @@
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+import torch.nn.functional as F
+from einops import rearrange, repeat
+import math
+from utils import grid_sample
+
+from freqfusion import FreqFusion
+
+#########################################
+
+class SepConv2d(torch.nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,act_layer=nn.ReLU):
+        super(SepConv2d, self).__init__()
+        self.depthwise = torch.nn.Conv2d(in_channels,
+                                         in_channels,
+                                         kernel_size=kernel_size,
+                                         stride=stride,
+                                         padding=padding,
+                                         dilation=dilation,
+                                         groups=in_channels)
+        self.pointwise = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1)
+        self.act_layer = act_layer() if act_layer is not None else nn.Identity()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+
+    def forward(self, x):
+        x = self.depthwise(x)
+        x = self.act_layer(x)
+        x = self.pointwise(x)
+        return x
+
+    def flops(self, H, W): 
+        flops = 0
+        flops += H*W*self.in_channels*self.kernel_size**2/self.stride**2
+        flops += H*W*self.in_channels*self.out_channels
+        return flops
+
+##########################################################################
+## Channel Attention Layer
+class CALayer(nn.Module):
+    def __init__(self, channel, reduction=16, bias=False):
+        super(CALayer, self).__init__()
+        # global average pooling: feature --> point
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        # feature channel downscale and upscale --> channel weight
+        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
+
+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)
+
+##########################################################################
+## Channel Attention Block (CAB)
+class CAB(nn.Module):
+    def __init__(self, n_feat, kernel_size, reduction, bias, act):
+        super(CAB, self).__init__()
+        modules_body = []
+        modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
+        modules_body.append(act)
+        modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
+
+        self.CA = CALayer(n_feat, reduction, bias=bias)
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x):
+        res = self.body(x)
+        res = self.CA(res)
+        res += x
+        return res
+
+#########################################
+######## Embedding for q,k,v ########
+class ConvProjection(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, kernel_size=3, q_stride=1, k_stride=1, v_stride=1, dropout = 0.,
+                 last_stage=False,bias=True):
+
+        super().__init__()
+
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        pad = (kernel_size - q_stride)//2
+        self.to_q = SepConv2d(dim, inner_dim, kernel_size, q_stride, pad, bias)
+        self.to_k = SepConv2d(dim, inner_dim, kernel_size, k_stride, pad, bias)
+        self.to_v = SepConv2d(dim, inner_dim, kernel_size, v_stride, pad, bias)
+
+    def forward(self, x, attn_kv=None):
+        b, n, c, h = *x.shape, self.heads
+        l = int(math.sqrt(n))
+        w = int(math.sqrt(n))
+
+        attn_kv = x if attn_kv is None else attn_kv
+        x = rearrange(x, 'b (l w) c -> b c l w', l=l, w=w)
+        attn_kv = rearrange(attn_kv, 'b (l w) c -> b c l w', l=l, w=w)
+        # print(attn_kv)
+        q = self.to_q(x)
+        q = rearrange(q, 'b (h d) l w -> b h (l w) d', h=h)
+        
+        k = self.to_k(attn_kv)
+        v = self.to_v(attn_kv)
+        k = rearrange(k, 'b (h d) l w -> b h (l w) d', h=h)
+        v = rearrange(v, 'b (h d) l w -> b h (l w) d', h=h)
+        return q,k,v    
+    
+    def flops(self, H, W): 
+        flops = 0
+        flops += self.to_q.flops(H, W)
+        flops += self.to_k.flops(H, W)
+        flops += self.to_v.flops(H, W)
+        return flops
+
+class LinearProjection(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0., bias=True):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.to_q = nn.Linear(dim, inner_dim, bias = bias)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = bias)
+        self.dim = dim
+        self.inner_dim = inner_dim
+
+    def forward(self, x, attn_kv=None):
+        B_, N, C = x.shape
+        attn_kv = x if attn_kv is None else attn_kv
+        q = self.to_q(x).reshape(B_, N, 1, self.heads, C // self.heads).permute(2, 0, 3, 1, 4).contiguous()
+        kv = self.to_kv(attn_kv).reshape(B_, N, 2, self.heads, C // self.heads).permute(2, 0, 3, 1, 4).contiguous()
+        q = q[0]
+        k, v = kv[0], kv[1] 
+        return q,k,v
+
+    def flops(self, H, W): 
+        flops = H*W*self.dim*self.inner_dim*3
+        return flops 
+
+class LinearProjection_Concat_kv(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0., bias=True):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = bias)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = bias)
+        self.dim = dim
+        self.inner_dim = inner_dim
+
+    def forward(self, x, attn_kv=None):
+        B_, N, C = x.shape
+        attn_kv = x if attn_kv is None else attn_kv
+        qkv_dec = self.to_qkv(x).reshape(B_, N, 3, self.heads, C // self.heads).permute(2, 0, 3, 1, 4).contiguous()
+        kv_enc = self.to_kv(attn_kv).reshape(B_, N, 2, self.heads, C // self.heads).permute(2, 0, 3, 1, 4).contiguous()
+        q, k_d, v_d = qkv_dec[0], qkv_dec[1], qkv_dec[2]  # make torchscript happy (cannot use tensor as tuple)
+        k_e, v_e = kv_enc[0], kv_enc[1] 
+        k = torch.cat((k_d,k_e),dim=2)
+        v = torch.cat((v_d,v_e),dim=2)
+        return q,k,v
+
+    def flops(self, H, W): 
+        flops = H*W*self.dim*self.inner_dim*5
+        return flops 
+
+#########################################
+
+########### SIA #############
+class WindowAttention(nn.Module):
+    def __init__(self, dim, win_size, num_heads, token_projection='linear', qkv_bias=True, qk_scale=None, attn_drop=0.,
+                 proj_drop=0., se_layer=False):
+
+        super().__init__()
+        self.dim = dim
+        self.win_size = win_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * win_size[0] - 1) * (2 * win_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.win_size[0])  # [0,...,Wh-1]
+        coords_w = torch.arange(self.win_size[1])  # [0,...,Ww-1]
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.win_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.win_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.win_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        # self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        if token_projection == 'conv':
+            self.qkv = ConvProjection(dim, num_heads, dim // num_heads, bias=qkv_bias)
+        elif token_projection == 'linear_concat':
+            self.qkv = LinearProjection_Concat_kv(dim, num_heads, dim // num_heads, bias=qkv_bias)
+        else:
+            self.qkv = LinearProjection(dim, num_heads, dim // num_heads, bias=qkv_bias)
+
+        self.token_projection = token_projection
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.ll = nn.Identity()
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.sigmoid = nn.Sigmoid()
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, dino_mat, point_feature, normal, attn_kv=None, mask=None):
+        B_, N, C = x.shape
+
+        dino_mat = dino_mat.unsqueeze(2)
+        normalizer = torch.sqrt((dino_mat @ dino_mat.transpose(-2, -1)).squeeze(-2)).detach()
+        normalizer = torch.clamp(normalizer, 1.0e-20, 1.0e10)
+        dino_mat = dino_mat.squeeze(2) / normalizer 
+        dino_mat_correlation_map = dino_mat @ dino_mat.transpose(-2, -1).contiguous()
+        dino_mat_correlation_map = torch.clamp(dino_mat_correlation_map, 0.0, 1.0e10)
+        dino_mat_correlation_map = torch.unsqueeze(dino_mat_correlation_map, dim=1)
+
+        point_feature = point_feature.unsqueeze(2)
+        Point = point_feature.repeat(1, 1, self.win_size[0] * self.win_size[1],1)
+        Point = Point - Point.transpose(-2, -3)
+        normal = normal.unsqueeze(2).repeat(1,1,self.win_size[0] * self.win_size[1],1)
+        # print(f'{Point.shape=}')
+        # print(f'{normal.shape=}')
+        Point = Point * normal
+        Point = torch.abs(torch.sum(Point, dim=3))
+
+        plane_correlation_map = 0.5 * (Point + Point.transpose(-1, -2))
+        plane_correlation_map = plane_correlation_map.unsqueeze(1)
+        plane_correlation_map = torch.exp(-plane_correlation_map)
+
+
+        q, k, v = self.qkv(x, attn_kv)
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1)) 
+
+        attn = dino_mat_correlation_map * attn
+        attn = plane_correlation_map * attn 
+
+
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+            self.win_size[0] * self.win_size[1], self.win_size[0] * self.win_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        ratio = attn.size(-1) // relative_position_bias.size(-1)
+        relative_position_bias = repeat(relative_position_bias, 'nH l c -> nH l (c d)', d=ratio)
+
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            mask = repeat(mask, 'nW m n -> nW m (n d)', d=ratio)
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N * ratio) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N * ratio)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+        x = (attn @ v).transpose(1, 2).contiguous().reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.ll(x)
+        x = self.proj_drop(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, win_size={self.win_size}, num_heads={self.num_heads}'
+
+    def flops(self, H, W):
+        # calculate flops for 1 window with token length of N
+        # print(N, self.dim)
+        flops = 0
+        N = self.win_size[0] * self.win_size[1]
+        nW = H * W / N
+        # qkv = self.qkv(x)
+        # flops += N * self.dim * 3 * self.dim
+        flops += self.qkv.flops(H, W)
+        # attn = (q @ k.transpose(-2, -1))
+        if self.token_projection != 'linear_concat':
+            flops += nW * self.num_heads * N * (self.dim // self.num_heads) * N
+            #  x = (attn @ v)
+            flops += nW * self.num_heads * N * N * (self.dim // self.num_heads)
+        else:
+            flops += nW * self.num_heads * N * (self.dim // self.num_heads) * N * 2
+            #  x = (attn @ v)
+            flops += nW * self.num_heads * N * N * 2 * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += nW * N * self.dim * self.dim
+        print("W-MSA:{%.2f}" % (flops / 1e9))
+        return flops
+
+
+#########################################
+########### feed-forward network #############
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+        self.in_features = in_features
+        self.hidden_features = hidden_features
+        self.out_features = out_features
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+    def flops(self, H, W):
+        flops = 0
+        # fc1
+        flops += H*W*self.in_features*self.hidden_features 
+        # fc2
+        flops += H*W*self.hidden_features*self.out_features
+        print("MLP:{%.2f}"%(flops/1e9))
+        return flops
+
+class LeFF(nn.Module):
+    def __init__(self, dim=32, hidden_dim=128, act_layer=nn.GELU,drop = 0.):
+        super().__init__()
+        self.linear1 = nn.Sequential(nn.Linear(dim, hidden_dim),
+                                act_layer())
+        self.dwconv = nn.Sequential(nn.Conv2d(hidden_dim,hidden_dim,groups=hidden_dim,kernel_size=3,stride=1,padding=1),
+                        act_layer())
+        self.linear2 = nn.Sequential(nn.Linear(hidden_dim, dim))
+        self.dim = dim
+        self.hidden_dim = hidden_dim
+
+    def forward(self, x, img_size=(128,128)):
+        # bs x hw x c
+        bs, hw, c = x.size()
+        # hh = int(math.sqrt(hw))
+        hh = img_size[0]
+        ww = img_size[1]
+
+        x = self.linear1(x)
+
+        # spatial restore
+        x = rearrange(x, ' b (h w) (c) -> b c h w ', h = hh, w = ww)
+        # bs,hidden_dim,32x32
+
+        x = self.dwconv(x)
+
+        # flaten
+        x = rearrange(x, ' b c h w -> b (h w) c', h = hh, w = ww)
+
+        x = self.linear2(x)
+
+        return x
+
+    def flops(self, H, W):
+        flops = 0
+        # fc1
+        flops += H*W*self.dim*self.hidden_dim 
+        # dwconv
+        flops += H*W*self.hidden_dim*3*3
+        # fc2
+        flops += H*W*self.hidden_dim*self.dim
+        print("LeFF:{%.2f}"%(flops/1e9))
+        return flops
+
+#########################################
+########### window operation#############
+def window_partition(x, win_size, dilation_rate=1):
+    B, H, W, C = x.shape
+    if dilation_rate !=1:
+        x = x.permute(0,3,1,2).contiguous() # B, C, H, W
+        assert type(dilation_rate) is int, 'dilation_rate should be a int'
+        x = F.unfold(x, kernel_size=win_size,dilation=dilation_rate,padding=4*(dilation_rate-1),stride=win_size) # B, C*Wh*Ww, H/Wh*W/Ww
+        windows = x.permute(0,2,1).contiguous().view(-1, C, win_size, win_size) # B' ,C ,Wh ,Ww
+        windows = windows.permute(0,2,3,1).contiguous() # B' ,Wh ,Ww ,C
+    else:
+        x = x.view(B, H // win_size, win_size, W // win_size, win_size, C)
+        windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, win_size, win_size, C) # B' ,Wh ,Ww ,C
+    return windows
+
+def window_reverse(windows, win_size, H, W, dilation_rate=1):
+    # B' ,Wh ,Ww ,C
+    B = int(windows.shape[0] / (H * W / win_size / win_size))
+    x = windows.view(B, H // win_size, W // win_size, win_size, win_size, -1)
+    if dilation_rate !=1:
+        x = windows.permute(0,5,3,4,1,2).contiguous() # B, C*Wh*Ww, H/Wh*W/Ww
+        x = F.fold(x, (H, W), kernel_size=win_size, dilation=dilation_rate, padding=4*(dilation_rate-1),stride=win_size)
+    else:
+        x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+#########################################
+# Downsample Block
+class Downsample(nn.Module):
+    def __init__(self, in_channel, out_channel):
+        super(Downsample, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channel, out_channel, kernel_size=4, stride=2, padding=1)
+            # nn.Conv2d(in_channel * 4, out_channel, kernel_size=3, padding=1)
+        )
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+
+    def forward(self, x, img_size=(128,128)):
+        B, L, C = x.shape
+        H = img_size[0]
+        W = img_size[1]
+        x = x.transpose(1, 2).contiguous().view(B, C, H, W)
+
+        out = self.conv(x).flatten(2).transpose(1,2).contiguous()  # B H*W C
+        return out
+
+    # def forward(self, x, img_size=(128,128)):
+    #     B, L, C = x.shape
+    #     H = img_size[0]
+    #     W = img_size[1]
+    #     x = x.transpose(1, 2).contiguous().view(B, C, H, W)
+
+    #     # new add
+    #     x = x.permute(0,2,3,1)
+
+    #     x0 = x[:, 0::2, 0::2, :]
+    #     x1 = x[:, 0::2, 1::2, :]
+    #     x2 = x[:, 1::2, 0::2, :]
+    #     x3 = x[:, 1::2, 1::2, :]
+    #     x = torch.cat([x0, x1, x2, x3], axis=-1)
+    #     x = x.permute(0,3,1,2)
+
+    #     out = self.conv(x).flatten(2).transpose(1,2).contiguous()  # B H*W C
+    #     return out
+
+    def flops(self, H, W):
+        flops = 0
+        # conv
+        flops += H/2*W/2*self.in_channel*self.out_channel*4*4
+        print("Downsample:{%.2f}"%(flops/1e9))
+        return flops
+
+# Upsample Block
+class Upsample(nn.Module):
+    def __init__(self, in_channel, out_channel):
+        super(Upsample, self).__init__()
+        self.deconv = nn.Sequential(
+            nn.ConvTranspose2d(in_channel, out_channel, kernel_size=2, stride=2)
+        )
+
+        # self.conv = nn.Sequential(
+        #     nn.Conv2d(in_channel, out_channel * 4, kernel_size=3, padding=1)
+        # )
+
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+
+    def forward(self, x, img_size=(128,128)):
+        B, L, C = x.shape
+        H = img_size[0]
+        W = img_size[1]
+        x = x.transpose(1, 2).contiguous().view(B, C, H, W)
+        out = self.deconv(x)
+
+        out = out.flatten(2).transpose(1,2).contiguous() # B H*W C
+        return out
+
+    # def forward(self, x, img_size=(128,128)):
+    #     B, L, C = x.shape
+    #     H = img_size[0]
+    #     W = img_size[1]
+    #     x = x.transpose(1, 2).contiguous().view(B, C, H, W)
+    #     out = self.conv(x)
+    #     # new add
+    #     pixel_shuffle = nn.PixelShuffle(2)
+    #     out = pixel_shuffle(out)
+
+    #     out = out.flatten(2).transpose(1,2).contiguous() # B H*W C
+    #     return out
+
+    def flops(self, H, W):
+        flops = 0
+        # conv
+        flops += H*2*W*2*self.in_channel*self.out_channel*2*2 
+        print("Upsample:{%.2f}"%(flops/1e9))
+        return flops
+
+# Input Projection
+class InputProj(nn.Module):
+    def __init__(self, in_channel=3, out_channel=64, kernel_size=3, stride=1, norm_layer=None,act_layer=nn.LeakyReLU):
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=stride, padding=kernel_size//2),
+            act_layer(inplace=True)
+        )
+        if norm_layer is not None:
+            self.norm = norm_layer(out_channel)
+        else:
+            self.norm = None
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2).contiguous()  # B H*W C
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+    def flops(self, H, W):
+        flops = 0
+        # conv
+        flops += H*W*self.in_channel*self.out_channel*3*3
+
+        if self.norm is not None:
+            flops += H*W*self.out_channel 
+        print("Input_proj:{%.2f}"%(flops/1e9))
+        return flops
+
+# Output Projection
+class OutputProj(nn.Module):
+    def __init__(self, in_channel=64, out_channel=3, kernel_size=3, stride=1, norm_layer=None,act_layer=None):
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=stride, padding=kernel_size//2),
+        )
+        if act_layer is not None:
+            self.proj.add_module(act_layer(inplace=True))
+        if norm_layer is not None:
+            self.norm = norm_layer(out_channel)
+        else:
+            self.norm = None
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+
+    def forward(self, x, img_size=(128,128)):
+        B, L, C = x.shape
+        H = img_size[0]
+        W = img_size[1]
+        # H = int(math.sqrt(L))
+        # W = int(math.sqrt(L))
+        x = x.transpose(1, 2).contiguous().view(B, C, H, W)
+        x = self.proj(x)
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+    def flops(self, H, W):
+        flops = 0
+        # conv
+        flops += H*W*self.in_channel*self.out_channel*3*3
+
+        if self.norm is not None:
+            flops += H*W*self.out_channel 
+        print("Output_proj:{%.2f}"%(flops/1e9))
+        return flops
+
+
+#########################################
+########### CA Transformer #############
+class CATransformerBlock(nn.Module):
+    def __init__(self, dim, input_resolution, num_heads, win_size=10, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, token_projection='linear', token_mlp='leff',
+                 se_layer=False):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.win_size = win_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        self.token_mlp = token_mlp
+        if min(self.input_resolution) <= self.win_size:
+            self.shift_size = 0
+            self.win_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.win_size, "shift_size must in 0-win_size"
+
+        self.norm1 = norm_layer(dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer,
+                       drop=drop) if token_mlp == 'ffn' else LeFF(dim, mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.CAB = CAB(dim, kernel_size=3, reduction=4, bias=False, act=nn.PReLU())
+
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"win_size={self.win_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
+
+    def forward(self, x, dino_mat, point, normal, mask=None, img_size=(128, 128)):
+        B, L, C = x.shape
+        H = img_size[0]
+        W = img_size[1]
+        assert L == W * H, \
+            f"Input image size ({H}*{W} doesn't match model ({L})."
+
+        shortcut = x
+        x = self.norm1(x)
+
+        # spatial restore
+        x = rearrange(x, ' b (h w) (c) -> b c h w ', h=H, w=W)
+        # bs,hidden_dim,32x32
+
+        x = self.CAB(x)
+
+        # flaten
+        x = rearrange(x, ' b c h w -> b (h w) c', h=H, w=W)
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x), img_size=img_size))
+        
+        return x
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # W-MSA/SW-MSA
+        flops += self.attn.flops(H, W)
+        # norm2
+        flops += self.dim * H * W
+        # mlp
+        flops += self.mlp.flops(H, W)
+        print("LeWin:{%.2f}" % (flops / 1e9))
+        return flops
+
+#########################################
+########### SIM Transformer #############
+class SIMTransformerBlock(nn.Module):
+    def __init__(self, dim, input_resolution, num_heads, win_size=10, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm,token_projection='linear',token_mlp='leff',se_layer=False):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.win_size = win_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        self.token_mlp = token_mlp
+        if min(self.input_resolution) <= self.win_size:
+            self.shift_size = 0
+            self.win_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.win_size, "shift_size must in 0-win_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, win_size=to_2tuple(self.win_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,
+            token_projection=token_projection,se_layer=se_layer)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,act_layer=act_layer, drop=drop) if token_mlp=='ffn' else LeFF(dim,mlp_hidden_dim,act_layer=act_layer, drop=drop)
+        self.CAB = CAB(dim, kernel_size=3, reduction=4, bias=False, act=nn.PReLU())
+
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"win_size={self.win_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
+
+    def forward(self, x, dino_mat, point, normal, mask=None, img_size = (128, 128)):
+        B, L, C = x.shape
+        H = img_size[0]
+        W = img_size[1]
+        assert L == W * H, \
+            f"Input image size ({H}*{W} doesn't match model ({L})."
+
+        C_dino_mat = dino_mat.shape[1]
+        C_point = point.shape[1]
+        C_normal = normal.shape[1]
+
+        if mask != None:
+            input_mask = F.interpolate(mask, size=(H,W)).permute(0,2,3,1).contiguous()
+            input_mask_windows = window_partition(input_mask, self.win_size) # nW, win_size, win_size, 1
+            attn_mask = input_mask_windows.view(-1, self.win_size * self.win_size) # nW, win_size*win_size
+            attn_mask = attn_mask.unsqueeze(2)*attn_mask.unsqueeze(1) # nW, win_size*win_size, win_size*win_size
+            attn_mask = attn_mask.masked_fill(attn_mask!=0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        else:
+            attn_mask = None
+
+        ## shift mask
+        if self.shift_size > 0:
+            # calculate attention mask for SW-MSA
+            shift_mask = torch.zeros((1, H, W, 1)).type_as(x)
+            h_slices = (slice(0, -self.win_size),
+                        slice(-self.win_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            w_slices = (slice(0, -self.win_size),
+                        slice(-self.win_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            cnt = 0
+            for h in h_slices:
+                for w in w_slices:
+                    shift_mask[:, h, w, :] = cnt
+                    cnt += 1
+            shift_mask_windows = window_partition(shift_mask, self.win_size)  # nW, win_size, win_size, 1
+            shift_mask_windows = shift_mask_windows.view(-1, self.win_size * self.win_size) # nW, win_size*win_size
+            shift_attn_mask = shift_mask_windows.unsqueeze(1) - shift_mask_windows.unsqueeze(2) # nW, win_size*win_size, win_size*win_size
+            shift_attn_mask = shift_attn_mask.masked_fill(shift_attn_mask != 0, float(-100.0)).masked_fill(shift_attn_mask == 0, float(0.0))
+            attn_mask = attn_mask + shift_attn_mask if attn_mask is not None else shift_attn_mask
+
+        shortcut = x
+        x = self.norm1(x)
+
+
+        x = x.view(B, H, W, C)
+        dino_mat = dino_mat.permute(0, 2, 3, 1).contiguous()
+        point = point.permute(0, 2, 3, 1).contiguous()
+        normal = normal.permute(0, 2, 3, 1).contiguous()
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            shifted_dino_mat = torch.roll(dino_mat, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            shifted_point = torch.roll(point, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            shifted_normal = torch.roll(normal, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+        else:
+            shifted_x = x
+            shifted_dino_mat = dino_mat
+            shifted_point = point
+            shifted_normal = normal
+
+        x_windows = window_partition(shifted_x, self.win_size)  # nW*B, win_size, win_size, C  N*C->C
+        x_windows = x_windows.view(-1, self.win_size * self.win_size, C)  # nW*B, win_size*win_size, C
+
+
+        dino_mat_windows = window_partition(shifted_dino_mat, self.win_size)  # nW*B, win_size, win_size, C  N*C->C
+        dino_mat_windows = dino_mat_windows.view(-1, self.win_size * self.win_size, C_dino_mat)  # nW*B, win_size*win_size, C
+
+        point_windows = window_partition(shifted_point, self.win_size)  # nW*B, win_size, win_size, C  N*C->C
+        point_windows = point_windows.view(-1, self.win_size * self.win_size, C_point)  # nW*B, win_size*win_size, C
+
+        normal_windows = window_partition(shifted_normal, self.win_size)  # nW*B, win_size, win_size, C  N*C->C
+        normal_windows = normal_windows.view(-1, self.win_size * self.win_size, C_normal)  # nW*B, win_size*win_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, dino_mat_windows, point_windows, normal_windows, mask=attn_mask)  # nW*B, win_size*win_size, C
+
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.win_size, self.win_size, C)
+
+
+        shifted_x = window_reverse(attn_windows, self.win_size, H, W)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+        x = x.view(B, H * W, C)
+
+        x = rearrange(x, ' b (h w) (c) -> b c h w ', h=H, w=W)
+        # bs,hidden_dim,32x32
+
+        x = self.CAB(x)
+
+        x = rearrange(x, ' b c h w -> b (h w) c', h=H, w=W)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x), img_size=img_size))
+        del attn_mask
+        return x
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # W-MSA/SW-MSA
+        flops += self.attn.flops(H, W)
+        # norm2
+        flops += self.dim * H * W
+        # mlp
+        flops += self.mlp.flops(H,W)
+        print("LeWin:{%.2f}"%(flops/1e9))
+        return flops
+
+
+#########################################
+########### Basic layer of ShadowFormer ################
+class BasicShadowFormer(nn.Module):
+    def __init__(self, dim, output_dim, input_resolution, depth, num_heads, win_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, use_checkpoint=False,
+                 token_projection='linear',token_mlp='ffn',se_layer=False,cab=False):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+        # build blocks
+        if cab:
+            self.blocks = nn.ModuleList([
+                CATransformerBlock(dim=dim, input_resolution=input_resolution,
+                                      num_heads=num_heads, win_size=win_size,
+                                      shift_size=0 if (i % 2 == 0) else win_size // 2,
+                                      mlp_ratio=mlp_ratio,
+                                      qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                      drop=drop, attn_drop=attn_drop,
+                                      drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                      norm_layer=norm_layer, token_projection=token_projection, token_mlp=token_mlp,
+                                      se_layer=se_layer)
+                for i in range(depth)])
+        else:
+            self.blocks = nn.ModuleList([
+                SIMTransformerBlock(dim=dim, input_resolution=input_resolution,
+                                     num_heads=num_heads, win_size=win_size,
+                                     shift_size=0 if (i % 2 == 0) else win_size // 2,
+                                     mlp_ratio=mlp_ratio,
+                                     qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                     drop=drop, attn_drop=attn_drop,
+                                     drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                     norm_layer=norm_layer,token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer)
+                for i in range(depth)])
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"    
+
+    def forward(self, x, dino_mat=None, point=None, normal=None, mask=None, img_size=(128,128)):
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x, dino_mat, point, normal, mask, img_size)
+        return x
+
+    def flops(self):
+        flops = 0
+        for blk in self.blocks:
+            flops += blk.flops()
+        return flops
+
+class ShadowFormer(nn.Module):
+    def __init__(self, img_size=256, in_chans=3,
+                embed_dim=32, depths=[2, 2, 2, 2, 2, 2, 2, 2, 2], num_heads=[1, 2, 4, 8, 16, 16, 8, 4, 2],
+                win_size=8, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                norm_layer=nn.LayerNorm, patch_norm=True,
+                use_checkpoint=False, token_projection='linear', token_mlp='leff', se_layer=True,
+                dowsample=Downsample, upsample=Upsample, **kwargs):
+        super().__init__()
+
+        self.num_enc_layers = len(depths)//2
+        self.num_dec_layers = len(depths)//2
+        self.embed_dim = embed_dim
+        self.patch_norm = patch_norm
+        self.mlp_ratio = mlp_ratio
+        self.token_projection = token_projection
+        self.mlp = token_mlp
+        self.win_size =win_size
+        self.reso = img_size
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        self.DINO_channel = 1024
+
+        # stochastic depth
+        enc_dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths[:self.num_enc_layers]))]
+        conv_dpr = [drop_path_rate]*depths[4]
+        dec_dpr = enc_dpr[::-1]
+
+        # build layers
+
+        # Input/Output
+        self.input_proj = InputProj(in_channel=4, out_channel=embed_dim, kernel_size=3, stride=1, act_layer=nn.LeakyReLU)
+        self.output_proj = OutputProj(in_channel=2*embed_dim, out_channel=in_chans, kernel_size=3, stride=1)
+
+        # Encoder
+        self.encoderlayer_0 = BasicShadowFormer(dim=embed_dim,
+                            output_dim=embed_dim,
+                            input_resolution=(img_size,
+                                                img_size),
+                            depth=depths[0],
+                            num_heads=num_heads[0],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=enc_dpr[sum(depths[:0]):sum(depths[:1])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer,cab=True)
+        self.dowsample_0 = dowsample(embed_dim, embed_dim*2)
+        self.encoderlayer_1 = BasicShadowFormer(dim=embed_dim*2,
+                            output_dim=embed_dim*2,
+                            input_resolution=(img_size // 2,
+                                                img_size // 2),
+                            depth=depths[1],
+                            num_heads=num_heads[1],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=enc_dpr[sum(depths[:1]):sum(depths[:2])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer, cab=True)
+        self.dowsample_1 = dowsample(embed_dim*2, embed_dim*4)
+        self.encoderlayer_2 = BasicShadowFormer(dim=embed_dim*4,
+                            output_dim=embed_dim*4,
+                            input_resolution=(img_size // (2 ** 2),
+                                                img_size // (2 ** 2)),
+                            depth=depths[2],
+                            num_heads=num_heads[2],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=enc_dpr[sum(depths[:2]):sum(depths[:3])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer)
+        self.dowsample_2 = dowsample(embed_dim*4, embed_dim*8)
+
+        # Bottleneck
+        channel_conv = embed_dim*16 
+        # channel_conv = embed_dim*8 if self.add_shadow_detect_dino_conact else embed_dim*4
+        self.conv = BasicShadowFormer(dim=channel_conv,
+                            output_dim=channel_conv,
+                            input_resolution=(img_size // (2 ** 3),
+                                                img_size // (2 ** 3)),
+                            depth=depths[4],
+                            num_heads=num_heads[4],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=conv_dpr,
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer)
+
+        # # Decoder
+        self.upsample_0 = upsample(channel_conv, embed_dim*4)
+        channel_0 =  embed_dim*8
+        self.decoderlayer_0 = BasicShadowFormer(dim=channel_0,
+                            output_dim=channel_0,
+                            input_resolution=(img_size // (2 ** 2),
+                                                img_size // (2 ** 2)),
+                            depth=depths[6],
+                            num_heads=num_heads[6],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=dec_dpr[sum(depths[5:6]):sum(depths[5:7])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer)
+        self.upsample_1 = upsample(channel_0, embed_dim*2)
+        channel_1 =  embed_dim*4
+        self.decoderlayer_1 = BasicShadowFormer(dim=channel_1,
+                            output_dim=channel_1,
+                            input_resolution=(img_size // 2,
+                                                img_size // 2),
+                            depth=depths[7],
+                            num_heads=num_heads[7],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=dec_dpr[sum(depths[5:7]):sum(depths[5:8])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer, cab=True)
+        self.upsample_2 = upsample(channel_1, embed_dim)
+        self.decoderlayer_2 = BasicShadowFormer(dim=embed_dim*2,
+                            output_dim=embed_dim*2,
+                            input_resolution=(img_size,
+                                                img_size),
+                            depth=depths[8],
+                            num_heads=num_heads[8],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=dec_dpr[sum(depths[5:8]):sum(depths[5:9])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer,cab=True)
+
+        self.Conv = nn.Conv2d(self.DINO_channel * 4, embed_dim * 8, kernel_size=1)
+        self.relu = nn.LeakyReLU()
+        self.apply(self._init_weights)
+        
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def extra_repr(self) -> str:
+        return f"embed_dim={self.embed_dim}, token_projection={self.token_projection}, token_mlp={self.mlp}, win_size={self.win_size}"
+
+    def forward(self, x, DINO_Mat_features=None, point=None, normal=None, mask=None):
+        point_feature=None
+        dino_mat =None
+        dino_mat1=None
+
+        self.img_size = torch.tensor((x.shape[2], x.shape[3]))
+        point_feature1 = grid_sample(point, self.img_size // 2)
+        point_feature2 = grid_sample(point, self.img_size // 4)
+        point_feature3 = grid_sample(point, self.img_size // 8)
+        normal1= grid_sample(normal, self.img_size // 2)
+        normal2= grid_sample(normal, self.img_size // 4)
+        normal3= grid_sample(normal, self.img_size // 8)
+
+        patch_features_0 = DINO_Mat_features[0]
+        patch_features_1 = DINO_Mat_features[1]
+        patch_features_2 = DINO_Mat_features[2]
+        patch_features_3 = DINO_Mat_features[3]
+        patch_feature_all = torch.cat((patch_features_0, patch_features_1,
+                                    patch_features_2, patch_features_3), dim=1)
+
+        # Get concatenate DINO Feature
+        dino_mat_cat = self.Conv(patch_feature_all)
+        dino_mat_cat = self.relu(dino_mat_cat)
+        B, C, W, H = dino_mat_cat.shape
+        dino_mat_cat_flat = dino_mat_cat.view(B, C, W * H).permute(0,2,1)
+
+
+        dino_mat2 = F.upsample_bilinear(DINO_Mat_features[-1], scale_factor=2)
+        dino_mat3 = DINO_Mat_features[-1]
+            
+        # RGBD
+        xi = torch.cat((x, point[:,2,:].unsqueeze(1)), dim=1)
+
+        y = self.input_proj(xi)
+        y = self.pos_drop(y)
+
+        # Encoder
+        self.img_size = (int(self.img_size[0]), int(self.img_size[1]))
+        conv0 = self.encoderlayer_0(y, dino_mat, point_feature, normal, mask, img_size = self.img_size)
+        pool0 = self.dowsample_0(conv0, img_size = self.img_size)
+
+        self.img_size = (int(self.img_size[0]/2), int(self.img_size[1]/2))
+        conv1 = self.encoderlayer_1(pool0, dino_mat1, point_feature1, normal1, img_size = self.img_size)
+        pool1 = self.dowsample_1(conv1, img_size = self.img_size)
+
+        self.img_size = (int(self.img_size[0] / 2), int(self.img_size[1] / 2))
+        conv2 = self.encoderlayer_2(pool1, dino_mat2, point_feature2, normal2, img_size = self.img_size)
+        pool2 = self.dowsample_2(conv2, img_size = self.img_size)
+
+        # Bottleneck
+        self.img_size = (int(self.img_size[0] / 2), int(self.img_size[1] / 2))
+        pool2 = torch.cat([pool2, dino_mat_cat_flat],-1)
+        conv3 = self.conv(pool2, dino_mat3, point_feature3, normal3, img_size = self.img_size)
+        print(f'{conv3.shape=}')
+
+        #Decoder
+        up0 = self.upsample_0(conv3, img_size = self.img_size)
+        self.img_size = (int(self.img_size[0] * 2), int(self.img_size[1] * 2))
+        print(f'{conv2.shape=}, {up0.shape=}')  # conv2.shape=torch.Size([1, 4096, 128]), up0.shape=torch.Size([1, 4096, 128])
+        deconv0 = torch.cat([up0,conv2],-1)   
+        deconv0 = self.decoderlayer_0(deconv0, dino_mat2, point_feature2, normal2, img_size = self.img_size)
+        print(f'{deconv0.shape=}')
+
+        up1 = self.upsample_1(deconv0, img_size = self.img_size)
+        self.img_size = (int(self.img_size[0] * 2), int(self.img_size[1] * 2))
+        print(f'{conv1.shape=}, {up1.shape=}')  # conv1.shape=torch.Size([1, 16384, 64]), up1.shape=torch.Size([1, 16384, 64])
+        deconv1 = torch.cat([up1,conv1],-1)
+        deconv1 = self.decoderlayer_1(deconv1, dino_mat1, point_feature1, normal1, img_size = self.img_size)
+        print(f'{deconv1.shape=}')
+
+        up2 = self.upsample_2(deconv1, img_size = self.img_size)
+        self.img_size = (int(self.img_size[0] * 2), int(self.img_size[1] * 2))
+        print(f'{conv0.shape=}, {up2.shape=}')  # conv0.shape=torch.Size([1, 65536, 32]), up2.shape=torch.Size([1, 65536, 32])
+        deconv2 = torch.cat([up2,conv0],-1)
+        deconv2 = self.decoderlayer_2(deconv2, dino_mat, point_feature, normal, mask, img_size = self.img_size)
+
+        # Output Projection
+        y = self.output_proj(deconv2, img_size = self.img_size) + x
+        return y
+
+
+
+class ShadowFormerFreq(nn.Module):
+    def __init__(self, img_size=256, in_chans=3,
+                embed_dim=32, depths=[2, 2, 2, 2, 2, 2, 2, 2, 2], num_heads=[1, 2, 4, 8, 16, 16, 8, 4, 2],
+                win_size=8, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                norm_layer=nn.LayerNorm, patch_norm=True,
+                use_checkpoint=False, token_projection='linear', token_mlp='leff', se_layer=True,
+                dowsample=Downsample, upsample=Upsample, **kwargs):
+        super().__init__()
+
+        self.num_enc_layers = len(depths)//2
+        self.num_dec_layers = len(depths)//2
+        self.embed_dim = embed_dim
+        self.patch_norm = patch_norm
+        self.mlp_ratio = mlp_ratio
+        self.token_projection = token_projection
+        self.mlp = token_mlp
+        self.win_size =win_size
+        self.reso = img_size
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        self.DINO_channel = 1024
+
+        # stochastic depth
+        enc_dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths[:self.num_enc_layers]))]
+        conv_dpr = [drop_path_rate]*depths[4]
+        dec_dpr = enc_dpr[::-1]
+
+        # build layers
+
+        # Input/Output
+        self.input_proj = InputProj(in_channel=4, out_channel=embed_dim, kernel_size=3, stride=1, act_layer=nn.LeakyReLU)
+        self.output_proj = OutputProj(in_channel=2*embed_dim, out_channel=in_chans, kernel_size=3, stride=1)
+
+        # Encoder
+        self.encoderlayer_0 = BasicShadowFormer(dim=embed_dim,
+                            output_dim=embed_dim,
+                            input_resolution=(img_size,
+                                                img_size),
+                            depth=depths[0],
+                            num_heads=num_heads[0],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=enc_dpr[sum(depths[:0]):sum(depths[:1])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer,cab=True)
+        self.dowsample_0 = dowsample(embed_dim, embed_dim*2)
+        self.encoderlayer_1 = BasicShadowFormer(dim=embed_dim*2,
+                            output_dim=embed_dim*2,
+                            input_resolution=(img_size // 2,
+                                                img_size // 2),
+                            depth=depths[1],
+                            num_heads=num_heads[1],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=enc_dpr[sum(depths[:1]):sum(depths[:2])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer, cab=True)
+        self.dowsample_1 = dowsample(embed_dim*2, embed_dim*4)
+        self.encoderlayer_2 = BasicShadowFormer(dim=embed_dim*4,
+                            output_dim=embed_dim*4,
+                            input_resolution=(img_size // (2 ** 2),
+                                                img_size // (2 ** 2)),
+                            depth=depths[2],
+                            num_heads=num_heads[2],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=enc_dpr[sum(depths[:2]):sum(depths[:3])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer)
+        self.dowsample_2 = dowsample(embed_dim*4, embed_dim*8)
+
+        # Bottleneck
+        channel_conv = embed_dim*16 
+        # channel_conv = embed_dim*8 if self.add_shadow_detect_dino_conact else embed_dim*4
+        self.conv = BasicShadowFormer(dim=channel_conv,
+                            output_dim=channel_conv,
+                            input_resolution=(img_size // (2 ** 3),
+                                                img_size // (2 ** 3)),
+                            depth=depths[4],
+                            num_heads=num_heads[4],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=conv_dpr,
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer)
+
+        # # Decoder
+        self.upsample_0 = upsample(channel_conv, embed_dim*4)
+        channel_0 =  embed_dim*8
+        self.decoderlayer_0 = BasicShadowFormer(dim=channel_0,
+                            output_dim=channel_0,
+                            input_resolution=(img_size // (2 ** 2),
+                                                img_size // (2 ** 2)),
+                            depth=depths[6],
+                            num_heads=num_heads[6],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=dec_dpr[sum(depths[5:6]):sum(depths[5:7])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer)
+        self.upsample_1 = upsample(channel_0, embed_dim*2)
+        channel_1 =  embed_dim*4
+        self.decoderlayer_1 = BasicShadowFormer(dim=channel_1,
+                            output_dim=channel_1,
+                            input_resolution=(img_size // 2,
+                                                img_size // 2),
+                            depth=depths[7],
+                            num_heads=num_heads[7],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=dec_dpr[sum(depths[5:7]):sum(depths[5:8])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer, cab=True)
+        self.upsample_2 = upsample(channel_1, embed_dim)
+        self.decoderlayer_2 = BasicShadowFormer(dim=embed_dim*2,
+                            output_dim=embed_dim*2,
+                            input_resolution=(img_size,
+                                                img_size),
+                            depth=depths[8],
+                            num_heads=num_heads[8],
+                            win_size=win_size,
+                            mlp_ratio=self.mlp_ratio,
+                            qkv_bias=qkv_bias, qk_scale=qk_scale,
+                            drop=drop_rate, attn_drop=attn_drop_rate,
+                            drop_path=dec_dpr[sum(depths[5:8]):sum(depths[5:9])],
+                            norm_layer=norm_layer,
+                            use_checkpoint=use_checkpoint,
+                            token_projection=token_projection,token_mlp=token_mlp,se_layer=se_layer,cab=True)
+
+        self.Conv = nn.Conv2d(self.DINO_channel * 4, embed_dim * 8, kernel_size=1)
+        self.relu = nn.LeakyReLU()
+        self.apply(self._init_weights)
+
+        self.freqfusion1 = FreqFusion(hr_channels=256,
+                                      lr_channels=512)
+    
+        self.freqfusion2 = FreqFusion(hr_channels=128,
+                                      lr_channels=256)     
+           
+        self.freqfusion3 = FreqFusion(hr_channels=64,
+                                      lr_channels=128)  
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def extra_repr(self) -> str:
+        return f"embed_dim={self.embed_dim}, token_projection={self.token_projection}, token_mlp={self.mlp}, win_size={self.win_size}"
+
+    def forward(self, x, DINO_Mat_features=None, point=None, normal=None, mask=None):
+        point_feature=None
+        dino_mat =None
+        dino_mat1=None
+
+        self.img_size = torch.tensor((x.shape[2], x.shape[3]))
+        point_feature1 = grid_sample(point, self.img_size // 2)
+        point_feature2 = grid_sample(point, self.img_size // 4)
+        point_feature3 = grid_sample(point, self.img_size // 8)
+        normal1= grid_sample(normal, self.img_size // 2)
+        normal2= grid_sample(normal, self.img_size // 4)
+        normal3= grid_sample(normal, self.img_size // 8)
+
+        patch_features_0 = DINO_Mat_features[0]
+        patch_features_1 = DINO_Mat_features[1]
+        patch_features_2 = DINO_Mat_features[2]
+        patch_features_3 = DINO_Mat_features[3]
+        patch_feature_all = torch.cat((patch_features_0, patch_features_1,
+                                    patch_features_2, patch_features_3), dim=1)
+
+        # Get concatenate DINO Feature
+        dino_mat_cat = self.Conv(patch_feature_all)
+        dino_mat_cat = self.relu(dino_mat_cat)
+        B, C, W, H = dino_mat_cat.shape
+        dino_mat_cat_flat = dino_mat_cat.view(B, C, W * H).permute(0,2,1)
+
+
+        dino_mat2 = F.upsample_bilinear(DINO_Mat_features[-1], scale_factor=2)
+        dino_mat3 = DINO_Mat_features[-1]
+            
+        # RGBD
+        xi = torch.cat((x, point[:,2,:].unsqueeze(1)), dim=1)
+
+        y = self.input_proj(xi)
+        y = self.pos_drop(y)
+
+        # Encoder
+        self.img_size = (int(self.img_size[0]), int(self.img_size[1]))
+        conv0 = self.encoderlayer_0(y, dino_mat, point_feature, normal, mask, img_size = self.img_size)
+        pool0 = self.dowsample_0(conv0, img_size = self.img_size)
+
+        self.img_size = (int(self.img_size[0]/2), int(self.img_size[1]/2))
+        conv1 = self.encoderlayer_1(pool0, dino_mat1, point_feature1, normal1, img_size = self.img_size)
+        pool1 = self.dowsample_1(conv1, img_size = self.img_size)
+
+        self.img_size = (int(self.img_size[0] / 2), int(self.img_size[1] / 2))
+        conv2 = self.encoderlayer_2(pool1, dino_mat2, point_feature2, normal2, img_size = self.img_size)
+        pool2 = self.dowsample_2(conv2, img_size = self.img_size)
+
+        # Bottleneck
+        self.img_size = (int(self.img_size[0] / 2), int(self.img_size[1] / 2))
+        pool2 = torch.cat([pool2, dino_mat_cat_flat],-1)
+        conv3 = self.conv(pool2, dino_mat3, point_feature3, normal3, img_size = self.img_size)
+        # print(f'{conv3.shape=}')  # conv3.shape=torch.Size([1, 1024, 512]) 1, 32, 32, 512
+        # conv3_B_C_H_W = conv3.view(conv3.shape[0], 32, 32, 512).permute(0, 3, 1, 2)
+        conv3_B_C_H_W = conv3.view(conv3.shape[0], int(conv3.shape[1]**0.5), int(conv3.shape[1]**0.5), 512).permute(0, 3, 1, 2)
+        # print(f'{conv3_B_C_H_W.shape=}')  # 1, 512, 32, 32
+
+        #Decoder
+        up0 = self.upsample_0(conv3, img_size = self.img_size)
+        self.img_size = (int(self.img_size[0] * 2), int(self.img_size[1] * 2))
+        # print(f'1.{conv2.shape=}, {up0.shape=}')  # conv2.shape=torch.Size([1, 4096, 128]), up0.shape=torch.Size([1, 4096, 128])
+        deconv0 = torch.cat([up0,conv2],-1)   
+        deconv0 = self.decoderlayer_0(deconv0, dino_mat2, point_feature2, normal2, img_size = self.img_size)
+        # print(f'1.{deconv0.shape=}')  # deconv0.shape=torch.Size([1, 4096, 256]) 1, 64, 64, 256
+        deconv0_B_C_H_W = deconv0.view(deconv0.shape[0], int(deconv0.shape[1]**0.5), int(deconv0.shape[1]**0.5), 256).permute(0, 3, 1, 2)
+        # print(f'1.{deconv0_B_C_H_W.shape=}')  # 1, 256, 64, 64
+
+        _, deconv0_B_C_H_W, lr_feat = self.freqfusion1(hr_feat=deconv0_B_C_H_W, lr_feat=conv3_B_C_H_W)  # 1, 256, 64, 64 & 1, 512, 32, 32
+        # print(f'1.{deconv0.shape=}, {lr_feat.shape=}')  # deconv0.shape=torch.Size([1, 256, 64, 64]), lr_feat.shape=torch.Size([1, 512, 64, 64])
+
+        deconv0 = deconv0_B_C_H_W.view(deconv0_B_C_H_W.shape[0], 256, -1).permute(0, 2, 1)
+
+
+        # print(f'1.{deconv0.shape=}')  # conv2.shape=torch.Size([1, 4096, 256]) 1, 64, 64, 256
+
+        deconv0_B_C_H_W = deconv0.view(deconv0.shape[0], int(deconv0.shape[1]**0.5), int(deconv0.shape[1]**0.5), 256).permute(0, 3, 1, 2)  # 1, 256, 64, 64
+        # print(f'2.{deconv0_B_C_H_W.shape=}')  # 1, 256, 64, 64
+
+        up1 = self.upsample_1(deconv0, img_size = self.img_size)
+        self.img_size = (int(self.img_size[0] * 2), int(self.img_size[1] * 2))
+        # print(f'2.{conv1.shape=}, {up1.shape=}')  # conv1.shape=torch.Size([1, 16384, 64]), up1.shape=torch.Size([1, 16384, 64]) 1, 128, 128, 64
+        deconv1 = torch.cat([up1,conv1],-1)  # 1, 16384, 128
+        deconv1 = self.decoderlayer_1(deconv1, dino_mat1, point_feature1, normal1, img_size = self.img_size)
+        # print(f'2.{deconv1.shape=}')  # 1, 16384, 128  1, 128, 128, 128
+        deconv1_B_C_H_W = deconv1.view(deconv1.shape[0], int(deconv1.shape[1]**0.5), int(deconv1.shape[1]**0.5), 128).permute(0, 3, 1, 2)
+        # print(f'2.{deconv1_B_C_H_W.shape=}')  # 1, 128, 128, 128
+
+        _, deconv1_B_C_H_W, lr_feat = self.freqfusion2(hr_feat=deconv1_B_C_H_W, lr_feat=deconv0_B_C_H_W)  # 1, 128, 128, 128 & 1, 256, 64, 64
+
+        # print(f'2.{deconv1_B_C_H_W.shape=}, {lr_feat.shape=}')  # hr_feat.shape=torch.Size([1, 128, 128, 128]), lr_feat.shape=torch.Size([1, 256, 128, 128])
+
+        deconv1 = deconv1_B_C_H_W.view(deconv1_B_C_H_W.shape[0], 128, -1).permute(0, 2, 1)
+        # print()
+        # print(f'3.{deconv1.shape=}')  # deconv1.shape=torch.Size([1, 16384, 128])  1, 128, 128, 128
+        
+        deconv1_B_C_H_W = deconv1.view(deconv1.shape[0], int(deconv1.shape[1]**0.5), int(deconv1.shape[1]**0.5), 128).permute(0, 3, 1, 2)  # 1, 128, 128, 128
+
+        up2 = self.upsample_2(deconv1, img_size = self.img_size)
+        self.img_size = (int(self.img_size[0] * 2), int(self.img_size[1] * 2))
+        # print(f'3.{conv0.shape=}, {up2.shape=}')  # conv0.shape=torch.Size([1, 65536, 32]), up2.shape=torch.Size([1, 65536, 32])  1, 256, 256, 32
+        deconv2 = torch.cat([up2,conv0],-1)  # 1, 256, 256, 64
+        # print(f'3.{deconv2.shape=}')  # 1, 65536, 64   1, 256, 256, 64
+        deconv2 = self.decoderlayer_2(deconv2, dino_mat, point_feature, normal, mask, img_size = self.img_size)
+        # print(f'3.{deconv2.shape=}')  # 1, 65536, 64   1, 256, 256, 64
+
+        deconv2_B_C_H_W = deconv2.view(deconv2.shape[0], int(deconv2.shape[1]**0.5), int(deconv2.shape[1]**0.5), 64).permute(0, 3, 1, 2)
+        # print(f'3.{deconv2_B_C_H_W.shape=}')
+
+        _, deconv2_B_C_H_W, lr_feat = self.freqfusion3(hr_feat=deconv2_B_C_H_W, lr_feat=deconv1_B_C_H_W)  # 1, 64, 256, 256 & 1, 128, 128, 128
+
+        # print('*'*5, f'3.{deconv2_B_C_H_W.shape=}, {lr_feat.shape=}')
+
+        deconv2 = deconv2_B_C_H_W.view(deconv2_B_C_H_W.shape[0], 64, -1).permute(0, 2, 1)
+        # print(f'4.{deconv2.shape=}')
+
+        # Output Projection
+        # print(f'4.{deconv2.shape=}')
+        # print(f'4.{x.shape=}')
+        y = self.output_proj(deconv2, img_size = self.img_size) + x
+        return y
+

requirements.txt

@@ -0,0 +1,233 @@
+absl-py==1.4.0
+addict==2.4.0
+aiofiles==23.2.1
+aiohttp==3.8.5
+aiosignal==1.3.1
+altair==5.1.0
+annotated-types==0.5.0
+ansi2html==1.8.0
+antlr4-python3-runtime==4.9.3
+anyio==3.7.1
+asttokens==2.4.1
+async-timeout==4.0.3
+attrs==23.1.0
+blinker==1.6.2
+boto3==1.28.62
+botocore==1.31.62
+bottle==0.12.25
+Brotli
+cachetools==5.3.1
+cattrs==23.1.2
+certifi 
+cffi 
+chainmap==1.0.3
+chardet==5.2.0
+charset-normalizer 
+click==8.1.7
+cloudpickle==2.2.1
+cmake==3.27.0
+colorama 
+combomethod==1.0.12
+comm==0.1.4
+common==0.1.2
+ConfigArgParse==1.7
+contourpy==1.1.0
+cryptography 
+cycler==0.11.0
+Cython==3.0.5
+dash==2.14.1
+dash-core-components==2.0.0
+dash-html-components==2.0.0
+dash-table==5.0.0
+data==0.4
+debugpy==1.8.1
+decorator==5.1.1
+dual==0.0.10
+dynamo3==0.4.10
+easydict==1.10
+efficientnet-pytorch==0.7.1
+einops==0.3.2
+exceptiongroup==1.1.3
+executing==2.0.1
+fastapi==0.103.0
+fastjsonschema==2.18.1
+ffmpy==0.3.1
+filelock==3.12.2
+Flask==2.3.3
+Flask-Cors==4.0.0
+flywheel==0.5.4
+fonttools==4.42.0
+frozenlist==1.4.0
+fsspec==2023.6.0
+funcsigs==1.0.2
+future 
+fvcore 
+google-auth==2.22.0
+google-auth-oauthlib==1.0.0
+gradio
+gradio_client
+grpcio==1.57.0
+h11==0.14.0
+h5py==3.9.0
+httpcore==0.17.3
+httpx==0.24.1
+huggingface-hub==0.16.4
+idna 
+imageio==2.31.1
+importlib-metadata==6.8.0
+importlib-resources==6.0.1
+iopath==0.1.9
+ipython==8.17.2
+ipywidgets==8.1.1
+itsdangerous==2.1.2
+jedi==0.19.1
+Jinja2==3.1.2
+jmespath==1.0.1
+joblib==1.3.2
+jsonschema==4.19.0
+jsonschema-specifications==2023.7.1
+jupyter_core==5.5.0
+jupyterlab-widgets==3.0.9
+kiwisolver==1.4.4
+kornia==0.7.0
+lazy_loader==0.3
+lightning-utilities==0.9.0
+lit==16.0.6
+Markdown==3.4.4
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib==3.3.4
+matplotlib-inline==0.1.6
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.0.4
+mypy-extensions==1.0.0
+natsort==8.4.0
+nbformat==5.7.0
+ndim
+nest-asyncio==1.5.8
+networkx==3.1
+ntplib==0.4.0
+nulltype==2.3.1
+numpy 
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.2.10.91
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-nccl-cu11==2.14.3
+nvidia-nvtx-cu11==11.7.91
+oauthlib==3.2.2
+omegaconf==2.3.0
+open3d==0.17.0
+opencv-python==4.8.0.74
+options==1.4.10
+orjson==3.9.5
+packaging==23.1
+
+pandas==2.0.3
+parso==0.8.3
+peewee==3.16.3
+pexpect==4.8.0
+Pillow==10.0.0
+platformdirs==3.10.0
+plotly==5.18.0
+portalocker 
+progressbar2==4.2.0
+prompt-toolkit==3.0.39
+protobuf==4.24.2
+prox==0.0.17
+ptflops==0.7.2.2
+ptyprocess==0.7.0
+pure-eval==0.2.2
+py-machineid==0.4.3
+pyasn1==0.5.0
+pyasn1-modules==0.3.0
+pybind11==2.11.1
+pycparser 
+pydantic==2.3.0
+pydantic_core==2.6.3
+
+pyDeprecate==0.3.2
+pydub==0.25.1
+Pygments==2.16.1
+PyNaCl==1.5.0
+pyOpenSSL 
+pyparsing==3.0.9
+pyquaternion==0.9.9
+pyre-extensions==0.0.29
+PySocks 
+python-dateutil==2.8.2
+python-geoip-python3==1.3
+
+python-multipart==0.0.6
+python-utils==3.7.0
+
+pytz==2023.3
+PyWavelets==1.4.1
+PyYAML==6.0.1
+referencing==0.30.2
+requests 
+requests-cache
+requests-oauthlib==1.3.1
+retrying==1.3.4
+rich==13.5.2
+rich-argparse==1.3.0
+rpds-py==0.10.0
+rsa==4.9
+s3transfer==0.7.0
+safetensors==0.3.1
+scikit-image==0.21.0
+scikit-learn==1.3.0
+scipy==1.11.1
+seaborn==0.12.2
+semantic-version==2.10.0
+six==1.12.0
+
+sniffio==1.3.0
+stack-data==0.6.3
+starlette==0.27.0
+sympy==1.12
+tabulate 
+tenacity==8.2.3
+tensorboard==2.14.0
+tensorboard-data-server==0.7.1
+
+termcolor 
+thop==0.1.1.post2209072238
+threadpoolctl==3.2.0
+tifffile==2023.7.18
+tight==0.1.0
+timm==0.9.5
+tomli==2.0.1
+tomli_w==1.0.0
+toolz==0.12.0
+gdown
+
+torchmetrics==1.1.1
+torchstat==0.0.7
+torchsummary==1.5.1
+
+tqdm==4.65.0
+traitlets==5.13.0
+triton==2.0.0
+typing-inspect==0.9.0
+typing_extensions 
+tzdata==2023.3
+url-normalize==1.4.3
+urllib3==1.26.16
+uvicorn==0.23.2
+wcwidth==0.2.9
+websockets==11.0.3
+Werkzeug==2.3.7
+widgetsnbextension==4.0.9
+wrapt==1.15.0
+x21
+xformers==0.0.20
+yacs 
+yarl==1.9.2
+zipp==3.16.2

run_test.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node 1 --master_port 29508 ./test_shadow.py --save_images

test_shadow.py

@@ -0,0 +1,271 @@
+import numpy as np
+import os
+import argparse
+from tqdm import tqdm
+from torch.utils.data.distributed import DistributedSampler
+import torch.nn as nn
+import torch
+from torch.utils.data import DataLoader
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.nn.functional as F
+import random
+# from utils.loader import get_validation_data
+from utils.loader import get_test_data
+import utils
+import cv2
+import torch.distributed as dist
+from skimage.metrics import peak_signal_noise_ratio as psnr_loss
+from skimage.metrics import structural_similarity as ssim_loss
+parser = argparse.ArgumentParser(description='RGB denoising evaluation on the validation set of SIDD')
+parser.add_argument('--input_dir', default='test_dir',
+    type=str, help='Directory of validation images')
+parser.add_argument('--result_dir', default='./output_dir',
+    type=str, help='Directory for results')
+parser.add_argument('--weights', default='ACVLab_shadow.pth'
+                    ,type=str, help='Path to weights')
+# parser.add_argument('--arch', default='ShadowFormer', type=str, help='arch')
+parser.add_argument('--arch', type=str, default='ShadowFormerFreq', help='archtechture')
+parser.add_argument('--batch_size', default=1, type=int, help='Batch size for dataloader')
+parser.add_argument('--save_images', action='store_true', default=False, help='Save denoised images in result directory')
+parser.add_argument('--cal_metrics', action='store_true', default=False, help='Measure denoised images with GT')
+parser.add_argument('--embed_dim', type=int, default=32, help='number of data loading workers')    
+parser.add_argument('--win_size', type=int, default=16, help='number of data loading workers')
+parser.add_argument('--token_projection', type=str, default='linear', help='linear/conv token projection')
+parser.add_argument('--token_mlp', type=str,default='leff', help='ffn/leff token mlp')
+
+parser.add_argument('--train_ps', type=int, default=256, help='patch size of training sample')
+parser.add_argument("--local-rank", type=int)
+
+args = parser.parse_args()
+
+local_rank = args.local_rank
+torch.cuda.set_device(local_rank)
+dist.init_process_group(backend='nccl')
+device = torch.device("cuda", local_rank)
+
+
+class SlidingWindowInference:
+    def __init__(self, window_size=512, overlap=64, img_multiple_of=64):
+        self.window_size = window_size
+        self.overlap = overlap
+        self.img_multiple_of = img_multiple_of
+        
+    def _pad_input(self, x, h_pad, w_pad):
+        """Handle padding using reflection padding"""
+        return F.pad(x, (0, w_pad, 0, h_pad), 'reflect')
+
+    def __call__(self, model, input_, point, normal, dino_net, device):
+        # Save original dimensions
+        original_height, original_width = input_.shape[2], input_.shape[3]
+        # print(f"Original size: {original_height}x{original_width}")
+        
+        # Calculate minimum dimensions needed (at least window_size and multiple of img_multiple_of)
+        H = max(self.window_size, 
+               ((original_height + self.img_multiple_of - 1) // self.img_multiple_of) * self.img_multiple_of)
+        W = max(self.window_size, 
+               ((original_width + self.img_multiple_of - 1) // self.img_multiple_of) * self.img_multiple_of)
+        # print(f"Target padded size: {H}x{W}")
+        
+        # Calculate required padding
+        padh = H - original_height
+        padw = W - original_width
+        # print(f"Padding: h={padh}, w={padw}")
+        
+        # Pad all inputs
+        input_pad = self._pad_input(input_, padh, padw)
+        point_pad = self._pad_input(point, padh, padw)
+        normal_pad = self._pad_input(normal, padh, padw)
+        
+        # If image was smaller than window_size, process it as a single window
+        if original_height <= self.window_size and original_width <= self.window_size:
+            # print("Image smaller than window size, processing as single padded window")
+            
+            # For DINO features
+            DINO_patch_size = 14
+            h_size = H * DINO_patch_size // 8
+            w_size = W * DINO_patch_size // 8
+            
+            UpSample_window = torch.nn.UpsamplingBilinear2d(size=(h_size, w_size))
+            
+            # Get DINO features
+            with torch.no_grad():
+                input_DINO = UpSample_window(input_pad)
+                dino_features = dino_net.module.get_intermediate_layers(input_DINO, 4, True)
+            
+            # Model inference
+            with torch.cuda.amp.autocast():
+                restored = model(input_pad, dino_features, point_pad, normal_pad)
+            
+            # Crop back to original size
+            output = restored[:, :, :original_height, :original_width]
+            return output
+        
+        # For larger images, proceed with sliding window approach
+        stride = self.window_size - self.overlap
+        h_steps = (H - self.window_size + stride - 1) // stride + 1
+        w_steps = (W - self.window_size + stride - 1) // stride + 1
+        # print(f"Steps: h={h_steps}, w={w_steps}")
+        
+        # Create output tensor and counter
+        output = torch.zeros_like(input_pad)
+        count = torch.zeros_like(input_pad)
+        
+        for h_idx in range(h_steps):
+            for w_idx in range(w_steps):
+                # Calculate current window position
+                h_start = min(h_idx * stride, H - self.window_size)
+                w_start = min(w_idx * stride, W - self.window_size)
+                h_end = h_start + self.window_size
+                w_end = w_start + self.window_size
+                
+                # Get current window
+                input_window = input_pad[:, :, h_start:h_end, w_start:w_end]
+                point_window = point_pad[:, :, h_start:h_end, w_start:w_end]
+                normal_window = normal_pad[:, :, h_start:h_end, w_start:w_end]
+                
+                # print(f"Processing window at ({h_idx}, {w_idx}): {input_window.shape}")
+                
+                # For DINO features
+                DINO_patch_size = 14
+                h_size = self.window_size * DINO_patch_size // 8
+                w_size = self.window_size * DINO_patch_size // 8
+                
+                UpSample_window = torch.nn.UpsamplingBilinear2d(size=(h_size, w_size))
+                
+                # Get DINO features
+                with torch.no_grad():
+                    input_DINO = UpSample_window(input_window)
+                    dino_features = dino_net.module.get_intermediate_layers(input_DINO, 4, True)
+                
+                # Model inference
+                with torch.cuda.amp.autocast():
+                    restored = model(input_window, dino_features, point_window, normal_window)
+                
+                # Create weight mask for smooth transition
+                weight = torch.ones_like(restored)
+                if self.overlap > 0:
+                    # Create gradual weights for overlap regions
+                    for i in range(self.overlap):
+                        ratio = i / self.overlap
+                        weight[:, :, i, :] *= ratio
+                        weight[:, :, -(i+1), :] *= ratio
+                        weight[:, :, :, i] *= ratio
+                        weight[:, :, :, -(i+1)] *= ratio
+                
+                # Accumulate results and weights
+                output[:, :, h_start:h_end, w_start:w_end] += restored * weight
+                count[:, :, h_start:h_end, w_start:w_end] += weight
+        
+        # Normalize output
+        output = output / (count + 1e-6)
+        
+        # Crop back to original size
+        output = output[:, :, :original_height, :original_width]
+        return output
+
+
+utils.mkdir(args.result_dir)
+
+# ######### Set Seeds ###########
+random.seed(1234)
+np.random.seed(1234)
+torch.manual_seed(1234)
+torch.cuda.manual_seed(1234)
+torch.cuda.manual_seed_all(1234)
+
+def worker_init_fn(worker_id):
+    random.seed(1234 + worker_id)
+
+g = torch.Generator()
+g.manual_seed(1234)
+
+torch.backends.cudnn.benchmark = True
+# torch.backends.cudnn.deterministic = True
+######### Model ###########
+model_restoration = utils.get_arch(args)
+model_restoration.to(device)
+model_restoration.eval()
+DINO_Net = torch.hub.load('./dinov2', 'dinov2_vitl14', source='local')
+DINO_Net.to(device)
+DINO_Net.eval()
+######### Load ###########
+utils.load_checkpoint(model_restoration, args.weights)
+print("===>Testing using weights: ", args.weights)
+
+######### DDP ###########
+
+model_restoration = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model_restoration).to(device)
+model_restoration = DDP(model_restoration, device_ids=[local_rank], output_device=local_rank)
+DINO_Net = DDP(DINO_Net, device_ids=[local_rank], output_device=local_rank)
+
+######### Test ###########
+img_multiple_of = 8 * args.win_size
+DINO_patch_size = 14
+
+def UpSample(img):
+    upsample = nn.UpsamplingBilinear2d(
+        size=((int)(img.shape[2] * (DINO_patch_size / 8)), 
+            (int)(img.shape[3] * (DINO_patch_size / 8))))
+    return upsample(img)
+
+img_options_train = {'patch_size':args.train_ps}
+test_dataset = get_test_data(args.input_dir, False)
+test_sampler = DistributedSampler(test_dataset, shuffle=False)
+test_loader = DataLoader(dataset=test_dataset, batch_size=1, num_workers=0, sampler=test_sampler, drop_last=False, worker_init_fn=worker_init_fn, generator=g)
+with torch.no_grad():
+    psnr_val_rgb_list = []
+    psnr_val_mask_list = []
+    ssim_val_rgb_list = []
+    rmse_val_rgb_list = []
+    for ii, data_test in enumerate(tqdm(test_loader), 0):
+            # rgb_gt = data_test[0].numpy().squeeze().transpose((1, 2, 0))
+            rgb_noisy = data_test[1].to(device)
+            point = data_test[2].to(device)
+            normal = data_test[3].to(device)
+            filenames = data_test[4]
+
+            # Pad the input if not_multiple_of win_size * 8
+            # height, width = rgb_noisy.shape[2], rgb_noisy.shape[3]
+            # H, W = ((height + img_multiple_of) // img_multiple_of) * img_multiple_of, (
+            #     (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
+            # rgb_noisy = F.pad(rgb_noisy, (0, padw, 0, padh), 'reflect')
+            # point = F.pad(point, (0, padw, 0, padh), 'reflect')
+            # normal = F.pad(normal, (0, padw, 0, padh), 'reflect')
+            # print(f'{rgb_noisy.shape=} {point.shape=} {normal.shape=}')
+            # UpSample_val = nn.UpsamplingBilinear2d(
+            #     size=((int)(rgb_noisy.shape[2] * (DINO_patch_size / 8)), 
+            #         (int)(rgb_noisy.shape[3] * (DINO_patch_size / 8))))
+            # with torch.cuda.amp.autocast():
+            #     # DINO_V2
+            #     input_DINO = UpSample_val(rgb_noisy)
+            #     dino_mat_features = DINO_Net.module.get_intermediate_layers(input_DINO, 4, True)
+            #     rgb_restored = model_restoration(rgb_noisy, dino_mat_features, point, normal)
+            sliding_window = SlidingWindowInference(
+                window_size=512,  # 與訓練相同的 patch size
+                overlap=64,       # 相應調整 overlap
+                img_multiple_of=8 * args.win_size
+            )
+
+            with torch.cuda.amp.autocast():
+                rgb_restored = sliding_window(
+                    model=model_restoration,
+                    input_=rgb_noisy,
+                    point=point,
+                    normal=normal,
+                    dino_net=DINO_Net,
+                    device=device
+                )
+
+        
+            rgb_restored = torch.clamp(rgb_restored, 0.0, 1.0)
+            # rgb_restored = rgb_restored[:, : ,:height, :width]
+            rgb_restored = torch.clamp(rgb_restored, 0, 1).cpu().numpy().squeeze().transpose((1, 2, 0))
+            
+
+            if args.save_images:
+                utils.save_img(rgb_restored * 255.0, os.path.join(args.result_dir, filenames[0]))
+
+

utils/__init__.py

@@ -0,0 +1,6 @@
+from .dir_utils import *
+from .dataset_utils import *
+from .image_utils import *
+from .model_utils import *
+from .shadow_mask_evaluate import *
+from .tta import *

utils/antialias.py

@@ -0,0 +1,125 @@
+# Copyright (c) 2019, Adobe Inc. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial-ShareAlike
+# 4.0 International Public License. To view a copy of this license, visit
+# https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode.
+
+
+
+########  https://github.com/adobe/antialiased-cnns/blob/master/models_lpf/__init__.py
+
+
+
+import torch
+import torch.nn.parallel
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Downsample(nn.Module):
+    def __init__(self, pad_type='reflect', filt_size=3, stride=2, channels=None, pad_off=0):
+        super(Downsample, self).__init__()
+        self.filt_size = filt_size
+        self.pad_off = pad_off
+        self.pad_sizes = [int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)), int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2))]
+        self.pad_sizes = [pad_size+pad_off for pad_size in self.pad_sizes]
+        self.stride = stride
+        self.off = int((self.stride-1)/2.)
+        self.channels = channels
+
+        # print('Filter size [%i]'%filt_size)
+        if(self.filt_size==1):
+            a = np.array([1.,])
+        elif(self.filt_size==2):
+            a = np.array([1., 1.])
+        elif(self.filt_size==3):
+            a = np.array([1., 2., 1.])
+        elif(self.filt_size==4):    
+            a = np.array([1., 3., 3., 1.])
+        elif(self.filt_size==5):    
+            a = np.array([1., 4., 6., 4., 1.])
+        elif(self.filt_size==6):    
+            a = np.array([1., 5., 10., 10., 5., 1.])
+        elif(self.filt_size==7):    
+            a = np.array([1., 6., 15., 20., 15., 6., 1.])
+
+        filt = torch.Tensor(a[:,None]*a[None,:])
+        filt = filt/torch.sum(filt)
+        self.register_buffer('filt', filt[None,None,:,:].repeat((self.channels,1,1,1)))
+
+        self.pad = get_pad_layer(pad_type)(self.pad_sizes)
+
+    def forward(self, inp):
+        if(self.filt_size==1):
+            if(self.pad_off==0):
+                return inp[:,:,::self.stride,::self.stride]    
+            else:
+                return self.pad(inp)[:,:,::self.stride,::self.stride]
+        else:
+            return F.conv2d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1])
+
+def get_pad_layer(pad_type):
+    if(pad_type in ['refl','reflect']):
+        PadLayer = nn.ReflectionPad2d
+    elif(pad_type in ['repl','replicate']):
+        PadLayer = nn.ReplicationPad2d
+    elif(pad_type=='zero'):
+        PadLayer = nn.ZeroPad2d
+    else:
+        print('Pad type [%s] not recognized'%pad_type)
+    return PadLayer
+
+
+class Downsample1D(nn.Module):
+    def __init__(self, pad_type='reflect', filt_size=3, stride=2, channels=None, pad_off=0):
+        super(Downsample1D, self).__init__()
+        self.filt_size = filt_size
+        self.pad_off = pad_off
+        self.pad_sizes = [int(1. * (filt_size - 1) / 2), int(np.ceil(1. * (filt_size - 1) / 2))]
+        self.pad_sizes = [pad_size + pad_off for pad_size in self.pad_sizes]
+        self.stride = stride
+        self.off = int((self.stride - 1) / 2.)
+        self.channels = channels
+
+        # print('Filter size [%i]' % filt_size)
+        if(self.filt_size == 1):
+            a = np.array([1., ])
+        elif(self.filt_size == 2):
+            a = np.array([1., 1.])
+        elif(self.filt_size == 3):
+            a = np.array([1., 2., 1.])
+        elif(self.filt_size == 4):
+            a = np.array([1., 3., 3., 1.])
+        elif(self.filt_size == 5):
+            a = np.array([1., 4., 6., 4., 1.])
+        elif(self.filt_size == 6):
+            a = np.array([1., 5., 10., 10., 5., 1.])
+        elif(self.filt_size == 7):
+            a = np.array([1., 6., 15., 20., 15., 6., 1.])
+
+        filt = torch.Tensor(a)
+        filt = filt / torch.sum(filt)
+        self.register_buffer('filt', filt[None, None, :].repeat((self.channels, 1, 1)))
+
+        self.pad = get_pad_layer_1d(pad_type)(self.pad_sizes)
+
+    def forward(self, inp):
+        if(self.filt_size == 1):
+            if(self.pad_off == 0):
+                return inp[:, :, ::self.stride]
+            else:
+                return self.pad(inp)[:, :, ::self.stride]
+        else:
+            return F.conv1d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1])
+
+
+def get_pad_layer_1d(pad_type):
+    if(pad_type in ['refl', 'reflect']):
+        PadLayer = nn.ReflectionPad1d
+    elif(pad_type in ['repl', 'replicate']):
+        PadLayer = nn.ReplicationPad1d
+    elif(pad_type == 'zero'):
+        PadLayer = nn.ZeroPad1d
+    else:
+        print('Pad type [%s] not recognized' % pad_type)
+    return PadLayer

utils/bundle_submissions.py

@@ -0,0 +1,104 @@
+ # Author: Tobias Plötz, TU Darmstadt ([email protected])
+
+ # This file is part of the implementation as described in the CVPR 2017 paper:
+ # Tobias Plötz and Stefan Roth, Benchmarking Denoising Algorithms with Real Photographs.
+ # Please see the file LICENSE.txt for the license governing this code.
+
+
+import numpy as np
+import scipy.io as sio
+import os
+import h5py
+
+def bundle_submissions_raw(submission_folder,session):
+    '''
+    Bundles submission data for raw denoising
+    submission_folder Folder where denoised images reside
+    Output is written to <submission_folder>/bundled/. Please submit
+    the content of this folder.
+    '''
+
+    out_folder = os.path.join(submission_folder, session)
+    # out_folder = os.path.join(submission_folder, "bundled/")
+    try:
+        os.mkdir(out_folder)
+    except:pass
+
+    israw = True
+    eval_version="1.0"
+
+    for i in range(50):
+        Idenoised = np.zeros((20,), dtype=np.object)
+        for bb in range(20):
+            filename = '%04d_%02d.mat'%(i+1,bb+1)
+            s = sio.loadmat(os.path.join(submission_folder,filename))
+            Idenoised_crop = s["Idenoised_crop"]
+            Idenoised[bb] = Idenoised_crop
+        filename = '%04d.mat'%(i+1)
+        sio.savemat(os.path.join(out_folder, filename),
+                    {"Idenoised": Idenoised,
+                     "israw": israw,
+                     "eval_version": eval_version},
+                    )
+
+def bundle_submissions_srgb(submission_folder,session):
+    '''
+    Bundles submission data for sRGB denoising
+    
+    submission_folder Folder where denoised images reside
+    Output is written to <submission_folder>/bundled/. Please submit
+    the content of this folder.
+    '''
+    out_folder = os.path.join(submission_folder, session)
+    # out_folder = os.path.join(submission_folder, "bundled/")
+    try:
+        os.mkdir(out_folder)
+    except:pass
+    israw = False
+    eval_version="1.0"
+
+    for i in range(50):
+        Idenoised = np.zeros((20,), dtype=np.object)
+        for bb in range(20):
+            filename = '%04d_%02d.mat'%(i+1,bb+1)
+            s = sio.loadmat(os.path.join(submission_folder,filename))
+            Idenoised_crop = s["Idenoised_crop"]
+            Idenoised[bb] = Idenoised_crop
+        filename = '%04d.mat'%(i+1)
+        sio.savemat(os.path.join(out_folder, filename),
+                    {"Idenoised": Idenoised,
+                     "israw": israw,
+                     "eval_version": eval_version},
+                    )
+
+
+
+def bundle_submissions_srgb_v1(submission_folder,session):
+    '''
+    Bundles submission data for sRGB denoising
+    
+    submission_folder Folder where denoised images reside
+    Output is written to <submission_folder>/bundled/. Please submit
+    the content of this folder.
+    '''
+    out_folder = os.path.join(submission_folder, session)
+    # out_folder = os.path.join(submission_folder, "bundled/")
+    try:
+        os.mkdir(out_folder)
+    except:pass
+    israw = False
+    eval_version="1.0"
+
+    for i in range(50):
+        Idenoised = np.zeros((20,), dtype=np.object)
+        for bb in range(20):
+            filename = '%04d_%d.mat'%(i+1,bb+1)
+            s = sio.loadmat(os.path.join(submission_folder,filename))
+            Idenoised_crop = s["Idenoised_crop"]
+            Idenoised[bb] = Idenoised_crop
+        filename = '%04d.mat'%(i+1)
+        sio.savemat(os.path.join(out_folder, filename),
+                    {"Idenoised": Idenoised,
+                     "israw": israw,
+                     "eval_version": eval_version},
+                    )

utils/dataset_utils.py

@@ -0,0 +1,46 @@
+import torch
+import os
+import torchvision.transforms as transforms
+
+
+class Augment_RGB_torch:
+### rotate and flip
+    def __init__(self, rotate=0):
+        self.rotate = rotate
+        pass
+    def transform0(self, torch_tensor):
+        return torch_tensor  
+
+    def transform1(self, torch_tensor):
+        H, W = torch_tensor.shape[1], torch_tensor.shape[2]
+        train_transform = transforms.Compose([
+        transforms.RandomRotation((self.rotate,self.rotate), interpolation=transforms.InterpolationMode.BILINEAR, expand=False),
+        transforms.Resize((int(H * 1.3), int(W * 1.3)), antialias=True),
+        # CenterCrop，if the size is larger than the original size, the excess will be filled with black
+        transforms.CenterCrop([H, W])
+        ])
+        return train_transform(torch_tensor)
+
+    def transform2(self, torch_tensor):
+        torch_tensor = torch.rot90(torch_tensor, k=1, dims=[-1,-2])
+        return torch_tensor
+    def transform3(self, torch_tensor):
+        torch_tensor = torch.rot90(torch_tensor, k=2, dims=[-1,-2])
+        return torch_tensor
+    def transform4(self, torch_tensor):
+        torch_tensor = torch.rot90(torch_tensor, k=3, dims=[-1,-2])
+        return torch_tensor
+    def transform5(self, torch_tensor):
+        torch_tensor = torch_tensor.flip(-2)
+        return torch_tensor
+    def transform6(self, torch_tensor):
+        torch_tensor = (torch.rot90(torch_tensor, k=1, dims=[-1,-2])).flip(-2)
+        return torch_tensor
+    def transform7(self, torch_tensor):
+        torch_tensor = (torch.rot90(torch_tensor, k=2, dims=[-1,-2])).flip(-2)
+        return torch_tensor
+    def transform8(self, torch_tensor):
+        torch_tensor = (torch.rot90(torch_tensor, k=3, dims=[-1,-2])).flip(-2)
+        return torch_tensor
+
+

utils/dir_utils.py

@@ -0,0 +1,22 @@
+import os
+from natsort import natsorted
+from glob import glob
+
+def mkdirs(paths):
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            mkdir(path)
+    else:
+        mkdir(paths)
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+def mknod(path):
+    if not os.path.exists(path):
+        os.mknod(path)
+
+def get_last_path(path, session):
+	x = natsorted(glob(os.path.join(path,'*%s'%session)))[-1]
+	return x

utils/image_utils.py

@@ -0,0 +1,204 @@
+import torch
+import numpy as np
+import pickle
+import cv2
+from skimage.color import rgb2lab
+import os
+import math
+os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
+
+def is_numpy_file(filename):
+    return any(filename.endswith(extension) for extension in [".npy"])
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in [".jpg"])
+
+def is_png_file(filename):
+    return any(filename.endswith(extension) for extension in [".png"])
+
+def is_pkl_file(filename):
+    return any(filename.endswith(extension) for extension in [".pkl"])
+
+def load_pkl(filename_):
+    with open(filename_, 'rb') as f:
+        ret_dict = pickle.load(f)
+    return ret_dict    
+
+def save_dict(dict_, filename_):
+    with open(filename_, 'wb') as f:
+        pickle.dump(dict_, f)    
+
+def load_npy(filepath):
+    img = np.load(filepath)
+    return img
+
+def load_SSAO(filepath):
+    img = cv2.imread(filepath)
+    # img = cv2.resize(img, [256, 256], interpolation=cv2.INTER_AREA)
+    img = img.astype(np.float32)
+    img = img/255.
+    return img
+def load_img(filepath):
+    img = cv2.cvtColor(cv2.imread(filepath), cv2.COLOR_BGR2RGB)
+    # img = cv2.resize(img, [256, 256], interpolation=cv2.INTER_AREA)
+    img = img.astype(np.float32)
+    img = img/255.
+    return img
+
+def load_val_img(filepath):
+    img = cv2.cvtColor(cv2.imread(filepath), cv2.COLOR_BGR2RGB)
+    # img = cv2.resize(img, [256, 256], interpolation=cv2.INTER_AREA)
+    resized_img = img.astype(np.float32)
+    resized_img = resized_img/255.
+    return resized_img
+
+def load_mask(filepath):
+    img = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
+    # img = cv2.resize(img, [256, 256], interpolation=cv2.INTER_AREA)
+    # img = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
+    # kernel = np.ones((8,8), np.uint8)
+    # erosion = cv2.erode(img, kernel, iterations=1)
+    # dilation = cv2.dilate(img, kernel, iterations=1)
+    # contour = dilation - erosion
+    img = img.astype(np.float32)
+    # contour = contour.astype(np.float32)
+    # contour = contour/255.
+    img = img/255.
+    return img
+
+def load_ssao(filepath):
+    img = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
+    # img = cv2.resize(img, [256, 256], interpolation=cv2.INTER_AREA)
+    img = img.astype(np.float32)
+    # contour = contour.astype(np.float32)
+    # contour = contour/255.
+    img = img/255.
+    return img
+
+def load_depth(filepath):
+    img = np.load(filepath)
+    # img = cv2.resize(img, [256, 256], interpolation=cv2.INTER_AREA)
+    # img = cv2.imread(filepath, cv2.IMREAD_UNCHANGED)
+    # img = img / 255
+    return img
+
+def load_normal(filepath):
+    img = np.load(filepath).transpose(1,2,0)
+    # img = cv2.resize(img, [256, 256], interpolation=cv2.INTER_AREA)
+    return img
+
+def load_val_mask(filepath):
+    img = cv2.imread(filepath, 0)
+    resized_img = img
+    # resized_img = cv2.resize(img, [256, 256], interpolation=cv2.INTER_AREA)
+    resized_img = resized_img.astype(np.float32)
+    resized_img = resized_img/255.
+    return resized_img
+
+def save_img(img, filepath):
+    cv2.imwrite(filepath, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+
+def myPSNR(tar_img, prd_img):
+    imdff = torch.clamp(prd_img,0,1) - torch.clamp(tar_img,0,1)
+    rmse = (imdff**2).mean().sqrt()
+    ps = 20*torch.log10(1/rmse)
+    return ps
+
+    # imdff = torch.clamp(prd_img,0,1) - torch.clamp(tar_img,0,1)
+    # rmse = (imdff**2).mean()
+    # ps = 10*torch.log10(1/rmse)
+    # return ps
+
+def batch_PSNR(img1, img2, average=True):
+    PSNR = []
+    for im1, im2 in zip(img1, img2):
+        psnr = myPSNR(im1, im2)
+        PSNR.append(psnr)
+    return sum(PSNR)/len(PSNR) if average else sum(PSNR)
+
+def tensor2im(input_image, imtype=np.uint8):
+    """"Converts a Tensor array into a numpy image array.
+    Parameters:
+        input_image (tensor) --  the input image tensor array
+        imtype (type)        --  the desired type of the converted numpy array
+    """
+    if not isinstance(input_image, np.ndarray):
+
+        if isinstance(input_image, torch.Tensor):  # get the data from a variable
+            image_tensor = input_image.data
+        else:
+            return input_image
+        image_numpy = image_tensor[0].cpu().float().numpy()  # convert it into a numpy array
+        if image_numpy.shape[0] == 1:  # grayscale to RGB
+            image_numpy = np.tile(image_numpy, (3, 1, 1))
+            # image_numpy = image_numpy.convert('L')
+
+        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling
+    else:  # if it is a numpy array, do nothing
+        image_numpy = input_image
+    # image_numpy =
+    return np.clip(image_numpy, 0, 255).astype(imtype)
+
+def calc_RMSE(real_img, fake_img):
+    # convert to LAB color space
+    real_lab = rgb2lab(real_img)
+    fake_lab = rgb2lab(fake_img)
+    return real_lab - fake_lab
+
+def tensor2uint(img):
+    img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy()
+    if img.ndim == 3:
+        img = np.transpose(img, (1, 2, 0))
+    return np.uint8((img*255.0).round())
+
+def imsave(img, img_path):
+    img = np.squeeze(img)
+    if img.ndim == 3:
+        img = img[:, :, [2, 1, 0]]
+    cv2.imwrite(img_path, img)
+
+def process_normal(normal):
+    normal = normal * 2.0 - 1.0
+    normal = normal[:,:,np.newaxis,:]
+    normalizer = np.sqrt(normal @ normal.transpose(0,1,3,2))
+    normalizer = np.squeeze(normalizer, axis=-2)
+    normalizer = np.clip(normalizer, 1.0e-20, 1.0e10)
+    normal = np.squeeze(normal, axis=-2)
+    normal = normal / normalizer 
+    return normal
+
+
+def depthToPoint(fov, depth):
+    # width = 512
+    # height = 512
+    height, width = depth.shape
+    fov_radians = np.deg2rad(fov)
+
+    focal_length = width / (2 * np.tan(fov_radians / 2))
+    fx = focal_length  
+    fy = focal_length 
+    cx = (width - 1) / 2.0  
+    cy = (height - 1) / 2.0  
+
+    x, y = np.meshgrid(range(width), range(height))
+    z = depth 
+    x_3d = (x - cx) * z / fx
+    y_3d = (y - cy) * z / fy
+    x_3d = x_3d.astype(np.float32)
+    y_3d = y_3d.astype(np.float32)
+
+
+    point_cloud = np.stack((x_3d, y_3d, z), axis=-1)
+
+    return point_cloud
+
+def grid_sample(input, img_size):
+    x = torch.linspace(-1, 1, img_size[0])
+    y = torch.linspace(-1, 1, img_size[1])
+    meshx, meshy = torch.meshgrid((x, y))
+    grid = torch.stack((meshy, meshx),2).unsqueeze(0).cuda()
+    grid = grid.repeat(input.shape[0],1,1,1)
+    input = torch.nn.functional.grid_sample(input, grid, mode="nearest", align_corners=False)
+    return input
+
+

utils/loader.py

@@ -0,0 +1,17 @@
+import os
+
+from dataset import DataLoaderTrain, DataLoaderVal, DataLoaderTest
+def get_training_data(rgb_dir, img_options, debug):
+    assert os.path.exists(rgb_dir)
+    return DataLoaderTrain(rgb_dir, img_options, None, debug)
+
+def get_validation_data(rgb_dir,  debug=False):
+    assert os.path.exists(rgb_dir)
+    return DataLoaderVal(rgb_dir, None, debug)
+
+def get_test_data(rgb_dir,  debug=False):
+    assert os.path.exists(rgb_dir)
+    return DataLoaderTest(rgb_dir, None, debug)
+
+
+

utils/misc.py

@@ -0,0 +1,124 @@
+import numpy as np
+import torchvision.utils as vutils
+
+
+def get_np_imgrid(array, nrow=3, padding=0, pad_value=0):
+    '''
+    achieves the same function of torchvision.utils.make_grid for
+    numpy array
+    '''
+    # assume every image has smae size
+    n, h, w, c = array.shape
+    row_num = n // nrow + (n % nrow != 0)
+    gh, gw = row_num*h + padding*(row_num-1), nrow*w + padding*(nrow - 1)
+    grid = np.ones((gh, gw, c), dtype=array.dtype) * pad_value
+    for i in range(n):
+        grow, gcol = i // nrow, i % nrow
+        off_y, off_x = grow * (h + padding), gcol * (w + padding)
+        grid[off_y : off_y + h, off_x : off_x + w] = array[i]
+    return grid
+
+
+def split_np_imgrid(imgrid, nimg, nrow, padding=0):
+    '''
+    reverse operation of make_grid.
+    args:
+        imgrid: HWC image grid
+        nimg: number of images in the grid
+        nrow: number of columns in image grid
+    return:
+        images: list, contains splitted images
+    '''
+    row_num = nimg // nrow + (nimg % nrow != 0)
+    gh, gw, _ = imgrid.shape
+    h, w = (gh - (row_num-1)*padding)//row_num, (gw - (nrow-1)*padding)//nrow
+    images = []
+    for gid in range(nimg):
+        grow, gcol = gid // nrow, gid % nrow 
+        off_i, off_j = grow * (h + padding), gcol * (w + padding)
+        images.append(imgrid[off_i:off_i+h, off_j:off_j+w])
+    return images
+
+
+class MDTableConvertor:
+    
+    def __init__(self, col_num):
+        self.col_num = col_num
+        
+    def _get_table_row(self, items):
+        row = ''
+        for item in items:
+            row += '| {:s} '.format(item)
+        row += '|\n'
+        return row
+
+    def convert(self, item_list, title=None):
+        '''
+        args: 
+            item_list: a list of items (str or can be converted to str)
+            that want to be presented in table.
+
+            title: None, or a list of strings. When set to None, empty title
+            row is used and column number is determined by col_num; Otherwise, 
+            it will be used as title row, its length will override col_num.
+
+        return: 
+            table: markdown table string.
+        '''
+        table = ''
+        if title: # not None or not []  both equal to true
+            col_num = len(title)
+            table += self._get_table_row(title)
+        else:
+            col_num=self.col_num
+            table += self._get_table_row([' ']*col_num) # empty title row
+        table += self._get_table_row(['-'] * col_num) # header spliter
+        for i in range(0, len(item_list), col_num):
+            table += self._get_table_row(item_list[i:i+col_num])
+        return table
+    
+
+def visual_dict_to_imgrid(visual_dict, col_num=4, padding=0):
+    '''
+    args:
+        visual_dict: a dictionary of images of the same size
+        col_num: number of columns in image grid
+        padding: number of padding pixels to seperate images
+    '''
+    im_names = []
+    im_tensors = []
+    for name, visual in visual_dict.items():
+        im_names.append(name)
+        im_tensors.append(visual)
+    im_grid = vutils.make_grid(im_tensors,
+                               nrow=col_num ,
+                               padding=0,
+                               pad_value=1.0)
+    layout = MDTableConvertor(col_num).convert(im_names)
+    
+    return im_grid, layout
+
+
+def count_parameters(model, trainable_only=False):
+    return sum(p.numel() for p in model.parameters())
+    
+    
+
+class WarmupExpLRScheduler(object):
+    def __init__(self, lr_start=1e-4, lr_max=4e-4, lr_min=5e-6, rampup_epochs=4, sustain_epochs=0, exp_decay=0.75):
+        self.lr_start = lr_start
+        self.lr_max = lr_max
+        self.lr_min = lr_min
+        self.rampup_epochs = rampup_epochs
+        self.sustain_epochs = sustain_epochs
+        self.exp_decay = exp_decay
+    
+    def __call__(self, epoch):
+        if epoch < self.rampup_epochs:
+            lr = (self.lr_max - self.lr_start) / self.rampup_epochs * epoch + self.lr_start
+        elif epoch < self.rampup_epochs + self.sustain_epochs:
+            lr = self.lr_max
+        else:
+            lr = (self.lr_max - self.lr_min) * self.exp_decay**(epoch - self.rampup_epochs - self.sustain_epochs) + self.lr_min
+        # print(lr)
+        return lr

utils/model_utils.py

@@ -0,0 +1,98 @@
+import torch
+import os
+from collections import OrderedDict
+
+def freeze(model):
+    for p in model.parameters():
+        p.requires_grad=False
+
+def unfreeze(model):
+    for p in model.parameters():
+        p.requires_grad=True
+
+def is_frozen(model):
+    x = [p.requires_grad for p in model.parameters()]
+    return not all(x)
+
+def save_checkpoint(model_dir, state, session):
+    epoch = state['epoch']
+    model_out_path = os.path.join(model_dir,"model_epoch_{}_{}.pth".format(epoch,session))
+    torch.save(state, model_out_path)
+
+def load_checkpoint(model, weights, strict=True):
+    checkpoint = torch.load(weights, map_location=torch.device('cpu'))
+    try:
+        state_dict = checkpoint["state_dict"]
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            new_state_dict[k] = v
+        model.load_state_dict(new_state_dict, strict=strict)
+    except:
+        state_dict = checkpoint["state_dict"]
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            name = k[7:] if 'module.' in k else k
+            new_state_dict[name] = v
+        model.load_state_dict(new_state_dict, strict=strict)
+
+def load_checkpoint_multigpu(model, weights):
+    checkpoint = torch.load(weights)
+    state_dict = checkpoint["state_dict"]
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        name = k[7:] 
+        new_state_dict[name] = v
+    model.load_state_dict(new_state_dict)
+
+def load_start_epoch(weights):
+    checkpoint = torch.load(weights,  map_location=torch.device('cpu'))
+    epoch = checkpoint["epoch"]
+    return epoch
+
+def load_optim(optimizer, weights):
+    checkpoint = torch.load(weights,  map_location=torch.device('cpu'))
+    optimizer.load_state_dict(checkpoint['optimizer'])
+    for p in optimizer.param_groups: lr = p['lr']
+    return lr
+
+def get_arch(opt):
+    from model import ShadowFormer, ShadowFormerFreq
+    arch = opt.arch
+
+    print('You choose '+arch+'...')
+    if arch == 'ShadowFormer':
+        model_restoration = ShadowFormer(img_size=opt.train_ps,embed_dim=opt.embed_dim,
+                                        win_size=opt.win_size,token_projection=opt.token_projection,
+                                        token_mlp=opt.token_mlp)
+    elif arch == 'ShadowFormerFreq':
+        model_restoration = ShadowFormerFreq(img_size=opt.train_ps,embed_dim=opt.embed_dim,
+                                        win_size=opt.win_size,token_projection=opt.token_projection,
+                                        token_mlp=opt.token_mlp)
+    else:
+        raise Exception("Arch error!")
+
+    return model_restoration
+
+
+def window_partition(x, win_size):
+    B, C, H, W = x.shape
+    x = x.permute(0,2,3,1)
+    x = x.reshape(B, H // win_size, win_size, W // win_size, win_size, C)
+    x = x.permute(0, 1, 3, 2, 4, 5).reshape(-1, win_size, win_size, C)
+    return x.permute(0,3,1,2)
+
+# def distributed_concat(var, num_total):
+#     var_list = [torch.zeros(1, dtype=var.dtype).cuda() for _ in range(torch.distributed.get_world_size())]
+#     torch.distributed.all_gather(var_list, var)
+#     # truncate the dummy elements added by SequentialDistributedSampler
+#     return var_list[:num_total]
+
+def distributed_concat(var, num_total):
+    # 確保 var 是一個 1D tensor (shape: [1])
+    var = var.view(1) if var.dim() == 0 else var
+
+    var_list = [torch.zeros_like(var).cuda() for _ in range(torch.distributed.get_world_size())]
+    torch.distributed.all_gather(var_list, var)
+
+    # truncate the dummy elements added by SequentialDistributedSampler
+    return var_list[:num_total]

utils/shadow_mask_evaluate.py

@@ -0,0 +1,141 @@
+import torch
+import numpy as np
+from collections import OrderedDict
+import pandas as pd
+import os
+from tqdm import tqdm
+import cv2
+from utils.misc import split_np_imgrid, get_np_imgrid
+
+
+def cal_ber(tn, tp, fn, fp):
+    return  0.5*(fp/(tn+fp) + fn/(fn+tp))
+
+def cal_acc(tn, tp, fn, fp):
+    return (tp + tn) / (tp + tn + fp + fn)
+
+
+def get_binary_classification_metrics(pred, gt, threshold=None):
+    if threshold is not None:
+        gt = (gt > threshold)
+        pred = (pred > threshold)
+    TP = np.logical_and(gt, pred).sum()
+    TN = np.logical_and(np.logical_not(gt), np.logical_not(pred)).sum()
+    FN = np.logical_and(gt, np.logical_not(pred)).sum()
+    FP = np.logical_and(np.logical_not(gt), pred).sum()
+    BER = cal_ber(TN, TP, FN, FP)
+    ACC = cal_acc(TN, TP, FN, FP)
+    return OrderedDict( [('TP', TP),
+                        ('TN', TN),
+                        ('FP', FP),
+                        ('FN', FN),
+                        ('BER', BER),
+                        ('ACC', ACC)]
+                      )
+
+
+def evaluate(res_root, pred_id, gt_id, nimg, nrow, threshold):
+    img_names  = os.listdir(res_root)
+    score_dict = OrderedDict()
+
+    for img_name in img_names:
+        im_grid_path = os.path.join(res_root, img_name)
+        im_grid = cv2.imread(im_grid_path)
+        ims = split_np_imgrid(im_grid, nimg, nrow)
+        pred = ims[pred_id]
+        gt = ims[gt_id]
+        score_dict[img_name] = get_binary_classification_metrics(pred,
+                                                                 gt,
+                                                                 threshold)
+            
+    df = pd.DataFrame(score_dict)
+    df['ave'] = df.mean(axis=1)
+
+    tn = df['ave']['TN']
+    tp = df['ave']['TP']
+    fn = df['ave']['FN']
+    fp = df['ave']['FP']
+
+    pos_err = (1 - tp / (tp + fn)) * 100
+    neg_err = (1 - tn / (tn + fp)) * 100
+    ber = (pos_err + neg_err) / 2
+    acc = (tn + tp) / (tn + tp + fn + fp)
+
+    return pos_err, neg_err, ber, acc, df
+
+
+
+
+
+
+
+###############################################
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, weight=1):
+        self.sum += val * weight
+        self.count += weight
+
+    def average(self):
+        if self.count == 0:
+            return 0
+        else:
+            return self.sum / self.count
+
+    def clear(self):
+        self.sum = 0
+        self.count = 0
+
+def compute_cm_torch(y_pred, y_label, n_class):
+    mask = (y_label >= 0) & (y_label < n_class)
+    hist = torch.bincount(n_class * y_label[mask] + y_pred[mask],
+                          minlength=n_class**2).reshape(n_class, n_class)
+    return hist
+
+class MyConfuseMatrixMeter(AverageMeter):
+    """More Clear Confusion Matrix Meter"""
+    def __init__(self, n_class):
+        super(MyConfuseMatrixMeter, self).__init__()
+        self.n_class = n_class
+
+    def update_cm(self, y_pred, y_label, weight=1):
+        y_label = y_label.type(torch.int64)
+        val = compute_cm_torch(y_pred=y_pred.flatten(), y_label=y_label.flatten(),
+                               n_class=self.n_class)
+        self.update(val, weight)
+
+    # def get_scores_binary(self):
+    #     assert self.n_class == 2, "this function can only be called for binary calssification problem"
+    #     tn, fp, fn, tp = self.sum.flatten()
+    #     eps = torch.finfo(torch.float32).eps
+    #     precision = tp / (tp + fp + eps)
+    #     recall = tp / (tp + fn + eps)
+    #     f1 = 2*recall*precision / (recall + precision + eps)
+    #     iou = tp / (tp + fn + fp + eps)
+    #     oa = (tp + tn) / (tp + tn + fn + fp + eps)
+    #     score_dict = {}
+    #     score_dict['precision'] = precision.item()
+    #     score_dict['recall'] = recall.item()
+    #     score_dict['f1'] = f1.item()
+    #     score_dict['iou'] = iou.item()
+    #     score_dict['oa'] = oa.item()
+    #     return score_dict
+    def get_scores_binary(self):
+        assert self.n_class == 2, "this function can only be called for binary calssification problem"
+        tn, fp, fn, tp = self.sum.flatten()
+        eps = torch.finfo(torch.float32).eps
+        pos_err = (1 - tp / (tp + fn + eps)) * 100
+        neg_err = (1 - tn / (tn + fp + eps)) * 100
+        ber = (pos_err + neg_err) / 2
+        acc = (tn + tp) / (tn + tp + fn + fp + eps)
+        score_dict = {}
+        score_dict['pos_err'] = pos_err
+        score_dict['neg_err'] = neg_err
+        score_dict['ber'] = ber
+        score_dict['acc'] = acc
+        return score_dict

utils/tta.py

@@ -0,0 +1,205 @@
+import torch
+
+
+class TestTimeAugmentation:
+    """Test-Time Augmentation for image restoration models"""
+    
+    def __init__(self, model, dino_net, device, use_flip=True, use_rot=True, use_multi_scale=False, scales=None):
+        """
+        Args:
+            model: The model to apply TTA to
+            dino_net: DINO feature extractor
+            device: Device to run inference on
+            use_flip: Whether to use horizontal and vertical flips
+            use_rot: Whether to use 90-degree rotations
+            use_multi_scale: Whether to use multi-scale testing
+            scales: List of scales to use for multi-scale testing, e.g. [0.8, 1.0, 1.2]
+        """
+        self.model = model
+        self.dino_net = dino_net
+        self.device = device
+        self.use_flip = use_flip
+        self.use_rot = use_rot
+        self.use_multi_scale = use_multi_scale
+        self.scales = scales or [1.0]
+        
+    def _apply_augmentation(self, image, point, normal, aug_type):
+        """Apply single augmentation to input images
+        
+        Args:
+            image: Input RGB image
+            point: Point map
+            normal: Normal map
+            aug_type: Augmentation type string (e.g., 'original', 'h_flip', etc.)
+            
+        Returns:
+            Augmented versions of image, point map and normal map
+        """
+        if aug_type == 'original':
+            return image, point, normal
+            
+        elif aug_type == 'h_flip':
+            # Horizontal flip
+            img_aug = torch.flip(image, dims=[3])
+            point_aug = torch.flip(point, dims=[3])
+            normal_aug = torch.flip(normal, dims=[3])
+            # For normal map, x direction needs to be flipped
+            normal_aug[:, 0, :, :] = -normal_aug[:, 0, :, :]
+            return img_aug, point_aug, normal_aug
+            
+        elif aug_type == 'v_flip':
+            # Vertical flip
+            img_aug = torch.flip(image, dims=[2])
+            point_aug = torch.flip(point, dims=[2])
+            normal_aug = torch.flip(normal, dims=[2])
+            # For normal map, y direction needs to be flipped
+            normal_aug[:, 1, :, :] = -normal_aug[:, 1, :, :]
+            return img_aug, point_aug, normal_aug
+            
+        elif aug_type == 'rot90':
+            # 90-degree rotation
+            img_aug = torch.rot90(image, k=1, dims=[2, 3])
+            point_aug = torch.rot90(point, k=1, dims=[2, 3])
+            normal_aug = torch.rot90(normal, k=1, dims=[2, 3])
+            # Swap x and y channels in normal map and negate x
+            normal_x = -normal_aug[:, 1, :, :].clone()
+            normal_y = normal_aug[:, 0, :, :].clone()
+            normal_aug[:, 0, :, :] = normal_x
+            normal_aug[:, 1, :, :] = normal_y
+            return img_aug, point_aug, normal_aug
+            
+        elif aug_type == 'rot180':
+            # 180-degree rotation
+            img_aug = torch.rot90(image, k=2, dims=[2, 3])
+            point_aug = torch.rot90(point, k=2, dims=[2, 3])
+            normal_aug = torch.rot90(normal, k=2, dims=[2, 3])
+            # For normal map, both x and y directions need to be flipped
+            normal_aug[:, 0, :, :] = -normal_aug[:, 0, :, :]
+            normal_aug[:, 1, :, :] = -normal_aug[:, 1, :, :]
+            return img_aug, point_aug, normal_aug
+            
+        elif aug_type == 'rot270':
+            # 270-degree rotation
+            img_aug = torch.rot90(image, k=3, dims=[2, 3])
+            point_aug = torch.rot90(point, k=3, dims=[2, 3])
+            normal_aug = torch.rot90(normal, k=3, dims=[2, 3])
+            # Swap x and y channels in normal map and negate y
+            normal_x = normal_aug[:, 1, :, :].clone()
+            normal_y = -normal_aug[:, 0, :, :].clone()
+            normal_aug[:, 0, :, :] = normal_x
+            normal_aug[:, 1, :, :] = normal_y
+            return img_aug, point_aug, normal_aug
+        
+        else:
+            raise ValueError(f"Unknown augmentation type: {aug_type}")
+    
+    def _reverse_augmentation(self, result, aug_type):
+        """Reverse the augmentation on the result
+        
+        Args:
+            result: Model output to reverse augmentation on
+            aug_type: Augmentation type string
+            
+        Returns:
+            De-augmented result
+        """
+        if aug_type == 'original':
+            return result
+            
+        elif aug_type == 'h_flip':
+            return torch.flip(result, dims=[3])
+            
+        elif aug_type == 'v_flip':
+            return torch.flip(result, dims=[2])
+            
+        elif aug_type == 'rot90':
+            return torch.rot90(result, k=3, dims=[2, 3])
+            
+        elif aug_type == 'rot180':
+            return torch.rot90(result, k=2, dims=[2, 3])
+            
+        elif aug_type == 'rot270':
+            return torch.rot90(result, k=1, dims=[2, 3])
+        
+        else:
+            raise ValueError(f"Unknown augmentation type: {aug_type}")
+    
+    def __call__(self, sliding_window, input_img, point, normal):
+        """
+        Apply TTA to the model and return ensemble result
+        
+        Args:
+            sliding_window: SlidingWindowInference class instance
+            input_img: Input RGB image [B, C, H, W]
+            point: Point map [B, C, H, W]
+            normal: Normal map [B, C, H, W]
+            
+        Returns:
+            Ensemble result with TTA [B, C, H, W]
+        """
+        # Define all augmentations to use
+        augmentations = ['original']
+        if self.use_flip:
+            augmentations.extend(['h_flip', 'v_flip'])
+        if self.use_rot:
+            augmentations.extend(['rot90', 'rot180', 'rot270'])
+        
+        # Initialize the result tensor
+        ensemble_result = torch.zeros_like(input_img)
+        ensemble_weight = 0.0
+        
+        # For each scale and augmentation
+        for scale in self.scales:
+            scale_weight = 1.0
+            if scale != 1.0:
+                # Resize inputs for multi-scale testing
+                h, w = input_img.shape[2], input_img.shape[3]
+                new_h, new_w = int(h * scale), int(w * scale)
+                
+                # Resize all inputs
+                resize_fn = torch.nn.functional.interpolate
+                input_img_scaled = resize_fn(input_img, size=(new_h, new_w), mode='bilinear', align_corners=False)
+                point_scaled = resize_fn(point, size=(new_h, new_w), mode='bilinear', align_corners=False)
+                normal_scaled = resize_fn(normal, size=(new_h, new_w), mode='bilinear', align_corners=False)
+                
+                # Normalize normal vectors after resizing
+                normal_norm = torch.sqrt(torch.sum(normal_scaled**2, dim=1, keepdim=True) + 1e-6)
+                normal_scaled = normal_scaled / normal_norm
+            else:
+                input_img_scaled = input_img
+                point_scaled = point
+                normal_scaled = normal
+            
+            # Apply each augmentation
+            for aug_type in augmentations:
+                # Apply augmentation
+                img_aug, point_aug, normal_aug = self._apply_augmentation(
+                    input_img_scaled, point_scaled, normal_scaled, aug_type
+                )
+                
+                # Run model inference with sliding window
+                with torch.cuda.amp.autocast():
+                    result_aug = sliding_window(
+                        model=self.model,
+                        input_=img_aug,
+                        point=point_aug,
+                        normal=normal_aug,
+                        dino_net=self.dino_net,
+                        device=self.device
+                    )
+                
+                # Reverse augmentation on the result
+                result_aug = self._reverse_augmentation(result_aug, aug_type)
+                
+                # Resize back to original size if using multi-scale
+                if scale != 1.0:
+                    result_aug = resize_fn(result_aug, size=(h, w), mode='bilinear', align_corners=False)
+                
+                # Add to ensemble
+                ensemble_result += result_aug * scale_weight
+                ensemble_weight += scale_weight
+        
+        # Average results
+        ensemble_result = ensemble_result / ensemble_weight
+        
+        return ensemble_result