def resnet3d_gn_ws(pretrained=True):
norm_cfg = dict(type='GN', num_groups=32, requires_grad=True)
backbone3d = dict(type='ResNet',
depth=50,
num_stages=4,
out_indices=(0, 1, 2, 3),
frozen_stages=1,
style='pytorch',
norm_cfg=norm_cfg,
conv_cfg=dict(type='ConvWS3'))
bb3 = build_backbone(backbone3d)
if pretrained:
backbone2d = dict(type='ResNet',
depth=50,
num_stages=4,
out_indices=(0, 1, 2, 3),
frozen_stages=1,
style='pytorch',
norm_cfg=norm_cfg,
conv_cfg=dict(type='ConvWS'))
bb2 = build_backbone(backbone2d)
bb2.init_weights('open-mmlab://jhu/resnet50_gn_ws')
bb2_weights = bb2.state_dict()
for w in list(bb2_weights.keys()):
try:
bb3.state_dict()[w].data.copy_(bb2_weights[w])
except RuntimeError:
tmp_bb3_weights = bb2_weights[w].unsqueeze(-1).repeat(
1, 1, 1, 1,
bb2.state_dict()[w].shape[-1])
bb3.state_dict()[w].data.copy_(tmp_bb3_weights /
tmp_bb3_weights.shape[-1])
return bb3
def __init__(self,
backbone,
neck=None,
head=None,
pretrains=None,
init_cfg=None,
frozen_modules=None,
train_cfg=None,
test_cfg=None):
super(SiamRPN, self).__init__(init_cfg)
if isinstance(pretrains, dict):
warnings.warn('DeprecationWarning: pretrains is deprecated, '
'please use "init_cfg" instead')
backbone_pretrain = pretrains.get('backbone', None)
if backbone_pretrain:
backbone.init_cfg = dict(type='Pretrained',
checkpoint=backbone_pretrain)
else:
backbone.init_cfg = None
self.backbone = build_backbone(backbone)
if neck is not None:
self.neck = build_neck(neck)
head = head.copy()
head.update(train_cfg=train_cfg.rpn, test_cfg=test_cfg.rpn)
self.head = build_head(head)
self.test_cfg = test_cfg
self.train_cfg = train_cfg
if frozen_modules is not None:
self.freeze_module(frozen_modules)
def __init__(self,
backbone,
rfp_steps=2,
rfp_sharing=False,
stage_with_rfp=(False, True, True, True),
**kwargs):
self.rfp_steps = rfp_steps
self.rfp_sharing = rfp_sharing
self.stage_with_rfp = stage_with_rfp
backbone["rfp"] = None
backbone["stage_with_rfp"] = stage_with_rfp
neck_out_channels = kwargs["neck"]["out_channels"]
if rfp_sharing:
backbone["rfp"] = neck_out_channels
super().__init__(backbone=backbone, **kwargs)
if not self.rfp_sharing:
backbone["rfp"] = neck_out_channels
self.rfp_modules = torch.nn.ModuleList()
for rfp_idx in range(1, rfp_steps):
rfp_module = builder.build_backbone(backbone)
rfp_module.init_weights(kwargs["pretrained"])
self.rfp_modules.append(rfp_module)
self.rfp_aspp = ASPP(neck_out_channels, neck_out_channels // 4)
self.rfp_weight = torch.nn.Conv2d(neck_out_channels,
1,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.rfp_weight.weight.data.fill_(0)
self.rfp_weight.bias.data.fill_(0)
def __init__(self,
backbone,
neck=None,
rpn_head=None,
roi_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(OBBTwoStageDetector, self).__init__()
self.backbone = build_backbone(backbone)
if neck is not None:
self.neck = build_neck(neck)
if rpn_head is not None:
rpn_train_cfg = train_cfg.rpn if train_cfg is not None else None
rpn_head_ = rpn_head.copy()
rpn_head_.update(train_cfg=rpn_train_cfg, test_cfg=test_cfg.rpn)
self.rpn_head = build_head(rpn_head_)
if roi_head is not None:
# update train and test cfg here for now
# TODO: refactor assigner & sampler
rcnn_train_cfg = train_cfg.rcnn if train_cfg is not None else None
roi_head.update(train_cfg=rcnn_train_cfg)
roi_head.update(test_cfg=test_cfg.rcnn)
self.roi_head = build_head(roi_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def __init__(self,
num_stages,
backbone,
neck=None,
shared_head=None,
rpn_head=None,
bbox_roi_extractor=None,
bbox_head=None,
count_head=None,
similar_head=None,
mask_roi_extractor=None,
mask_head=None,
pretrained=None,
train_cfg=None,
test_cfg=None):
super(CaSe_CascadeRCNN, self).__init__()
self.num_stages = num_stages
self.backbone = builder.build_backbone(backbone)
if neck is not None:
self.neck = builder.build_neck(neck)
if shared_head is not None:
self.shared_head = builder.build_shared_head(shared_head)
if rpn_head is not None:
self.rpn_head = builder.build_head(rpn_head)
if count_head is not None:
self.count_head = builder.build_head(count_head)
if similar_head is not None:
self.similar_head = builder.build_head(similar_head)
if bbox_head is not None:
self.bbox_roi_extractor = nn.ModuleList()
self.bbox_head = nn.ModuleList()
if not isinstance(bbox_roi_extractor, list):
bbox_roi_extractor = [
bbox_roi_extractor for _ in range(num_stages)
]
if not isinstance(bbox_head, list):
bbox_head = [bbox_head for _ in range(num_stages)]
assert len(bbox_roi_extractor) == len(bbox_head) == self.num_stages
for roi_extractor, head in zip(bbox_roi_extractor, bbox_head):
self.bbox_roi_extractor.append(
builder.build_roi_extractor(roi_extractor))
self.bbox_head.append(builder.build_head(head))
if mask_head is not None:
self.mask_roi_extractor = builder.build_roi_extractor(
mask_roi_extractor)
self.mask_head = builder.build_head(mask_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
def __init__(self,
preprocessor=None,
backbone=None,
encoder=None,
decoder=None,
loss=None,
label_convertor=None,
train_cfg=None,
test_cfg=None,
max_seq_len=40,
pretrained=None,
init_cfg=None):
super().__init__(init_cfg=init_cfg)
# Label convertor (str2tensor, tensor2str)
assert label_convertor is not None
label_convertor.update(max_seq_len=max_seq_len)
self.label_convertor = build_convertor(label_convertor)
# Preprocessor module, e.g., TPS
self.preprocessor = None
if preprocessor is not None:
self.preprocessor = build_preprocessor(preprocessor)
# Backbone
assert backbone is not None
self.backbone = build_backbone(backbone)
# Encoder module
self.encoder = None
if encoder is not None:
self.encoder = build_encoder(encoder)
# Decoder module
assert decoder is not None
decoder.update(num_classes=self.label_convertor.num_classes())
decoder.update(start_idx=self.label_convertor.start_idx)
decoder.update(padding_idx=self.label_convertor.padding_idx)
decoder.update(max_seq_len=max_seq_len)
self.decoder = build_decoder(decoder)
# Loss
assert loss is not None
loss.update(ignore_index=self.label_convertor.padding_idx)
self.loss = build_loss(loss)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.max_seq_len = max_seq_len
if pretrained is not None:
warnings.warn('DeprecationWarning: pretrained is a deprecated \
key, please consider using init_cfg')
self.init_cfg = dict(type='Pretrained', checkpoint=pretrained)
def __init__(self,
backbone,
neck=None,
shared_head=None,
rpn_head=None,
bbox_roi_extractor=None,
bbox_head=None,
mask_roi_extractor=None,
mask_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None,
seg_cfg=None,
cap_cfg=None,
seg_scales=[0.5, 0.75],
):
super(EncoderDecoder, self).__init__()
self.backbone = builder.build_backbone(backbone)
if neck is not None:
self.neck = builder.build_neck(neck)
if shared_head is not None:
self.shared_head = builder.build_shared_head(shared_head)
if rpn_head is not None:
self.rpn_head = builder.build_head(rpn_head)
if bbox_head is not None:
self.bbox_roi_extractor = builder.build_roi_extractor(
bbox_roi_extractor)
self.bbox_head = builder.build_head(bbox_head)
if mask_head is not None:
if mask_roi_extractor is not None:
self.mask_roi_extractor = builder.build_roi_extractor(
mask_roi_extractor)
self.share_roi_extractor = False
else:
self.share_roi_extractor = True
self.mask_roi_extractor = self.bbox_roi_extractor
self.mask_head = builder.build_head(mask_head)
if seg_cfg is not None:
self.seg_decoder = DeeplabDecoder(**seg_cfg)
self.seg_scales = seg_scales
if cap_cfg is not None:
self.cap_decoder = CapDecoder(**cap_cfg)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def __init__(self,
preprocessor=None,
backbone=None,
encoder=None,
decoder=None,
loss=None,
label_convertor=None,
train_cfg=None,
test_cfg=None,
max_seq_len=40,
pretrained=None):
super().__init__()
# Label convertor (str2tensor, tensor2str)
assert label_convertor is not None
label_convertor.update(max_seq_len=max_seq_len)
self.label_convertor = build_convertor(label_convertor)
# Preprocessor module, e.g., TPS
self.preprocessor = None
if preprocessor is not None:
self.preprocessor = build_preprocessor(preprocessor)
# Backbone
assert backbone is not None
self.backbone = build_backbone(backbone)
# Encoder module
self.encoder = None
if encoder is not None:
self.encoder = build_encoder(encoder)
# Decoder module
assert decoder is not None
decoder.update(num_classes=self.label_convertor.num_classes())
decoder.update(start_idx=self.label_convertor.start_idx)
decoder.update(padding_idx=self.label_convertor.padding_idx)
decoder.update(max_seq_len=max_seq_len)
self.decoder = build_decoder(decoder)
# Loss
assert loss is not None
loss.update(ignore_index=self.label_convertor.padding_idx)
self.loss = build_loss(loss)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.max_seq_len = max_seq_len
self.init_weights(pretrained=pretrained)
def __init__(
self,
backbone,
neck=None,
shared_head=None,
rpn_head=None,
bbox_roi_extractor=None,
bbox_head=None,
mask_roi_extractor=None,
mask_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None,
seg_cfg=None,
cap_cfg=None,
):
super(CapModel, self).__init__()
self.backbone = builder.build_backbone(backbone)
# if neck is not None:
# self.neck = builder.build_neck(neck)
# if shared_head is not None:
# self.shared_head = builder.build_shared_head(shared_head)
# if rpn_head is not None:
# self.rpn_head = builder.build_head(rpn_head)
# if bbox_head is not None:
# self.bbox_roi_extractor = builder.build_roi_extractor(
# bbox_roi_extractor)
# self.bbox_head = builder.build_head(bbox_head)
# if mask_head is not None:
# if mask_roi_extractor is not None:
# self.mask_roi_extractor = builder.build_roi_extractor(
# mask_roi_extractor)
# self.share_roi_extractor = False
# else:
# self.share_roi_extractor = True
# self.mask_roi_extractor = self.bbox_roi_extractor
# self.mask_head = builder.build_head(mask_head)
self.cap_decoder = CapDecoder(**cap_cfg)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def __init__(self,
backbone,
neck,
rpn_head,
train_cfg,
test_cfg,
pretrained=None):
super(QG_RPN, self).__init__()
self.backbone = builder.build_backbone(backbone)
self.neck = builder.build_neck(neck) if neck is not None else None
self.rpn_head = builder.build_head(rpn_head)
self.rpn_modulator = RPN_Modulator()
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def run(self):
"""Runs validation only with the network."""
# Get the checkpoint file
print('loading checkpoint file ...')
cp = torch.load(self.cfg.work_dir + '/latest.pth')
print('done')
print('loading state dictionary ...')
# Initialize network first as separate modules so we can access WFCOS
backbone = build_backbone(self.cfg.model.backbone).cuda()
neck = build_neck(self.cfg.model.neck).cuda()
head = build_head(self.cfg.model.bbox_head).cuda()
# Load the state dicts
backbone_state = OrderedDict()
neck_state = OrderedDict()
head_state = OrderedDict()
for key in cp['state_dict'].keys():
if 'backbone' in key:
backbone_state[key.split('.', 1)[1]] = cp['state_dict'][key]
elif 'neck' in key:
neck_state[key.split('.', 1)[1]] = cp['state_dict'][key]
elif 'bbox_head' in key:
head_state[key.split('.', 1)[1]] = cp['state_dict'][key]
backbone.load_state_dict(backbone_state)
neck.load_state_dict(neck_state)
head.load_state_dict(head_state)
# Set to eval mode
backbone.eval()
neck.eval()
head.eval()
print('done')
print('starting inference validation run ...')
for i, (img, cls) in enumerate(self.loader):
out = backbone(img)
out = neck(out)
out = head(out)
img_metas = [{'img_shape': (640, 800), 'scale_factor': 1}]
bboxes = head.get_bboxes(out[0], out[1], out[2], img_metas,
self.cfg.test_cfg)
pass
print('done')
def __init__(self,
backbone,
neck,
rpn_head,
train_cfg,
test_cfg,
pretrained=None):
super(OBBRPN, self).__init__()
self.backbone = build_backbone(backbone)
self.neck = build_neck(neck) if neck is not None else None
rpn_train_cfg = train_cfg.rpn if train_cfg is not None else None
rpn_head.update(train_cfg=rpn_train_cfg)
rpn_head.update(test_cfg=test_cfg.rpn)
self.rpn_head = build_head(rpn_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def __init__(self,
backbone,
neck=None,
bbox_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(OBBSingleStageDetector, self).__init__()
self.backbone = build_backbone(backbone)
if neck is not None:
self.neck = build_neck(neck)
bbox_head.update(train_cfg=train_cfg)
bbox_head.update(test_cfg=test_cfg)
self.bbox_head = build_head(bbox_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def __init__(self,
preprocessor=None,
backbone=None,
neck=None,
head=None,
loss=None,
label_convertor=None,
train_cfg=None,
test_cfg=None,
pretrained=None,
init_cfg=None):
super().__init__(init_cfg=init_cfg)
# Label_convertor
assert label_convertor is not None
self.label_convertor = build_convertor(label_convertor)
# Preprocessor module, e.g., TPS
self.preprocessor = None
if preprocessor is not None:
self.preprocessor = build_preprocessor(preprocessor)
# Backbone
assert backbone is not None
self.backbone = build_backbone(backbone)
# Neck
assert neck is not None
self.neck = build_neck(neck)
# Head
assert head is not None
head.update(num_classes=self.label_convertor.num_classes())
self.head = build_head(head)
# Loss
assert loss is not None
self.loss = build_loss(loss)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
if pretrained is not None:
warnings.warn('DeprecationWarning: pretrained is a deprecated \
key, please consider using init_cfg')
self.init_cfg = dict(type='Pretrained', checkpoint=pretrained)
def __init__(self,
backbone,
neck,
rpn_head,
train_cfg,
test_cfg,
bbox_roi_extractor=None,
pretrained=None):
super(newRPN, self).__init__()
self.backbone = builder.build_backbone(backbone)
self.neck = builder.build_neck(neck) if neck is not None else None
self.rpn_head = builder.build_rpn_head(rpn_head)
if bbox_roi_extractor is not None:
self.bbox_roi_extractor = builder.build_roi_extractor(
bbox_roi_extractor)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def __init__(self,
num_stages,
backbone,
neck,
rpn_head,
train_cfg,
test_cfg,
pretrained=None):
super(CascadeRPN, self).__init__()
assert num_stages == len(rpn_head)
self.num_stages = num_stages
self.backbone = builder.build_backbone(backbone)
self.neck = builder.build_neck(neck) if neck is not None else None
self.rpn_head = nn.ModuleList()
for head in rpn_head:
self.rpn_head.append(builder.build_head(head))
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def __init__(self,
backbone,
neck=None,
shared_head=None,
rpn_head=None,
bbox_roi_extractor=None,
bbox_head=None,
count_head=None,
similar_head=None,
mask_roi_extractor=None,
mask_head=None,
pretrained=None,
train_cfg=None,
test_cfg=None):
super(CaSe, self).__init__()
self.backbone = builder.build_backbone(backbone)
if neck is not None:
self.neck = builder.build_neck(neck)
if shared_head is not None:
self.shared_head = builder.build_shared_head(shared_head)
if rpn_head is not None:
self.rpn_head = builder.build_head(rpn_head)
if count_head is not None:
self.count_head = builder.build_head(count_head)
if similar_head is not None:
self.similar_head = builder.build_head(similar_head)
if bbox_head is not None:
self.bbox_roi_extractor = builder.build_roi_extractor(
bbox_roi_extractor)
self.bbox_head = builder.build_head(bbox_head)
if mask_head is not None:
self.mask_roi_extractor = builder.build_roi_extractor(
mask_roi_extractor)
self.mask_head = builder.build_head(mask_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
def __init__(self,
preprocessor=None,
backbone=None,
neck=None,
head=None,
loss=None,
label_convertor=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super().__init__()
# Label_convertor
assert label_convertor is not None
self.label_convertor = build_convertor(label_convertor)
# Preprocessor module, e.g., TPS
self.preprocessor = None
if preprocessor is not None:
self.preprocessor = build_preprocessor(preprocessor)
# Backbone
assert backbone is not None
self.backbone = build_backbone(backbone)
# Neck
assert neck is not None
self.neck = build_neck(neck)
# Head
assert head is not None
head.update(num_classes=self.label_convertor.num_classes())
self.head = build_head(head)
# Loss
assert loss is not None
self.loss = build_loss(loss)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def __init__(self,
backbone,
neck=None,
mask_head=None,
shape_transform_module=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super().__init__()
self.backbone = builder.build_backbone(backbone)
if neck is not None:
self.neck = builder.build_neck(neck)
self.mask_head = builder.build_head(mask_head)
if shape_transform_module is not None:
self.shape_transform_module = build_roi_extractor(
shape_transform_module)
else:
self.shape_transform_module = None
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
def __init__(self,
backbone,
neck=None,
head=None,
init_cfg=None,
frozen_modules=None,
train_cfg=None,
test_cfg=None):
super(Stark, self).__init__(init_cfg)
self.backbone = build_backbone(backbone)
self.neck = build_neck(neck)
self.head = build_head(head)
self.test_cfg = test_cfg
self.train_cfg = train_cfg
# Set the update interval
self.update_intervals = self.test_cfg['update_intervals']
self.num_extra_template = len(self.update_intervals)
if frozen_modules is not None:
self.freeze_module(frozen_modules)
def __init__(self,
pretrains=None,
backbone=None,
neck=None,
head=None,
frozen_modules=None,
train_cfg=None,
test_cfg=None):
super(SiamRPN, self).__init__()
self.backbone = build_backbone(backbone)
if neck is not None:
self.neck = build_neck(neck)
head = head.copy()
head.update(train_cfg=train_cfg.rpn, test_cfg=test_cfg.rpn)
self.head = build_head(head)
self.test_cfg = test_cfg
self.train_cfg = train_cfg
self.init_weights(pretrains)
if frozen_modules is not None:
self.freeze_module(frozen_modules)
def __init__(self,
backbone,
neck=None,
rpn_head=None,
bbox_roi_extractor=None,
action_head=None,
bbox_head=None,
mask_roi_extractor=None,
mask_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(TwoStageDetector, self).__init__()
self.backbone = builder.build_backbone(backbone)
if neck is not None:
self.neck = builder.build_neck(neck)
else:
raise NotImplementedError
if rpn_head is not None:
self.rpn_head = builder.build_rpn_head(rpn_head)
if action_head is not None:
self.action_head = builder_.build_actionhead(action_head)
if bbox_head is not None:
self.bbox_roi_extractor = builder.build_roi_extractor(
bbox_roi_extractor)
self.bbox_head = builder.build_bbox_head(bbox_head)
if mask_head is not None:
self.mask_roi_extractor = builder.build_roi_extractor(
mask_roi_extractor)
self.mask_head = builder.build_mask_head(mask_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.init_weights(pretrained=pretrained)
