ObjectDetectio-watermarking/model/layer.py

import torch


class cbs(torch.nn.Module):
    def __init__(self, in_, out_, kernel_size, stride):
        super().__init__()
        self.conv = torch.nn.Conv2d(in_, out_, kernel_size=kernel_size, stride=stride, padding=(kernel_size - 1) // 2,
                                    bias=False)
        self.bn = torch.nn.BatchNorm2d(out_, eps=0.001, momentum=0.03)
        self.silu = torch.nn.SiLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.silu(x)
        return x


class concat(torch.nn.Module):
    def __init__(self, dim=1):
        super().__init__()
        self.concat = torch.concat
        self.dim = dim

    def forward(self, x):
        x = self.concat(x, dim=self.dim)
        return x


class residual(torch.nn.Module):  # in_->in_，len->len
    def __init__(self, in_, config=None):
        super().__init__()
        if not config:  # 正常版本
            self.cbs0 = cbs(in_, in_, kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, in_, kernel_size=3, stride=1)
        else:  # 剪枝版本。len(config) = 2
            self.cbs0 = cbs(in_, config[0], kernel_size=1, stride=1)
            self.cbs1 = cbs(config[0], config[1], kernel_size=3, stride=1)

    def forward(self, x):
        x0 = self.cbs0(x)
        x0 = self.cbs1(x0)
        return x + x0


