A

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

__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"]

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))

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