def __init__(self, device, input_size=None, num_classes=20, trainable=False, conf_thresh=0.01, nms_thresh=0.5, hr=False):
super(myYOLO, self).__init__()
self.device = device
self.num_classes = num_classes
self.trainable = trainable
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.stride = 32
self.grid_cell = self.create_grid(input_size)
self.input_size = input_size
self.scale = np.array([[[input_size[1], input_size[0], input_size[1], input_size[0]]]])
self.scale_torch = torch.tensor(self.scale.copy(), device=device).float()
# we use resnet18 as backbone
self.backbone = resnet18(pretrained=True)
# neck
self.SPP = nn.Sequential(
Conv(512, 256, k=1),
SPP(),
BottleneckCSP(256*4, 512, n=1, shortcut=False)
)
self.SAM = SAM(512)
self.conv_set = BottleneckCSP(512, 512, n=3, shortcut=False)
self.pred = nn.Conv2d(512, 1 + self.num_classes + 4, 1)
# pred = nn.Conv2d(512, 1 + 20 + 4, 1)
"""
def __init__(self,
device,
input_size=None,
num_classes=20,
trainable=False,
conf_thresh=0.01,
nms_thresh=0.5,
hr=False):
super(myYOLO, self).__init__()
self.device = device
self.num_classes = num_classes
self.trainable = trainable
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.stride = 32
self.grid_cell = self.create_grid(input_size)
self.input_size = input_size
self.scale = np.array(
[[[input_size[1], input_size[0], input_size[1], input_size[0]]]])
self.scale_torch = torch.tensor(self.scale.copy(),
device=device).float()
# we use resnet18 as backbone
self.backbone = resnet18(pretrained=True)
self.SPP = SPP(512, 512)
self.SAM = SAM(512)
self.conv_set = nn.Sequential(
Conv2d(512, 256, 1, leakyReLU=True),
Conv2d(256, 512, 3, padding=1, leakyReLU=True),
Conv2d(512, 256, 1, leakyReLU=True),
Conv2d(256, 512, 3, padding=1, leakyReLU=True),
)
self.pred = nn.Conv2d(512, 1 + self.num_classes + 4, 1)
def __init__(self, device, input_size=None, num_classes=2, trainable=False, conf_thresh=0.05, nms_thresh=0.2, anchor_size=None, hr=False):
super(YOLOv2slim, 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.anchor_size = torch.tensor(anchor_size)
self.anchor_number = len(anchor_size)
self.stride = 32
self.grid_cell, self.all_anchor_wh = self.create_grid(input_size)
self.scale = np.array([[[input_size[1], input_size[0], input_size[1], input_size[0]]]])
self.scale_torch = torch.tensor(self.scale.copy(), device=device).float()
# backbone darknet-tiny
self.backbone = darknet_tiny(pretrained=trainable, hr=hr)
# detection head
self.convsets_1 = nn.Sequential(
SPP(),
Conv2d(512*4, 256, 1, leakyReLU=True),
Conv2d(256, 512, 3, 1, leakyReLU=True)
)
self.route_layer = Conv2d(256, 32, 1, leakyReLU=True)
self.reorg = reorg_layer(stride=2)
self.convsets_2 = Conv2d(640, 512, 3, 1, leakyReLU=True)
# prediction layer
self.pred_ = nn.Conv2d(512, self.anchor_number*(1 + 4 + self.num_classes), 1)
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,
device,
input_size=None,
num_classes=2,
trainable=False,
conf_thresh=0.05,
nms_thresh=0.20,
anchor_size=None,
hr=False):
super(YOLOv3tiny, 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 = [16, 32]
self.anchor_size = torch.tensor(anchor_size).view(
2,
len(anchor_size) // 2, 2)
self.anchor_number = self.anchor_size.size(1)
self.grid_cell, self.stride_tensor, self.all_anchors_wh = self.create_grid(
input_size)
self.scale = np.array(
[[[input_size[1], input_size[0], input_size[1], input_size[0]]]])
self.scale_torch = torch.tensor(self.scale.copy(),
device=device).float()
# backbone darknet-light
self.backbone = darknet_light(pretrained=trainable, hr=hr)
# s = 32
self.conv_set_2 = nn.Sequential(
SPP(), Conv2d(1024 * 4, 256, 1, leakyReLU=True))
self.conv_1x1_2 = Conv2d(256, 128, 1, leakyReLU=True)
self.extra_conv_2 = Conv2d(256, 512, 3, padding=1, leakyReLU=True)
self.pred_2 = nn.Conv2d(
512, self.anchor_number * (1 + 4 + self.num_classes), 1)
# s = 16
self.conv_set_1 = Conv2d(384, 256, 3, padding=1, leakyReLU=True)
self.pred_1 = nn.Conv2d(
256, self.anchor_number * (1 + 4 + self.num_classes), 1)
def __init__(self,
device,
input_size=None,
trainable=False,
num_classes=None,
conf_thresh=0.3,
nms_thresh=0.45,
topk=100,
use_nms=False,
hr=False):
super(CenterNet, self).__init__()
self.device = device
self.input_size = input_size
self.trainable = trainable
self.num_classes = num_classes
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.use_nms = use_nms
self.stride = 4
self.topk = topk
self.grid_cell = self.create_grid(input_size)
self.scale = np.array(
[[[input_size, input_size, input_size, input_size]]])
self.scale_torch = torch.tensor(self.scale.copy(),
device=self.device).float()
self.backbone = resnet18(pretrained=trainable)
self.smooth5 = nn.Sequential(SPP(), Conv2d(512 * 4, 256, ksize=1),
Conv2d(256, 512, ksize=3, padding=1))
self.deconv5 = DeConv2d(512, 256, ksize=4, stride=2) # 32 -> 16
self.deconv4 = DeConv2d(256, 256, ksize=4, stride=2) # 16 -> 8
self.deconv3 = DeConv2d(256, 256, ksize=4, stride=2) # 8 -> 4
self.cls_pred = nn.Sequential(
Conv2d(256, 64, ksize=3, padding=1, leakyReLU=True),
nn.Conv2d(64, self.num_classes, kernel_size=1))
self.txty_pred = nn.Sequential(
Conv2d(256, 64, ksize=3, padding=1, leakyReLU=True),
nn.Conv2d(64, 2, kernel_size=1))
self.twth_pred = nn.Sequential(
Conv2d(256, 64, ksize=3, padding=1, leakyReLU=True),
nn.Conv2d(64, 2, kernel_size=1))
def __init__(self, device, input_size=None, trainable=False, num_classes=None, conf_thresh=0.05, nms_thresh=0.45, topk=100, use_nms=False, hr=False):
super(CenterNet, self).__init__()
self.device = device
self.input_size = input_size
self.trainable = trainable
self.num_classes = num_classes
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.stride = 4
self.topk = topk
self.use_nms = use_nms
self.grid_cell = self.create_grid(input_size)
self.scale = np.array([[[input_size[0], input_size[1], input_size[0], input_size[1]]]])
self.scale_torch = torch.tensor(self.scale.copy(), device=self.device).float()
# backbone
self.backbone = resnet18(pretrained=trainable)
# neck
self.spp = nn.Sequential(
Conv(512, 256, k=1),
SPP(),
BottleneckCSP(256*4, 512, n=1, shortcut=False)
)
# head
self.deconv5 = DeConv(512, 256, ksize=4, stride=2) # 32 -> 16
self.deconv4 = DeConv(256, 256, ksize=4, stride=2) # 16 -> 8
self.deconv3 = DeConv(256, 256, ksize=4, stride=2) # 8 -> 4
self.cls_pred = nn.Sequential(
Conv(256, 64, k=3, p=1),
nn.Conv2d(64, self.num_classes, kernel_size=1)
)
self.txty_pred = nn.Sequential(
Conv(256, 64, k=3, p=1),
nn.Conv2d(64, 2, kernel_size=1)
)
self.twth_pred = nn.Sequential(
Conv(256, 64, k=3, p=1),
nn.Conv2d(64, 2, kernel_size=1)
)
def __init__(self, device, input_size=None, trainable=False, conf_thresh=0.01, nms_thresh=0.3, anchor_size=None, hr=False):
super(TinyYOLAF, self).__init__()
self.device = device
self.input_size = input_size
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.anchor_number = self.anchor_size.size(1)
self.grid_cell, self.stride_tensor, self.all_anchors_wh = self.create_grid(input_size)
self.scale = np.array([[[input_size, input_size, input_size, input_size]]])
self.scale_torch = torch.tensor(self.scale.copy(), device=device).float()
# backbone darknet-tiny
self.backbone = darknet_tiny(pretrained=trainable, hr=hr)
# s = 32
self.conv_set_3 = nn.Sequential(
SPP(),
Conv2d(512*4, 256, 1, leakyReLU=True),
Conv2d(256, 512, 3, padding=1, leakyReLU=True)
)
self.conv_1x1_3 = Conv2d(512, 128, 1, leakyReLU=True)
self.pred_3 = nn.Conv2d(512, self.anchor_number*(2 + 4), 1)
# s = 16
self.conv_set_2 = nn.Sequential(
Conv2d(384, 128, 1, leakyReLU=True),
Conv2d(128, 256, 3, padding=1, leakyReLU=True)
)
self.conv_1x1_2 = Conv2d(256, 64, 1, leakyReLU=True)
self.pred_2 = nn.Conv2d(256, self.anchor_number*(2 + 4), 1)
# s = 8
self.conv_set_1 = nn.Sequential(
Conv2d(192, 64, 1, leakyReLU=True),
Conv2d(64, 128, 3, padding=1, leakyReLU=True)
)
self.pred_1 = nn.Conv2d(128, self.anchor_number*(2 + 4), 1)
def __init__(self,
device,
input_size=None,
trainable=False,
num_classes=None,
backbone='r18',
conf_thresh=0.05,
nms_thresh=0.45,
topk=100,
use_nms=False,
hr=False):
super(CenterNet, self).__init__()
self.device = device
self.input_size = input_size
self.trainable = trainable
self.num_classes = num_classes
self.bk = backbone
self.conf_thresh = conf_thresh
self.nms_thresh = nms_thresh
self.stride = 4
self.topk = topk
self.use_nms = use_nms
self.grid_cell = self.create_grid(input_size)
self.scale = np.array(
[[[input_size[0], input_size[1], input_size[0], input_size[1]]]])
self.scale_torch = torch.tensor(self.scale.copy(),
device=self.device).float()
# backbone
if self.bk == 'r18':
print("Use backbone : resnet-18")
self.backbone = resnet18(pretrained=trainable)
c = 512
elif self.bk == 'r34':
print("Use backbone : resnet-34")
self.backbone = resnet34(pretrained=trainable)
c = 512
elif self.bk == 'r50':
print("Use backbone : resnet-50")
self.backbone = resnet50(pretrained=trainable)
c = 2048
elif self.bk == 'r101':
print("Use backbone : resnet-101")
self.backbone = resnet101(pretrained=trainable)
c = 2048
else:
print("Only support r18, r34, r50, r101 as backbone !!")
exit()
# neck
self.spp = nn.Sequential(
Conv(c, 256, k=1), SPP(),
BottleneckCSP(256 * 4, c, n=1, shortcut=False))
# head
self.deconv5 = DeConv(c, 256, ksize=4, stride=2) # 32 -> 16
self.deconv4 = DeConv(256, 256, ksize=4, stride=2) # 16 -> 8
self.deconv3 = DeConv(256, 256, ksize=4, stride=2) # 8 -> 4
self.cls_pred = nn.Sequential(
Conv(256, 64, k=3, p=1),
nn.Conv2d(64, self.num_classes, kernel_size=1))
self.txty_pred = nn.Sequential(Conv(256, 64, k=3, p=1),
nn.Conv2d(64, 2, kernel_size=1))
self.twth_pred = nn.Sequential(Conv(256, 64, k=3, p=1),
nn.Conv2d(64, 2, kernel_size=1))
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(YOLOv3Spp, 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 darknet-53 (optional: darknet-19)
self.backbone = darknet53(pretrained=trainable, hr=hr)
# s = 32
self.conv_set_3 = nn.Sequential(SPP(), Conv(1024 * 4, 512, k=1),
Conv(512, 1024, k=3, p=1),
Conv(1024, 512, k=1),
Conv(512, 1024, k=3, p=1),
Conv(1024, 512, k=1))
self.conv_1x1_3 = Conv(512, 256, k=1)
self.extra_conv_3 = Conv(512, 1024, k=3, p=1)
self.pred_3 = nn.Conv2d(1024,
self.num_anchors * (1 + 4 + self.num_classes),
kernel_size=1)
# s = 16
self.conv_set_2 = nn.Sequential(Conv(768, 256, k=1),
Conv(256, 512, k=3, p=1),
Conv(512, 256, k=1),
Conv(256, 512, k=3, p=1),
Conv(512, 256, k=1))
self.conv_1x1_2 = Conv(256, 128, k=1)
self.extra_conv_2 = Conv(256, 512, k=3, p=1)
self.pred_2 = nn.Conv2d(512,
self.num_anchors * (1 + 4 + self.num_classes),
kernel_size=1)
# s = 8
self.conv_set_1 = nn.Sequential(Conv(384, 128, k=1),
Conv(128, 256, k=3, p=1),
Conv(256, 128, k=1),
Conv(128, 256, k=3, p=1),
Conv(256, 128, k=1))
self.extra_conv_1 = Conv(128, 256, k=3, p=1)
self.pred_1 = nn.Conv2d(256,
self.num_anchors * (1 + 4 + self.num_classes),
kernel_size=1)