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)