From 1135b49f4f4c90c8bb37a42f194ba09629b554f1 Mon Sep 17 00:00:00 2001 From: "baizhou.chen" Date: Thu, 16 May 2024 10:13:54 +0800 Subject: [PATCH] A --- LICENSE | 2 +- __init__.py | 11 ++ anyline.py | 54 ++++++++ checkpoints/autodownload | 0 teed_module.py | 268 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 334 insertions(+), 1 deletion(-) create mode 100644 __init__.py create mode 100644 anyline.py create mode 100644 checkpoints/autodownload create mode 100644 teed_module.py diff --git a/LICENSE b/LICENSE index e8a2ee7..b5e395e 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2024 TheMisto +Copyright (c) 2024 TheMisto.ai(深圳混合元组科技有限公司) Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..fb89ccd --- /dev/null +++ b/__init__.py @@ -0,0 +1,11 @@ +from .anyline import AnyLine + +NODE_CLASS_MAPPINGS = { + "AnyLinePreprocessor": AnyLine +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "AnyLinePreprocessor": "TheMisto Anyline" +} +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/anyline.py b/anyline.py new file mode 100644 index 0000000..8e37567 --- /dev/null +++ b/anyline.py @@ -0,0 +1,54 @@ +import torch +import numpy as np +import kornia as kn + +class AnyLine: + + @classmethod + def INPUT_TYPES(s): + + return { + # TODO need to be other thing + "required": { + "E-mail":("STRING",{ + "default": "Enter your resign email（输入你注册的email地址）" + }), + "image": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("anyline_image",) + + FUNCTION = "get_anyline" + CATEGORY = "TheMisto/image/preprocessor" + + def get_anyline(self, email, image): + return (image,) + +NODE_CLASS_MAPPINGS = { + "AnyLine": AnyLine +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "AnyLine": "TheMisto Anyline" +} + + +if __name__ == "__main__": + # loader + checkpoint_path = "" + model = TED().to(device) + model.load_state_dict(torch.load(checkpoint_path, map_location=device)) + model.eval() + + with torch.no_grad(): + images = sample_batched['images'].to(device) + preds = model(images, single_test=resize_input) + + img_height, img_width = img_shape[0].item(), img_shape[1].item() + image_vis = kn.utils.tensor_to_image(torch.sigmoid(tensor_image)) + image_vis = (255.0 * (1.0 - image_vis)).astype(np.uint8) + + image_vis = cv2.resize(image_vis, (img_width, img_height)) \ No newline at end of file diff --git a/checkpoints/autodownload b/checkpoints/autodownload new file mode 100644 index 0000000..e69de29 diff --git a/teed_module.py b/teed_module.py new file mode 100644 index 0000000..c4ac74a --- /dev/null +++ b/teed_module.py @@ -0,0 +1,268 @@ +# This module is strongly reference from https://arxiv.org/abs/2308.06468 +# Thanks for the great work and the magic activation function +import torch +import torch.nn as nn +import torch.nn.functional as F + +@torch.jit.script +def mish_(input): + return input * torch.tanh(F.softplus(input)) + +@torch.jit.script +def Fsmish(input): + return input * torch.tanh(torch.log(1+torch.sigmoid(input))) + +class Mish(nn.Module): + """ + Applies the mish function element-wise: + mish(x) = x * tanh(softplus(x)) = x * tanh(ln(1 + exp(x))) + Shape: + - Input: (N, *) where * means, any number of additional + dimensions + - Output: (N, *), same shape as the input + Reference: https://pytorch.org/docs/stable/generated/torch.nn.Mish.html + """ + def __init__(self): + """ + Init method. + """ + super().__init__() + + def forward(self, input): + """ + Forward pass of the function. + """ + if torch.__version__ >= "1.9": + return F.mish(input) + else: + return mish_(input) + +class Smish(nn.Module): + """ + the same thing as Mish + """ + def __init__(self): + """ + Init method. + """ + super().__init__() + + def forward(self, input): + """ + Forward pass of the function. + """ + return Fsmish(input) + + + +def weight_init(m): + if isinstance(m, (nn.Conv2d,)): + torch.nn.init.xavier_normal_(m.weight, gain=1.0) + + if m.bias is not None: + torch.nn.init.zeros_(m.bias) + + # for fusion layer + if isinstance(m, (nn.ConvTranspose2d,)): + torch.nn.init.xavier_normal_(m.weight, gain=1.0) + if m.bias is not None: + torch.nn.init.zeros_(m.bias) + +class CoFusion(nn.Module): + # from LDC + + def __init__(self, in_ch, out_ch): + super(CoFusion, self).__init__() + self.conv1 = nn.Conv2d(in_ch, 32, kernel_size=3,stride=1, padding=1) # before 64 + self.conv3= nn.Conv2d(32, out_ch, kernel_size=3, + stride=1, padding=1)# before 64 instead of 32 + self.relu = nn.ReLU() + self.norm_layer1 = nn.GroupNorm(4, 32) # before 64 + + def forward(self, x): + # fusecat = torch.cat(x, dim=1) + attn = self.relu(self.norm_layer1(self.conv1(x))) + attn = F.softmax(self.conv3(attn), dim=1) + return ((x * attn).sum(1)).unsqueeze(1) + +class CoFusion2(nn.Module): + def __init__(self, in_ch, out_ch): + super(CoFusion2, self).__init__() + self.conv1 = nn.Conv2d(in_ch, 32, kernel_size=3,stride=1, padding=1) + self.conv3 = nn.Conv2d(32, out_ch, kernel_size=3,stride=1, padding=1) + self.smish = Smish() + +def forward(self, x): + attn = self.conv1(self.smish(x)) + attn = self.conv3(self.smish(attn)) + return ((x * attn).sum(1)).unsqueeze(1) + +class DoubleFusion(nn.Module): + def __init__(self, in_ch, out_ch): + super(DoubleFusion, self).__init__() + self.DWconv1 = nn.Conv2d(in_ch, in_ch*8, kernel_size=3,stride=1, padding=1, groups=in_ch) + self.PSconv1 = nn.PixelShuffle(1) + self.DWconv2 = nn.Conv2d(24, 24*1, kernel_size=3,stride=1, padding=1,groups=24) + self.AF= Smish() + + + def forward(self, x): + attn = self.PSconv1(self.DWconv1(self.AF(x))) + attn2 = self.PSconv1(self.DWconv2(self.AF(attn))) + return Fsmish(((attn2 +attn).sum(1)).unsqueeze(1)) + +class _DenseLayer(nn.Sequential): + def __init__(self, input_features, out_features): + super(_DenseLayer, self).__init__() + + self.add_module('conv1', nn.Conv2d(input_features, out_features, + kernel_size=3, stride=1, padding=2, bias=True)), + self.add_module('smish1', Smish()), + self.add_module('conv2', nn.Conv2d(out_features, out_features, + kernel_size=3, stride=1, bias=True)) + def forward(self, x): + x1, x2 = x + new_features = super(_DenseLayer, self).forward(Fsmish(x1)) # F.relu() + return 0.5 * (new_features + x2), x2 + +class _DenseBlock(nn.Sequential): + def __init__(self, num_layers, input_features, out_features): + super(_DenseBlock, self).__init__() + for i in range(num_layers): + layer = _DenseLayer(input_features, out_features) + self.add_module('denselayer%d' % (i + 1), layer) + input_features = out_features + +class UpConvBlock(nn.Module): + def __init__(self, in_features, up_scale): + super(UpConvBlock, self).__init__() + self.up_factor = 2 + self.constant_features = 16 + + layers = self.make_deconv_layers(in_features, up_scale) + assert layers is not None, layers + self.features = nn.Sequential(*layers) + + def make_deconv_layers(self, in_features, up_scale): + layers = [] + all_pads=[0,0,1,3,7] + for i in range(up_scale): + kernel_size = 2 ** up_scale + pad = all_pads[up_scale] # kernel_size-1 + out_features = self.compute_out_features(i, up_scale) + layers.append(nn.Conv2d(in_features, out_features, 1)) + layers.append(Smish()) + layers.append(nn.ConvTranspose2d( + out_features, out_features, kernel_size, stride=2, padding=pad)) + in_features = out_features + return layers + + def compute_out_features(self, idx, up_scale): + return 1 if idx == up_scale - 1 else self.constant_features + + def forward(self, x): + return self.features(x) + + +class SingleConvBlock(nn.Module): + def __init__(self, in_features, out_features, stride, use_ac=False): + super(SingleConvBlock, self).__init__() + # self.use_bn = use_bs + self.use_ac=use_ac + self.conv = nn.Conv2d(in_features, out_features, 1, stride=stride, + bias=True) + if self.use_ac: + self.smish = Smish() + + def forward(self, x): + x = self.conv(x) + if self.use_ac: + return self.smish(x) + else: + return x + +class DoubleConvBlock(nn.Module): + def __init__(self, in_features, mid_features, + out_features=None, + stride=1, + use_act=True): + super(DoubleConvBlock, self).__init__() + + self.use_act = use_act + if out_features is None: + out_features = mid_features + self.conv1 = nn.Conv2d(in_features, mid_features, + 3, padding=1, stride=stride) + self.conv2 = nn.Conv2d(mid_features, out_features, 3, padding=1) + self.smish= Smish()#nn.ReLU(inplace=True) + + def forward(self, x): + x = self.conv1(x) + x = self.smish(x) + x = self.conv2(x) + if self.use_act: + x = self.smish(x) + return x + +class MistoTEED(nn.Module): + def __init__(self): + super(MistoTEED, self).__init__() + self.block_1 = DoubleConvBlock(3, 16, 16, stride=2,) + self.block_2 = DoubleConvBlock(16, 32, use_act=False) + self.dblock_3 = _DenseBlock(1, 32, 48) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.side_1 = SingleConvBlock(16, 32, 2) + self.pre_dense_3 = SingleConvBlock(32, 48, 1) + self.up_block_1 = UpConvBlock(16, 1) + self.up_block_2 = UpConvBlock(32, 1) + self.up_block_3 = UpConvBlock(48, 2) + self.block_cat = DoubleFusion(3,3) + self.apply(weight_init) + + def slice(self, tensor, slice_shape): + t_shape = tensor.shape + img_h, img_w = slice_shape + if img_w!=t_shape[-1] or img_h!=t_shape[2]: + new_tensor = F.interpolate( + tensor, size=(img_h, img_w), mode='bicubic',align_corners=False) + else: + new_tensor=tensor + return new_tensor + def resize_input(self,tensor): + t_shape = tensor.shape + if t_shape[2] % 8 != 0 or t_shape[3] % 8 != 0: + img_w= ((t_shape[3]// 8) + 1) * 8 + img_h = ((t_shape[2] // 8) + 1) * 8 + new_tensor = F.interpolate( + tensor, size=(img_h, img_w), mode='bicubic', align_corners=False) + else: + new_tensor = tensor + return new_tensor + + def crop_bdcn(data1, h, w, crop_h, crop_w): + _, _, h1, w1 = data1.size() + assert (h <= h1 and w <= w1) + data = data1[:, :, crop_h:crop_h + h, crop_w:crop_w + w] + return data + + + def forward(self, x,is_eval=False): + assert x.ndim == 4, x.shape + block_1 = self.block_1(x) + block_1_side = self.side_1(block_1) + block_2 = self.block_2(block_1) + block_2_down = self.maxpool(block_2) + block_2_add = block_2_down + block_1_side + block_3_pre_dense = self.pre_dense_3(block_2_down) + block_3, _ = self.dblock_3([block_2_add, block_3_pre_dense]) + out_1 = self.up_block_1(block_1) + out_2 = self.up_block_2(block_2) + out_3 = self.up_block_3(block_3) + results = [out_1, out_2, out_3] + block_cat = torch.cat(results, dim=1) + block_cat = self.block_cat(block_cat) + results.append(block_cat) + if is_eval: + pass + else: + return results \ No newline at end of file