def __init__(self):
super(PointNet, self).__init__()
# downsample 1
DownSample = PointNetSetAbstraction
self.ds1 = DownSample(npoint=358,
radius=0.2,
nsample=64,
in_channel=6 + 0,
mlp=[64, 64, 128],
group_all=False)
# downsample 2
self.ds2 = DownSample(npoint=128,
radius=0.4,
nsample=64,
in_channel=128 + 3,
mlp=[128, 128, 256],
group_all=False)
# downsample 3
self.ds3 = DownSample(npoint=None,
radius=None,
nsample=None,
in_channel=256 + 3,
mlp=[256, 512, 1024],
group_all=True)
self.overlap_atte = MutiHeadAttention(n_head=4,
in_channel=256,
qk_channel=64,
v_channel=64,
out_channel=256,
mid_channel=1024)
self.up3 = UpSample(in_channel=1280, mlp=[256, 256])
self.up2 = UpSample(in_channel=384, mlp=[256, 128])
self.up1 = UpSample(in_channel=128 + 6 + 0, mlp=[512, 256, 256])
def __init__(self):
super(PointNetMsg, self).__init__()
DownSample = PointNetSetAbstractionMsg
self.ds1 = DownSample(358, [0.1, 0.2, 0.4], [32, 64, 128],
3 + 3 + 162,
[[32, 32, 64], [64, 64, 128], [64, 96, 128]],
bn=False)
# downsample 2
self.ds2 = DownSample(128, [0.4, 0.8], [64, 128],
128 + 128 + 64,
[[128, 128, 256], [128, 196, 256]],
bn=False)
# downsample 3
self.ds3 = PointNetSetAbstraction(npoint=None,
radius=None,
nsample=None,
in_channel=512 + 3,
mlp=[256, 512, 1024],
group_all=True,
bn=False)
# self.overlap_atte = MutiHeadAttention(n_head=4, in_channel=256, qk_channel=64, v_channel=64, out_channel=256,
# mid_channel=1024)
self.up3 = UpSample(in_channel=1024 + 256 + 256,
mlp=[256, 256],
bn=False)
self.up2 = UpSample(in_channel=256 + 128 + 128 + 64,
mlp=[256, 128],
bn=False)
self.up1 = UpSample(in_channel=128 + 6 + 3 + 162,
mlp=[128, 256],
bn=False)
def __init__(self):
super(SelfAggregation, self).__init__()
# downsample 1
DownSample = PointNetSetAbstraction
self.ds1 = DownSample(npoint=358,
radius=0.2,
nsample=64,
in_channel=162 + 3,
mlp=[1024, 512, 512],
group_all=False,
cat_xyz=True)
# downsample 2
self.ds2 = DownSample(npoint=128,
radius=0.4,
nsample=64,
in_channel=512 + 3,
mlp=[512, 512, 256],
group_all=False,
cat_xyz=True)
# downsample 3
self.ds3 = DownSample(npoint=None,
radius=None,
nsample=None,
in_channel=256 + 3,
mlp=[256, 512, 1024],
group_all=True,
cat_xyz=True)
self.up3 = UpSample(in_channel=1280, mlp=[512, 512])
self.up2 = UpSample(in_channel=512 + 512, mlp=[512, 256])
self.up1 = UpSample(in_channel=256 + 162, mlp=[512, 256, 128])
def __init__(self,
device,
input_size=None,
num_classes=20,
trainable=False,
conf_thresh=0.001,
nms_thresh=0.50,
anchor_size=None,
hr=False):
super(YOLOv4, self).__init__()
self.device = device
self.input_size = input_size
self.num_classes = num_classes
self.trainable = trainable
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.stride = [8, 16, 32]
self.anchor_size = torch.tensor(anchor_size).view(
3,
len(anchor_size) // 3, 2)
self.num_anchors = self.anchor_size.size(1)
self.grid_cell, self.stride_tensor, self.all_anchors_wh = self.create_grid(
input_size)
# backbone cspdarknet-53
self.backbone = cspdarknet53(pretrained=trainable, hr=hr)
# self.backbone = darknet53(pretrained=trainable, hr=hr)
# neck
self.spp = nn.Sequential(
Conv(1024, 512, k=1), SPP(),
BottleneckCSP(512 * 4, 1024, n=3, shortcut=False))
# self.neck = DilateEncoder(1024, 1024)
# head
self.head_conv_0 = Conv(1024, 512, k=1) # 10
self.head_upsample_0 = UpSample(scale_factor=2)
self.head_csp_0 = BottleneckCSP(512 + 512, 512, n=3, shortcut=False)
# P3/8-small
self.head_conv_1 = Conv(512, 256, k=1) # 14
self.head_upsample_1 = UpSample(scale_factor=2)
self.head_csp_1 = BottleneckCSP(256 + 256, 256, n=3, shortcut=False)
# P4/16-medium
self.head_conv_2 = Conv(256, 256, k=3, p=1, s=2)
self.head_csp_2 = BottleneckCSP(256 + 256, 512, n=3, shortcut=False)
# P8/32-large
self.head_conv_3 = Conv(512, 512, k=3, p=1, s=2)
self.head_csp_3 = BottleneckCSP(512 + 512, 1024, n=3, shortcut=False)
# det conv
self.head_det_1 = nn.Conv2d(
256, self.num_anchors * (1 + self.num_classes + 4), 1)
self.head_det_2 = nn.Conv2d(
512, self.num_anchors * (1 + self.num_classes + 4), 1)
self.head_det_3 = nn.Conv2d(
1024, self.num_anchors * (1 + self.num_classes + 4), 1)
def __init__(self, in_channels, n_classes, bilinear=True):
super(Unet, self).__init__()
self.in_channels = in_channels
self.n_classes = n_classes
self.bilinear = bilinear
self.conv1 = Conv_BN_Relu(self.in_channels, 64)
self.down1 = DownSample(in_channels=64, out_channels=128)
self.down2 = DownSample(in_channels=128, out_channels=256)
self.down3 = DownSample(in_channels=256, out_channels=512)
self.down4 = DownSample(in_channels=512, out_channels=512)
self.up1 = UpSample(in_channels=1024, out_channels=256)
self.up2 = UpSample(in_channels=512, out_channels=128)
self.up3 = UpSample(in_channels=256, out_channels=64)
self.up4 = UpSample(in_channels=128, out_channels=64)
self.output = Output_Layer_Conv(64, self.n_classes)