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)