import os
from pathlib import Path
from types import SimpleNamespace

import cv2
import numpy as np
import torch
from PIL import Image

import torch.nn.functional  as F 

def bilinear_preprocessing(warped_img, point_positions, img_size):
    """
    Utility function that preprocesss an image.
    preprocess warped_img based on the 2D grid point_positions with a size img_size.
    Args:
        warped_img  :       torch.Tensor of shape BxCxHxW (dtype float)
        point_positions:    torch.Tensor of shape Bx2xGhxGw (dtype float)
        img_size:           tuple of int [w, h]
    """
    upsampled_grid = F.interpolate(
        point_positions, size=(img_size[1], img_size[0]), mode="bilinear", align_corners=True
    )
    preprocessed_img = F.grid_sample(warped_img, upsampled_grid.transpose(1, 2).transpose(2, 3), align_corners=True)
    return preprocessed_img

def tensor_to_cv2image_mask(tensor, remove_padding=True):
    image = tensor.numpy()
    image = image.transpose((1, 2, 0))
    image = image * 255
    if remove_padding:
        image_ = np.sum(image, -1)
        image_h = np.sum(image_, 1)
        if 0 in image_h:
            h_border = np.min(np.where(image_h == 0)[0])
        else:
            h_border = image.shape[0]
        image_w = np.sum(image_, 0)
        if 0 in image_w:
            w_border = np.min(np.where(image_w == 0)[0])
        else:
            w_border = image.shape[1]
        image = image[:h_border, :w_border]
    image = image.astype(np.uint8)
    return image

