def __init__(self,
with_avg_pool=False,
with_reg=True,
with_cls=True,
roi_feat_size=7,
in_channels=256,
num_classes=80,
bbox_coder=dict(target_means=[0., 0., 0., 0.],
target_stds=[0.1, 0.1, 0.2, 0.2]),
reg_class_agnostic=False,
reg_decoded_bbox=False,
loss_cls=dict(type='CrossEntropyLoss',
use_sigmoid=False,
loss_weight=1.0),
loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
loss_weight=1.0)):
super(BBoxHead, self).__init__()
assert with_cls or with_reg
self.with_avg_pool = with_avg_pool
self.with_cls = with_cls
self.with_reg = with_reg
self.roi_feat_size = _pair(roi_feat_size)
self.roi_feat_area = self.roi_feat_size[0] * self.roi_feat_size[1]
self.in_channels = in_channels
self.num_classes = num_classes
self.reg_class_agnostic = reg_class_agnostic
self.reg_decoded_bbox = reg_decoded_bbox
if bbox_coder is not None:
self.bbox_coder = DeltaXYWHBBoxCoder(**bbox_coder)
else:
self.bbox_coder = None
self.loss_cls = build_module_from_registers(loss_cls,
module_name='loss')
self.loss_bbox = build_module_from_registers(loss_bbox,
module_name='loss')
in_channels = self.in_channels
if self.with_avg_pool:
self.avg_pool = nn.AvgPool2d(self.roi_feat_size)
else:
in_channels *= self.roi_feat_area
if self.with_cls:
# adding background class
self.fc_cls = nn.Linear(in_channels, num_classes + 1)
if self.with_reg:
out_dim_reg = 4 if reg_class_agnostic else 4 * num_classes
self.fc_reg = nn.Linear(in_channels, out_dim_reg)
self.debug_imgs = None
def __init__(
self,
direction='AtoB',
pretrained=None,
pretrained_modules=None,
generator_a=dict(),
generator_b=dict(),
discriminator_a=dict(),
discriminator_b=dict(),
loss_gan=dict(),
loss_cycle_A=dict(),
loss_cycle_B=dict(),
loss_idt=dict(),
train_cfg=None, # TODO: pool_size in train_cfg
test_cfg=None,
**kwargs):
super(CycleGAN, self).__init__(pretrained=pretrained,
pretrained_modules=pretrained_modules)
self.direction = direction
self.netG_A = build_module_from_registers(generator_a,
module_name='generator')
self.netG_B = build_module_from_registers(generator_b,
module_name='generator')
self.mode = 'train'
if self.mode == 'train':
self.netD_A = build_module_from_registers(
discriminator_a, module_name='discriminator')
self.netD_B = build_module_from_registers(
discriminator_b, module_name='discriminator')
if loss_idt is not None and loss_idt.get(
'loss_weight') is not None:
input_nc = generator_a['in_channels']
output_nc = generator_a['out_channels']
assert input_nc == output_nc
pool_size = train_cfg['pool_size']
self.fake_A_pool = ImagePool(pool_size)
self.fake_B_pool = ImagePool(pool_size)
self.loss_gan = build_module_from_registers(loss_gan,
module_name='loss')
self.loss_cycle_A = build_module_from_registers(loss_cycle_A,
module_name='loss')
self.loss_cycle_B = build_module_from_registers(loss_cycle_B,
module_name='loss')
self.loss_idt = build_module_from_registers(loss_idt,
module_name='loss')
self.init_weights()
def __init__(self,
num_classes,
in_channels,
loss_cls=dict(type='CrossEntropyLoss'),
spatial_type='avg',
dropout_ratio=0.8,
init_std=0.01):
super().__init__()
self.spatial_type = spatial_type
self.dropout_ratio = dropout_ratio
self.init_std = init_std
if self.dropout_ratio != 0:
self.dropout = nn.Dropout(p=self.dropout_ratio)
else:
self.dropout = None
self.fc_cls = nn.Linear(in_channels, num_classes)
if self.spatial_type == 'avg':
self.avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
else:
self.avg_pool = None
self.loss_cls = build_module_from_registers(loss_cls,
module_name='loss')
def build_model(self, algorithm, train_cfg, test_cfg):
sub_cfg = dict(train_cfg=train_cfg, test_cfg=test_cfg)
_model = build_module_from_registers(algorithm,
module_name='algorithm',
sub_cfg=sub_cfg)
model = PCLDataParallel(_model.cuda(self.gpu_ids[0]),
device_ids=self.gpu_ids)
return model
def __init__(self,
pretrained=None,
pretrained_modules=None,
down_ratio=4,
num_classes=1,
frame_rate=30,
track_buffer=30,
conf_thres=0.6,
max_objs=128,
min_box_area=200,
backbone=None,
head=None,
filter=dict(type='KalmanFilter'),
train_cfg=None,
test_cfg=None):
super(FairMOT, self).__init__(pretrained=pretrained,
pretrained_modules=pretrained_modules)
self.backbone = build_module_from_registers(backbone,
module_name='backbone')
self.head = build_module_from_registers(head, module_name='head')
self.init_weights()
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.num_classes = num_classes
self.down_ratio = down_ratio
self.max_objs = max_objs
self.min_box_area = min_box_area
self.det_head_cfg = head['det_head']
self.filter = build_module_from_registers(filter, module_name='filter')
self.tracked_stracks = [] # type: list[STrack]
self.lost_stracks = [] # type: list[STrack]
self.removed_stracks = [] # type: list[STrack]
self.frame_id = 0
self.det_thres = conf_thres
self.buffer_size = int(frame_rate / 30.0 * track_buffer)
self.max_time_lost = self.buffer_size
self.max_per_image = 128
def __init__(self,
bbox_roi_extractor=dict(roi_layer=dict(type='RoIAlign',
out_size=7,
sample_num=0),
out_channels=256,
featmap_strides=[4, 8, 16, 32]),
bbox_head=dict(
type='Shared2FCBBoxHead',
in_channels=256,
fc_out_channels=1024,
roi_feat_size=7,
num_classes=80,
bbox_coder=dict(target_means=[0., 0., 0., 0.],
target_std=[0.1, 0.1, 0.1, 0.1]),
reg_class_agnostic=False,
loss_cls=dict(type='CrossEntropyLoss',
use_sigmoid=False,
loss_weight=-1.0),
loss_bbox=dict(type='L1Loss', loss_weight=-1.0)),
shared_head=None,
assigner=dict(pos_iou_thr=0.5,
neg_iou_thr=0.5,
min_pos_iou=0.5,
match_low_quality=False,
ignore_iof_thr=-1),
sampler=dict(num=512,
pos_fraction=0.5,
neg_pos_ub=-1,
add_gt_as_proposal=True),
pos_weight=-1,
train_cfg=None,
test_cfg=None,
**kwargs):
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.shared_head = shared_head
self.pos_weight = pos_weight
self.bbox_head = build_module_from_registers(bbox_head,
module_name='head')
if bbox_roi_extractor is not None:
self.bbox_roi_extractor = SingleRoIExtractor(**bbox_roi_extractor)
else:
self.bbox_roi_extractor = None
if assigner is not None:
self.bbox_assigner = MaxIoUAssigner(**assigner)
else:
self.bbox_assigner = None
if sampler is not None:
self.bbox_sampler = RandomSampler(**sampler)
else:
self.bbox_sampler = None
def __init__(self,
pretrained,
resample_rate=8,
speed_ratio=8,
channel_ratio=8,
slow_pathway=dict(
type='ResNet3dPathway',
depth=50,
pretrained=None,
lateral=True,
conv1_kernel=(1, 7, 7),
dilations=(1, 1, 1, 1),
conv1_stride_t=1,
pool1_stride_t=1,
inflate=(0, 0, 1, 1)),
fast_pathway=dict(
type='ResNet3dPathway',
depth=50,
pretrained=None,
lateral=False,
base_channels=8,
conv1_kernel=(5, 7, 7),
conv1_stride_t=1,
pool1_stride_t=1)):
super().__init__()
self.pretrained = pretrained
self.resample_rate = resample_rate
self.speed_ratio = speed_ratio
self.channel_ratio = channel_ratio
if slow_pathway['lateral']:
slow_pathway['speed_ratio'] = speed_ratio
slow_pathway['channel_ratio'] = channel_ratio
self.slow_path = build_module_from_registers(
slow_pathway, module_name='backbone')
self.fast_path = build_module_from_registers(
fast_pathway, module_name='backbone')
def __init__(self, transforms):
assert isinstance(transforms, Sequence)
self.transforms = []
for transform in transforms:
if isinstance(transform, dict):
self.transforms.append(
build_module_from_registers(transform,
module_name='pipeline'))
elif callable(transform):
self.transforms.append(transform)
else:
raise TypeError(
f'transform must be callable or a dict, but got {type(transform)}'
)
def __init__(self,
pretrained=None,
pretrained_modules=None,
backbone=None,
head=None,
auxiliary_head=None,
train_cfg=None,
test_cfg=None,
**kwargs):
super(FCN, self).__init__(pretrained=pretrained,
pretrained_modules=pretrained_modules)
self.backbone = build_module_from_registers(backbone,
module_name='backbone')
self.head = build_module_from_registers(head, module_name='head')
self.auxiliary_head = build_module_from_registers(auxiliary_head,
module_name='head')
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights()
def build_datasets(self, data, runtime_cfg, train_cfg, test_cfg):
datasets = {}
sub_cfg = dict(train_cfg=train_cfg, test_cfg=test_cfg)
for task in self.tasks:
if 'test' in task:
mode = 'test'
else:
mode = task
if data.get(mode, None) is None:
continue
config = data[mode]
datasets[task] = build_module_from_registers(config,
module_name='data',
sub_cfg=sub_cfg)
return datasets
def __init__(self,
channels,
first_level,
final_kernel=1,
head_conv=256,
reid_dim=512,
reid_nID=14455,
id_loss=None,
**kwargs):
super(FairMOTReIDHead, self).__init__()
self.head_conv = head_conv
self.first_level = first_level
self.reid_nID = reid_nID
self.embedding_dim = reid_dim
self.embedding_scale = math.sqrt(2) * math.log(self.reid_nID - 1)
if self.head_conv > 0 and self.embedding_dim is not None:
self.id = nn.Sequential(
nn.Conv2d(channels[self.first_level],
head_conv,
kernel_size=3,
padding=1,
bias=True), nn.ReLU(inplace=True),
nn.Conv2d(self.head_conv,
self.embedding_dim,
kernel_size=final_kernel,
stride=1,
padding=final_kernel // 2,
bias=True))
elif self.head_conv <= 0 and self.embedding_dim is not None:
self.id = nn.Conv2d(channels[self.first_level],
self.embedding_dim,
kernel_size=final_kernel,
stride=1,
padding=final_kernel // 2,
bias=True)
self.id_loss = build_module_from_registers(id_loss, module_name='loss')
self.id_classifier = nn.Linear(self.embedding_dim, self.reid_nID)
def __init__(self,
channels,
first_level,
final_kernel=1,
head_conv=256,
num_classes=1,
cat_spec_wh=False,
reg_offset=True,
offset=2,
hm_loss=None,
wh_loss=None,
offset_loss=None,
**kwargs):
super(FairMOTDetHead, self).__init__()
self.num_classes = num_classes
self.head_conv = head_conv
self.first_level = first_level
self.hm_loss = hm_loss
self.wh_loss = wh_loss
self.offset_loss = offset_loss
self.hm_channels = self.num_classes
if not cat_spec_wh:
self.wh_channels = 2
else:
self.wh_channels = 2 * self.num_classes
if reg_offset:
self.offset_channels = offset
else:
self.offset_channels = None
self.heads = dict(hm=self.hm_channels,
wh=self.wh_channels,
offset=self.offset_channels)
for head, classes in self.heads.items():
if self.head_conv > 0 and classes is not None:
fc = nn.Sequential(
nn.Conv2d(channels[self.first_level],
head_conv,
kernel_size=3,
padding=1,
bias=True), nn.ReLU(inplace=True),
nn.Conv2d(self.head_conv,
classes,
kernel_size=final_kernel,
stride=1,
padding=final_kernel // 2,
bias=True))
elif self.head_conv <= 0 and classes is not None:
fc = nn.Conv2d(channels[self.first_level],
classes,
kernel_size=final_kernel,
stride=1,
padding=final_kernel // 2,
bias=True)
setattr(self, head, fc)
self.hm_loss = build_module_from_registers(hm_loss, module_name='loss')
self.wh_loss = build_module_from_registers(wh_loss, module_name='loss')
self.offset_loss = build_module_from_registers(offset_loss,
module_name='loss')
def __init__(self,
num_classes=1,
in_channels=0,
feat_channels=256,
anchor_generator=dict(scales=[8, 16, 32],
ratios=[0.5, 1.0, 2.0],
strides=[4, 8, 16, 32, 64]),
box_coder=dict(target_means=[.0, .0, .0, .0],
target_stds=[1.0, 1.0, 1.0, 1.0]),
sampler=dict(num=256,
pos_fraction=-0.5,
neg_pos_ub=-1,
add_gt_as_proposal=False),
assigner=dict(pos_iou_thr=0.7,
neg_iou_thr=0.3,
min_pos_iou=0.3,
match_low_quality=True,
ignore_iof_thr=-1),
allowed_border=-1,
pos_weight=-1,
reg_decoded_bbox=False,
background_label=0,
loss_cls=dict(),
loss_bbox=dict(),
train_cfg=None,
test_cfg=None,
**kwargs):
super(RPNHead, self).__init__()
self.num_classes = num_classes
self.in_channels = in_channels
self.feat_channels = feat_channels
self.allowed_border = allowed_border
self.pos_weight = pos_weight
assert loss_cls is not None
self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
if self.use_sigmoid_cls:
self.cls_out_channels = num_classes
else:
self.cls_out_channels = num_classes + 1
self.sampling = loss_cls['type'] not in ['FocalLoss', 'GHMC']
self.reg_decoded_bbox = reg_decoded_bbox
self.background_label = (num_classes if background_label is None else
background_label)
assert (self.background_label == 0
or self.background_label == num_classes)
if box_coder is not None:
self.bbox_coder = DeltaXYWHBBoxCoder(**box_coder)
else:
self.bbox_coder = None
self.loss_cls = build_module_from_registers(loss_cls,
module_name='loss')
self.loss_bbox = build_module_from_registers(loss_bbox,
module_name='loss')
self.train_cfg = train_cfg
self.test_cfg = test_cfg
if assigner is not None:
self.assigner = MaxIoUAssigner(**assigner)
else:
self.assigner = None
if self.sampling and sampler is not None:
self.sampler = RandomSampler(**sampler)
else:
self.sampler = PseudoSampler()
if anchor_generator is not None:
self.anchor_generator = AnchorGenerator(**anchor_generator)
self.num_anchors = self.anchor_generator.num_base_anchors[0]
else:
self.anchor_generator = None
self.rpn_conv = nn.Conv2d(self.in_channels,
self.feat_channels,
3,
padding=1)
self.rpn_cls = nn.Conv2d(self.feat_channels,
self.num_anchors * self.cls_out_channels, 1)
self.rpn_reg = nn.Conv2d(self.feat_channels, self.num_anchors * 4, 1)
def __init__(self,
temporal_scale,
prop_boundary_ratio,
num_samples,
num_samples_per_bin,
feat_dim,
hidden_dim_1d=256,
hidden_dim_2d=128,
hidden_dim_3d=512,
tem_loss=dict(),
pem_reg_loss=dict(),
pem_cls_loss=dict(),
**kwargs):
super(BMNHead, self).__init__()
self.tscale = temporal_scale
self.temporal_gap = 1. / temporal_scale
self.pbr = prop_boundary_ratio
self.num_samples = num_samples
self.num_samples_per_bin = num_samples_per_bin
self.feat_dim = feat_dim
self.hidden_dim_1d = hidden_dim_1d
self.hidden_dim_2d = hidden_dim_2d
self.hidden_dim_3d = hidden_dim_3d
self.tem_loss = build_module_from_registers(tem_loss,
module_name='loss')
self.pem_reg_loss = build_module_from_registers(pem_reg_loss,
module_name='loss')
self.pem_cls_loss = build_module_from_registers(pem_cls_loss,
module_name='loss')
self.x_1d_b = nn.Sequential(
nn.Conv1d(self.feat_dim,
self.hidden_dim_1d,
kernel_size=3,
padding=1,
groups=4), nn.ReLU(inplace=True),
nn.Conv1d(self.hidden_dim_1d,
self.hidden_dim_1d,
kernel_size=3,
padding=1,
groups=4), nn.ReLU(inplace=True))
# Temporal Evaluation Module
self.x_1d_s = nn.Sequential(
nn.Conv1d(self.hidden_dim_1d,
self.hidden_dim_1d,
kernel_size=3,
padding=1,
groups=4), nn.ReLU(inplace=True),
nn.Conv1d(self.hidden_dim_1d, 1, kernel_size=1), nn.Sigmoid())
self.x_1d_e = nn.Sequential(
nn.Conv1d(self.hidden_dim_1d,
self.hidden_dim_1d,
kernel_size=3,
padding=1,
groups=4), nn.ReLU(inplace=True),
nn.Conv1d(self.hidden_dim_1d, 1, kernel_size=1), nn.Sigmoid())
# Proposal Evaluation Module
self.x_1d_p = nn.Sequential(
nn.Conv1d(self.hidden_dim_1d,
self.hidden_dim_2d,
kernel_size=3,
padding=1), nn.ReLU(inplace=True))
self.x_3d_p = nn.Sequential(
nn.Conv3d(self.hidden_dim_2d,
self.hidden_dim_3d,
kernel_size=(self.num_sample, 1, 1),
stride=(self.num_sample, 1, 1)), nn.ReLU(inplace=True))
self.x_2d_p = nn.Sequential(
nn.Conv2d(self.hidden_dim_3d, self.hidden_dim_2d, kernel_size=1),
nn.ReLU(inplace=True),
nn.Conv2d(self.hidden_dim_2d,
self.hidden_dim_2d,
kernel_size=3,
padding=1), nn.ReLU(inplace=True),
nn.Conv2d(self.hidden_dim_2d,
self.hidden_dim_2d,
kernel_size=3,
padding=1), nn.ReLU(inplace=True),
nn.Conv2d(self.hidden_dim_2d, 2, kernel_size=1), nn.Sigmoid())
self._get_interp1d_mask()
self.bm_mask = self._get_bm_mask()
def __init__(self, upsample=True, loss=None):
super(FCRNHead, self).__init__()
self.upsample = upsample
self.loss_func = build_module_from_registers(loss, module_name='loss')
if self.upsample:
self.up_layer = nn.Upsample((228, 304), mode='bilinear')
def __init__(self,
in_channels=0,
head_conv=0,
num_classes=0,
output_stride=4,
K=7,
max_output_objects=0,
hm_loss=None,
mov_loss=None,
wh_loss=None):
super(MOCHead, self).__init__()
assert head_conv > 0
head_info = dict(hm=num_classes, mov=2 * K, wh=2 * K)
self.hm = nn.Sequential(
nn.Conv2d(K * in_channels,
head_conv,
kernel_size=3,
padding=1,
bias=True), nn.ReLU(inplace=True),
nn.Conv2d(head_conv,
head_info['hm'],
kernel_size=1,
stride=1,
padding=0,
bias=True))
self.mov = nn.Sequential(
nn.Conv2d(K * in_channels,
head_conv,
kernel_size=3,
padding=1,
bias=True), nn.ReLU(inplace=True),
nn.Conv2d(head_conv,
head_info['mov'],
kernel_size=1,
stride=1,
padding=0,
bias=True))
self.wh = nn.Sequential(
nn.Conv2d(in_channels,
head_conv,
kernel_size=3,
padding=1,
bias=True), nn.ReLU(inplace=True),
nn.Conv2d(head_conv,
head_info['wh'] // K,
kernel_size=1,
stride=1,
padding=0,
bias=True))
self.hm_loss = build_module_from_registers(hm_loss, module_name='loss')
self.wh_loss = build_module_from_registers(wh_loss, module_name='loss')
self.mov_loss = build_module_from_registers(mov_loss,
module_name='loss')
self.num_classes = num_classes
self.output_stride = output_stride
self.N = max_output_objects
self.K = K
def build_modules(self):
kwargs = self.module_cfg['kwargs']
for key, value in kwargs.items():
module = build_module_from_registers(value, module_name=key)
setattr(self, key, module)
