From d26d854ddaab365b632731faa28a950494686aac Mon Sep 17 00:00:00 2001 From: routineLife1 <1315176066@qq.com> Date: Sat, 9 Nov 2024 18:05:53 +0800 Subject: [PATCH] feat: support DAT (#22) * implement DAT. * feat: support DAT * fix: Do not use mutable data structures for argument defaults --------- Co-authored-by: Tohrusky <65994850+Tohrusky@users.noreply.github.com> --- ccrestoration/arch/__init__.py | 1 + ccrestoration/arch/dat_arch.py | 991 +++++++++++++++++++++++++++++ ccrestoration/config/__init__.py | 1 + ccrestoration/config/dat_config.py | 182 ++++++ ccrestoration/model/__init__.py | 1 + ccrestoration/model/dat_model.py | 48 ++ ccrestoration/type/arch.py | 1 + ccrestoration/type/config.py | 16 + ccrestoration/type/model.py | 1 + tests/test_dat.py | 71 +++ 10 files changed, 1313 insertions(+) create mode 100644 ccrestoration/arch/dat_arch.py create mode 100644 ccrestoration/config/dat_config.py create mode 100644 ccrestoration/model/dat_model.py create mode 100644 tests/test_dat.py diff --git a/ccrestoration/arch/__init__.py b/ccrestoration/arch/__init__.py index 908ab3e..596d872 100644 --- a/ccrestoration/arch/__init__.py +++ b/ccrestoration/arch/__init__.py @@ -13,3 +13,4 @@ from ccrestoration.arch.iconvsr_arch import IconVSR # noqa from ccrestoration.arch.msrswvsr_arch import MSRSWVSR # noqa from ccrestoration.arch.scunet_arch import SCUNet # noqa +from ccrestoration.arch.dat_arch import DAT # noqa diff --git a/ccrestoration/arch/dat_arch.py b/ccrestoration/arch/dat_arch.py new file mode 100644 index 0000000..a199975 --- /dev/null +++ b/ccrestoration/arch/dat_arch.py @@ -0,0 +1,991 @@ +# type: ignore +import numpy as np +import torch +from einops import rearrange +from einops.layers.torch import Rearrange +from torch import nn +from torch.nn import functional as F +from torch.nn.init import trunc_normal_ +from torch.utils import checkpoint + +from ccrestoration.arch import ARCH_REGISTRY +from ccrestoration.arch.arch_util import DropPath, Upsample +from ccrestoration.type import ArchType + + +@ARCH_REGISTRY.register(name=ArchType.DAT) +class DAT(nn.Module): + """Dual Aggregation Transformer + Args: + img_size (int): Input image size. Default: 64 + in_chans (int): Number of input image channels. Default: 3 + embed_dim (int): Patch embedding dimension. Default: 180 + depths (tuple(int)): Depth of each residual group (number of DATB in each RG). + split_size (tuple(int)): Height and Width of spatial window. + num_heads (tuple(int)): Number of attention heads in different residual groups. + expansion_factor (float): Ratio of ffn hidden dim to embedding dim. Default: 4 + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. Default: None + drop_rate (float): Dropout rate. Default: 0 + attn_drop_rate (float): Attention dropout rate. Default: 0 + drop_path_rate (float): Stochastic depth rate. Default: 0.1 + act_layer (nn.Module): Activation layer. Default: nn.GELU + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm + use_chk (bool): Whether to use checkpointing to save memory. + upscale: Upscale factor. 2/3/4 for image SR + img_range: Image range. 1. or 255. + resi_connection: The convolutional block before residual connection. '1conv'/'3conv' + """ + + def __init__( + self, + img_size=64, + in_chans=3, + embed_dim=180, + split_size=[2, 4], # noqa + depth=[2, 2, 2, 2], # noqa + num_heads=[2, 2, 2, 2], # noqa + expansion_factor=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.1, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + use_chk=False, + upscale=2, + img_range=1.0, + resi_connection="1conv", + upsampler="pixelshuffle", + **kwargs, + ): + super().__init__() + + num_in_ch = in_chans + num_out_ch = in_chans + num_feat = 64 + self.img_range = img_range + if in_chans == 3: + rgb_mean = (0.5, 0.5, 0.5) + self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1) + else: + self.mean = torch.zeros(1, 1, 1, 1) + self.upscale = upscale + self.upsampler = upsampler + + # ------------------------- 1, Shallow Feature Extraction ------------------------- # + self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1) + + # ------------------------- 2, Deep Feature Extraction ------------------------- # + self.num_layers = len(depth) + self.use_chk = use_chk + self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models + heads = num_heads + + self.before_RG = nn.Sequential(Rearrange("b c h w -> b (h w) c"), nn.LayerNorm(embed_dim)) + + curr_dim = embed_dim + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, np.sum(depth))] # stochastic depth decay rule + + self.layers = nn.ModuleList() + for i in range(self.num_layers): + layer = ResidualGroup( + dim=embed_dim, + num_heads=heads[i], + reso=img_size, + split_size=split_size, + expansion_factor=expansion_factor, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_paths=dpr[sum(depth[:i]) : sum(depth[: i + 1])], + act_layer=act_layer, + norm_layer=norm_layer, + depth=depth[i], + use_chk=use_chk, + resi_connection=resi_connection, + rg_idx=i, + ) + self.layers.append(layer) + + self.norm = norm_layer(curr_dim) + # build the last conv layer in deep feature extraction + if resi_connection == "1conv": + self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1) + elif resi_connection == "3conv": + # to save parameters and memory + self.conv_after_body = nn.Sequential( + nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1), + nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0), + nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1), + ) + + # ------------------------- 3, Reconstruction ------------------------- # + if self.upsampler == "pixelshuffle": + # for classical SR + self.conv_before_upsample = nn.Sequential( + nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True) + ) + self.upsample = Upsample(upscale, num_feat) + self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) + elif self.upsampler == "pixelshuffledirect": + # for lightweight SR (to save parameters) + self.upsample = UpsampleOneStep(upscale, embed_dim, num_out_ch, (img_size, img_size)) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d)): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def forward_features(self, x): + _, _, H, W = x.shape + x_size = [H, W] + x = self.before_RG(x) + for layer in self.layers: + x = layer(x, x_size) + x = self.norm(x) + x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W) + + return x + + def forward(self, x): + """Input: x: (B, C, H, W)""" + self.mean = self.mean.type_as(x) + x = (x - self.mean) * self.img_range + + if self.upsampler == "pixelshuffle": + # for image SR + x = self.conv_first(x) + x = self.conv_after_body(self.forward_features(x)) + x + x = self.conv_before_upsample(x) + x = self.conv_last(self.upsample(x)) + elif self.upsampler == "pixelshuffledirect": + # for lightweight SR + x = self.conv_first(x) + x = self.conv_after_body(self.forward_features(x)) + x + x = self.upsample(x) + + x = x / self.img_range + self.mean + return x + + +def img2windows(img, H_sp, W_sp): + """Input: Image (B, C, H, W) + Output: Window Partition (B', N, C) + """ + B, C, H, W = img.shape + img_reshape = img.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp) + img_perm = img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, H_sp * W_sp, C) + return img_perm + + +def windows2img(img_splits_hw, H_sp, W_sp, H, W): + """Input: Window Partition (B', N, C) + Output: Image (B, H, W, C) + """ + B = int(img_splits_hw.shape[0] / (H * W / H_sp / W_sp)) + + img = img_splits_hw.view(B, H // H_sp, W // W_sp, H_sp, W_sp, -1) + img = img.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return img + + +class SpatialGate(nn.Module): + """Spatial-Gate. + + Args: + ---- + dim (int): Half of input channels. + + """ + + def __init__(self, dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.conv = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim) # DW Conv + + def forward(self, x, H, W): + # Split + x1, x2 = x.chunk(2, dim=-1) + B, N, C = x.shape + x2 = ( + self.conv(self.norm(x2).transpose(1, 2).contiguous().view(B, C // 2, H, W)) + .flatten(2) + .transpose(-1, -2) + .contiguous() + ) + + return x1 * x2 + + +class SGFN(nn.Module): + """Spatial-Gate Feed-Forward Network. + + Args: + ---- + in_features (int): Number of input channels. + hidden_features (int | None): Number of hidden channels. Default: None + out_features (int | None): Number of output channels. Default: None + act_layer (nn.Module): Activation layer. Default: nn.GELU + drop (float): Dropout rate. Default: 0.0 + + """ + + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + drop=0.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.sg = SpatialGate(hidden_features // 2) + self.fc2 = nn.Linear(hidden_features // 2, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x, H, W): + """Input: x: (B, H*W, C), H, W + Output: x: (B, H*W, C) + """ + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + + x = self.sg(x, H, W) + x = self.drop(x) + + x = self.fc2(x) + x = self.drop(x) + return x + + +class DynamicPosBias(nn.Module): + # The implementation builds on Crossformer code https://github.com/cheerss/CrossFormer/blob/main/models/crossformer.py + + """Dynamic Relative Position Bias. + + Args: + ---- + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + residual (bool): If True, use residual strage to connect conv. + + """ + + def __init__(self, dim, num_heads, residual): + super().__init__() + self.residual = residual + self.num_heads = num_heads + self.pos_dim = dim // 4 + self.pos_proj = nn.Linear(2, self.pos_dim) + self.pos1 = nn.Sequential( + nn.LayerNorm(self.pos_dim), + nn.ReLU(inplace=True), + nn.Linear(self.pos_dim, self.pos_dim), + ) + self.pos2 = nn.Sequential( + nn.LayerNorm(self.pos_dim), + nn.ReLU(inplace=True), + nn.Linear(self.pos_dim, self.pos_dim), + ) + self.pos3 = nn.Sequential( + nn.LayerNorm(self.pos_dim), + nn.ReLU(inplace=True), + nn.Linear(self.pos_dim, self.num_heads), + ) + + def forward(self, biases): + if self.residual: + pos = self.pos_proj(biases) # 2Gh-1 * 2Gw-1, heads + pos = pos + self.pos1(pos) + pos = pos + self.pos2(pos) + pos = self.pos3(pos) + else: + pos = self.pos3(self.pos2(self.pos1(self.pos_proj(biases)))) + return pos + + +class Spatial_Attention(nn.Module): + """Spatial Window Self-Attention. + It supports rectangle window (containing square window). + + Args: + ---- + dim (int): Number of input channels. + idx (int): The indentix of window. (0/1) + split_size (tuple(int)): Height and Width of spatial window. + dim_out (int | None): The dimension of the attention output. Default: None + num_heads (int): Number of attention heads. Default: 6 + attn_drop (float): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float): Dropout ratio of output. Default: 0.0 + qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set + position_bias (bool): The dynamic relative position bias. Default: True + + """ + + def __init__( + self, + dim, + idx, + split_size=[8, 8], # noqa + dim_out=None, + num_heads=6, + attn_drop=0.0, + proj_drop=0.0, + qk_scale=None, + position_bias=True, + ): + super().__init__() + self.dim = dim + self.dim_out = dim_out or dim + self.split_size = split_size + self.num_heads = num_heads + self.idx = idx + self.position_bias = position_bias + + head_dim = dim // num_heads + self.scale = qk_scale or head_dim**-0.5 + + if idx == 0: + H_sp, W_sp = self.split_size[0], self.split_size[1] + elif idx == 1: + W_sp, H_sp = self.split_size[0], self.split_size[1] + else: + print("ERROR MODE", idx) + exit(0) + self.H_sp = H_sp + self.W_sp = W_sp + + if self.position_bias: + self.pos = DynamicPosBias(self.dim // 4, self.num_heads, residual=False) + # generate mother-set + position_bias_h = torch.arange(1 - self.H_sp, self.H_sp) + position_bias_w = torch.arange(1 - self.W_sp, self.W_sp) + biases = torch.stack(torch.meshgrid([position_bias_h, position_bias_w], indexing="ij")) + biases = biases.flatten(1).transpose(0, 1).contiguous().float() + self.register_buffer("rpe_biases", biases) + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.H_sp) + coords_w = torch.arange(self.W_sp) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij")) + coords_flatten = torch.flatten(coords, 1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += self.H_sp - 1 + relative_coords[:, :, 1] += self.W_sp - 1 + relative_coords[:, :, 0] *= 2 * self.W_sp - 1 + relative_position_index = relative_coords.sum(-1) + self.register_buffer("relative_position_index", relative_position_index) + + self.attn_drop = nn.Dropout(attn_drop) + + def im2win(self, x, H, W): + B, N, C = x.shape + x = x.transpose(-2, -1).contiguous().view(B, C, H, W) + x = img2windows(x, self.H_sp, self.W_sp) + x = x.reshape(-1, self.H_sp * self.W_sp, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).contiguous() + return x + + def forward(self, qkv, H, W, mask=None): + """Input: qkv: (B, 3*L, C), H, W, mask: (B, N, N), N is the window size + Output: x (B, H, W, C) + """ + q, k, v = qkv[0], qkv[1], qkv[2] + + B, L, C = q.shape + assert L == H * W, "flatten img_tokens has wrong size" + + # partition the q,k,v, image to window + q = self.im2win(q, H, W) + k = self.im2win(k, H, W) + v = self.im2win(v, H, W) + + q = q * self.scale + attn = q @ k.transpose(-2, -1) # B head N C @ B head C N --> B head N N + + # calculate drpe + if self.position_bias: + pos = self.pos(self.rpe_biases) + # select position bias + relative_position_bias = pos[self.relative_position_index.view(-1)].view( + self.H_sp * self.W_sp, self.H_sp * self.W_sp, -1 + ) + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + attn = attn + relative_position_bias.unsqueeze(0) + + N = attn.shape[3] + + # use mask for shift window + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + + attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype) + attn = self.attn_drop(attn) + + x = attn @ v + x = x.transpose(1, 2).reshape(-1, self.H_sp * self.W_sp, C) # B head N N @ B head N C + + # merge the window, window to image + x = windows2img(x, self.H_sp, self.W_sp, H, W) # B H' W' C + + return x + + +class Axial_Spatial_Attention(nn.Module): + """Axial Spatial Self-Attention + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. Default: 6 + split_size (tuple(int)): Height and Width of spatial window. + shift_size (tuple(int)): Shift size for spatial window. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. + drop (float): Dropout rate. Default: 0.0 + attn_drop (float): Attention dropout rate. Default: 0.0 + rg_idx (int): The indentix of Residual Group (RG) + b_idx (int): The indentix of Block in each RG + """ + + def __init__( + self, + dim, + num_heads, + reso=64, + split_size=[8, 8], # noqa + shift_size=[1, 2], # noqa + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + rg_idx=0, + b_idx=0, + ): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.split_size = split_size + self.shift_size = shift_size + self.b_idx = b_idx + self.rg_idx = rg_idx + self.patches_resolution = reso + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + + assert 0 <= self.shift_size[0] < self.split_size[0], "shift_size must in 0-split_size0" + assert 0 <= self.shift_size[1] < self.split_size[1], "shift_size must in 0-split_size1" + + self.branch_num = 2 + + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(drop) + + self.attns = nn.ModuleList( + [ + Spatial_Attention( + dim // 2, + idx=i, + split_size=split_size, + num_heads=num_heads // 2, + dim_out=dim // 2, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + position_bias=True, + ) + for i in range(self.branch_num) + ] + ) + + if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or ( + self.rg_idx % 2 != 0 and self.b_idx % 4 == 0 + ): + attn_mask = self.calculate_mask(self.patches_resolution, self.patches_resolution) + self.register_buffer("attn_mask_0", attn_mask[0]) + self.register_buffer("attn_mask_1", attn_mask[1]) + else: + attn_mask = None + self.register_buffer("attn_mask_0", None) + self.register_buffer("attn_mask_1", None) + + self.dwconv = nn.Sequential( + nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim), + nn.BatchNorm2d(dim), + nn.GELU(), + ) + self.channel_interaction = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(dim, dim // 8, kernel_size=1), + nn.BatchNorm2d(dim // 8), + nn.GELU(), + nn.Conv2d(dim // 8, dim, kernel_size=1), + ) + self.spatial_interaction = nn.Sequential( + nn.Conv2d(dim, dim // 16, kernel_size=1), + nn.BatchNorm2d(dim // 16), + nn.GELU(), + nn.Conv2d(dim // 16, 1, kernel_size=1), + ) + + def calculate_mask(self, H, W): + # The implementation builds on Swin Transformer code https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py + # calculate attention mask for shift window + img_mask_0 = torch.zeros((1, H, W, 1)) # 1 H W 1 idx=0 + img_mask_1 = torch.zeros((1, H, W, 1)) # 1 H W 1 idx=1 + h_slices_0 = ( + slice(0, -self.split_size[0]), + slice(-self.split_size[0], -self.shift_size[0]), + slice(-self.shift_size[0], None), + ) + w_slices_0 = ( + slice(0, -self.split_size[1]), + slice(-self.split_size[1], -self.shift_size[1]), + slice(-self.shift_size[1], None), + ) + + h_slices_1 = ( + slice(0, -self.split_size[1]), + slice(-self.split_size[1], -self.shift_size[1]), + slice(-self.shift_size[1], None), + ) + w_slices_1 = ( + slice(0, -self.split_size[0]), + slice(-self.split_size[0], -self.shift_size[0]), + slice(-self.shift_size[0], None), + ) + cnt = 0 + for h in h_slices_0: + for w in w_slices_0: + img_mask_0[:, h, w, :] = cnt + cnt += 1 + cnt = 0 + for h in h_slices_1: + for w in w_slices_1: + img_mask_1[:, h, w, :] = cnt + cnt += 1 + + # calculate mask for window-0 + img_mask_0 = img_mask_0.view( + 1, + H // self.split_size[0], + self.split_size[0], + W // self.split_size[1], + self.split_size[1], + 1, + ) + img_mask_0 = ( + img_mask_0.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.split_size[0], self.split_size[1], 1) + ) # nW, sw[0], sw[1], 1 + mask_windows_0 = img_mask_0.view(-1, self.split_size[0] * self.split_size[1]) + attn_mask_0 = mask_windows_0.unsqueeze(1) - mask_windows_0.unsqueeze(2) + attn_mask_0 = attn_mask_0.masked_fill(attn_mask_0 != 0, -100.0).masked_fill(attn_mask_0 == 0, 0.0) + + # calculate mask for window-1 + img_mask_1 = img_mask_1.view( + 1, + H // self.split_size[1], + self.split_size[1], + W // self.split_size[0], + self.split_size[0], + 1, + ) + img_mask_1 = ( + img_mask_1.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.split_size[1], self.split_size[0], 1) + ) # nW, sw[1], sw[0], 1 + mask_windows_1 = img_mask_1.view(-1, self.split_size[1] * self.split_size[0]) + attn_mask_1 = mask_windows_1.unsqueeze(1) - mask_windows_1.unsqueeze(2) + attn_mask_1 = attn_mask_1.masked_fill(attn_mask_1 != 0, -100.0).masked_fill(attn_mask_1 == 0, 0.0) + + return attn_mask_0, attn_mask_1 + + def forward(self, x, H, W): + """Input: x: (B, H*W, C), H, W + Output: x: (B, H*W, C) + """ + B, L, C = x.shape + assert L == H * W, "flatten img_tokens has wrong size" + + qkv = self.qkv(x).reshape(B, -1, 3, C).permute(2, 0, 1, 3) # 3, B, HW, C + # V without partition + v = qkv[2].transpose(-2, -1).contiguous().view(B, C, H, W) + + # image padding + max_split_size = max(self.split_size[0], self.split_size[1]) + pad_l = pad_t = 0 + pad_r = (max_split_size - W % max_split_size) % max_split_size + pad_b = (max_split_size - H % max_split_size) % max_split_size + + qkv = qkv.reshape(3 * B, H, W, C).permute(0, 3, 1, 2) # 3B C H W + qkv = F.pad(qkv, (pad_l, pad_r, pad_t, pad_b)).reshape(3, B, C, -1).transpose(-2, -1) # l r t b + _H = pad_b + H + _W = pad_r + W + _L = _H * _W + + # window-0 and window-1 on split channels [C/2, C/2]; for square windows (e.g., 8x8), window-0 and window-1 can be merged + # shift in block: (0, 4, 8, ...), (2, 6, 10, ...), (0, 4, 8, ...), (2, 6, 10, ...), ... + if (self.rg_idx % 2 == 0 and self.b_idx > 0 and (self.b_idx - 2) % 4 == 0) or ( + self.rg_idx % 2 != 0 and self.b_idx % 4 == 0 + ): + qkv = qkv.view(3, B, _H, _W, C) + qkv_0 = torch.roll( + qkv[:, :, :, :, : C // 2], + shifts=(-self.shift_size[0], -self.shift_size[1]), + dims=(2, 3), + ) + qkv_0 = qkv_0.view(3, B, _L, C // 2) + qkv_1 = torch.roll( + qkv[:, :, :, :, C // 2 :], + shifts=(-self.shift_size[1], -self.shift_size[0]), + dims=(2, 3), + ) + qkv_1 = qkv_1.view(3, B, _L, C // 2) + + if self.patches_resolution != _H or self.patches_resolution != _W: + mask_tmp = self.calculate_mask(_H, _W) + x1_shift = self.attns[0](qkv_0, _H, _W, mask=mask_tmp[0].to(x.device)) + x2_shift = self.attns[1](qkv_1, _H, _W, mask=mask_tmp[1].to(x.device)) + else: + x1_shift = self.attns[0](qkv_0, _H, _W, mask=self.attn_mask_0) + x2_shift = self.attns[1](qkv_1, _H, _W, mask=self.attn_mask_1) + + x1 = torch.roll(x1_shift, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2)) + x2 = torch.roll(x2_shift, shifts=(self.shift_size[1], self.shift_size[0]), dims=(1, 2)) + x1 = x1[:, :H, :W, :].reshape(B, L, C // 2) + x2 = x2[:, :H, :W, :].reshape(B, L, C // 2) + # attention output + attened_x = torch.cat([x1, x2], dim=2) + + else: + x1 = self.attns[0](qkv[:, :, :, : C // 2], _H, _W)[:, :H, :W, :].reshape(B, L, C // 2) + x2 = self.attns[1](qkv[:, :, :, C // 2 :], _H, _W)[:, :H, :W, :].reshape(B, L, C // 2) + # attention output + attened_x = torch.cat([x1, x2], dim=2) + + # convolution output + conv_x = self.dwconv(v) + + # Adaptive Interaction Module (AIM) + # C-Map (before sigmoid) + channel_map = self.channel_interaction(conv_x).permute(0, 2, 3, 1).contiguous().view(B, 1, C) + # S-Map (before sigmoid) + attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W) + spatial_map = self.spatial_interaction(attention_reshape) + + # C-I + attened_x = attened_x * torch.sigmoid(channel_map) + # S-I + conv_x = torch.sigmoid(spatial_map) * conv_x + conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(B, L, C) + + x = attened_x + conv_x + + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class Axial_Channel_Attention(nn.Module): + # The implementation builds on XCiT code https://github.com/facebookresearch/xcit + + """Axial Channel Self-Attention + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. Default: 6 + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. + attn_drop (float): Attention dropout rate. Default: 0.0 + drop_path (float): Stochastic depth rate. Default: 0.0 + """ + + def __init__( + self, + dim, + num_heads=8, + qkv_bias=False, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + ): + super().__init__() + self.num_heads = num_heads + self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1)) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.dwconv = nn.Sequential( + nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim), + nn.BatchNorm2d(dim), + nn.GELU(), + ) + self.channel_interaction = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Conv2d(dim, dim // 8, kernel_size=1), + nn.BatchNorm2d(dim // 8), + nn.GELU(), + nn.Conv2d(dim // 8, dim, kernel_size=1), + ) + self.spatial_interaction = nn.Sequential( + nn.Conv2d(dim, dim // 16, kernel_size=1), + nn.BatchNorm2d(dim // 16), + nn.GELU(), + nn.Conv2d(dim // 16, 1, kernel_size=1), + ) + + def forward(self, x, H, W): + """Input: x: (B, H*W, C), H, W + Output: x: (B, H*W, C) + """ + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) + qkv = qkv.permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] + + q = q.transpose(-2, -1) + k = k.transpose(-2, -1) + v = v.transpose(-2, -1) + + v_ = v.reshape(B, C, N).contiguous().view(B, C, H, W) + + q = torch.nn.functional.normalize(q, dim=-1) + k = torch.nn.functional.normalize(k, dim=-1) + + attn = (q @ k.transpose(-2, -1)) * self.temperature + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + # attention output + attened_x = (attn @ v).permute(0, 3, 1, 2).reshape(B, N, C) + + # convolution output + conv_x = self.dwconv(v_) + + # Adaptive Interaction Module (AIM) + # C-Map (before sigmoid) + attention_reshape = attened_x.transpose(-2, -1).contiguous().view(B, C, H, W) + channel_map = self.channel_interaction(attention_reshape) + # S-Map (before sigmoid) + spatial_map = self.spatial_interaction(conv_x).permute(0, 2, 3, 1).contiguous().view(B, N, 1) + + # S-I + attened_x = attened_x * torch.sigmoid(spatial_map) + # C-I + conv_x = conv_x * torch.sigmoid(channel_map) + conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(B, N, C) + + x = attened_x + conv_x + + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class DATB(nn.Module): + def __init__( + self, + dim, + num_heads, + reso=64, + split_size=[2, 4], # noqa + shift_size=[1, 2], # noqa + expansion_factor=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + rg_idx=0, + b_idx=0, + ): + super().__init__() + + self.norm1 = norm_layer(dim) + + if b_idx % 2 == 0: + # DSTB + self.attn = Axial_Spatial_Attention( + dim, + num_heads=num_heads, + reso=reso, + split_size=split_size, + shift_size=shift_size, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + rg_idx=rg_idx, + b_idx=b_idx, + ) + else: + # DCTB + self.attn = Axial_Channel_Attention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + ffn_hidden_dim = int(dim * expansion_factor) + self.ffn = SGFN( + in_features=dim, + hidden_features=ffn_hidden_dim, + out_features=dim, + act_layer=act_layer, + ) + self.norm2 = norm_layer(dim) + + def forward(self, x, x_size): + """Input: x: (B, H*W, C), x_size: (H, W) + Output: x: (B, H*W, C) + """ + H, W = x_size + x = x + self.drop_path(self.attn(self.norm1(x), H, W)) + x = x + self.drop_path(self.ffn(self.norm2(x), H, W)) + + return x + + +class ResidualGroup(nn.Module): + """ResidualGroup + Args: + dim (int): Number of input channels. + reso (int): Input resolution. + num_heads (int): Number of attention heads. + split_size (tuple(int)): Height and Width of spatial window. + expansion_factor (float): Ratio of ffn hidden dim to embedding dim. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None): Override default qk scale of head_dim ** -0.5 if set. Default: None + drop (float): Dropout rate. Default: 0 + attn_drop(float): Attention dropout rate. Default: 0 + drop_paths (float | None): Stochastic depth rate. + act_layer (nn.Module): Activation layer. Default: nn.GELU + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm + depth (int): Number of dual aggregation Transformer blocks in residual group. + use_chk (bool): Whether to use checkpointing to save memory. + resi_connection: The convolutional block before residual connection. '1conv'/'3conv' + """ + + def __init__( + self, + dim, + reso, + num_heads, + split_size=[2, 4], # noqa + expansion_factor=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_paths=None, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + depth=2, + use_chk=False, + resi_connection="1conv", + rg_idx=0, + ): + super().__init__() + self.use_chk = use_chk + self.reso = reso + + self.blocks = nn.ModuleList( + [ + DATB( + dim=dim, + num_heads=num_heads, + reso=reso, + split_size=split_size, + shift_size=[split_size[0] // 2, split_size[1] // 2], + expansion_factor=expansion_factor, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_paths[i], + act_layer=act_layer, + norm_layer=norm_layer, + rg_idx=rg_idx, + b_idx=i, + ) + for i in range(depth) + ] + ) + + if resi_connection == "1conv": + self.conv = nn.Conv2d(dim, dim, 3, 1, 1) + elif resi_connection == "3conv": + self.conv = nn.Sequential( + nn.Conv2d(dim, dim // 4, 3, 1, 1), + nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(dim // 4, dim // 4, 1, 1, 0), + nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(dim // 4, dim, 3, 1, 1), + ) + + def forward(self, x, x_size): + """Input: x: (B, H*W, C), x_size: (H, W) + Output: x: (B, H*W, C) + """ + H, W = x_size + res = x + for blk in self.blocks: + if self.use_chk: + x = checkpoint.checkpoint(blk, x, x_size, use_reentrant=False) + else: + x = blk(x, x_size) + x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W) + x = self.conv(x) + x = rearrange(x, "b c h w -> b (h w) c") + x = res + x + + return x + + +class UpsampleOneStep(nn.Sequential): + """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle) + Used in lightweight SR to save parameters. + + Args: + ---- + scale (int): Scale factor. Supported scales: 2^n and 3. + num_feat (int): Channel number of intermediate features. + + """ + + def __init__(self, scale, num_feat, num_out_ch, input_resolution=None): + self.num_feat = num_feat + self.input_resolution = input_resolution + m = [] + m.append(nn.Conv2d(num_feat, (scale**2) * num_out_ch, 3, 1, 1)) + m.append(nn.PixelShuffle(scale)) + super(UpsampleOneStep, self).__init__(*m) + + def flops(self): + h, w = self.input_resolution + flops = h * w * self.num_feat * 3 * 9 + return flops diff --git a/ccrestoration/config/__init__.py b/ccrestoration/config/__init__.py index 46a12bc..c67ead7 100644 --- a/ccrestoration/config/__init__.py +++ b/ccrestoration/config/__init__.py @@ -12,3 +12,4 @@ from ccrestoration.config.iconvsr_config import IconVSRConfig # noqa from ccrestoration.config.animesr_config import AnimeSRConfig # noqa from ccrestoration.config.scunet_config import SCUNetConfig # noqa +from ccrestoration.config.dat_config import DATConfig # noqa diff --git a/ccrestoration/config/dat_config.py b/ccrestoration/config/dat_config.py new file mode 100644 index 0000000..e84b425 --- /dev/null +++ b/ccrestoration/config/dat_config.py @@ -0,0 +1,182 @@ +from typing import Any, List, Optional, Tuple, Union + +from pydantic import field_validator +from torch import nn + +from ccrestoration.config import CONFIG_REGISTRY +from ccrestoration.type import ArchType, BaseConfig, ConfigType, ModelType + + +class DATConfig(BaseConfig): + arch: Union[ArchType, str] = ArchType.DAT + model: Union[ModelType, str] = ModelType.DAT + scale: int = 2 + in_chans: int = 3 + img_size: Union[int, Tuple[int, ...]] = 64 + img_range: float = 1.0 + split_size: Union[List[int], Tuple[int, ...]] = [2, 4] # noqa + depth: Union[List[int], Tuple[int, ...]] = [6, 6, 6, 6, 6, 6] # noqa + embed_dim: int = 180 + num_heads: Union[List[int], Tuple[int, ...]] = [6, 6, 6, 6, 6, 6] # noqa + expansion_factor: float = 4.0 + resi_connection: str = "1conv" + qkv_bias: bool = True + qk_scale: Optional[float] = None + drop_rate: float = 0.0 + attn_drop_rate: float = 0.0 + drop_path_rate: float = 0.1 + act_layer: Any = nn.GELU + norm_layer: Any = nn.LayerNorm + use_chk: bool = False + upsampler: str = "pixelshuffle" + + @field_validator("scale") + def scale_match(cls, v: int) -> int: + if v not in [1, 2, 3, 4, 8]: + raise ValueError("scale factor must be one of [1, 2, 3, 4, 8]") + return v + + @field_validator("upsampler") + def upsampler_match(cls, v: str) -> str: + if v not in ["pixelshuffle", "pixelshuffledirect", "nearest+conv", ""]: + raise ValueError("Upsampler must be one of ['pixelshuffle','pixelshuffledirect','nearest+conv', '']") + return v + + @field_validator("resi_connection") + def resi_connection_match(cls, v: str) -> str: + if v not in ["1conv", "3conv"]: + raise ValueError("Residual connection must be one of ['1conv', '3conv']") + return v + + +DATConfigs = [ + # official models + # dat_s + DATConfig( + name=ConfigType.DAT_S_2x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_S_2x.pth", + hash="160330dd8a40b141e12713ca9bfde09e36a03c533b455965f157d023672cb794", + scale=2, + split_size=[8, 16], + expansion_factor=2, + ), + DATConfig( + name=ConfigType.DAT_S_3x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_S_3x.pth", + hash="d1446d3eb2fbaad472c6fd6c2ac03a3467265c5f93822a741b09838b80f18b62", + scale=3, + split_size=[8, 16], + expansion_factor=2, + ), + DATConfig( + name=ConfigType.DAT_S_4x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_S_4x.pth", + hash="2ba8a7cbe2fd88f3499d08d1b353fa17a98b815832f4426a7144a0ef9f3bfcf7", + scale=4, + split_size=[8, 16], + expansion_factor=2, + ), + # dat_m + DATConfig( + name=ConfigType.DAT_2x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_2x.pth", + hash="7760aa96e4ee77e29d4f89c3a4486200042e019461fdb8aa286f49aa00b89b51", + scale=2, + split_size=[8, 32], + expansion_factor=4, + ), + DATConfig( + name=ConfigType.DAT_3x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_3x.pth", + hash="581973e02c06f90d4eb90acf743ec9604f56f3c2c6f9e1e2c2b38ded1f80d197", + scale=3, + split_size=[8, 32], + expansion_factor=4, + ), + DATConfig( + name=ConfigType.DAT_4x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_4x.pth", + hash="391a6ce69899dff5ea3214557e9d585608254579217169faf3d4c353caff049e", + scale=4, + split_size=[8, 32], + expansion_factor=4, + ), + # dat_2 + DATConfig( + name=ConfigType.DAT_2_2x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_2_2x.pth", + hash="aea2c76996c2b3e7cb034380e07738608bd59cc34c667331df7269e4b670ac18", + scale=2, + split_size=[8, 32], + expansion_factor=2, + ), + DATConfig( + name=ConfigType.DAT_2_3x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_2_3x.pth", + hash="e19fbb2e6addf5cecf90472937fabdd905853b79c4ef807b4dd184c30bb22a28", + scale=3, + split_size=[8, 32], + expansion_factor=2, + ), + DATConfig( + name=ConfigType.DAT_2_4x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_2_4x.pth", + hash="05b5c17bb5d1939ec0ec6b9368368d82d8c45b80c134e370f798efec0aeec395", + scale=4, + split_size=[8, 32], + expansion_factor=2, + ), + # dat_light + DATConfig( + name=ConfigType.DAT_light_2x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_light_2x.pth", + hash="3888bf067c1a3790adb64210e108690859f4900ffd158fe39a2cfe057ae8300c", + scale=2, + split_size=[8, 32], + expansion_factor=2, + depth=[18], + embed_dim=60, + num_heads=[6], + resi_connection="3conv", + upsampler="pixelshuffledirect", + ), + DATConfig( + name=ConfigType.DAT_light_3x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_light_3x.pth", + hash="dfb99deab865db05771080c73c50a06b67c9499c3a27ca64ab3729fa573d5cfe", + scale=3, + split_size=[8, 32], + expansion_factor=2, + depth=[18], + embed_dim=60, + num_heads=[6], + resi_connection="3conv", + upsampler="pixelshuffledirect", + ), + DATConfig( + name=ConfigType.DAT_light_4x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_light_4x.pth", + hash="e18223fc41500e72e6bd576f2a576eae0818f58fa8cc8646a2233f5faf7d0f74", + scale=4, + split_size=[8, 32], + expansion_factor=2, + depth=[18], + embed_dim=60, + num_heads=[6], + resi_connection="3conv", + upsampler="pixelshuffledirect", + ), + # community models + DATConfig( + name=ConfigType.DAT_APISR_GAN_generator_4x, + url="https://github.com/TensoRaws/ccrestoration/releases/download/model_zoo/DAT_APISR_GAN_generator_4x.pth", + hash="cc625c4ec5242a57e46f0941a3c36e3a731dcd859750c5f72a5251045b1e6d72", + scale=4, + split_size=[8, 16], + expansion_factor=2, + upsampler="pixelshuffledirect", + ), +] + +for cfg in DATConfigs: + CONFIG_REGISTRY.register(cfg) diff --git a/ccrestoration/model/__init__.py b/ccrestoration/model/__init__.py index cd8a4ab..1b99c92 100644 --- a/ccrestoration/model/__init__.py +++ b/ccrestoration/model/__init__.py @@ -16,3 +16,4 @@ from ccrestoration.model.iconvsr_model import IconVSRModel # noqa from ccrestoration.model.animesr_model import AnimeSRModel # noqa from ccrestoration.model.scunet_model import SCUNetModel # noqa +from ccrestoration.model.dat_model import DATModel # noqa diff --git a/ccrestoration/model/dat_model.py b/ccrestoration/model/dat_model.py new file mode 100644 index 0000000..b531bb1 --- /dev/null +++ b/ccrestoration/model/dat_model.py @@ -0,0 +1,48 @@ +from typing import Any + +from ccrestoration.arch import DAT +from ccrestoration.config import DATConfig +from ccrestoration.model import MODEL_REGISTRY +from ccrestoration.model.sr_base_model import SRBaseModel +from ccrestoration.type import ModelType + + +@MODEL_REGISTRY.register(name=ModelType.DAT) +class DATModel(SRBaseModel): + def load_model(self) -> Any: + cfg: DATConfig = self.config + state_dict = self.get_state_dict() + + if "params_ema" in state_dict: + state_dict = state_dict["params_ema"] + elif "params" in state_dict: + state_dict = state_dict["params"] + elif "model_state_dict" in state_dict: + # For APISR's model + state_dict = state_dict["model_state_dict"] + + model = DAT( + img_size=cfg.img_size, + in_chans=cfg.in_chans, + embed_dim=cfg.embed_dim, + split_size=cfg.split_size, + depth=cfg.depth, + num_heads=cfg.num_heads, + qkv_bias=cfg.qkv_bias, + qk_scale=cfg.qk_scale, + drop_rate=cfg.drop_rate, + attn_drop_rate=cfg.attn_drop_rate, + drop_path_rate=cfg.drop_path_rate, + act_layer=cfg.act_layer, + norm_layer=cfg.norm_layer, + use_chk=cfg.use_chk, + upscale=cfg.scale, + img_range=cfg.img_range, + upsampler=cfg.upsampler, + resi_connection=cfg.resi_connection, + expansion_factor=cfg.expansion_factor, + ) + + model.load_state_dict(state_dict) + model.eval().to(self.device) + return model diff --git a/ccrestoration/type/arch.py b/ccrestoration/type/arch.py index f49633f..dff1c6a 100644 --- a/ccrestoration/type/arch.py +++ b/ccrestoration/type/arch.py @@ -11,6 +11,7 @@ class ArchType(str, Enum): EDSR = "EDSR" SWINIR = "SWINIR" SCUNET = "SCUNET" + DAT = "DAT" # ------------------------------------- Auxiliary Network ---------------------------------------------------------- diff --git a/ccrestoration/type/config.py b/ccrestoration/type/config.py index 0d64345..e338d7f 100644 --- a/ccrestoration/type/config.py +++ b/ccrestoration/type/config.py @@ -56,6 +56,22 @@ class ConfigType(str, Enum): SCUNet_color_real_psnr_1x = "SCUNet_color_real_psnr_1x.pth" SCUNet_color_real_gan_1x = "SCUNet_color_real_gan_1x.pth" + # DAT + DAT_S_2x = "DAT_S_2x.pth" + DAT_S_3x = "DAT_S_3x.pth" + DAT_S_4x = "DAT_S_4x.pth" + DAT_2x = "DAT_2x.pth" + DAT_3x = "DAT_3x.pth" + DAT_4x = "DAT_4x.pth" + DAT_2_2x = "DAT_2_2x.pth" + DAT_2_3x = "DAT_2_3x.pth" + DAT_2_4x = "DAT_2_4x.pth" + DAT_light_2x = "DAT_light_2x.pth" + DAT_light_3x = "DAT_light_3x.pth" + DAT_light_4x = "DAT_light_4x.pth" + + DAT_APISR_GAN_generator_4x = "DAT_APISR_GAN_generator_4x.pth" + # ------------------------------------- Auxiliary Network ---------------------------------------------------------- # SpyNet diff --git a/ccrestoration/type/model.py b/ccrestoration/type/model.py index 9a842bf..0915d7f 100644 --- a/ccrestoration/type/model.py +++ b/ccrestoration/type/model.py @@ -10,6 +10,7 @@ class ModelType(str, Enum): EDSR = "EDSR" SwinIR = "SwinIR" SCUNet = "SCUNet" + DAT = "DAT" # ------------------------------------- Auxiliary Network ---------------------------------------------------------- diff --git a/tests/test_dat.py b/tests/test_dat.py new file mode 100644 index 0000000..85a59e8 --- /dev/null +++ b/tests/test_dat.py @@ -0,0 +1,71 @@ +import os + +import cv2 +import pytest + +from ccrestoration import AutoConfig, AutoModel, BaseConfig, ConfigType +from ccrestoration.model import SRBaseModel + +from .util import ASSETS_PATH, calculate_image_similarity, compare_image_size, get_device, load_image + + +class Test_DAT: + def test_official_light(self) -> None: + img1 = load_image() + + for k in [ + ConfigType.DAT_light_2x, + ConfigType.DAT_light_3x, + ConfigType.DAT_light_4x, + ]: + print(f"Testing {k}") + cfg: BaseConfig = AutoConfig.from_pretrained(k) + model: SRBaseModel = AutoModel.from_config(config=cfg, fp16=False, device=get_device()) + print(model.device) + + img2 = model.inference_image(img1) + cv2.imwrite(str(ASSETS_PATH / f"test_{k}_out.jpg"), img2) + + assert calculate_image_similarity(img1, img2) + assert compare_image_size(img1, img2, cfg.scale) + + @pytest.mark.skipif(os.environ.get("GITHUB_ACTIONS") == "true", reason="Skip on CI test") + def test_official(self) -> None: + img1 = load_image() + + for k in [ + ConfigType.DAT_S_2x, + ConfigType.DAT_S_3x, + ConfigType.DAT_S_4x, + ConfigType.DAT_2x, + ConfigType.DAT_3x, + ConfigType.DAT_4x, + ConfigType.DAT_2_2x, + ConfigType.DAT_2_3x, + ConfigType.DAT_2_4x, + ]: + print(f"Testing {k}") + cfg: BaseConfig = AutoConfig.from_pretrained(k) + model: SRBaseModel = AutoModel.from_config(config=cfg, fp16=False, device=get_device()) + print(model.device) + + img2 = model.inference_image(img1) + cv2.imwrite(str(ASSETS_PATH / f"test_{k}_out.jpg"), img2) + + assert calculate_image_similarity(img1, img2) + assert compare_image_size(img1, img2, cfg.scale) + + def test_custom(self) -> None: + img1 = load_image() + + for k in [ConfigType.DAT_APISR_GAN_generator_4x]: + print(f"Testing {k}") + cfg: BaseConfig = AutoConfig.from_pretrained(k) + model: SRBaseModel = AutoModel.from_config(config=cfg, fp16=False, device=get_device()) + print(model.device) + + img2 = model.inference_image(img1) + cv2.imwrite(str(ASSETS_PATH / f"test_{k}_out.jpg"), img2) + + assert calculate_image_similarity(img1, img2) + assert compare_image_size(img1, img2, cfg.scale)