# Embedded preprocessor model code.
_EMBEDDED_MODEL_SOURCES = [('models.block',
  'import torch.nn as nn\n'
  '\n'
  '\n'
  'def build_lateral_connection(input_dim, output_dim):\n'
  '    return nn.Sequential(\n'
  '        nn.Conv2d(input_dim, input_dim, 1, 1, 0),\n'
  '        nn.Conv2d(input_dim, input_dim*2, 3, 1, 1),\n'
  '        nn.Conv2d(input_dim*2, input_dim*2, 3, 1, 1),\n'
  '        nn.Conv2d(input_dim*2, output_dim, 1, 1, 0)\n'
  '    )\n'
  '\n'
  '\n'
  'class ConvWithActivation(nn.Module):\n'
  '    def __init__(self, conv_type, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, '
  "groups=1, bias=True, activation='relu'):\n"
  '        super(ConvWithActivation, self).__init__()\n'
  "        if conv_type == 'conv':\n"
  '            conv_func = nn.Conv2d \n'
  "        elif conv_type == 'deconv':\n"
  '            conv_func = nn.ConvTranspose2d\n'
  '        self.conv2d = conv_func(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)\n'
  '        self.conv2d = nn.utils.spectral_norm(self.conv2d)\n'
  '        self.activation = get_activation(activation)\n'
  '\n'
  '        for m in self.modules():\n'
  '            if isinstance(m, conv_func):\n'
  '                nn.init.kaiming_normal_(m.weight)\n'
  '        \n'
  '    def forward(self, x):\n'
  '        x = self.conv2d(x)\n'
  '        x = self.activation(x)\n'
  '        return x\n'
  '\n'
  '    \n'
  'def get_activation(type):\n'
  "    if type == 'leaky relu':\n"
  '        return nn.LeakyReLU(0.2, inplace=True)\n'
  "    elif type == 'relu':\n"
  '        return nn.ReLU(inplace=True)\n'
  "    elif type == 'sigmoid':\n"
  '        return nn.Sigmoid()\n'
  '    else:\n'
  '        raise NotImplementedError'),
 ('models.seg',
  'import torch\n'
  'import torch.nn as nn\n'
  'import torch.nn.functional as F\n'
  'import numpy as np\n'
  '\n'
  '\n'
  'class sobel_net(nn.Module):\n'
  '    def __init__(self):\n'
  '        super().__init__()\n'
  '        self.conv_opx = nn.Conv2d(1, 1, 3, bias=False)\n'
  '        self.conv_opy = nn.Conv2d(1, 1, 3, bias=False)\n'
  "        sobel_kernelx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype='float32').reshape((1, 1, 3, 3))\n"
  "        sobel_kernely = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype='float32').reshape((1, 1, 3, 3))\n"
  '        self.conv_opx.weight.data = torch.from_numpy(sobel_kernelx)\n'
  '        self.conv_opy.weight.data = torch.from_numpy(sobel_kernely)\n'
  '\n'
  '        for p in self.parameters():\n'
  '            p.requires_grad = False\n'
  '\n'
  '    def forward(self, im):  # input rgb\n'
  '        x = (0.299 * im[:, 0, :, :] + 0.587 * im[:, 1, :, :] + 0.114 * im[:, 2, :, :]).unsqueeze(1)  # rgb2gray\n'
  '        gradx = self.conv_opx(x)\n'
  '        grady = self.conv_opy(x)\n'
  '\n'
  '        x = (gradx ** 2 + grady ** 2) ** 0.5\n'
  '        x = (x - x.min()) / (x.max() - x.min())\n'
  '        x = F.pad(x, (1, 1, 1, 1))\n'
  '\n'
  '        x = torch.cat([im, x], dim=1)\n'
  '        return x\n'
  '\n'
  '\n'
  'class REBNCONV(nn.Module):\n'
  '    def __init__(self, in_ch=3, out_ch=3, dirate=1):\n'
  '        super(REBNCONV, self).__init__()\n'
  '\n'
  '        self.conv_s1 = nn.Conv2d(in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate)\n'
  '        self.bn_s1 = nn.BatchNorm2d(out_ch)\n'
  '        self.relu_s1 = nn.ReLU(inplace=True)\n'
  '\n'
  '    def forward(self, x):\n'
  '        hx = x\n'
  '        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))\n'
  '\n'
  '        return xout\n'
  '\n'
  '\n'
  "## upsample tensor 'src' to have the same spatial size with tensor 'tar'\n"
  'def _upsample_like(src, tar):\n'
  "    src = F.interpolate(src, size=tar.shape[2:], mode='bilinear', align_corners=False)\n"
  '\n'
  '    return src\n'
  '\n'
  '\n'
  '### RSU-7 ###\n'
  'class RSU7(nn.Module):  # UNet07DRES(nn.Module):\n'
  '\n'
  '    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):\n'
  '        super(RSU7, self).__init__()\n'
  '\n'
  '        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)\n'
  '\n'
  '        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)\n'
  '        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '        self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '\n'
  '        self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)\n'
  '\n'
  '        self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)\n'
  '\n'
  '    def forward(self, x):\n'
  '        hx = x\n'
  '        hxin = self.rebnconvin(hx)\n'
  '\n'
  '        hx1 = self.rebnconv1(hxin)\n'
  '        hx = self.pool1(hx1)\n'
  '\n'
  '        hx2 = self.rebnconv2(hx)\n'
  '        hx = self.pool2(hx2)\n'
  '\n'
  '        hx3 = self.rebnconv3(hx)\n'
  '        hx = self.pool3(hx3)\n'
  '\n'
  '        hx4 = self.rebnconv4(hx)\n'
  '        hx = self.pool4(hx4)\n'
  '\n'
  '        hx5 = self.rebnconv5(hx)\n'
  '        hx = self.pool5(hx5)\n'
  '\n'
  '        hx6 = self.rebnconv6(hx)\n'
  '\n'
  '        hx7 = self.rebnconv7(hx6)\n'
  '\n'
  '        hx6d = self.rebnconv6d(torch.cat((hx7, hx6), 1))\n'
  '        hx6dup = _upsample_like(hx6d, hx5)\n'
  '\n'
  '        hx5d = self.rebnconv5d(torch.cat((hx6dup, hx5), 1))\n'
  '        hx5dup = _upsample_like(hx5d, hx4)\n'
  '\n'
  '        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))\n'
  '        hx4dup = _upsample_like(hx4d, hx3)\n'
  '\n'
  '        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))\n'
  '        hx3dup = _upsample_like(hx3d, hx2)\n'
  '\n'
  '        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))\n'
  '        hx2dup = _upsample_like(hx2d, hx1)\n'
  '\n'
  '        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))\n'
  '\n'
  '        return hx1d + hxin\n'
  '\n'
  '\n'
  '### RSU-6 ###\n'
  'class RSU6(nn.Module):  # UNet06DRES(nn.Module):\n'
  '\n'
  '    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):\n'
  '        super(RSU6, self).__init__()\n'
  '\n'
  '        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)\n'
  '\n'
  '        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)\n'
  '        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '\n'
  '        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)\n'
  '\n'
  '        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)\n'
  '\n'
  '    def forward(self, x):\n'
  '        hx = x\n'
  '\n'
  '        hxin = self.rebnconvin(hx)\n'
  '\n'
  '        hx1 = self.rebnconv1(hxin)\n'
  '        hx = self.pool1(hx1)\n'
  '\n'
  '        hx2 = self.rebnconv2(hx)\n'
  '        hx = self.pool2(hx2)\n'
  '\n'
  '        hx3 = self.rebnconv3(hx)\n'
  '        hx = self.pool3(hx3)\n'
  '\n'
  '        hx4 = self.rebnconv4(hx)\n'
  '        hx = self.pool4(hx4)\n'
  '\n'
  '        hx5 = self.rebnconv5(hx)\n'
  '\n'
  '        hx6 = self.rebnconv6(hx5)\n'
  '\n'
  '        hx5d = self.rebnconv5d(torch.cat((hx6, hx5), 1))\n'
  '        hx5dup = _upsample_like(hx5d, hx4)\n'
  '\n'
  '        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))\n'
  '        hx4dup = _upsample_like(hx4d, hx3)\n'
  '\n'
  '        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))\n'
  '        hx3dup = _upsample_like(hx3d, hx2)\n'
  '\n'
  '        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))\n'
  '        hx2dup = _upsample_like(hx2d, hx1)\n'
  '\n'
  '        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))\n'
  '\n'
  '        return hx1d + hxin\n'
  '\n'
  '\n'
  '### RSU-5 ###\n'
  'class RSU5(nn.Module):  # UNet05DRES(nn.Module):\n'
  '\n'
  '    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):\n'
  '        super(RSU5, self).__init__()\n'
  '\n'
  '        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)\n'
  '\n'
  '        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)\n'
  '        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '\n'
  '        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)\n'
  '\n'
  '        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)\n'
  '\n'
  '    def forward(self, x):\n'
  '        hx = x\n'
  '\n'
  '        hxin = self.rebnconvin(hx)\n'
  '\n'
  '        hx1 = self.rebnconv1(hxin)\n'
  '        hx = self.pool1(hx1)\n'
  '\n'
  '        hx2 = self.rebnconv2(hx)\n'
  '        hx = self.pool2(hx2)\n'
  '\n'
  '        hx3 = self.rebnconv3(hx)\n'
  '        hx = self.pool3(hx3)\n'
  '\n'
  '        hx4 = self.rebnconv4(hx)\n'
  '\n'
  '        hx5 = self.rebnconv5(hx4)\n'
  '\n'
  '        hx4d = self.rebnconv4d(torch.cat((hx5, hx4), 1))\n'
  '        hx4dup = _upsample_like(hx4d, hx3)\n'
  '\n'
  '        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))\n'
  '        hx3dup = _upsample_like(hx3d, hx2)\n'
  '\n'
  '        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))\n'
  '        hx2dup = _upsample_like(hx2d, hx1)\n'
  '\n'
  '        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))\n'
  '\n'
  '        return hx1d + hxin\n'
  '\n'
  '\n'
  '### RSU-4 ###\n'
  'class RSU4(nn.Module):  # UNet04DRES(nn.Module):\n'
  '\n'
  '    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):\n'
  '        super(RSU4, self).__init__()\n'
  '\n'
  '        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)\n'
  '\n'
  '        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)\n'
  '        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)\n'
  '\n'
  '        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)\n'
  '\n'
  '        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)\n'
  '        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)\n'
  '\n'
  '    def forward(self, x):\n'
  '        hx = x\n'
  '\n'
  '        hxin = self.rebnconvin(hx)\n'
  '\n'
  '        hx1 = self.rebnconv1(hxin)\n'
  '        hx = self.pool1(hx1)\n'
  '\n'
  '        hx2 = self.rebnconv2(hx)\n'
  '        hx = self.pool2(hx2)\n'
  '\n'
  '        hx3 = self.rebnconv3(hx)\n'
  '\n'
  '        hx4 = self.rebnconv4(hx3)\n'
  '\n'
  '        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))\n'
  '        hx3dup = _upsample_like(hx3d, hx2)\n'
  '\n'
  '        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))\n'
  '        hx2dup = _upsample_like(hx2d, hx1)\n'
  '\n'
  '        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))\n'
  '\n'
  '        return hx1d + hxin\n'
  '\n'
  '\n'
  '### RSU-4F ###\n'
  'class RSU4F(nn.Module):  # UNet04FRES(nn.Module):\n'
  '\n'
  '    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):\n'
  '        super(RSU4F, self).__init__()\n'
  '\n'
  '        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)\n'
  '\n'
  '        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)\n'
  '        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)\n'
  '        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)\n'
  '\n'
  '        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)\n'
  '\n'
  '        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)\n'
  '        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)\n'
  '        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)\n'
  '\n'
  '    def forward(self, x):\n'
  '        hx = x\n'
  '\n'
  '        hxin = self.rebnconvin(hx)\n'
  '\n'
  '        hx1 = self.rebnconv1(hxin)\n'
  '        hx2 = self.rebnconv2(hx1)\n'
  '        hx3 = self.rebnconv3(hx2)\n'
  '\n'
  '        hx4 = self.rebnconv4(hx3)\n'
  '\n'
  '        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))\n'
  '        hx2d = self.rebnconv2d(torch.cat((hx3d, hx2), 1))\n'
  '        hx1d = self.rebnconv1d(torch.cat((hx2d, hx1), 1))\n'
  '\n'
  '        return hx1d + hxin\n'
  '\n'
  '\n'
  '##### U^2-Net ####\n'
  'class U2NET(nn.Module):\n'
  '\n'
  '    def __init__(self, in_ch=3, out_ch=1):\n'
  '        super(U2NET, self).__init__()\n'
  '        self.edge = sobel_net()\n'
  '\n'
  '        self.stage1 = RSU7(in_ch, 32, 64)\n'
  '        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.stage2 = RSU6(64, 32, 128)\n'
  '        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.stage3 = RSU5(128, 64, 256)\n'
  '        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.stage4 = RSU4(256, 128, 512)\n'
  '        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.stage5 = RSU4F(512, 256, 512)\n'
  '        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.stage6 = RSU4F(512, 256, 512)\n'
  '\n'
  '        # decoder\n'
  '        self.stage5d = RSU4F(1024, 256, 512)\n'
  '        self.stage4d = RSU4(1024, 128, 256)\n'
  '        self.stage3d = RSU5(512, 64, 128)\n'
  '        self.stage2d = RSU6(256, 32, 64)\n'
  '        self.stage1d = RSU7(128, 16, 64)\n'
  '\n'
  '        self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)\n'
  '        self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)\n'
  '        self.side3 = nn.Conv2d(128, out_ch, 3, padding=1)\n'
  '        self.side4 = nn.Conv2d(256, out_ch, 3, padding=1)\n'
  '        self.side5 = nn.Conv2d(512, out_ch, 3, padding=1)\n'
  '        self.side6 = nn.Conv2d(512, out_ch, 3, padding=1)\n'
  '\n'
  '        self.outconv = nn.Conv2d(6, out_ch, 1)\n'
  '\n'
  '    def forward(self, x):\n'
  '        x = self.edge(x)\n'
  '        hx = x\n'
  '\n'
  '        # stage 1\n'
  '        hx1 = self.stage1(hx)\n'
  '        hx = self.pool12(hx1)\n'
  '\n'
  '        # stage 2\n'
  '        hx2 = self.stage2(hx)\n'
  '        hx = self.pool23(hx2)\n'
  '\n'
  '        # stage 3\n'
  '        hx3 = self.stage3(hx)\n'
  '        hx = self.pool34(hx3)\n'
  '\n'
  '        # stage 4\n'
  '        hx4 = self.stage4(hx)\n'
  '        hx = self.pool45(hx4)\n'
  '\n'
  '        # stage 5\n'
  '        hx5 = self.stage5(hx)\n'
  '        hx = self.pool56(hx5)\n'
  '\n'
  '        # stage 6\n'
  '        hx6 = self.stage6(hx)\n'
  '        hx6up = _upsample_like(hx6, hx5)\n'
  '\n'
  '        # -------------------- decoder --------------------\n'
  '        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))\n'
  '        hx5dup = _upsample_like(hx5d, hx4)\n'
  '\n'
  '        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))\n'
  '        hx4dup = _upsample_like(hx4d, hx3)\n'
  '\n'
  '        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))\n'
  '        hx3dup = _upsample_like(hx3d, hx2)\n'
  '\n'
  '        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))\n'
  '        hx2dup = _upsample_like(hx2d, hx1)\n'
  '\n'
  '        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))\n'
  '\n'
  '        # side output\n'
  '        d1 = self.side1(hx1d)\n'
  '\n'
  '        d2 = self.side2(hx2d)\n'
  '        d2 = _upsample_like(d2, d1)\n'
  '\n'
  '        d3 = self.side3(hx3d)\n'
  '        d3 = _upsample_like(d3, d1)\n'
  '\n'
  '        d4 = self.side4(hx4d)\n'
  '        d4 = _upsample_like(d4, d1)\n'
  '\n'
  '        d5 = self.side5(hx5d)\n'
  '        d5 = _upsample_like(d5, d1)\n'
  '\n'
  '        d6 = self.side6(hx6)\n'
  '        d6 = _upsample_like(d6, d1)\n'
  '\n'
  '        d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))\n'
  '\n'
  '        return torch.sigmoid(d0), torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid(\n'
  '            d4), torch.sigmoid(d5), torch.sigmoid(d6)\n'
  '\n'
  '\n'
  '### U^2-Net small ###\n'
  'class U2NETP(nn.Module):\n'
  '\n'
  '    def __init__(self, in_ch=3, out_ch=1):\n'
  '        super(U2NETP, self).__init__()\n'
  '\n'
  '        self.stage1 = RSU7(in_ch, 16, 64)\n'
  '        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.stage2 = RSU6(64, 16, 64)\n'
  '        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.stage3 = RSU5(64, 16, 64)\n'
  '        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.stage4 = RSU4(64, 16, 64)\n'
  '        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.stage5 = RSU4F(64, 16, 64)\n'
  '        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)\n'
  '\n'
  '        self.stage6 = RSU4F(64, 16, 64)\n'
  '\n'
  '        # decoder\n'
  '        self.stage5d = RSU4F(128, 16, 64)\n'
  '        self.stage4d = RSU4(128, 16, 64)\n'
  '        self.stage3d = RSU5(128, 16, 64)\n'
  '        self.stage2d = RSU6(128, 16, 64)\n'
  '        self.stage1d = RSU7(128, 16, 64)\n'
  '\n'
  '        self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)\n'
  '        self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)\n'
  '        self.side3 = nn.Conv2d(64, out_ch, 3, padding=1)\n'
  '        self.side4 = nn.Conv2d(64, out_ch, 3, padding=1)\n'
  '        self.side5 = nn.Conv2d(64, out_ch, 3, padding=1)\n'
  '        self.side6 = nn.Conv2d(64, out_ch, 3, padding=1)\n'
  '\n'
  '        self.outconv = nn.Conv2d(6, out_ch, 1)\n'
  '\n'
  '    def forward(self, x):\n'
  '        hx = x\n'
  '\n'
  '        # stage 1\n'
  '        hx1 = self.stage1(hx)\n'
  '        hx = self.pool12(hx1)\n'
  '\n'
  '        # stage 2\n'
  '        hx2 = self.stage2(hx)\n'
  '        hx = self.pool23(hx2)\n'
  '\n'
  '        # stage 3\n'
  '        hx3 = self.stage3(hx)\n'
  '        hx = self.pool34(hx3)\n'
  '\n'
  '        # stage 4\n'
  '        hx4 = self.stage4(hx)\n'
  '        hx = self.pool45(hx4)\n'
  '\n'
  '        # stage 5\n'
  '        hx5 = self.stage5(hx)\n'
  '        hx = self.pool56(hx5)\n'
  '\n'
  '        # stage 6\n'
  '        hx6 = self.stage6(hx)\n'
  '        hx6up = _upsample_like(hx6, hx5)\n'
  '\n'
  '        # decoder\n'
  '        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))\n'
  '        hx5dup = _upsample_like(hx5d, hx4)\n'
  '\n'
  '        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))\n'
  '        hx4dup = _upsample_like(hx4d, hx3)\n'
  '\n'
  '        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))\n'
  '        hx3dup = _upsample_like(hx3d, hx2)\n'
  '\n'
  '        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))\n'
  '        hx2dup = _upsample_like(hx2d, hx1)\n'
  '\n'
  '        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))\n'
  '\n'
  '        # side output\n'
  '        d1 = self.side1(hx1d)\n'
  '\n'
  '        d2 = self.side2(hx2d)\n'
  '        d2 = _upsample_like(d2, d1)\n'
  '\n'
  '        d3 = self.side3(hx3d)\n'
  '        d3 = _upsample_like(d3, d1)\n'
  '\n'
  '        d4 = self.side4(hx4d)\n'
  '        d4 = _upsample_like(d4, d1)\n'
  '\n'
  '        d5 = self.side5(hx5d)\n'
  '        d5 = _upsample_like(d5, d1)\n'
  '\n'
  '        d6 = self.side6(hx6)\n'
  '        d6 = _upsample_like(d6, d1)\n'
  '\n'
  '        d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1))\n'
  '\n'
  '        return torch.sigmoid(d0), torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid(\n'
  '            d4), torch.sigmoid(d5), torch.sigmoid(d6)\n'
  '\n'
  '\n'
  'def get_parameter_number(net):\n'
  '    total_num = sum(p.numel() for p in net.parameters())\n'
  '    trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)\n'
  "    return {'Total': total_num, 'Trainable': trainable_num}\n"
  '\n'
  '\n'
  "if __name__ == '__main__':\n"
  '    net = U2NET(4, 1).cuda()\n'
  '    print(get_parameter_number(net))  # 69090500 加attention后69442032\n'
  '    with torch.no_grad():\n'
  '        inputs = torch.zeros(1, 3, 256, 256).cuda()\n'
  '        outs = net(inputs)\n'
  '        print(outs[0].shape)  # torch.Size([2, 3, 256, 256]) torch.Size([2, 2, 256, 256])\n'),
 ('models.encoder.swin_transformer_v2',
  '"""\n'
  'Swin TransformerV2, modified from\n'
  'https://github.com/SwinTransformer/Swin-Transformer-Object-Detection\n'
  '"""\n'
  '\n'
  'import torch\n'
  'import torch.nn as nn\n'
  'import torch.nn.functional as F\n'
  'import torch.utils.checkpoint as checkpoint\n'
  'import numpy as np\n'
  '# Local implementations of the layer utilities used by this model.\n'
  'def to_2tuple(x):\n'
  '    # Accept scalars (including floats), as well as iterable pairs.\n'
  '    try:\n'
  '        return tuple(x)\n'
  '    except TypeError:\n'
  '        return (x, x)\n'
  '\n'
  'def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):\n'
  '    return nn.init.trunc_normal_(tensor, mean=mean, std=std, a=a, b=b)\n'
  '\n'
  'class DropPath(nn.Module):\n'
  '    def __init__(self, drop_prob=0., scale_by_keep=True):\n'
  '        super().__init__()\n'
  '        self.drop_prob = float(drop_prob)\n'
  '        self.scale_by_keep = scale_by_keep\n'
  '\n'
  '    def forward(self, x):\n'
  '        if self.drop_prob == 0. or not self.training:\n'
  '            return x\n'
  '        keep_prob = 1. - self.drop_prob\n'
  '        shape = (x.shape[0],) + (1,) * (x.ndim - 1)\n'
  '        random_tensor = x.new_empty(shape).bernoulli_(keep_prob)\n'
  '        if self.scale_by_keep:\n'
  '            random_tensor.div_(keep_prob)\n'
  '        return x * random_tensor\n'
  '\n'
  '\n'
  'class Mlp(nn.Module):\n'
  '    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n'
  '        super().__init__()\n'
  '        out_features = out_features or in_features\n'
  '        hidden_features = hidden_features or in_features\n'
  '        self.fc1 = nn.Linear(in_features, hidden_features)\n'
  '        self.act = act_layer()\n'
  '        self.fc2 = nn.Linear(hidden_features, out_features)\n'
  '        self.drop = nn.Dropout(drop)\n'
  '\n'
  '    def forward(self, x):\n'
  '        x = self.fc1(x)\n'
  '        x = self.act(x)\n'
  '        x = self.drop(x)\n'
  '        x = self.fc2(x)\n'
  '        x = self.drop(x)\n'
  '        return x\n'
  '\n'
  '\n'
  'def window_partition(x, window_size):\n'
  '    """\n'
  '    Args:\n'
  '        x: (B, H, W, C)\n'
  '        window_size (int): window size\n'
  '    Returns:\n'
  '        windows: (num_windows*B, window_size, window_size, C)\n'
  '    """\n'
  '    B, H, W, C = x.shape\n'
  '    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)\n'
  '    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)\n'
  '    return windows\n'
  '\n'
  '\n'
  'def window_reverse(windows, window_size, H, W):\n'
  '    """\n'
  '    Args:\n'
  '        windows: (num_windows*B, window_size, window_size, C)\n'
  '        window_size (int): Window size\n'
  '        H (int): Height of image\n'
  '        W (int): Width of image\n'
  '    Returns:\n'
  '        x: (B, H, W, C)\n'
  '    """\n'
  '    B = int(windows.shape[0] / (H * W / window_size / window_size))\n'
  '    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)\n'
  '    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)\n'
  '    return x\n'
  '\n'
  '\n'
  'class WindowAttention(nn.Module):\n'
  '    r""" Window based multi-head self attention (W-MSA) module with relative position bias.\n'
  '    It supports both of shifted and non-shifted window.\n'
  '    Args:\n'
  '        dim (int): Number of input channels.\n'
  '        window_size (tuple[int]): The height and width of the window.\n'
  '        num_heads (int): Number of attention heads.\n'
  '        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True\n'
  '        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0\n'
  '        proj_drop (float, optional): Dropout ratio of output. Default: 0.0\n'
  '        pretrained_window_size (tuple[int]): The height and width of the window in pre-training.\n'
  '    """\n'
  '\n'
  '    def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.,\n'
  '                 pretrained_window_size=[0, 0]):\n'
  '\n'
  '        super().__init__()\n'
  '        self.dim = dim\n'
  '        self.window_size = window_size  # Wh, Ww\n'
  '        self.pretrained_window_size = pretrained_window_size\n'
  '        self.num_heads = num_heads\n'
  '\n'
  '        self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)\n'
  '\n'
  '        # mlp to generate continuous relative position bias\n'
  '        self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True),\n'
  '                                     nn.ReLU(inplace=True),\n'
  '                                     nn.Linear(512, num_heads, bias=False))\n'
  '\n'
  '        # get relative_coords_table\n'
  '        relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32)\n'
  '        relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32)\n'
  '        relative_coords_table = torch.stack(\n'
  '            torch.meshgrid([relative_coords_h,\n'
  '                            relative_coords_w])).permute(1, 2, 0).contiguous().unsqueeze(0)  # 1, 2*Wh-1, 2*Ww-1, '
  '2\n'
  '        if pretrained_window_size[0] > 0:\n'
  '            relative_coords_table[:, :, :, 0] /= (pretrained_window_size[0] - 1)\n'
  '            relative_coords_table[:, :, :, 1] /= (pretrained_window_size[1] - 1)\n'
  '        else:\n'
  '            relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1)\n'
  '            relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1)\n'
  '        relative_coords_table *= 8  # normalize to -8, 8\n'
  '        relative_coords_table = torch.sign(relative_coords_table) * torch.log2(\n'
  '            torch.abs(relative_coords_table) + 1.0) / np.log2(8)\n'
  '\n'
  '        self.register_buffer("relative_coords_table", relative_coords_table)\n'
  '\n'
  '        # get pair-wise relative position index for each token inside the window\n'
  '        coords_h = torch.arange(self.window_size[0])\n'
  '        coords_w = torch.arange(self.window_size[1])\n'
  '        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww\n'
  '        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww\n'
  '        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww\n'
  '        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2\n'
  '        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0\n'
  '        relative_coords[:, :, 1] += self.window_size[1] - 1\n'
  '        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1\n'
  '        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww\n'
  '        self.register_buffer("relative_position_index", relative_position_index)\n'
  '\n'
  '        self.qkv = nn.Linear(dim, dim * 3, bias=False)\n'
  '        if qkv_bias:\n'
  '            self.q_bias = nn.Parameter(torch.zeros(dim))\n'
  '            self.v_bias = nn.Parameter(torch.zeros(dim))\n'
  '        else:\n'
  '            self.q_bias = None\n'
  '            self.v_bias = None\n'
  '        self.attn_drop = nn.Dropout(attn_drop)\n'
  '        self.proj = nn.Linear(dim, dim)\n'
  '        self.proj_drop = nn.Dropout(proj_drop)\n'
  '        self.softmax = nn.Softmax(dim=-1)\n'
  '\n'
  '    def forward(self, x, mask=None):\n'
  '        """\n'
  '        Args:\n'
  '            x: input features with shape of (num_windows*B, N, C)\n'
  '            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None\n'
  '        """\n'
  '        B_, N, C = x.shape\n'
  '        qkv_bias = None\n'
  '        if self.q_bias is not None:\n'
  '            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))\n'
  '        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)\n'
  '        qkv = qkv.reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)\n'
  '        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)\n'
  '\n'
  '        # cosine attention\n'
  '        attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))\n'
  '        logit_scale = torch.clamp(self.logit_scale, max=torch.log(torch.tensor(1. / '
  '0.01).to(self.logit_scale.device))).exp()\n'
  '        attn = attn * logit_scale\n'
  '\n'
  '        relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads)\n'
  '        relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(\n'
  '            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # '
  'Wh*Ww,Wh*Ww,nH\n'
  '        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww\n'
  '        relative_position_bias = 16 * torch.sigmoid(relative_position_bias)\n'
  '        attn = attn + relative_position_bias.unsqueeze(0)\n'
  '\n'
  '        if mask is not None:\n'
  '            nW = mask.shape[0]\n'
  '            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)\n'
  '            attn = attn.view(-1, self.num_heads, N, N)\n'
  '            attn = self.softmax(attn)\n'
  '        else:\n'
  '            attn = self.softmax(attn)\n'
  '\n'
  '        attn = self.attn_drop(attn)\n'
  '\n'
  '        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)\n'
  '        x = self.proj(x)\n'
  '        x = self.proj_drop(x)\n'
  '        return x\n'
  '\n'
  '\n'
  'class SwinTransformerBlock(nn.Module):\n'
  '    r""" Swin Transformer Block.\n'
  '    Args:\n'
  '        dim (int): Number of input channels.\n'
  '        num_heads (int): Number of attention heads.\n'
  '        window_size (int): Window size.\n'
  '        shift_size (int): Shift size for SW-MSA.\n'
  '        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n'
  '        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n'
  '        drop (float, optional): Dropout rate. Default: 0.0\n'
  '        attn_drop (float, optional): Attention dropout rate. Default: 0.0\n'
  '        drop_path (float, optional): Stochastic depth rate. Default: 0.0\n'
  '        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU\n'
  '        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm\n'
  '        pretrained_window_size (int): Window size in pre-training.\n'
  '    """\n'
  '\n'
  '    def __init__(self, dim, num_heads, window_size=7, shift_size=0,\n'
  '                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., drop_path=0.,\n'
  '                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, pretrained_window_size=0):\n'
  '        super().__init__()\n'
  '        self.dim = dim\n'
  '        self.num_heads = num_heads\n'
  '        self.window_size = window_size\n'
  '        self.shift_size = shift_size\n'
  '        self.mlp_ratio = mlp_ratio\n'
  '\n'
  '        self.norm1 = norm_layer(dim)\n'
  '        self.attn = WindowAttention(\n'
  '            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,\n'
  '            qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop,\n'
  '            pretrained_window_size=to_2tuple(pretrained_window_size))\n'
  '\n'
  '        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n'
  '        self.norm2 = norm_layer(dim)\n'
  '        mlp_hidden_dim = int(dim * mlp_ratio)\n'
  '        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n'
  '        \n'
  '        self.H = None\n'
  '        self.W = None\n'
  '\n'
  '    def forward(self, x, mask_matrix):\n'
  '        B, L, C = x.shape\n'
  '        H, W = self.H, self.W\n'
  '        assert L == H * W, "input feature has wrong size"\n'
  '\n'
  '        shortcut = x\n'
  '        x = x.view(B, H, W, C)\n'
  '        \n'
  '        # pad feature maps to multiples of window size\n'
  '        pad_l = pad_t = 0\n'
  '        pad_r = (self.window_size - W % self.window_size) % self.window_size\n'
  '        pad_b = (self.window_size - H % self.window_size) % self.window_size\n'
  '        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))\n'
  '        _, Hp, Wp, _ = x.shape\n'
  '        \n'
  '        # cyclic shift\n'
  '        if self.shift_size > 0:\n'
  '            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))\n'
  '            attn_mask = mask_matrix\n'
  '        else:\n'
  '            shifted_x = x\n'
  '            attn_mask = None\n'
  '\n'
  '        # partition windows\n'
  '        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C\n'
  '        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C\n'
  '\n'
  '        # W-MSA/SW-MSA\n'
  '        attn_windows = self.attn(x_windows, mask=attn_mask)  # nW*B, window_size*window_size, C\n'
  '\n'
  '        # merge windows\n'
  '        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)\n'
  "        shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp)  # B H' W' C\n"
  '\n'
  '        # reverse cyclic shift\n'
  '        if self.shift_size > 0:\n'
  '            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))\n'
  '        else:\n'
  '            x = shifted_x\n'
  '        \n'
  '        if pad_r > 0 or pad_b > 0:\n'
  '            x = x[:, :H, :W, :].contiguous()\n'
  '\n'
  '        x = x.view(B, H * W, C)\n'
  '        x = shortcut + self.drop_path(self.norm1(x))\n'
  '\n'
  '        # FFN\n'
  '        x = x + self.drop_path(self.norm2(self.mlp(x)))\n'
  '\n'
  '        return x\n'
  '\n'
  '\n'
  'class PatchMerging(nn.Module):\n'
  '    r""" Patch Merging Layer.\n'
  '    Args:\n'
  '        dim (int): Number of input channels.\n'
  '        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm\n'
  '    """\n'
  '\n'
  '    def __init__(self, dim, norm_layer=nn.LayerNorm):\n'
  '        super().__init__()\n'
  '        self.dim = dim\n'
  '        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)\n'
  '        self.norm = norm_layer(2 * dim)\n'
  '\n'
  '    def forward(self, x, H, W):\n'
  '        """ Forward function.\n'
  '        Args:\n'
  '            x: Input feature, tensor size (B, H*W, C).\n'
  '            H, W: Spatial resolution of the input feature.\n'
  '        """\n'
  '        B, L, C = x.shape\n'
  '        assert L == H * W, "input feature has wrong size"\n'
  '        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."\n'
  '        \n'
  '        x = x.view(B, H, W, C)\n'
  '\n'
  '        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C\n'
  '        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C\n'
  '        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C\n'
  '        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C\n'
  '        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C\n'
  '        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C\n'
  '\n'
  '        x = self.reduction(x)\n'
  '        x = self.norm(x)\n'
  '\n'
  '        return x\n'
  '\n'
  '\n'
  'class BasicLayer(nn.Module):\n'
  '    """ A basic Swin Transformer layer for one stage.\n'
  '    Args:\n'
  '        dim (int): Number of input channels.\n'
  '        depth (int): Number of blocks.\n'
  '        num_heads (int): Number of attention heads.\n'
  '        window_size (int): Local window size.\n'
  '        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n'
  '        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n'
  '        drop (float, optional): Dropout rate. Default: 0.0\n'
  '        attn_drop (float, optional): Attention dropout rate. Default: 0.0\n'
  '        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0\n'
  '        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n'
  '        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None\n'
  '        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.\n'
  '        pretrained_window_size (int): Local window size in pre-training.\n'
  '    """\n'
  '\n'
  '    def __init__(self, dim, depth, num_heads, window_size,\n'
  '                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0.,\n'
  '                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,\n'
  '                 pretrained_window_size=0):\n'
  '\n'
  '        super().__init__()\n'
  '        self.window_size = window_size\n'
  '        self.shift_size = window_size // 2\n'
  '        self.dim = dim\n'
  '        self.depth = depth\n'
  '        self.use_checkpoint = use_checkpoint\n'
  '\n'
  '        # build blocks\n'
  '        self.blocks = nn.ModuleList([\n'
  '            SwinTransformerBlock(dim=dim,\n'
  '                                 num_heads=num_heads, window_size=window_size,\n'
  '                                 shift_size=0 if (i % 2 == 0) else window_size // 2,\n'
  '                                 mlp_ratio=mlp_ratio,\n'
  '                                 qkv_bias=qkv_bias,\n'
  '                                 drop=drop, attn_drop=attn_drop,\n'
  '                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,\n'
  '                                 norm_layer=norm_layer,\n'
  '                                 pretrained_window_size=pretrained_window_size)\n'
  '            for i in range(depth)])\n'
  '\n'
  '        # patch merging layer\n'
  '        if downsample is not None:\n'
  '            self.downsample = downsample(dim=dim, norm_layer=norm_layer)\n'
  '        else:\n'
  '            self.downsample = None\n'
  '\n'
  '    def forward(self, x, H, W):\n'
  '        \n'
  '        # calculate attention mask for SW-MSA\n'
  '        Hp = int(np.ceil(H / self.window_size)) * self.window_size\n'
  '        Wp = int(np.ceil(W / self.window_size)) * self.window_size\n'
  '        img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # 1 Hp Wp 1\n'
  '        h_slices = (slice(0, -self.window_size),\n'
  '                    slice(-self.window_size, -self.shift_size),\n'
  '                    slice(-self.shift_size, None))\n'
  '        w_slices = (slice(0, -self.window_size),\n'
  '                    slice(-self.window_size, -self.shift_size),\n'
  '                    slice(-self.shift_size, None))\n'
  '        cnt = 0\n'
  '        for h in h_slices:\n'
  '            for w in w_slices:\n'
  '                img_mask[:, h, w, :] = cnt\n'
  '                cnt += 1\n'
  '\n'
  '        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1\n'
  '        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)\n'
  '        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)\n'
  '        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))\n'
  '        \n'
  '        for blk in self.blocks:\n'
  '            blk.H, blk.W = H, W\n'
  '            if self.use_checkpoint:\n'
  '                x = checkpoint.checkpoint(blk, x, attn_mask)\n'
  '            else:\n'
  '                x = blk(x, attn_mask)\n'
  '        if self.downsample is not None:\n'
  '            x_down = self.downsample(x, H, W)\n'
  '            Wh, Ww = (H + 1) // 2, (W + 1) // 2\n'
  '            return x, H, W, x_down, Wh, Ww\n'
  '        else:\n'
  '            return x, H, W, x, H, W\n'
  '\n'
  '    def _init_respostnorm(self):\n'
  '        for blk in self.blocks:\n'
  '            nn.init.constant_(blk.norm1.bias, 0)\n'
  '            nn.init.constant_(blk.norm1.weight, 0)\n'
  '            nn.init.constant_(blk.norm2.bias, 0)\n'
  '            nn.init.constant_(blk.norm2.weight, 0)\n'
  '\n'
  '\n'
  'class PatchEmbed(nn.Module):\n'
  '    r""" Image to Patch Embedding\n'
  '    Args:\n'
  '        patch_size (int): Patch token size. Default: 4.\n'
  '        in_chans (int): Number of input image channels. Default: 3.\n'
  '        embed_dim (int): Number of linear projection output channels. Default: 96.\n'
  '        norm_layer (nn.Module, optional): Normalization layer. Default: None\n'
  '    """\n'
  '\n'
  '    def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):\n'
  '        super().__init__()\n'
  '        patch_size = to_2tuple(patch_size)\n'
  '        self.patch_size = patch_size\n'
  '\n'
  '        self.in_chans = in_chans\n'
  '        self.embed_dim = embed_dim\n'
  '\n'
  '        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)\n'
  '        if norm_layer is not None:\n'
  '            self.norm = norm_layer(embed_dim)\n'
  '        else:\n'
  '            self.norm = None\n'
  '\n'
  '    def forward(self, x):\n'
  '        """Forward function."""\n'
  '        # padding\n'
  '        _, _, H, W = x.size()\n'
  '        if W % self.patch_size[1] != 0:\n'
  '            x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))\n'
  '        if H % self.patch_size[0] != 0:\n'
  '            x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))\n'
  '\n'
  '        x = self.proj(x)  # B C Wh Ww\n'
  '        if self.norm is not None:\n'
  '            Wh, Ww = x.size(2), x.size(3)\n'
  '            x = x.flatten(2).transpose(1, 2)\n'
  '            x = self.norm(x)\n'
  '            x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)\n'
  '\n'
  '        return x\n'
  '\n'
  '\n'
  'class SwinTransformerV2(nn.Module):\n'
  '    r""" Swin Transformer\n'
  '        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -\n'
  '          https://arxiv.org/pdf/2103.14030\n'
  '    Args:\n'
  '        patch_size (int | tuple(int)): Patch size. Default: 4\n'
  '        in_chans (int): Number of input image channels. Default: 3\n'
  '        embed_dim (int): Patch embedding dimension. Default: 96\n'
  '        depths (tuple(int)): Depth of each Swin Transformer layer.\n'
  '        num_heads (tuple(int)): Number of attention heads in different layers.\n'
  '        window_size (int): Window size. Default: 7\n'
  '        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4\n'
  '        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True\n'
  '        drop_rate (float): Dropout rate. Default: 0\n'
  '        attn_drop_rate (float): Attention dropout rate. Default: 0\n'
  '        drop_path_rate (float): Stochastic depth rate. Default: 0.1\n'
  '        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.\n'
  '        patch_norm (bool): If True, add normalization after patch embedding. Default: True\n'
  '        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False\n'
  '        pretrained_window_sizes (tuple(int)): Pretrained window sizes of each layer.\n'
  '    """\n'
  '\n'
  '    def __init__(self, patch_size=4, in_chans=3,\n'
  '                 embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],\n'
  '                 window_size=7, mlp_ratio=4., qkv_bias=True,\n'
  '                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,\n'
  '                 norm_layer=nn.LayerNorm, patch_norm=True,\n'
  '                 use_checkpoint=False, pretrained_window_sizes=[0, 0, 0, 0],\n'
  '                 frozen_stages=-1, out_features=None, **kwargs):\n'
  '        super().__init__()\n'
  '\n'
  '        self.num_layers = len(depths)\n'
  '        self.embed_dim = embed_dim\n'
  '        self.patch_norm = patch_norm\n'
  '        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))\n'
  '        self.mlp_ratio = mlp_ratio\n'
  '        \n'
  '        self.frozen_stages = frozen_stages\n'
  '        self._out_features = out_features\n'
  '\n'
  '        # split image into non-overlapping patches\n'
  '        self.patch_embed = PatchEmbed(\n'
  '            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,\n'
  '            norm_layer=norm_layer if self.patch_norm else None)\n'
  '\n'
  '        self.pos_drop = nn.Dropout(p=drop_rate)\n'
  '\n'
  '        # stochastic depth\n'
  '        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule\n'
  '        \n'
  '        # build layers\n'
  '        self.layers = nn.ModuleList()\n'
  '        for i_layer in range(self.num_layers):\n'
  '            layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer),\n'
  '                               num_heads=num_heads[i_layer],\n'
  '                               depth=depths[i_layer],\n'
  '                               window_size=window_size,\n'
  '                               mlp_ratio=self.mlp_ratio,\n'
  '                               qkv_bias=qkv_bias,\n'
  '                               drop=drop_rate, attn_drop=attn_drop_rate,\n'
  '                               drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],\n'
  '                               norm_layer=norm_layer,\n'
  '                               downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,\n'
  '                               use_checkpoint=use_checkpoint,\n'
  '                               pretrained_window_size=pretrained_window_sizes[i_layer])\n'
  '            self.layers.append(layer)\n'
  '\n'
  '        self.num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]\n'
  '        \n'
  '        # add a norm layer for each output\n'
  '        for i_layer in range(self.num_layers):\n'
  "            stage = f'stage{i_layer+1}'\n"
  '            if stage in self._out_features:\n'
  '                layer = norm_layer(self.num_features[i_layer])\n'
  "                layer_name = f'norm{i_layer}'\n"
  '                self.add_module(layer_name, layer)\n'
  '\n'
  '        self.apply(self._init_weights)\n'
  '        for bly in self.layers:\n'
  '            bly._init_respostnorm()\n'
  '\n'
  '    def _init_weights(self, m):\n'
  '        if isinstance(m, nn.Linear):\n'
  '            trunc_normal_(m.weight, std=.02)\n'
  '            if isinstance(m, nn.Linear) and m.bias is not None:\n'
  '                nn.init.constant_(m.bias, 0)\n'
  '        elif isinstance(m, nn.LayerNorm):\n'
  '            nn.init.constant_(m.bias, 0)\n'
  '            nn.init.constant_(m.weight, 1.0)\n'
  '            \n'
  '    def _freeze_stages(self):\n'
  '        if self.frozen_stages >= 0:\n'
  '            self.patch_embed.eval()\n'
  '            for param in self.patch_embed.parameters():\n'
  '                param.requires_grad = False\n'
  '\n'
  '        if self.frozen_stages >= 2:\n'
  '            self.pos_drop.eval()\n'
  '            for i in range(0, self.frozen_stages - 1):\n'
  '                m = self.layers[i]\n'
  '                m.eval()\n'
  '                for param in m.parameters():\n'
  '                    param.requires_grad = False\n'
  '\n'
  '    @torch.jit.ignore\n'
  '    def no_weight_decay(self):\n'
  "        return {'absolute_pos_embed'}\n"
  '\n'
  '    @torch.jit.ignore\n'
  '    def no_weight_decay_keywords(self):\n'
  '        return {"cpb_mlp", "logit_scale", \'relative_position_bias_table\'}\n'
  '\n'
  '    def forward(self, x):\n'
  '        """Forward function."""\n'
  '        x = self.patch_embed(x)\n'
  '\n'
  '        Wh, Ww = x.size(2), x.size(3)\n'
  '        x = x.flatten(2).transpose(1, 2)\n'
  '        x = self.pos_drop(x)\n'
  '\n'
  '        outs = []\n'
  '        for i in range(self.num_layers):\n'
  '            layer = self.layers[i]\n'
  '            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)\n'
  "            name = f'stage{i+1}'\n"
  '            if name in self._out_features:\n'
  "                norm_layer = getattr(self, f'norm{i}')\n"
  '                x_out = norm_layer(x_out)\n'
  '                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()\n'
  '                outs.append(out)\n'
  '\n'
  '        return outs\n'
  '\n'
  '\n'
  'class SwinTransformerV2ForSimMIM(SwinTransformerV2):\n'
  '    def __init__(self, **kwargs):\n'
  '        super().__init__(**kwargs)\n'
  '        self.mask_token = nn.Parameter(torch.zeros(1, self.embed_dim, 1, 1))\n'
  '        trunc_normal_(self.mask_token, mean=0., std=.02)\n'
  '    \n'
  '    def forward(self, x, mask):\n'
  '        """Forward function."""\n'
  '        x = self.patch_embed(x)\n'
  '\n'
  '        B, _, Wh, Ww = x.shape\n'
  '        mask_tokens = self.mask_token.expand(B, -1, Wh, Ww)\n'
  '        mask = mask.unsqueeze(1)\n'
  '        x = x * (1. - mask) + mask_tokens * mask\n'
  '\n'
  '        x = x.flatten(2).transpose(1, 2)\n'
  '        x = self.pos_drop(x)\n'
  '\n'
  '        outs = []\n'
  '        for i in range(self.num_layers):\n'
  '            layer = self.layers[i]\n'
  '            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)\n'
  "            name = f'stage{i+1}'\n"
  '            if name in self._out_features:\n'
  "                norm_layer = getattr(self, f'norm{i}')\n"
  '                x_out = norm_layer(x_out)\n'
  '                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()\n'
  '                outs.append(out)\n'
  '\n'
  '        return outs\n'
  '\n'
  '\n'
  'def build_swinv2_encoder(depths, embed_dim, num_heads, drop_path_rate, pretrained_ws, window_size, use_checkpoint, '
  'encoder_type):\n'
  "    if encoder_type == 'SwinTransformerV2':\n"
  '        encoder_cls = SwinTransformerV2\n'
  "    elif encoder_type == 'SwinTransformerV2ForSimMIM':\n"
  '        encoder_cls = SwinTransformerV2ForSimMIM\n'
  '    else:\n'
  "        raise ValueError(f'encoder {encoder_type} not supported')\n"
  '\n'
  '    return encoder_cls(\n'
  '        patch_size=4,\n'
  '        in_chans=3,\n'
  '        embed_dim=embed_dim,\n'
  '        depths=depths, \n'
  '        num_heads=num_heads,\n'
  '        window_size=window_size,\n'
  '        mlp_ratio=4,\n'
  '        qkv_bias=True,\n'
  '        qk_scale=None,\n'
  '        drop_rate=0.,\n'
  '        attn_drop_rate=0.,\n'
  '        drop_path_rate=drop_path_rate,\n'
  '        patch_norm=True,\n'
  '        pretrained_window_sizes=[pretrained_ws] * len(depths),\n'
  '        frozen_stages=-1,\n'
  '        out_features=["stage1", "stage2", "stage3", "stage4", "stage5"],\n'
  '        use_checkpoint=use_checkpoint,\n'
  '    )'),
 ('models.encoder.swinv2_encoder',
  '# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\n'
  '"""\n'
  'Backbone modules.\n'
  '"""\n'
  'import torch\n'
  'import torch.nn as nn\n'
  '\n'
  'from .swin_transformer_v2 import build_swinv2_encoder\n'
  '\n'
  '\n'
  'class SwinV2Encoder(nn.Module):\n'
  '    def __init__(self, train_backbone, weight_path, embed_dim, depths, num_heads,\n'
  '        drop_path_rate, pretrained_ws, window_size, use_checkpoint, encoder_type):\n'
  '        super(SwinV2Encoder, self).__init__()\n'
  '        self.backbone = build_swinv2_encoder(\n'
  '                            embed_dim=embed_dim, \n'
  '                            depths=depths, \n'
  '                            num_heads=num_heads, \n'
  '                            drop_path_rate=drop_path_rate, \n'
  '                            pretrained_ws=pretrained_ws, \n'
  '                            window_size=window_size, \n'
  '                            use_checkpoint=use_checkpoint,\n'
  '                            encoder_type=encoder_type)\n'
  '        \n'
  '        if not train_backbone:\n'
  '            for name, parameter in self.backbone.named_parameters():\n'
  '                parameter.requires_grad_(False)\n'
  '\n'
  "        # if is_main_process() and weight_path != '':\n"
  '        #     self.load_pretrained_weights(weight_path)\n'
  '        self.num_channels = embed_dim * 8\n'
  '        self.encoder_type = encoder_type\n'
  '    \n'
  '    def forward(self, input, mask=None):\n'
  "        if self.encoder_type == 'SwinTransformerV2':\n"
  '            feats = self.backbone(input)\n'
  "        elif self.encoder_type == 'SwinTransformerV2ForSimMIM':\n"
  '            feats = self.backbone(input, mask)\n'
  '        return feats\n'
  '\n'
  '    def load_pretrained_weights(self, pth_path):\n'
  '        model_dict = self.backbone.state_dict()\n'
  "        pth_dict = torch.load(pth_path, map_location='cpu')\n"
  "        if 'model' in pth_dict:\n"
  "            pth_dict = pth_dict['model']\n"
  '\n'
  '        loaded_keys = []\n'
  '        ignore_keys = []\n'
  '        for model_key in model_dict.keys():\n'
  "            if 'relative_coords_table' in model_key or \\\n"
  "               'relative_position_index' in model_key:\n"
  '                ignore_keys.append(model_key)\n'
  '                continue\n'
  '            if model_key in pth_dict.keys():\n'
  '                model_dict[model_key] = pth_dict[model_key]\n'
  '                loaded_keys.append(model_key)\n'
  "            elif 'cpb_mlp' in model_key:\n"
  "                model_key_rn = model_key.replace('cpb_mlp', 'rpe_mlp')\n"
  '                if model_key_rn in pth_dict.keys():\n'
  '                    model_dict[model_key] = pth_dict[model_key_rn]\n'
  '                    loaded_keys.append(model_key)\n'
  '                    loaded_keys.append(model_key_rn)\n'
  '\n'
  '        missing_keys = [ele for ele in model_dict.keys() if not ele in loaded_keys and not ele in ignore_keys]\n'
  '        unexpected_keys = [ele for ele in pth_dict.keys() if not ele in loaded_keys and not ele in ignore_keys]\n'
  '\n'
  "        print(f'Load pretrained SwinTransformer weights from {pth_path}')\n"
  "        print('Loaded keys: ', loaded_keys)\n"
  "        print('Missing keys: ', missing_keys)\n"
  "        print('Unexpected keys: ', unexpected_keys)\n"
  "        print('Ignored keys:', ignore_keys)\n"
  '\n'
  '        self.backbone.load_state_dict(model_dict)'),
 ('models.encoder',
  'from .swinv2_encoder import SwinV2Encoder\n'
  '\n'
  'def build_encoder(args):\n'
  '    train_encoder = args.lr_encoder_ratio > 0\n'
  "    if args.encoder == 'swinv2':\n"
  "        encoder_type = 'SwinTransformerV2'   \n"
  '        encoder = SwinV2Encoder(train_encoder, args.pretrained_encoder, args.swin_enc_embed_dim, '
  'args.swin_enc_depths, \n'
  '            args.swin_enc_num_heads, args.swin_enc_drop_path_rate, args.swin_enc_pretrained_ws, '
  'args.swin_enc_window_size,\n'
  '            args.swin_use_checkpoint or args.swin_enc_use_checkpoint, encoder_type)\n'
  '    else:\n'
  '        raise NotImplementedError\n'
  '    return encoder'),
 ('models.decoder.swinv2_decoder',
  'import torch\n'
  'import numpy as np\n'
  'import torch.nn as nn\n'
  'import torch.utils.checkpoint as checkpoint\n'
  '\n'
  'from collections.abc import Iterable\n'
  'def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):\n'
  '    return nn.init.trunc_normal_(tensor, mean=mean, std=std, a=a, b=b)\n'
  '\n'
  'from models.block import build_lateral_connection, ConvWithActivation\n'
  'from models.encoder.swin_transformer_v2 import SwinTransformerBlock, window_partition\n'
  '\n'
  'class PatchSplit(nn.Module):\n'
  '    """ Patch Merging Layer\n'
  '    Args:\n'
  '        dim (int): Number of input channels.\n'
  '        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm\n'
  '    """\n'
  '    def __init__(self, dim, norm_layer=nn.LayerNorm):\n'
  '        super(PatchSplit, self).__init__()\n'
  '        self.dim = dim\n'
  '        self.upsample = nn.Linear(dim // 4, dim // 2, bias=False)\n'
  '        self.norm = norm_layer(dim // 2)\n'
  '\n'
  '    def forward(self, x, H, W):\n'
  '        B, L, C = x.shape\n'
  '        assert(L == H * W and C % 4 == 0)\n'
  '\n'
  '        x = x.reshape(B, H, W, 4, C//4)\n'
  '        x = x[:, :, :, [0, 2, 1, 3], :]\n'
  '        x = x.reshape(B, H, W, 2, 2, C//4)\n'
  '        x = x.permute(0, 1, 3, 2, 4, 5)\n'
  '        x = x.reshape(B, H * 2 * W * 2, C//4)\n'
  '        \n'
  '        x = self.upsample(x)\n'
  '        x = self.norm(x) # Swin V2 Post Norm\n'
  '\n'
  '        return x\n'
  '\n'
  '\n'
  'class BasicLayer(nn.Module):\n'
  '    """ A basic Swin Transformer layer for one stage.\n'
  '    Args:\n'
  '        dim (int): Number of input channels.\n'
  '        depth (int): Number of blocks.\n'
  '        num_heads (int): Number of attention heads.\n'
  '        window_size (int): Local window size.\n'
  '        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n'
  '        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n'
  '        drop (float, optional): Dropout rate. Default: 0.0\n'
  '        attn_drop (float, optional): Attention dropout rate. Default: 0.0\n'
  '        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0\n'
  '        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n'
  '        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None\n'
  '        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.\n'
  '        pretrained_window_size (int): Local window size in pre-training.\n'
  '    """\n'
  '\n'
  '    def __init__(self, dim, depth, num_heads, window_size,\n'
  '                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0.,\n'
  '                 drop_path=0., norm_layer=nn.LayerNorm, upsample=None, use_checkpoint=False,\n'
  '                 pretrained_window_size=0):\n'
  '\n'
  '        super().__init__()\n'
  '        self.window_size = window_size\n'
  '        self.shift_size = window_size // 2\n'
  '        self.dim = dim\n'
  '        self.depth = depth\n'
  '        self.use_checkpoint = use_checkpoint\n'
  '\n'
  '        # build blocks\n'
  '        self.blocks = nn.ModuleList([\n'
  '            SwinTransformerBlock(dim=dim,\n'
  '                                 num_heads=num_heads, window_size=window_size,\n'
  '                                 shift_size=0 if (i % 2 == 0) else window_size // 2,\n'
  '                                 mlp_ratio=mlp_ratio,\n'
  '                                 qkv_bias=qkv_bias,\n'
  '                                 drop=drop, attn_drop=attn_drop,\n'
  '                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,\n'
  '                                 norm_layer=norm_layer,\n'
  '                                 pretrained_window_size=pretrained_window_size)\n'
  '            for i in range(depth)])\n'
  '\n'
  '        # patch merging layer\n'
  '        if upsample is not None:\n'
  '            self.upsample = upsample(dim=dim, norm_layer=norm_layer)\n'
  '        else:\n'
  '            self.upsample = None\n'
  '\n'
  '    def forward(self, x, H, W):\n'
  '        \n'
  '        # calculate attention mask for SW-MSA\n'
  '        Hp = int(np.ceil(H / self.window_size)) * self.window_size\n'
  '        Wp = int(np.ceil(W / self.window_size)) * self.window_size\n'
  '        img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # 1 Hp Wp 1\n'
  '        h_slices = (slice(0, -self.window_size),\n'
  '                    slice(-self.window_size, -self.shift_size),\n'
  '                    slice(-self.shift_size, None))\n'
  '        w_slices = (slice(0, -self.window_size),\n'
  '                    slice(-self.window_size, -self.shift_size),\n'
  '                    slice(-self.shift_size, None))\n'
  '        cnt = 0\n'
  '        for h in h_slices:\n'
  '            for w in w_slices:\n'
  '                img_mask[:, h, w, :] = cnt\n'
  '                cnt += 1\n'
  '\n'
  '        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1\n'
  '        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)\n'
  '        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)\n'
  '        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))\n'
  '        \n'
  '        for blk in self.blocks:\n'
  '            blk.H, blk.W = H, W\n'
  '            if self.use_checkpoint:\n'
  '                x = checkpoint.checkpoint(blk, x, attn_mask)\n'
  '            else:\n'
  '                x = blk(x, attn_mask)\n'
  '\n'
  '        if self.upsample is not None:\n'
  '            x_up = self.upsample(x, H, W)\n'
  '            Wh, Ww = H * 2, W * 2\n'
  '            return x, H, W, x_up, Wh, Ww\n'
  '        else:\n'
  '            return x, H, W, x, H, W\n'
  '\n'
  '    def _init_respostnorm(self):\n'
  '        for blk in self.blocks:\n'
  '            nn.init.constant_(blk.norm1.bias, 0)\n'
  '            nn.init.constant_(blk.norm1.weight, 0)\n'
  '            nn.init.constant_(blk.norm2.bias, 0)\n'
  '            nn.init.constant_(blk.norm2.weight, 0)\n'
  '\n'
  '\n'
  'class SwinTransformerV2Decoder(nn.Module):\n'
  '    r""" Swin Transformer\n'
  '        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -\n'
  '          https://arxiv.org/pdf/2103.14030\n'
  '    Args:\n'
  '        patch_size (int | tuple(int)): Patch size. Default: 4\n'
  '        in_chans (int): Number of input image channels. Default: 3\n'
  '        embed_dim (int): Patch embedding dimension. Default: 96\n'
  '        depths (tuple(int)): Depth of each Swin Transformer layer.\n'
  '        num_heads (tuple(int)): Number of attention heads in different layers.\n'
  '        window_size (int): Window size. Default: 7\n'
  '        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4\n'
  '        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True\n'
  '        drop_rate (float): Dropout rate. Default: 0\n'
  '        attn_drop_rate (float): Attention dropout rate. Default: 0\n'
  '        drop_path_rate (float): Stochastic depth rate. Default: 0.1\n'
  '        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.\n'
  '        patch_norm (bool): If True, add normalization after patch embedding. Default: True\n'
  '        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False\n'
  '        pretrained_window_sizes (tuple(int)): Pretrained window sizes of each layer.\n'
  '    """\n'
  '\n'
  '    def __init__(self, \n'
  '                 encoder_dim=768,\n'
  '                 embed_dim=768,\n'
  '                 depths=[2, 6, 2, 2, 2],\n'
  '                 num_heads=[24, 12, 6, 3, 2],\n'
  '                 window_size=7, \n'
  '                 mlp_ratio=4., \n'
  '                 qkv_bias=True,\n'
  '                 drop_rate=0., \n'
  '                 attn_drop_rate=0., \n'
  '                 drop_path_rate=0.1,\n'
  '                 norm_layer=nn.LayerNorm, \n'
  '                 patch_norm=True,\n'
  '                 use_checkpoint=False, \n'
  '                 pretrained_window_sizes=[0, 0, 0, 0, 0],\n'
  '                 frozen_stages=-1,\n'
  '                 skip_stages=None,\n'
  '                 intermediate_erase_stages=None,\n'
  '                 mask_stage=None,\n'
  '                 pred_mask=True):\n'
  '        super(SwinTransformerV2Decoder, self).__init__()\n'
  '        \n'
  '        self.num_layers = len(depths)\n'
  '        self.embed_dim = embed_dim\n'
  '        self.patch_norm = patch_norm\n'
  '        self.mlp_ratio = mlp_ratio\n'
  '        self.frozen_stages = frozen_stages\n'
  '        self.pred_mask = pred_mask\n'
  '        self.skip_stages = skip_stages\n'
  '        self.intermediate_erase_stages = intermediate_erase_stages\n'
  '        self.mask_stage = mask_stage\n'
  '\n'
  '        self.pos_drop = nn.Dropout(p=drop_rate)\n'
  '\n'
  '        # stochastic depth\n'
  '        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule\n'
  '        \n'
  '        # build layers\n'
  '        self.layers = nn.ModuleList()\n'
  '        for i_layer in range(self.num_layers):\n'
  '            layer = BasicLayer(dim=int(embed_dim * 0.5 ** i_layer),\n'
  '                               num_heads=num_heads[i_layer],\n'
  '                               depth=depths[i_layer],\n'
  '                               window_size=window_size,\n'
  '                               mlp_ratio=self.mlp_ratio,\n'
  '                               qkv_bias=qkv_bias,\n'
  '                               drop=drop_rate, attn_drop=attn_drop_rate,\n'
  '                               drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],\n'
  '                               norm_layer=norm_layer,\n'
  '                               upsample=PatchSplit,\n'
  '                               use_checkpoint=use_checkpoint,\n'
  '                               pretrained_window_size=pretrained_window_sizes[i_layer])\n'
  '            self.layers.append(layer)\n'
  '            \n'
  '        num_features = [int(embed_dim * 0.5 ** (i + 1)) for i in range(self.num_layers)]\n'
  '        self.num_features = num_features\n'
  '\n'
  '        self.pred_conv = nn.Sequential(\n'
  '            nn.Conv2d(num_features[self.num_layers - 1], 128, 3, 2, 1),   # 512 → 256\n'
  '            nn.BatchNorm2d(128),\n'
  '            nn.ReLU(inplace=True),\n'
  '\n'
  '            nn.Conv2d(128, 256, 3, 2, 1), # 256 → 128\n'
  '            nn.BatchNorm2d(256),\n'
  '            nn.ReLU(inplace=True),\n'
  '\n'
  '            nn.Conv2d(256, 256, 3, 1, padding=2, dilation=2),  # 保 128\n'
  '            nn.BatchNorm2d(256),\n'
  '            nn.ReLU(inplace=True),\n'
  '\n'
  '            nn.Conv2d(256, 192, 3, 2, 1), # 128 → 64\n'
  '            nn.BatchNorm2d(192),\n'
  '            nn.ReLU(inplace=True),\n'
  '\n'
  '            nn.Conv2d(192, 2, 3, 2, 1)    # 64 → 32\n'
  '        )\n'
  '\n'
  '        if not self.skip_stages is None:\n'
  '            self.lateral_connection_list = nn.ModuleList([\n'
  '                build_lateral_connection(\n'
  '                    int(encoder_dim * 0.5 ** (layer_idx + 1)), \n'
  '                    num_features[layer_idx] \n'
  '                )\n'
  '                for layer_idx in to_layer_idx(self.skip_stages)\n'
  '            ])\n'
  '        \n'
  '        self.intermediate_convs_ = nn.ModuleList([\n'
  '            nn.Sequential(\n'
  '                nn.Conv2d(96, 192, kernel_size=3, stride=2, padding=1),  # -> 64x64\n'
  '                nn.BatchNorm2d(192),\n'
  '                nn.ReLU(inplace=True),\n'
  '\n'
  '                nn.Conv2d(192, 192, kernel_size=3, stride=2, padding=1), # -> 32x32\n'
  '                nn.BatchNorm2d(192),\n'
  '                nn.ReLU(inplace=True),\n'
  '\n'
  '                nn.Conv2d(192, 2, kernel_size=1)  # -> 2 x 32 x 32\n'
  '            ),\n'
  '            nn.Sequential(\n'
  '                nn.Conv2d(48, 96, kernel_size=3, stride=2, padding=1),   # -> 128x128\n'
  '                nn.BatchNorm2d(96),\n'
  '                nn.ReLU(inplace=True),\n'
  '\n'
  '                nn.Conv2d(96, 192, kernel_size=3, stride=2, padding=1),  # -> 64x64\n'
  '                nn.BatchNorm2d(192),\n'
  '                nn.ReLU(inplace=True),\n'
  '\n'
  '                nn.Conv2d(192, 192, kernel_size=3, stride=2, padding=1), # -> 32x32\n'
  '                nn.BatchNorm2d(192),\n'
  '                nn.ReLU(inplace=True),\n'
  '\n'
  '                nn.Conv2d(192, 2, kernel_size=1)  # -> 2 x 32 x 32\n'
  '            )\n'
  '        ])        \n'
  '\n'
  '        # if not self.mask_stage is None:\n'
  '        mask_layer_idx = to_layer_idx(self.mask_stage)\n'
  '        self.mask_conv = nn.Sequential(\n'
  "            ConvWithActivation('deconv', num_features[mask_layer_idx], 64, 3, 2, 1),\n"
  '            nn.Conv2d(64, 1, 3, 1, 1))\n'
  '\n'
  '        # add a norm layer for each output\n'
  "        self.output_stages = self.intermediate_erase_stages + [f'stage{self.num_layers + 1}']\n"
  '        if self.pred_mask:\n'
  '            self.output_stages = self.output_stages + [self.mask_stage]\n'
  '        self.output_stages = list(set(self.output_stages))\n'
  '        self.output_stages.sort()\n'
  '        self.output_idx = to_layer_idx(self.output_stages)\n'
  '        for layer_idx in self.output_idx:\n'
  '            layer = norm_layer(num_features[layer_idx])\n'
  "            layer_name = f'norm{layer_idx}'\n"
  '            self.add_module(layer_name, layer)\n'
  '\n'
  '        self.apply(self._init_weights)\n'
  '        for bly in self.layers:\n'
  '            bly._init_respostnorm()\n'
  '\n'
  '    def _init_weights(self, m):\n'
  '        if isinstance(m, nn.Linear):\n'
  '            trunc_normal_(m.weight, std=.02)\n'
  '            if isinstance(m, nn.Linear) and m.bias is not None:\n'
  '                nn.init.constant_(m.bias, 0)\n'
  '        elif isinstance(m, nn.LayerNorm):\n'
  '            nn.init.constant_(m.bias, 0)\n'
  '            nn.init.constant_(m.weight, 1.0)\n'
  '            \n'
  '    def _freeze_stages(self):\n'
  '        if self.frozen_stages >= 0:\n'
  '            self.patch_embed.eval()\n'
  '            for param in self.patch_embed.parameters():\n'
  '                param.requires_grad = False\n'
  '\n'
  '        if self.frozen_stages >= 2:\n'
  '            self.pos_drop.eval()\n'
  '            for i in range(0, self.frozen_stages - 1):\n'
  '                m = self.layers[i]\n'
  '                m.eval()\n'
  '                for param in m.parameters():\n'
  '                    param.requires_grad = False\n'
  '\n'
  '    @torch.jit.ignore\n'
  '    def no_weight_decay(self):\n'
  "        return {'absolute_pos_embed'}\n"
  '\n'
  '    @torch.jit.ignore\n'
  '    def no_weight_decay_keywords(self):\n'
  '        return {"cpb_mlp", "logit_scale", \'relative_position_bias_table\'}\n'
  '\n'
  '    def forward(self, x, skip_features):\n'
  '        """Forward function."""\n'
  '\n'
  '        Wh, Ww = x.size(2), x.size(3)\n'
  '        x = x.flatten(2).transpose(1, 2)\n'
  '        x = self.pos_drop(x)\n'
  '\n'
  '        outputs = []\n'
  '        mask = None\n'
  '        for i in range(self.num_layers):\n'
  '            layer = self.layers[i]\n'
  '            _, _, _, x, Wh, Ww = layer(x, Wh, Ww)\n'
  '            \n'
  "            name = f'stage{i+2}'\n"
  '\n'
  '            if name in self.skip_stages:\n'
  '                skip_feature = self.lateral_connection_list[self.skip_stages.index(name)](skip_features[-i-2])\n'
  '                skip_feature = skip_feature.flatten(2).transpose(1, 2)\n'
  '                x = x + skip_feature\n'
  '\n'
  '            if name in self.output_stages:\n'
  "                norm_layer = getattr(self, f'norm{i}')\n"
  '                x_out = norm_layer(x)\n'
  '                x_out = x_out.view(-1, Wh, Ww, self.num_features[i]).permute(0, 3, 1, 2).contiguous()\n'
  '\n'
  '            if name in self.intermediate_erase_stages:\n'
  '                # inter_conv = self.intermediate_convs[self.intermediate_erase_stages.index(name)]\n'
  '                inter_conv = self.intermediate_convs_[self.intermediate_erase_stages.index(name)]\n'
  '                inter_erase_output = inter_conv(x_out)\n'
  '                outputs.append(inter_erase_output)\n'
  '\n'
  '            if name == self.mask_stage:\n'
  '                mask = self.mask_conv(x_out)\n'
  '\n'
  '        erase_output = self.pred_conv(x_out)\n'
  '        outputs.append(erase_output)\n'
  '        return outputs, mask\n'
  '\n'
  'def to_layer_idx(stages):\n'
  '    if isinstance(stages, str):\n'
  "        return int(stages.replace('stage', '')) - 2\n"
  '    elif isinstance(stages, Iterable):\n'
  "        return [int(stage.replace('stage', '')) - 2 for stage in stages]\n"
  '\n'
  '\n'
  'def build_swin_v2_decoder(args):\n'
  '    decoder = SwinTransformerV2Decoder(\n'
  '        window_size=args.swin_dec_window_size, \n'
  '        mlp_ratio=4., \n'
  '        qkv_bias=True,\n'
  '        drop_rate=0., \n'
  '        attn_drop_rate=0., \n'
  '        drop_path_rate=args.swin_dec_drop_path_rate,\n'
  '        norm_layer=nn.LayerNorm, \n'
  '        patch_norm=True,\n'
  '        use_checkpoint=args.swin_use_checkpoint or args.swin_dec_use_checkpoint, \n'
  '        pretrained_window_sizes=[args.swin_dec_pretrained_ws] * len(args.swin_dec_depths),\n'
  '        frozen_stages=-1,\n'
  "        skip_stages=['stage2', 'stage3', 'stage4'],\n"
  "        # intermediate_erase_stages=['stage4', 'stage5'] if args.intermediate_erase else [],\n"
  "        intermediate_erase_stages=['stage4', 'stage5'] ,\n"
  "        mask_stage='stage5' if args.pred_mask else None,\n"
  '        pred_mask=args.pred_mask,\n'
  '        embed_dim=args.swin_enc_embed_dim * 8,\n'
  '        depths=args.swin_dec_depths,\n'
  '        num_heads=args.swin_dec_num_heads,\n'
  '        encoder_dim=args.swin_enc_embed_dim * 8,\n'
  '    )\n'
  '\n'
  '    if args.pretrained_decoder:\n'
  '        decoder = load_pretrained_decoder(decoder, args.pretrained_decoder)\n'
  '    \n'
  '    return decoder\n'
  '\n'
  'def load_pretrained_decoder(decoder, pth_path):\n'
  '    decoder_dict = decoder.state_dict()\n'
  "    pth_dict = torch.load(pth_path, map_location='cpu')\n"
  "    if 'model' in pth_dict:\n"
  "        pth_dict = pth_dict['model']\n"
  '\n'
  '    loaded_keys = []\n'
  '    ignore_keys = []\n'
  '    for dec_key in decoder_dict.keys():\n'
  "        if 'relative_coords_table' in dec_key or \\\n"
  "           'relative_position_index' in dec_key:\n"
  '            ignore_keys.append(dec_key)\n'
  '            continue\n'
  '\n'
  '        if dec_key in pth_dict.keys():\n'
  '            decoder_dict[dec_key] = pth_dict[dec_key]\n'
  '            loaded_keys.append(dec_key)\n'
  '        \n'
  '    missing_keys = [ele for ele in decoder_dict.keys() if not ele in loaded_keys and not ele in ignore_keys]\n'
  '    unexpected_keys = [ele for ele in pth_dict.keys() if not ele in loaded_keys and not ele in ignore_keys]\n'
  '\n'
  "    print(f'Load pretrained SwinTransformer Decoder weights from {pth_path}')\n"
  "    print('Loaded keys: ', loaded_keys)\n"
  "    print('Missing keys: ', missing_keys)\n"
  "    print('Unexpected keys: ', unexpected_keys)\n"
  "    print('Ignored keys:', ignore_keys)\n"
  '\n'
  '    decoder.load_state_dict(decoder_dict)\n'
  '\n'
  '    return decoder'),
 ('models.decoder',
  'from .swinv2_decoder import build_swin_v2_decoder\n'
  '\n'
  'def build_decoder(args):\n'
  "    if args.decoder == 'swinv2':\n"
  '        return build_swin_v2_decoder(args)\n'
  '    raise NotImplementedError\n'),
 ('models.build_seg',
  'import torch\n'
  'import torch.nn as nn\n'
  'from models.seg import U2NETP\n'
  '\n'
  'class Net(nn.Module):\n'
  '    def __init__(self):\n'
  '        super(Net, self).__init__()\n'
  '        self.msk = U2NETP(3, 1)\n'
  '\n'
  '    def forward(self, x):\n'
  '        msk, _1,_2,_3,_4,_5,_6 = self.msk(x)\n'
  '        return msk\n'
  '\n'
  'def reload_seg_model(model, path=""):\n'
  '    if not bool(path):\n'
  '        return model\n'
  '    else:\n'
  '        model_dict = model.state_dict()\n'
  "        # pretrained_dict = torch.load(path, map_location='cuda:0')\n"
  "        pretrained_dict = torch.load(path, map_location='cpu')\n"
  '        print(len(pretrained_dict.keys()))\n'
  '        pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict}\n'
  '        print(len(pretrained_dict.keys()))\n'
  '        model_dict.update(pretrained_dict)\n'
  '        model.load_state_dict(model_dict)\n'
  '\n'
  '        return model\n'
  '\n'
  'def build(args):\n'
  '    model = Net()\n'
  '    device = torch.device(args.device)\n'
  '    model = model.to(device)\n'
  '\n'
  '    if args.distributed:\n'
  '        model = torch.nn.parallel.DistributedDataParallel(\n'
  '            model,\n'
  '            device_ids=[args.gpu],\n'
  '            # find_unused_parameters=True\n'
  '        )\n'
  '        \n'
  '    return model'),
 ('models.viteraser',
  'import torch\n'
  'import torch.nn as nn\n'
  'import torch.nn.functional as F\n'
  '\n'
  'from .decoder import build_decoder\n'
  'from .encoder import build_encoder\n'
  '\n'
  'class SegHead(nn.Module):\n'
  '    def __init__(self, in_channel, encoder_stride):\n'
  '        super(SegHead, self).__init__()\n'
  '        self.decoder = nn.Sequential(\n'
  '            nn.Conv2d(\n'
  '                in_channels=in_channel,\n'
  '                out_channels=encoder_stride ** 2, kernel_size=1),\n'
  '            nn.PixelShuffle(encoder_stride)\n'
  '        )\n'
  '\n'
  '    def forward(self, feature):\n'
  '        seg_res = self.decoder(feature)\n'
  '        return seg_res\n'
  '\n'
  'class pointHead(nn.Module):\n'
  '    def __init__(self, in_channel, encoder_stride):\n'
  '        super(pointHead, self).__init__()\n'
  '        self.decoder = nn.Sequential(\n'
  '            nn.Conv2d(\n'
  '                in_channels=in_channel,\n'
  '                out_channels=2 *(encoder_stride ** 2), kernel_size=1),\n'
  '            nn.PixelShuffle(encoder_stride)\n'
  '        )\n'
  '\n'
  '    def forward(self, feature):\n'
  '        seg_res = self.decoder(feature)\n'
  '        return seg_res\n'
  '\n'
  'class ViTEraser(nn.Module):\n'
  '    # def __init__(self, encoder, decoder, vgg16):\n'
  '    def __init__(self, encoder, decoder):\n'
  '        super(ViTEraser, self).__init__()\n'
  '        self.encoder = encoder\n'
  '        self.decoder = decoder\n'
  '        self.pixel_embed = nn.Linear(encoder.num_channels, decoder.embed_dim)\n'
  '\n'
  '        # tiny\n'
  '        self.seg_mask = SegHead(in_channel=96 * 8,  encoder_stride=32) \n'
  '        self.seg_point_2 = pointHead(in_channel=96 * 8,  encoder_stride=2) \n'
  '\n'
  '        # base\n'
  '        # self.seg_mask = SegHead(in_channel=1024,  encoder_stride=32) \n'
  '        # self.seg_point_2 = pointHead(in_channel=1024,  encoder_stride=2) \n'
  '\n'
  '    def forward(self, images):        \n'
  '        enc_ms_feats = self.encoder(images)\n'
  '\n'
  '        pred_mask_coarse = self.seg_mask(enc_ms_feats[-1])\n'
  '        pred_point_coarse = self.seg_point_2(enc_ms_feats[-1])\n'
  '        \n'
  '        enc_feat = self.pixel_embed(\n'
  '                enc_ms_feats[-1].permute(0, 2, 3, 1)\n'
  '            ).permute(0, 3, 1, 2)\n'
  '\n'
  '        outputs, pred_mask = self.decoder(enc_feat, enc_ms_feats)\n'
  '        \n'
  '        if not self.training:\n'
  '            return outputs[-1], pred_mask\n'
  '\n'
  '        return {\n'
  "            'outputs': outputs,\n"
  '            "pred_point_coarse": pred_point_coarse,\n'
  '            "pred_mask_coarse" : pred_mask_coarse,\n'
  "            'pred_mask': pred_mask\n"
  '        }\n'
  '    \n'
  'class MLP(nn.Module):\n'
  '    """ Very simple multi-layer perceptron (also called FFN)"""\n'
  '\n'
  '    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):\n'
  '        super().__init__()\n'
  '        self.num_layers = num_layers\n'
  '        h = [hidden_dim] * (num_layers - 1)\n'
  '        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))\n'
  '\n'
  '    def forward(self, x):\n'
  '        for i, layer in enumerate(self.layers):\n'
  '            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)\n'
  '        return x\n'
  '\n'
  '\n'
  'def load_pretrained_model(model, weight_path, ignore_encoder=False):\n'
  "    weight = torch.load(weight_path, map_location='cpu')['model']\n"
  '    model_dict = model.state_dict()\n'
  '\n'
  '    loaded_keys = []\n'
  '    ignore_keys = []\n'
  '    for k, v in weight.items():\n'
  "        if 'relative_coords_table' in k or \\\n"
  "               'relative_position_index' in k:\n"
  '            ignore_keys.append(k)\n'
  '            continue \n'
  '\n'
  "        if ignore_encoder and k.startswith('encoder.'):\n"
  '            ignore_keys.append(k)\n'
  '            continue\n'
  '        \n'
  '        if k in model_dict.keys():\n'
  '            model_dict[k] = v\n'
  '            loaded_keys.append(k)\n'
  '        else:\n'
  '            ignore_keys.append(k)\n'
  '        \n'
  '    model.load_state_dict(model_dict)\n'
  "    print(f'Load Model from {weight_path}')\n"
  "    print('Loaded keys:', loaded_keys)\n"
  "    print('Ignored keys:', ignore_keys)\n"
  '    return model\n'
  '\n'
  '\n'
  'def build(args):\n'
  '    encoder = build_encoder(args)\n'
  '    decoder = build_decoder(args)\n'
  '    \n'
  '    model = ViTEraser(\n'
  '        encoder=encoder, \n'
  '        decoder=decoder)\n'
  '\n'
  '    if args.pretrained_model:\n'
  '        model = load_pretrained_model(model, args.pretrained_model, args.load_pretrain_ignore_encoder)\n'
  '\n'
  '    device = torch.device(args.device)\n'
  '    model = model.to(device)\n'
  '\n'
  '    if args.distributed:\n'
  '        model = torch.nn.parallel.DistributedDataParallel(\n'
  '            model,\n'
  '            device_ids=[args.gpu],\n'
  '            # find_unused_parameters=True\n'
  '        )\n'
  '        \n'
  '    return model'),
 ('models',
  'from .viteraser import build as build_viteraser_model\n'
  'from .build_seg import build as build_seg_model\n'
  '\n'
  'def build_model2(args):\n'
  '    return build_viteraser_model(args)\n'
  '\n'
  'def build_model1(args):\n'
  '    return build_seg_model(args)')]


