Upload 7 files

GeoTr.py

@@ -0,0 +1,233 @@
+from extractor import BasicEncoder
+from position_encoding import build_position_encoding
+
+import argparse
+import numpy as np
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+import copy
+from typing import Optional
+
+
+class attnLayer(nn.Module):
+    def __init__(self, d_model, nhead=8, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn_list = nn.ModuleList([copy.deepcopy(nn.MultiheadAttention(d_model, nhead, dropout=dropout)) for i in range(2)])
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2_list = nn.ModuleList([copy.deepcopy(nn.LayerNorm(d_model)) for i in range(2)])
+
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2_list = nn.ModuleList([copy.deepcopy(nn.Dropout(dropout)) for i in range(2)])
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, tgt, memory_list, tgt_mask=None, memory_mask=None,
+                     tgt_key_padding_mask=None, memory_key_padding_mask=None,
+                     pos=None, memory_pos=None):
+        q = k = self.with_pos_embed(tgt, pos)
+        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+        for memory, multihead_attn, norm2, dropout2, m_pos in zip(memory_list, self.multihead_attn_list, self.norm2_list, self.dropout2_list, memory_pos):
+            tgt2 = multihead_attn(query=self.with_pos_embed(tgt, pos),
+                                       key=self.with_pos_embed(memory, m_pos),
+                                       value=memory, attn_mask=memory_mask,
+                                       key_padding_mask=memory_key_padding_mask)[0]
+            tgt = tgt + dropout2(tgt2)
+            tgt = norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+    def forward_pre(self, tgt, memory, tgt_mask=None, memory_mask=None,
+                     tgt_key_padding_mask=None, memory_key_padding_mask=None,
+                     pos=None, memory_pos=None):
+        tgt2 = self.norm1(tgt)
+        q = k = self.with_pos_embed(tgt2, pos)
+        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, pos),
+                                   key=self.with_pos_embed(memory, memory_pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(self, tgt, memory_list, tgt_mask=None, memory_mask=None,
+                     tgt_key_padding_mask=None, memory_key_padding_mask=None,
+                     pos=None, memory_pos=None):
+        if self.normalize_before:
+            return self.forward_pre(tgt, memory_list, tgt_mask, memory_mask,
+                                    tgt_key_padding_mask, memory_key_padding_mask, pos, memory_pos)
+        return self.forward_post(tgt, memory_list, tgt_mask, memory_mask,
+                                 tgt_key_padding_mask, memory_key_padding_mask, pos, memory_pos)
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
+
+
+class TransDecoder(nn.Module):
+    def __init__(self, num_attn_layers, hidden_dim=128):
+        super(TransDecoder, self).__init__()
+        attn_layer = attnLayer(hidden_dim)
+        self.layers = _get_clones(attn_layer, num_attn_layers)
+        self.position_embedding = build_position_encoding(hidden_dim)
+
+    def forward(self, imgf, query_embed):
+        pos = self.position_embedding(torch.ones(imgf.shape[0], imgf.shape[2], imgf.shape[3]).bool().cuda())  # torch.Size([1, 128, 36, 36])
+        
+        bs, c, h, w = imgf.shape
+        imgf = imgf.flatten(2).permute(2, 0, 1)
+        query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)
+        pos = pos.flatten(2).permute(2, 0, 1)
+
+        for layer in self.layers:
+            query_embed = layer(query_embed, [imgf], pos=pos, memory_pos=[pos, pos])
+        query_embed = query_embed.permute(1, 2, 0).reshape(bs, c, h, w)
+
+        return query_embed
+
+
+class TransEncoder(nn.Module):
+    def __init__(self, num_attn_layers, hidden_dim=128):
+        super(TransEncoder, self).__init__()
+        attn_layer = attnLayer(hidden_dim)
+        self.layers = _get_clones(attn_layer, num_attn_layers)
+        self.position_embedding = build_position_encoding(hidden_dim)
+
+    def forward(self, imgf):
+        pos = self.position_embedding(torch.ones(imgf.shape[0], imgf.shape[2], imgf.shape[3]).bool().cuda())  # torch.Size([1, 128, 36, 36])
+        bs, c, h, w = imgf.shape
+        imgf = imgf.flatten(2).permute(2, 0, 1)  
+        pos = pos.flatten(2).permute(2, 0, 1)
+
+        for layer in self.layers:
+            imgf = layer(imgf, [imgf], pos=pos, memory_pos=[pos, pos])
+        imgf = imgf.permute(1, 2, 0).reshape(bs, c, h, w)
+
+        return imgf
+
+
+class FlowHead(nn.Module):
+    def __init__(self, input_dim=128, hidden_dim=256):
+        super(FlowHead, self).__init__()
+        self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
+        self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        return self.conv2(self.relu(self.conv1(x)))
+
+
+class UpdateBlock(nn.Module):
+    def __init__(self, hidden_dim=128):
+        super(UpdateBlock, self).__init__()
+        self.flow_head = FlowHead(hidden_dim, hidden_dim=256)
+        self.mask = nn.Sequential(
+            nn.Conv2d(hidden_dim, 256, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 64*9, 1, padding=0))
+
+    def forward(self, imgf, coords1):
+        mask = .25 * self.mask(imgf)  # scale mask to balence gradients
+        dflow = self.flow_head(imgf)
+        coords1 = coords1 + dflow
+
+        return mask, coords1
+
+
+def coords_grid(batch, ht, wd):
+    coords = torch.meshgrid(torch.arange(ht), torch.arange(wd))
+    coords = torch.stack(coords[::-1], dim=0).float()
+    return coords[None].repeat(batch, 1, 1, 1)
+
+
+def upflow8(flow, mode='bilinear'):
+    new_size = (8 * flow.shape[2], 8 * flow.shape[3])
+    return  8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)
+
+
+class GeoTr(nn.Module):
+    def __init__(self, num_attn_layers):
+        super(GeoTr, self).__init__()
+        self.num_attn_layers = num_attn_layers
+
+        self.hidden_dim = hdim = 256
+
+        self.fnet = BasicEncoder(output_dim=hdim, norm_fn='instance')
+
+        self.TransEncoder = TransEncoder(self.num_attn_layers, hidden_dim=hdim)
+        self.TransDecoder = TransDecoder(self.num_attn_layers, hidden_dim=hdim)
+        self.query_embed = nn.Embedding(1296, self.hidden_dim)
+        
+        self.update_block = UpdateBlock(self.hidden_dim)
+                                    
+    def initialize_flow(self, img):
+        N, C, H, W = img.shape
+        coodslar = coords_grid(N, H, W).to(img.device)
+        coords0 = coords_grid(N, H // 8, W // 8).to(img.device)
+        coords1 = coords_grid(N, H // 8, W // 8).to(img.device)
+
+        return coodslar, coords0, coords1
+
+    def upsample_flow(self, flow, mask):
+        N, _, H, W = flow.shape
+        mask = mask.view(N, 1, 9, 8, 8, H, W)
+        mask = torch.softmax(mask, dim=2) 
+
+        up_flow = F.unfold(8 * flow, [3, 3], padding=1) 
+        up_flow = up_flow.view(N, 2, 9, 1, 1, H, W) 
+
+        up_flow = torch.sum(mask * up_flow, dim=2)
+        up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
+        
+        return up_flow.reshape(N, 2, 8 * H, 8 * W)
+
+    def forward(self, image1):
+        fmap = self.fnet(image1)
+        fmap = torch.relu(fmap)
+        
+        fmap = self.TransEncoder(fmap)
+        fmap = self.TransDecoder(fmap, self.query_embed.weight)  
+
+        # convex upsample baesd on fmap
+        coodslar, coords0, coords1 = self.initialize_flow(image1)
+        coords1 = coords1.detach()
+        mask, coords1 = self.update_block(fmap, coords1)
+        flow_up = self.upsample_flow(coords1 - coords0, mask)
+        bm_up = coodslar + flow_up
+
+        return bm_up 

IllTr.py

@@ -0,0 +1,284 @@
+import torch
+import torch.nn as nn
+from torch.functional import Tensor
+from torch.nn.modules.activation import Tanhshrink
+from timm.models.layers import trunc_normal_
+from functools import partial
+
+
+class Ffn(nn.Module):
+    # feed forward network layer after attention
+    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)
+
+    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
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        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)
+
+    def forward(self, x, task_embed=None, level=0):
+        N, L, D = x.shape
+        qkv = self.qkv(x).reshape(N, L, 3, self.num_heads, D // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        # for decoder's task_embedding of different levels of attention layers
+        if task_embed != None:
+            _N, _H, _L, _D = q.shape
+            task_embed = task_embed.reshape(1, _H, _L, _D)
+            if level == 1:
+                q += task_embed
+                k += task_embed
+            if level == 2:
+                q += task_embed
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(N, L, D)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, dim, num_heads, ffn_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.norm2 = norm_layer(dim)
+        ffn_hidden_dim = int(dim * ffn_ratio)
+        self.ffn = Ffn(in_features=dim, hidden_features=ffn_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, dim, num_heads, ffn_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn1 = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.norm2 = norm_layer(dim)
+        self.attn2 = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.norm3 = norm_layer(dim)
+        ffn_hidden_dim = int(dim * ffn_ratio)
+        self.ffn = Ffn(in_features=dim, hidden_features=ffn_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, task_embed):
+        x = x + self.attn1(self.norm1(x), task_embed=task_embed, level=1)
+        x = x + self.attn2(self.norm2(x), task_embed=task_embed, level=2)
+        x = x + self.ffn(self.norm3(x))
+        return x
+
+
+class ResBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels, channels, kernel_size=5, stride=1,
+                               padding=2, bias=False)
+        self.bn1 = nn.InstanceNorm2d(channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(channels, channels, kernel_size=5, stride=1,
+                               padding=2, bias=False)
+        self.bn2 = nn.InstanceNorm2d(channels)
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Head(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(Head, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1,
+                               padding=1, bias=False)
+        self.bn1 = nn.InstanceNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.resblock = ResBlock(out_channels)
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.resblock(out)
+
+        return out
+
+
+class PatchEmbed(nn.Module):
+    """ Feature to Patch Embedding
+    input : N C H W
+    output: N num_patch P^2*C
+    """
+    def __init__(self, patch_size=1, in_channels=64):
+        super().__init__()
+        self.patch_size = patch_size
+        self.dim = self.patch_size ** 2 * in_channels
+
+    def forward(self, x):
+        N, C, H, W = ori_shape = x.shape
+
+        p = self.patch_size
+        num_patches = (H // p) * (W // p)
+        out = torch.zeros((N, num_patches, self.dim)).to(x.device)
+        i, j = 0, 0
+        for k in range(num_patches):
+            if i + p > W:
+                i = 0
+                j += p
+            out[:, k, :] = x[:, :, i:i + p, j:j + p].flatten(1)
+            i += p
+        return out, ori_shape
+
+
+class DePatchEmbed(nn.Module):
+    """ Patch Embedding to Feature
+    input : N num_patch P^2*C
+    output: N C H W
+    """
+    def __init__(self, patch_size=1, in_channels=64):
+        super().__init__()
+        self.patch_size = patch_size
+        self.num_patches = None
+        self.dim = self.patch_size ** 2 * in_channels
+
+    def forward(self, x, ori_shape):
+        N, num_patches, dim = x.shape
+        _, C, H, W = ori_shape
+        p = self.patch_size
+        out = torch.zeros(ori_shape).to(x.device)
+        i, j = 0, 0
+        for k in range(num_patches):
+            if i + p > W:
+                i = 0
+                j += p
+            out[:, :, i:i + p, j:j + p] = x[:, k, :].reshape(N, C, p, p)
+            i += p
+        return out
+
+
+class Tail(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(Tail, self).__init__()
+        self.output = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+
+    def forward(self, x):
+        out = self.output(x)
+        return out
+
+
+class IllTr_Net(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+
+    def __init__(self, patch_size=1, in_channels=3, mid_channels=16, num_classes=1000, depth=12,
+                 num_heads=8, ffn_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 norm_layer=nn.LayerNorm):
+        super(IllTr_Net, self).__init__()
+
+        self.num_classes = num_classes
+        self.embed_dim = patch_size * patch_size * mid_channels
+        self.head = Head(in_channels, mid_channels)
+        self.patch_embedding = PatchEmbed(patch_size=patch_size, in_channels=mid_channels)
+        self.embed_dim = self.patch_embedding.dim
+        if self.embed_dim % num_heads != 0:
+            raise RuntimeError("Embedding dim must be devided by numbers of heads")
+
+        self.pos_embed = nn.Parameter(torch.zeros(1, (128 // patch_size) ** 2, self.embed_dim))
+        self.task_embed = nn.Parameter(torch.zeros(6, 1, (128 // patch_size) ** 2, self.embed_dim))
+
+        self.encoder = nn.ModuleList([
+            EncoderLayer(
+                dim=self.embed_dim, num_heads=num_heads, ffn_ratio=ffn_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, norm_layer=norm_layer)
+            for _ in range(depth)])
+        self.decoder = nn.ModuleList([
+            DecoderLayer(
+                dim=self.embed_dim, num_heads=num_heads, ffn_ratio=ffn_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, norm_layer=norm_layer)
+            for _ in range(depth)])
+
+        self.de_patch_embedding = DePatchEmbed(patch_size=patch_size, in_channels=mid_channels)
+        # tail
+        self.tail = Tail(int(mid_channels), in_channels)
+        
+        self.acf = nn.Hardtanh(0,1)
+
+        trunc_normal_(self.pos_embed, std=.02)
+        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)
+
+    def forward(self, x):
+        x = self.head(x)
+        x, ori_shape = self.patch_embedding(x)
+        x = x + self.pos_embed[:, :x.shape[1]]
+        
+        for blk in self.encoder:
+            x = blk(x)
+
+        for blk in self.decoder:
+            x = blk(x, self.task_embed[0, :, :x.shape[1]])
+      
+        x = self.de_patch_embedding(x, ori_shape)
+        x = self.tail(x)
+        
+        x = self.acf(x)
+        return x
+
+
+def IllTr(**kwargs):
+    model = IllTr_Net(
+        patch_size=4, depth=6, num_heads=8, ffn_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs)
+
+    return model

demo.py

@@ -0,0 +1,178 @@
+#origin
+
+from seg import U2NETP
+from GeoTr import GeoTr
+from IllTr import IllTr
+from inference_ill import rec_ill
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import skimage.io as io
+import numpy as np
+import cv2
+import glob
+import os
+from PIL import Image
+import argparse
+import warnings
+warnings.filterwarnings('ignore')
+
+
+
+
+
+import gradio as gr
+
+
+class GeoTr_Seg(nn.Module):
+    def __init__(self):
+        super(GeoTr_Seg, self).__init__()
+        self.msk = U2NETP(3, 1)
+        self.GeoTr = GeoTr(num_attn_layers=6)
+        
+    def forward(self, x):
+        msk, _1,_2,_3,_4,_5,_6 = self.msk(x)
+        msk = (msk > 0.5).float()
+        x = msk * x
+
+        bm = self.GeoTr(x)
+        bm = (2 * (bm / 286.8) - 1) * 0.99
+        
+        return bm
+        
+
+def reload_model(model, path=""):
+    if not bool(path):
+        return model
+    else:
+        model_dict = model.state_dict()
+        pretrained_dict = torch.load(path, map_location='cuda:0')
+        print(len(pretrained_dict.keys()))
+        pretrained_dict = {k[7:]: v for k, v in pretrained_dict.items() if k[7:] in model_dict}
+        print(len(pretrained_dict.keys()))
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+
+        return model
+        
+
+def reload_segmodel(model, path=""):
+    if not bool(path):
+        return model
+    else:
+        model_dict = model.state_dict()
+        pretrained_dict = torch.load(path, map_location='cuda:0')
+        print(len(pretrained_dict.keys()))
+        pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict}
+        print(len(pretrained_dict.keys()))
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+
+        return model
+        
+
+def rec(opt):
+    # print(torch.__version__) # 1.5.1
+    img_list = os.listdir(opt.distorrted_path)  # distorted images list
+
+    if not os.path.exists(opt.gsave_path):  # create save path
+        os.mkdir(opt.gsave_path)
+    if not os.path.exists(opt.isave_path):  # create save path
+        os.mkdir(opt.isave_path)
+    
+    GeoTr_Seg_model = GeoTr_Seg().cuda()
+    # reload segmentation model
+    reload_segmodel(GeoTr_Seg_model.msk, opt.Seg_path)
+    # reload geometric unwarping model
+    reload_model(GeoTr_Seg_model.GeoTr, opt.GeoTr_path)
+    
+    IllTr_model = IllTr().cuda()
+    # reload illumination rectification model
+    reload_model(IllTr_model, opt.IllTr_path)
+    
+    # To eval mode
+    GeoTr_Seg_model.eval()
+    IllTr_model.eval()
+  
+    for img_path in img_list:
+        name = img_path.split('.')[-2]  # image name
+
+        img_path = opt.distorrted_path + img_path  # read image and to tensor
+        im_ori = np.array(Image.open(img_path))[:, :, :3] / 255. 
+        h, w, _ = im_ori.shape
+        im = cv2.resize(im_ori, (288, 288))
+        im = im.transpose(2, 0, 1)
+        im = torch.from_numpy(im).float().unsqueeze(0)
+        
+        with torch.no_grad():
+            # geometric unwarping
+            bm = GeoTr_Seg_model(im.cuda())
+            bm = bm.cpu()
+            bm0 = cv2.resize(bm[0, 0].numpy(), (w, h))  # x flow
+            bm1 = cv2.resize(bm[0, 1].numpy(), (w, h))  # y flow
+            bm0 = cv2.blur(bm0, (3, 3))
+            bm1 = cv2.blur(bm1, (3, 3))
+            lbl = torch.from_numpy(np.stack([bm0, bm1], axis=2)).unsqueeze(0)  # h * w * 2
+            
+            out = F.grid_sample(torch.from_numpy(im_ori).permute(2,0,1).unsqueeze(0).float(), lbl, align_corners=True)
+            img_geo = ((out[0]*255).permute(1, 2, 0).numpy())[:,:,::-1].astype(np.uint8)
+            cv2.imwrite(opt.gsave_path + name + '_geo' + '.png', img_geo)  # save
+            
+            # illumination rectification
+            if opt.ill_rec:
+                ill_savep = opt.isave_path + name + '_ill' + '.png'
+                rec_ill(IllTr_model, img_geo, saveRecPath=ill_savep)
+        
+        print('Done: ', img_path)
+        
+
+
+
+
+
+def process_image(input_image):
+    GeoTr_Seg_model = GeoTr_Seg().cuda()
+    reload_segmodel(GeoTr_Seg_model.msk, './model_pretrained/seg.pth')
+    reload_model(GeoTr_Seg_model.GeoTr, './model_pretrained/geotr.pth')
+
+    IllTr_model = IllTr().cuda()
+    reload_model(IllTr_model, './model_pretrained/illtr.pth')
+
+    GeoTr_Seg_model.eval()
+    IllTr_model.eval()
+
+    im_ori = np.array(input_image)[:, :, :3] / 255.
+    h, w, _ = im_ori.shape
+    im = cv2.resize(im_ori, (288, 288))
+    im = im.transpose(2, 0, 1)
+    im = torch.from_numpy(im).float().unsqueeze(0)
+
+    with torch.no_grad():
+        bm = GeoTr_Seg_model(im.cuda())
+        bm = bm.cpu()
+        bm0 = cv2.resize(bm[0, 0].numpy(), (w, h))
+        bm1 = cv2.resize(bm[0, 1].numpy(), (w, h))
+        bm0 = cv2.blur(bm0, (3, 3))
+        bm1 = cv2.blur(bm1, (3, 3))
+        lbl = torch.from_numpy(np.stack([bm0, bm1], axis=2)).unsqueeze(0)
+
+        out = F.grid_sample(torch.from_numpy(im_ori).permute(2, 0, 1).unsqueeze(0).float(), lbl, align_corners=True)
+        img_geo = ((out[0] * 255).permute(1, 2, 0).numpy()).astype(np.uint8)
+        
+        ill_rec=False
+
+        if ill_rec:
+            img_ill = rec_ill(IllTr_model, img_geo)
+            return Image.fromarray(img_ill)
+        else:
+            return Image.fromarray(img_geo)
+
+# Define Gradio interface
+input_image = gr.inputs.Image()
+output_image = gr.outputs.Image(type='pil')
+
+
+iface = gr.Interface(fn=process_image, inputs=input_image, outputs=output_image, title="Image Correction")
+iface.launch(server_port=1234, server_name="0.0.0.0")
+

extractor.py

@@ -0,0 +1,115 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn='group', stride=1):
+        super(ResidualBlock, self).__init__()
+  
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+        
+        elif norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(planes)
+            self.norm2 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.BatchNorm2d(planes)
+        
+        elif norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(planes)
+            self.norm2 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            if not stride == 1:
+                self.norm3 = nn.Sequential()
+
+        if stride == 1:
+            self.downsample = None
+        
+        else:    
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
+
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+
+class BasicEncoder(nn.Module):
+    def __init__(self, output_dim=128, norm_fn='batch'):
+        super(BasicEncoder, self).__init__()
+        self.norm_fn = norm_fn
+
+        if self.norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
+            
+        elif self.norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(64)
+
+        elif self.norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(64)
+
+        elif self.norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.in_planes = 64
+        self.layer1 = self._make_layer(64,  stride=1)
+        self.layer2 = self._make_layer(128, stride=2)
+        self.layer3 = self._make_layer(192, stride=2)
+
+        # output convolution
+        self.conv2 = nn.Conv2d(192, output_dim, kernel_size=1)
+  
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+        
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+
+        x = self.conv2(x)
+
+        return x

position_encoding.py

@@ -0,0 +1,111 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Various positional encodings for the transformer.
+"""
+import math
+import torch
+from torch import nn
+from typing import List
+from typing import Optional
+from torch import Tensor
+
+
+class NestedTensor(object):
+    def __init__(self, tensors, mask: Optional[Tensor]):
+        self.tensors = tensors
+        self.mask = mask
+
+    def to(self, device):
+        # type: (Device) -> NestedTensor # noqa
+        cast_tensor = self.tensors.to(device)
+        mask = self.mask
+        if mask is not None:
+            assert mask is not None
+            cast_mask = mask.to(device)
+        else:
+            cast_mask = None
+        return NestedTensor(cast_tensor, cast_mask)
+
+    def decompose(self):
+        return self.tensors, self.mask
+
+    def __repr__(self):
+        return str(self.tensors)
+
+
+class PositionEmbeddingSine(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images.
+    """
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, mask):
+        assert mask is not None
+        y_embed = mask.cumsum(1, dtype=torch.float32)
+        x_embed = mask.cumsum(2, dtype=torch.float32)
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32).cuda()
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+        # print(pos.shape)
+        return pos
+
+
+class PositionEmbeddingLearned(nn.Module):
+    """
+    Absolute pos embedding, learned.
+    """
+    def __init__(self, num_pos_feats=256):
+        super().__init__()
+        self.row_embed = nn.Embedding(50, num_pos_feats)
+        self.col_embed = nn.Embedding(50, num_pos_feats)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.uniform_(self.row_embed.weight)
+        nn.init.uniform_(self.col_embed.weight)
+
+    def forward(self, tensor_list: NestedTensor):
+        x = tensor_list.tensors
+        h, w = x.shape[-2:]
+        i = torch.arange(w, device=x.device)
+        j = torch.arange(h, device=x.device)
+        x_emb = self.col_embed(i)
+        y_emb = self.row_embed(j)
+        pos = torch.cat([
+            x_emb.unsqueeze(0).repeat(h, 1, 1),
+            y_emb.unsqueeze(1).repeat(1, w, 1),
+        ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)
+        return pos
+
+def build_position_encoding(hidden_dim=512, position_embedding='sine'):
+    N_steps = hidden_dim // 2
+    if position_embedding in ('v2', 'sine'):
+        position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
+    elif position_embedding in ('v3', 'learned'):
+        position_embedding = PositionEmbeddingLearned(N_steps)
+    else:
+        raise ValueError(f"not supported {position_embedding}")
+
+    return position_embedding
+
+

requirements.txt

@@ -0,0 +1,7 @@
+numpy
+opencv_python
+Pillow
+scikit_image 
+thop 
+torch 
+gradio

seg.py

@@ -0,0 +1,567 @@
+import torch
+import torch.nn as nn
+from torchvision import models
+import torch.nn.functional as F
+import numpy as np
+
+
+class sobel_net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv_opx = nn.Conv2d(1, 1, 3, bias=False)
+        self.conv_opy = nn.Conv2d(1, 1, 3, bias=False)
+        sobel_kernelx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype='float32').reshape((1, 1, 3, 3))
+        sobel_kernely = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype='float32').reshape((1, 1, 3, 3))
+        self.conv_opx.weight.data = torch.from_numpy(sobel_kernelx)
+        self.conv_opy.weight.data = torch.from_numpy(sobel_kernely)
+
+        for p in self.parameters():
+            p.requires_grad = False
+
+    def forward(self, im):  # input rgb
+        x = (0.299 * im[:, 0, :, :] + 0.587 * im[:, 1, :, :] + 0.114 * im[:, 2, :, :]).unsqueeze(1)  # rgb2gray
+        gradx = self.conv_opx(x)
+        grady = self.conv_opy(x)
+
+        x = (gradx ** 2 + grady ** 2) ** 0.5
+        x = (x - x.min()) / (x.max() - x.min())
+        x = F.pad(x, (1, 1, 1, 1))
+
+        x = torch.cat([im, x], dim=1)
+        return x
+
+
+class REBNCONV(nn.Module):
+    def __init__(self, in_ch=3, out_ch=3, dirate=1):
+        super(REBNCONV, self).__init__()
+
+        self.conv_s1 = nn.Conv2d(in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate)
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src, tar):
+    src = F.interpolate(src, size=tar.shape[2:], mode='bilinear', align_corners=False)
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):  # UNet07DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU7, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d = self.rebnconv6d(torch.cat((hx7, hx6), 1))
+        hx6dup = _upsample_like(hx6d, hx5)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6dup, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-6 ###
+class RSU6(nn.Module):  # UNet06DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-5 ###
+class RSU5(nn.Module):  # UNet05DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-4 ###
+class RSU4(nn.Module):  # UNet04DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-4F ###
+class RSU4F(nn.Module):  # UNet04FRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d, hx2), 1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d, hx1), 1))
+
+        return hx1d + hxin
+
+
+##### U^2-Net ####
+class U2NET(nn.Module):
+
+    def __init__(self, in_ch=3, out_ch=1):
+        super(U2NET, self).__init__()
+        self.edge = sobel_net()
+
+        self.stage1 = RSU7(in_ch, 32, 64)
+        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage2 = RSU6(64, 32, 128)
+        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage3 = RSU5(128, 64, 256)
+        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage4 = RSU4(256, 128, 512)
+        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage5 = RSU4F(512, 256, 512)
+        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage6 = RSU4F(512, 256, 512)
+
+        # decoder
+        self.stage5d = RSU4F(1024, 256, 512)
+        self.stage4d = RSU4(1024, 128, 256)
+        self.stage3d = RSU5(512, 64, 128)
+        self.stage2d = RSU6(256, 32, 64)
+        self.stage1d = RSU7(128, 16, 64)
+
+        self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side3 = nn.Conv2d(128, out_ch, 3, padding=1)
+        self.side4 = nn.Conv2d(256, out_ch, 3, padding=1)
+        self.side5 = nn.Conv2d(512, out_ch, 3, padding=1)
+        self.side6 = nn.Conv2d(512, out_ch, 3, padding=1)
+
+        self.outconv = nn.Conv2d(6, out_ch, 1)
+
+    def forward(self, x):
+        x = self.edge(x)
+        hx = x
+
+        # stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        # stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        # stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        # stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        # stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        # stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6, hx5)
+
+        # -------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+        # side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2, d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3, d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4, d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5, d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6, d1)
+
+        d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
+
+        return torch.sigmoid(d0), torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid(
+            d4), torch.sigmoid(d5), torch.sigmoid(d6)
+
+
+### U^2-Net small ###
+class U2NETP(nn.Module):
+
+    def __init__(self, in_ch=3, out_ch=1):
+        super(U2NETP, self).__init__()
+
+        self.stage1 = RSU7(in_ch, 16, 64)
+        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage2 = RSU6(64, 16, 64)
+        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage3 = RSU5(64, 16, 64)
+        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage4 = RSU4(64, 16, 64)
+        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage5 = RSU4F(64, 16, 64)
+        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage6 = RSU4F(64, 16, 64)
+
+        # decoder
+        self.stage5d = RSU4F(128, 16, 64)
+        self.stage4d = RSU4(128, 16, 64)
+        self.stage3d = RSU5(128, 16, 64)
+        self.stage2d = RSU6(128, 16, 64)
+        self.stage1d = RSU7(128, 16, 64)
+
+        self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side3 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side4 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side5 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side6 = nn.Conv2d(64, out_ch, 3, padding=1)
+
+        self.outconv = nn.Conv2d(6, out_ch, 1)
+
+    def forward(self, x):
+        hx = x
+
+        # stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        # stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        # stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        # stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        # stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        # stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6, hx5)
+
+        # decoder
+        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+        # side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2, d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3, d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4, d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5, d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6, d1)
+
+        d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))
+
+        return torch.sigmoid(d0), torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid(
+            d4), torch.sigmoid(d5), torch.sigmoid(d6)
+
+
+def get_parameter_number(net):
+    total_num = sum(p.numel() for p in net.parameters())
+    trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)
+    return {'Total': total_num, 'Trainable': trainable_num}
+
+
+if __name__ == '__main__':
+    net = U2NET(4, 1).cuda()
+    print(get_parameter_number(net))  # 69090500 加attention后69442032
+    with torch.no_grad():
+        inputs = torch.zeros(1, 3, 256, 256).cuda()
+        outs = net(inputs)
+        print(outs[0].shape)  # torch.Size([2, 3, 256, 256]) torch.Size([2, 2, 256, 256])