Spaces:

zenosai
/

MonkeyOCRv2

Running on Zero

MonkeyOCRv2 / modeling /modeling_preprocessor.py

zenosai

update MonkeyOCRv2 Space

46de218 8 days ago

92.6 kB

	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_dim2, input_dim2, 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, 2Wh-1, 2Ww-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, WhWw, WhWw\n'
	' relative_coords = relative_coords.permute(1, 2, 0).contiguous() # WhWw, WhWw, 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) # WhWw, WhWw\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, WhWw, WhWw) 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) # '
	'WhWw,WhWw,nH\n'
	' relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, WhWw, WhWw\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) # nWB, window_sizewindow_size, C\n'
	'\n'
	' # W-MSA/SW-MSA\n'
	' attn_windows = self.attn(x_windows, mask=attn_mask) # nWB, window_sizewindow_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/2W/2 4C\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