def _install_embedded_models():
    import sys
    import types

    package_names = {"models", "models.encoder", "models.decoder"}
    for name, _source in _EMBEDDED_MODEL_SOURCES:
        module = types.ModuleType(name)
        module.__file__ = f"<embedded {name}>"
        module.__loader__ = None
        if name in package_names:
            module.__path__ = []
            module.__package__ = name
        else:
            module.__package__ = name.rpartition(".")[0]
        module.__embedded_preprocessor_model__ = True
        sys.modules[name] = module

    for name, source in _EMBEDDED_MODEL_SOURCES:
        module = sys.modules[name]
        exec(compile(source, module.__file__, "exec"), module.__dict__)

    return sys.modules["models"].build_model2, sys.modules["models"].build_model1


build_model2, build_model1 = _install_embedded_models()

def _default_preprocessor_args(device: str):
    return SimpleNamespace(
        resume="",
        print_freq=5,
        save_interval=2,
        lr=1e-4,
        lr_encoder_ratio=0.2,
        batch_size=1,
        weight_decay=1e-4,
        epochs=250,
        warmup_min_lr=0.0001,
        min_lr=0.000001,
        warmup_epochs=10,
        milestones=[80],
        segmim_finetune=False,
        eval=True,
        output_dir="",
        device=device,
        seed=42,
        clip_max_norm=0,
        layer_decay=0.75,
        pretrained_model="",
        load_pretrain_ignore_encoder=False,
        pretrained_encoder="",
        pretrained_decoder="",
        pretrained_vgg16="",
        encoder="swinv2",
        decoder="swinv2",
        swin_dec_depths=[2, 6, 2, 2, 2],
        swin_dec_num_heads=[24, 12, 6, 3, 2],
        swin_dec_window_size=16,
        swin_dec_drop_path_rate=0.2,
        swin_dec_pretrained_ws=8,
        swin_enc_depths=[2, 2, 6, 2],
        swin_enc_num_heads=[3, 6, 12, 24],
        swin_enc_drop_path_rate=0.2,
        swin_enc_embed_dim=96,
        swin_enc_pretrained_ws=8,
        swin_enc_window_size=16,
        pred_mask=True,
        intermediate_erase=True,
        swin_use_checkpoint=False,
        swin_enc_use_checkpoint=False,
        swin_dec_use_checkpoint=False,
        distributed=False,
        gpu=0,
    )


