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,
vertex_head,
polyrnn_head,
loss_vertex=dict(type='GaussianFocalLoss',
alpha=2.0,
gamma=4.0,
loss_weight=1),
loss_polygon=dict(type='CrossEntropyLoss',
use_mask=False,
loss_weight=1.0),
loss_offset=dict(type='L1Loss', loss_weight=1.0),
loss_type=0,
params=dict(dt_threshold=2, radius=1),
):
super(PolygonHead, self).__init__()
self.vertex_head = build_head(vertex_head)
polyrnn_head['loss_type'] = loss_type
self.polyrnn_head = build_head(polyrnn_head)
self.loss_vertex_cfg = loss_vertex
self.loss_polygon_cfg = loss_polygon
self.loss_vertex = build_loss(loss_vertex)
self.loss_polygon = build_loss(loss_polygon)
self.loss_offset = build_loss(loss_offset)
self.loss_type = loss_type
self.params = params
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,
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,
num_query=1,
transformer=None,
positional_encoding=dict(
type='SinePositionalEncoding',
num_feats=128,
normalize=True),
bbox_head=None,
cls_head=None,
loss_cls=dict(
type='CrossEntropyLoss',
use_sigmoid=False,
loss_weight=1.0,
),
loss_bbox=dict(type='L1Loss', loss_weight=5.0),
loss_iou=dict(type='GIoULoss', loss_weight=2.0),
train_cfg=None,
test_cfg=None,
init_cfg=None,
frozen_modules=None,
**kwargs):
super(StarkHead, self).__init__(init_cfg=init_cfg)
self.transformer = build_transformer(transformer)
self.positional_encoding = build_positional_encoding(
positional_encoding)
assert bbox_head is not None
self.bbox_head = build_head(bbox_head)
if cls_head is None:
# the stage-1 training
self.loss_bbox = build_loss(loss_bbox)
self.loss_iou = build_loss(loss_iou)
self.cls_head = None
else:
# the stage-2 training
self.cls_head = build_head(cls_head)
self.loss_cls = build_loss(loss_cls)
self.embed_dims = self.transformer.embed_dims
self.num_query = num_query
self.query_embedding = nn.Embedding(self.num_query, self.embed_dims)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.fp16_enabled = False
if frozen_modules is not None:
assert isinstance(frozen_modules, list)
for module in frozen_modules:
m = getattr(self, module)
# TODO: Study the influence of freezing BN running_mean and
# running_variance of `frozen_modules` in the 2nd stage train.
# The official code doesn't freeze these.
for param in m.parameters():
param.requires_grad = False
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_mask_head(self, mask_roi_extractor, mask_head):
"""Initialize ``mask_head``"""
if mask_roi_extractor is not None:
self.mask_roi_extractor = 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 = build_head(mask_head)
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,
backbone,
neck=None,
bbox_head=None,
mask_feat_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(SingleStageSegDetector, self).__init__()
self.backbone = build_backbone(backbone)
if neck is not None:
self.neck = build_neck(neck)
if mask_feat_head is not None:
self.mask_feat_head = build_head(mask_feat_head)
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,
heads,
feat_channels=256,
align_type='AlignConv',
train_cfg=None,
test_cfg=None):
super(S2AHead, self).__init__()
# Skip cls branch to speedup inference speed
if isinstance(test_cfg['skip_cls'], list):
skip_cls = test_cfg['skip_cls']
test_cfg.pop('skip_cls')
assert len(skip_cls) == len(heads)
test_cfg = [deepcopy(test_cfg) for _ in range(len(heads))]
for skip, cfg in zip(skip_cls, test_cfg):
cfg['skip_cls'] = skip
else:
test_cfg = [test_cfg for _ in heads]
self.heads = nn.ModuleList()
for i, head in enumerate(heads):
if train_cfg is not None:
head.update(train_cfg=train_cfg[i])
head.update(test_cfg=test_cfg[i])
head_module = build_head(head)
if i == 0:
self.anchor_generator = head_module.anchor_generator
self.num_anchors = head_module.num_anchors
else:
head_module.anchor_generator = self.anchor_generator
head_module.num_anchors = self.num_anchors
self.heads.append(head_module)
self.num_stages = len(self.heads)
self.num_classes = self.heads[-1].num_classes
assert self.num_stages >= 2
if isinstance(align_type, str):
self.align_type = [align_type for _ in range(len(self.heads) - 1)]
else:
assert len(align_type) == len(self.heads) - 1
self.align_type = align_type
self.feat_channels = feat_channels
self.bbox_type = 'obb'
self.reg_dim = get_bbox_dim(self.bbox_type)
self._init_layers()
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,
# neck=None,
# bbox_head=None,
cfg=None,
train_cfg=None,
test_cfg=None,
pretrained=None):
super(SingleStageDetector, self).__init__()
# self.backbone = builder.build_backbone(backbone)
if 'neck' in cfg:
self.neck = builder.build_neck(cfg['neck'])
self.bbox_head = builder.build_head(cfg['bbox_head'])
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(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,
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_bbox_head(self, bbox_roi_extractor, bbox_head):
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(self.num_stages)
]
if not isinstance(bbox_head, list):
bbox_head = [bbox_head for _ in range(self.num_stages)]
assert len(bbox_roi_extractor) == len(bbox_head) == self.num_stages
last_type = ['hbb', 'obb', 'poly']
for roi_extractor, head in zip(bbox_roi_extractor, bbox_head):
self.bbox_roi_extractor.append(build_roi_extractor(roi_extractor))
self.bbox_head.append(build_head(head))
assert self.bbox_head[-1].start_bbox_type in last_type
last_type = [self.bbox_head[-1].end_bbox_type]
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,
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_bbox_head(self, bbox_roi_extractor, bbox_head):
"""Initialize ``bbox_head``"""
self.bbox_roi_extractor = build_roi_extractor(bbox_roi_extractor)
self.bbox_head = build_head(bbox_head)