class c3(torch.nn.Module):  # in_->out_，len->len
    def __init__(self, in_, out_, n, config=None):
        super().__init__()
        if not config:  # 正常版本
            self.cbs0 = cbs(in_, in_ // 2, kernel_size=1, stride=1)
            self.sequential1 = torch.nn.Sequential(*(residual(in_ // 2) for _ in range(n)))
            self.cbs2 = cbs(in_, in_ // 2, kernel_size=1, stride=1)
            self.concat3 = concat(dim=1)
            self.cbs4 = cbs(in_, out_, kernel_size=1, stride=1)
        else:  # 剪枝版本。len(config) = 3 + 2 * n
            self.cbs0 = cbs(in_, config[0], kernel_size=1, stride=1)
            self.sequential1 = torch.nn.Sequential(
                *(residual(config[0 + 2 * _] if _ == 0 else config[1 + 2 * _] + config[2 + 2 * _],
                           config[1 + 2 * _:3 + 2 * _]) for _ in range(n)))
            self.cbs2 = cbs(config[0], config[1 + 2 * n], kernel_size=1, stride=1)
            self.concat3 = concat(dim=1)
            self.cbs4 = cbs(config[0] + config[2 * n - 1] + config[2 * n] + config[1 + 2 * n], config[2 + 2 * n],
                            kernel_size=1, stride=1)

    def forward(self, x):
        x0 = self.cbs0(x)
        x1 = self.sequential1(x0)
        x1 = x0 + x1
        x2 = self.cbs2(x)
        x = self.concat3([x1, x2])
        x = self.cbs4(x)
        return x


class elan(torch.nn.Module):  # in_->out_，len->len
    def __init__(self, in_, out_, n, config=None):
        super().__init__()
        if not config:  # 正常版本
            self.cbs0 = cbs(in_, out_ // 4, kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, out_ // 4, kernel_size=1, stride=1)
            self.sequential2 = torch.nn.Sequential(
                *(cbs(out_ // 4, out_ // 4, kernel_size=3, stride=1) for _ in range(n)))
            self.sequential3 = torch.nn.Sequential(
                *(cbs(out_ // 4, out_ // 4, kernel_size=3, stride=1) for _ in range(n)))
            self.concat4 = concat()
            self.cbs5 = cbs(out_, out_, kernel_size=1, stride=1)
        else:  # 剪枝版本。len(config) = 3 + 2 * n
            self.cbs0 = cbs(in_, config[0], kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, config[1], kernel_size=1, stride=1)
            self.sequential2 = torch.nn.Sequential(
                *(cbs(config[1 + _], config[2 + _], kernel_size=3, stride=1) for _ in range(n)))
            self.sequential3 = torch.nn.Sequential(
                *(cbs(config[1 + n + _], config[2 + n + _], kernel_size=3, stride=1) for _ in range(n)))
            self.concat4 = concat()
            self.cbs5 = cbs(config[0] + config[1] + config[1 + n] + config[1 + 2 * n], config[2 + 2 * n],
                            kernel_size=1, stride=1)

    def forward(self, x):
        x0 = self.cbs0(x)
        x1 = self.cbs1(x)
        x2 = self.sequential2(x1)
        x3 = self.sequential3(x2)
        x = self.concat4([x0, x1, x2, x3])
        x = self.cbs5(x)
        return x


class elan_h(torch.nn.Module):  # in_->out_，len->len
    def __init__(self, in_, out_, config=None):
        super().__init__()
        if not config:  # 正常版本
            self.cbs0 = cbs(in_, in_ // 2, kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, in_ // 2, kernel_size=1, stride=1)
            self.cbs2 = cbs(in_ // 2, in_ // 4, kernel_size=3, stride=1)
            self.cbs3 = cbs(in_ // 4, in_ // 4, kernel_size=3, stride=1)
            self.cbs4 = cbs(in_ // 4, in_ // 4, kernel_size=3, stride=1)
            self.cbs5 = cbs(in_ // 4, in_ // 4, kernel_size=3, stride=1)
            self.concat6 = concat()
            self.cbs7 = cbs(2 * in_, out_, kernel_size=1, stride=1)
        else:  # 剪枝版本。len(config) = 7
            self.cbs0 = cbs(in_, config[0], kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, config[1], kernel_size=1, stride=1)
            self.cbs2 = cbs(config[1], config[2], kernel_size=3, stride=1)
            self.cbs3 = cbs(config[2], config[3], kernel_size=3, stride=1)
            self.cbs4 = cbs(config[3], config[4], kernel_size=3, stride=1)
            self.cbs5 = cbs(config[4], config[5], kernel_size=3, stride=1)
            self.concat6 = concat()
            self.cbs7 = cbs(config[0] + config[1] + config[2] + config[3] + config[4] + config[5], config[6],
                            kernel_size=1, stride=1)

    def forward(self, x):
        x0 = self.cbs0(x)
        x1 = self.cbs1(x)
        x2 = self.cbs2(x1)
        x3 = self.cbs3(x2)
        x4 = self.cbs4(x3)
        x5 = self.cbs5(x4)
        x = self.concat6([x0, x1, x2, x3, x4, x5])
        x = self.cbs7(x)
        return x


class mp(torch.nn.Module):  # in_->out_，len->len//2
    def __init__(self, in_, out_, config=None):
        super().__init__()
        if not config:  # 正常版本
            self.maxpool0 = torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1)
            self.cbs1 = cbs(in_, out_ // 2, 1, 1)
            self.cbs2 = cbs(in_, out_ // 2, 1, 1)
            self.cbs3 = cbs(out_ // 2, out_ // 2, 3, 2)
            self.concat4 = concat(dim=1)
        else:  # 剪枝版本。len(config) = 3
            self.maxpool0 = torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1)
            self.cbs1 = cbs(in_, config[0], 1, 1)
            self.cbs2 = cbs(in_, config[1], 1, 1)
            self.cbs3 = cbs(config[1], config[2], 3, 2)
            self.concat4 = concat(dim=1)

    def forward(self, x):
        x0 = self.maxpool0(x)
        x0 = self.cbs1(x0)
        x1 = self.cbs2(x)
        x1 = self.cbs3(x1)
        x = self.concat4([x0, x1])
        return x


class sppf(torch.nn.Module):  # in_->out_，len->len
    def __init__(self, in_, out_, config=None):
        super().__init__()
        if not config:  # 正常版本
            self.cbs0 = cbs(in_, in_ // 2, kernel_size=1, stride=1)
            self.MaxPool2d1 = torch.nn.MaxPool2d(kernel_size=5, stride=1, padding=2, dilation=1)
            self.MaxPool2d2 = torch.nn.MaxPool2d(kernel_size=9, stride=1, padding=4, dilation=1)
            self.MaxPool2d3 = torch.nn.MaxPool2d(kernel_size=13, stride=1, padding=6, dilation=1)
            self.concat4 = concat(dim=1)
            self.cbs5 = cbs(2 * in_, out_, kernel_size=1, stride=1)
        else:  # 剪枝版本。len(config) = 2
            self.cbs0 = cbs(in_, config[0], kernel_size=1, stride=1)
            self.MaxPool2d1 = torch.nn.MaxPool2d(kernel_size=5, stride=1, padding=2, dilation=1)
            self.MaxPool2d2 = torch.nn.MaxPool2d(kernel_size=9, stride=1, padding=4, dilation=1)
            self.MaxPool2d3 = torch.nn.MaxPool2d(kernel_size=13, stride=1, padding=6, dilation=1)
            self.concat4 = concat(dim=1)
            self.cbs5 = cbs(4 * config[0], config[1], kernel_size=1, stride=1)

    def forward(self, x):
        x = self.cbs0(x)
        x0 = self.MaxPool2d1(x)
        x1 = self.MaxPool2d2(x0)
        x2 = self.MaxPool2d3(x1)
        x = self.concat4([x, x0, x1, x2])
        x = self.cbs5(x)
        return x


class sppcspc(torch.nn.Module):  # in_->out_，len->len
    def __init__(self, in_, out_, config=None):
        super().__init__()
        if not config:  # 正常版本
            self.cbs0 = cbs(in_, in_ // 2, kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, in_ // 2, kernel_size=1, stride=1)
            self.cbs2 = cbs(in_ // 2, in_ // 2, kernel_size=3, stride=1)
            self.cbs3 = cbs(in_ // 2, in_ // 2, kernel_size=1, stride=1)
            self.MaxPool2d4 = torch.nn.MaxPool2d(kernel_size=5, stride=1, padding=2, dilation=1)
            self.MaxPool2d5 = torch.nn.MaxPool2d(kernel_size=9, stride=1, padding=4, dilation=1)
            self.MaxPool2d6 = torch.nn.MaxPool2d(kernel_size=13, stride=1, padding=6, dilation=1)
            self.concat7 = concat(dim=1)
            self.cbs8 = cbs(2 * in_, in_ // 2, kernel_size=1, stride=1)
            self.cbs9 = cbs(in_ // 2, in_ // 2, kernel_size=3, stride=1)
            self.concat10 = concat(dim=1)
            self.cbs11 = cbs(in_, out_, kernel_size=1, stride=1)
        else:  # 剪枝版本。len(config) = 7
            self.cbs0 = cbs(in_, config[0], kernel_size=1, stride=1)
            self.cbs1 = cbs(in_, config[1], kernel_size=1, stride=1)
            self.cbs2 = cbs(config[1], config[2], kernel_size=3, stride=1)
            self.cbs3 = cbs(config[2], config[3], kernel_size=1, stride=1)
            self.MaxPool2d4 = torch.nn.MaxPool2d(kernel_size=5, stride=1, padding=2, dilation=1)
            self.MaxPool2d5 = torch.nn.MaxPool2d(kernel_size=9, stride=1, padding=4, dilation=1)
            self.MaxPool2d6 = torch.nn.MaxPool2d(kernel_size=13, stride=1, padding=6, dilation=1)
            self.concat7 = concat(dim=1)
            self.cbs8 = cbs(4 * config[3], config[4], kernel_size=1, stride=1)
            self.cbs9 = cbs(config[4], config[5], kernel_size=3, stride=1)
            self.concat10 = concat(dim=1)
            self.cbs11 = cbs(config[0] + config[5], config[6], kernel_size=1, stride=1)

    def forward(self, x):
        x0 = self.cbs0(x)
        x1 = self.cbs1(x)
        x1 = self.cbs2(x1)
        x1 = self.cbs3(x1)
        x4 = self.MaxPool2d4(x1)
        x5 = self.MaxPool2d5(x1)
        x6 = self.MaxPool2d6(x1)
        x = self.concat7([x1, x4, x5, x6])
        x = self.cbs8(x)
        x = self.cbs9(x)
        x = self.concat10([x, x0])
        x = self.cbs11(x)
        return x


class head(torch.nn.Module):  # in_->(batch, 3, output_size, output_size, 5+output_class))，len->len
    def __init__(self, in_, output_size, output_class):
        super().__init__()
        self.output_size = output_size
        self.output_class = output_class
        self.output = torch.nn.Conv2d(in_, 3 * (5 + output_class), kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        x = self.output(x).reshape(-1, 3, self.output_size, self.output_size, 5 + self.output_class)  # 变形
        return x


# 参考yolox
class split_head(torch.nn.Module):  # in_->(batch, 3, output_size, output_size, 5+output_class))，len->len
    def __init__(self, in_, output_size, output_class, config=None):
        super().__init__()
        self.output_size = output_size
        self.output_class = output_class
        if not config:  # 正常版本
            out_ = 3 * (5 + self.output_class)
            self.cbs0 = cbs(in_, out_, kernel_size=1, stride=1)
            self.cbs1 = cbs(out_, out_, kernel_size=3, stride=1)
            self.cbs2 = cbs(out_, out_, kernel_size=3, stride=1)
            self.cbs3 = cbs(out_, out_, kernel_size=3, stride=1)
            self.cbs4 = cbs(out_, out_, kernel_size=3, stride=1)
            self.Conv2d5 = torch.nn.Conv2d(out_, 12, kernel_size=1, stride=1, padding=0)
            self.Conv2d6 = torch.nn.Conv2d(out_, 3, kernel_size=1, stride=1, padding=0)
            self.Conv2d7 = torch.nn.Conv2d(out_, 3 * self.output_class, kernel_size=1, stride=1, padding=0)
            self.concat8 = concat(4)
        else:  # 剪枝版本。len(config) = 8
            self.cbs0 = cbs(in_, config[0], kernel_size=1, stride=1)
            self.cbs1 = cbs(config[0], config[1], kernel_size=1, stride=1)
            self.cbs2 = cbs(config[1], config[2], kernel_size=1, stride=1)
            self.cbs3 = cbs(config[0], config[3], kernel_size=1, stride=1)
            self.cbs4 = cbs(config[3], config[4], kernel_size=1, stride=1)
            self.Conv2d5 = torch.nn.Conv2d(config[5], 12, kernel_size=1, stride=1, padding=0)
            self.Conv2d6 = torch.nn.Conv2d(config[6], 3, kernel_size=1, stride=1, padding=0)
            self.Conv2d7 = torch.nn.Conv2d(config[7], 3 * self.output_class, kernel_size=1, stride=1, padding=0)
            self.concat8 = concat(4)

    def forward(self, x):
        x = self.cbs0(x)
        x0 = self.cbs1(x)
        x0 = self.cbs2(x0)
        x1 = self.cbs3(x)
        x1 = self.cbs4(x1)
        x2 = self.Conv2d5(x0).reshape(-1, 3, self.output_size, self.output_size, 4)  # 变形
        x3 = self.Conv2d6(x0).reshape(-1, 3, self.output_size, self.output_size, 1)  # 变形
        x4 = self.Conv2d7(x1).reshape(-1, 3, self.output_size, self.output_size, self.output_class)  # 变形
        x = self.concat8([x2, x3, x4])
        return x


class image_deal(torch.nn.Module):  # 归一化
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = x / 255
        x = x.permute(0, 3, 1, 2)
        return x


class decode(torch.nn.Module):  # (Cx,Cy,w,h,confidence...)原始输出->(Cx,Cy,w,h,confidence...)真实坐标
    def __init__(self, input_size):
        super().__init__()
        self.stride = (8, 16, 32)
        output_size = [int(input_size // i) for i in self.stride]
        self.anchor = (((12, 16), (19, 36), (40, 28)), ((36, 75), (76, 55), (72, 146)),
                       ((142, 110), (192, 243), (459, 401)))
        self.grid = [0, 0, 0]
        for i in range(3):
            self.grid[i] = torch.arange(output_size[i])
        self.frame_sigmoid = torch.nn.Sigmoid()

    def forward(self, output):
        device = output[0].device
        # 遍历每一个大层
        for i in range(3):
            self.grid[i] = self.grid[i].to(device)  # 放到对应的设备上
            # 中心坐标[0-1]->[-0.5-1.5]->[-0.5*stride-80/40/20.5*stride]
            output[i] = self.frame_sigmoid(output[i])  # 边框输出归一化
            output[i][..., 0] = (2 * output[i][..., 0] - 0.5 + self.grid[i].unsqueeze(1)) * self.stride[i]
            output[i][..., 1] = (2 * output[i][..., 1] - 0.5 + self.grid[i]) * self.stride[i]
            # 遍历每一个大层中的小层
            for j in range(3):
                output[i][:, j, ..., 2] = 4 * output[i][:, j, ..., 2] ** 2 * self.anchor[i][j][0]  # [0-1]->[0-4*anchor]
                output[i][:, j, ..., 3] = 4 * output[i][:, j, ..., 3] ** 2 * self.anchor[i][j][1]  # [0-1]->[0-4*anchor]
        return output


class deploy(torch.nn.Module):
    def __init__(self, model, input_size):
        super().__init__()
        self.image_deal = image_deal()
        self.model = model
        self.decode = decode(input_size)

    def forward(self, x):
        x = self.image_deal(x)
        x = self.model(x)
        x = self.decode(x)
        return x