def _load_state_dict(model, checkpoint_path: str):
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    if isinstance(checkpoint, dict) and "model" in checkpoint:
        checkpoint = checkpoint["model"]
    model.load_state_dict(checkpoint, strict=False)


class Preprocessor:
    def __init__(self, model_path: str, device: str | None = None, batch_size: int = 16):
        self.model_path = Path(model_path)
        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        self.batch_size = max(1, int(batch_size))
        args = _default_preprocessor_args(str(self.device))

        self.model1 = build_model1(args)
        self.model2 = build_model2(args)

        _load_state_dict(self.model1, str(self.model_path / "preprocessor1.pth"))
        _load_state_dict(self.model2, str(self.model_path / "preprocessor2.pth"))

        self.model1.eval()
        self.model2.eval()

    @torch.no_grad()
    def preprocess_image(self, image: Image.Image):
        return self.preprocess_images([image])[0]

    @torch.no_grad()
    def preprocess_images(self, images: list[Image.Image], batch_size: int | None = None):
        if not images:
            return []

        batch_size = max(1, int(batch_size or self.batch_size))
        preprocessed_images = []
        img_size = [512, 512]

        for start in range(0, len(images), batch_size):
            batch_images = images[start:start + batch_size]
            img_arrays = [
                np.asarray(image.convert("RGB")).astype(np.float32) / 255.0
                for image in batch_images
            ]
            inp_batch = torch.stack([
                torch.from_numpy(cv2.resize(img, img_size).transpose(2, 0, 1))
                for img in img_arrays
            ], dim=0).to(self.device)

            pred_mask_batch = self.model1(inp_batch)
            pred_mask_01_batch = (pred_mask_batch > 0.8).float()

            largest_masks = []
            for mask_tensor in pred_mask_01_batch:
                mask = mask_tensor.squeeze().cpu().numpy().astype(np.uint8)
                num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
                if num_labels > 1:
                    largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
                    largest_mask = (labels == largest_label).astype(np.uint8)
                else:
                    largest_mask = mask
                largest_masks.append(torch.from_numpy(largest_mask).float().unsqueeze(0))

            largest_mask_batch = torch.stack(largest_masks, dim=0).to(self.device)
            outputs_batch, _ = self.model2(inp_batch * largest_mask_batch)

            for img, point_positions in zip(img_arrays, outputs_batch):
                size = img.shape[:2][::-1]
                preprocessed = bilinear_preprocessing(
                    warped_img=torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).to(self.device),
                    point_positions=point_positions.unsqueeze(0),
                    img_size=tuple(size),
                )
                preprocessed = (preprocessed[0].detach().cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8)
                preprocessed_images.append(Image.fromarray(preprocessed).convert("RGB"))

            del inp_batch, pred_mask_batch, pred_mask_01_batch, largest_mask_batch, outputs_batch
            if self.device.type == "cuda":
                torch.cuda.empty_cache()

        return preprocessed_images

    def preprocess_doc(self, doc: dict):
        return {**doc, "images": self.preprocess_images(doc["images"])}

    def preprocess_docs(self, docs: list[dict], batch_size: int | None = None):
        if not docs:
            return []

        all_images = []
        doc_image_counts = []
        for doc in docs:
            images = doc["images"]
            doc_image_counts.append(len(images))
            all_images.extend(images)

        all_preprocessed = self.preprocess_images(all_images, batch_size=batch_size)

        preprocessed_docs = []
        offset = 0
        for doc, count in zip(docs, doc_image_counts):
            preprocessed_docs.append({**doc, "images": all_preprocessed[offset:offset + count]})
            offset += count

        return preprocessed_docs