class SegmentationNetwork(nn.Module): def __init__(self, backbone=None, output_channels=2, backbone_output_channels=512): super().__init__() if not backbone: b, _ = remove_backbone_head(resnet18(pretrained=False)) backbone = b # ResNet produces 512-length outputs backbone.out_channels = backbone_output_channels self.segmentation_network = FasterRCNN(backbone, num_classes=output_channels) def forward(self, x, boxes=None): x = self.segmentation_network(x, boxes) return x def infer(self, x): self.eval() self.segmentation_network.eval() x = self.segmentation_network(x) return convert_bounding_box_inference(x)
def fasterrcnn_resnetxx_fpnxx(cfg): backbone = resnet.__dict__[cfg['backbone_name']]( pretrained=cfg['backbone_pretrained'], norm_layer=misc_nn_ops.FrozenBatchNorm2d) # freeze layers for name, parameter in backbone.named_parameters(): if 'layer2' not in name and 'layer3' not in name and 'layer4' not in name: parameter.requires_grad_(False) # wrapper backbone with fpn return_layers = cfg['fpn']['return_layers'] in_channels_stage2 = backbone.inplanes // 8 in_channels_list = [in_channels_stage2 * 2**i for i in range(len(return_layers))] out_channels = cfg['fpn']['out_channels'] backbone_fpn = BackboneWithFPN(backbone, return_layers, in_channels_list, out_channels) anchor_generator = AnchorGenerator(**cfg['anchor_generator']) # print(anchor_generator.num_anchors_per_location()) roi_pooler = MultiScaleRoIAlign(**cfg['box_roi_pool']) model = FasterRCNN(backbone_fpn, num_classes=cfg['num_classes'], rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) if os.path.exists(cfg['fasterrcnn_pretrained']): state_dict = torch.load(cfg['fasterrcnn_pretrained']) model.load_state_dict(state_dict) return model
def __init__(self, h_g, h_l, g, k, s, c): super(glimpse_network, self).__init__() self.retina = retina(g, k, s) self.feature_extractor = nn.Sequential( *list(resnet50(pretrained=True).children())[:-1]) #detection feature backbone = torchvision.models.mobilenet_v2(pretrained=True).features backbone.out_channels = 1280 anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), )) roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], output_size=7, sampling_ratio=2) self.detection_model = FasterRCNN( backbone, num_classes=2, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) self.detection_model.eval() # glimpse layer D_in = k * g * g * c D_in = 2048 self.fc1 = nn.Linear(D_in, h_g) # location layer D_in = 2 self.fc2 = nn.Linear(D_in, h_l) self.fc3 = nn.Linear(h_g, h_g + h_l) self.fc4 = nn.Linear(h_l, h_g + h_l)
def get_model( backbone_name="resnet50", detector_name="fasterrcnn", trainable_layers=3, model_ckpt=None, ): """Constructs a fasterrcnn or maskrcnn detector with the given backbone""" num_classes = 2 # 1 class (wheat) + background if model_ckpt: # backbone = resnet_fpn_backbone('resnet101', True) backbone = timm_resnet_fpn_backbone(backbone_name, False, trainable_layers) else: backbone = timm_resnet_fpn_backbone(backbone_name, True, trainable_layers) if detector_name == "fasterrcnn": model = FasterRCNN(backbone, num_classes) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor( in_features, num_classes) elif detector_name == "maskrcnn": model = MaskRCNN(backbone, num_classes) in_features_mask = ( model.roi_heads.mask_predictor.conv5_mask.in_channels) hidden_layer = 256 model.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, hidden_layer, num_classes) else: raise Exception(f"{detector_name} is not supported") if model_ckpt is not None: model.load_state_dict(torch.load(model_ckpt)["model_state_dict"]) print("loaded ckpt") return model
def start_evaluation(test_data_loader: DataLoader, model: FasterRCNN, device: str, epoch: int, logger: Logger, args: argparse.Namespace) -> Tensor: """ Evaluate the model with the test set :param test_data_loader: Data loader for test data: :param model: Model that is being tested :param device: Device for the computation :param epoch: Current epoch :param logger: Logger for logging handling :param args: Arguments :return: """ logger.info(f'Start evaluation after {epoch} epochs') model.eval() scores = [] for idx, result in enumerate(test_data_loader): images = list(image.to(device) for image in result[0]) targets = result[1] with torch.set_grad_enabled(False): outputs = model(images) for output_idx, element in enumerate(outputs): predicted_labels = element['labels'] true_labels = targets[output_idx]['labels'] if len(element['scores']) != 0: scores.append(torch.mean(element['scores'])) if idx % args.print_status: logger.info( f'Scores {element["scores"]} \n' f'Labels predicted: {predicted_labels} Groundtruth labels: {true_labels}' ) avg_score = torch.mean(torch.Tensor(scores)) return avg_score
def train_one_epoch(train_data_loader: DataLoader, model: FasterRCNN, device: str, logger: Logger, optimizer: SGD, args: argparse.Namespace) -> Tuple[float, float]: """ Training the model for one epoch :param train_data_loader: Dataloader for training data :param model: FasterRCNN model :param device: Device for the computation :param logger: Logger for logging handling :param optimizer: Optmizer for the training :param args: Arguments :return: """ start = time.time() model.train() loss_per_iteration = [] for idx, result in enumerate(train_data_loader): images = list(image.to(device) for image in result[0]) targets = result[1] targets = [{k: v.to(device) for k, v in t.items()} for t in targets] optimizer.zero_grad() with torch.set_grad_enabled(True): loss_dict = model(images, targets) losses = sum(loss for loss in loss_dict.values()) losses.backward() optimizer.step() loss_per_iteration.append(losses) loss_doc_str = "".join("{}:{} ".format(key, val) for key, val in loss_dict.items()) if idx % args.print_status == 0: logger.info(f'Iteration: [{idx}/{len(train_data_loader)}]\t' f'Loss: {losses} \t' f'Loss_dict: {loss_doc_str}') epoch_time = time.time() - start return loss_per_iteration, epoch_time
def demo(): # load a pre-trained model for classification and return # only the features backbone = torchvision.models.mobilenet_v2(pretrained=True).features # FasterRCNN needs to know the number of # output channels in a backbone. For mobilenet_v2, it's 1280 # so we need to add it here backbone.out_channels = 1280 # let's make the RPN generate 5 x 3 anchors per spatial # location, with 5 different sizes and 3 different aspect # ratios. We have a Tuple[Tuple[int]] because each feature # map could potentially have different sizes and # aspect ratios anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), )) # let's define what are the feature maps that we will # use to perform the region of interest cropping, as well as # the size of the crop after rescaling. # if your backbone returns a Tensor, featmap_names is expected to # be [0]. More generally, the backbone should return an # OrderedDict[Tensor], and in featmap_names you can choose which # feature maps to use. roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=7, sampling_ratio=2) # put the pieces together inside a FasterRCNN model model = FasterRCNN(backbone, num_classes=2, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) model.eval() x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)] predictions = model(x) print(predictions)
class ResNet50_FasterRCNN: def __init__(self, pretrained=False): # Building our FasterRCNN model for objects detection backbone = resnet_fpn_backbone('resnet50', pretrained=pretrained) num_classes = 4 + 1 anchor_generator = AnchorGenerator(sizes=(40, 60, 150, 200, 250), aspect_ratios=(0.7, 1.0, 1.3)) self.model = FRCNN(backbone, num_classes=num_classes, rpn_anchor_generator=anchor_generator) def train(self): self.model.train() def to(self, device): self.model.to(device) def eval(self): self.model.eval() def parameters(self): return self.model.parameters() def get_state_dict(self): return self.model.state_dict() def set_state_dict(self, state_dict): self.model.load_state_dict(state_dict) def fit_batch(self, images, target): return self.model(images, target) def predict_batch(self, images): return self.model(images)
def __init__(self, backbone_name: str, pretrained: bool = True, finetune: bool = True, num_classes: int = 2): self.__pretrained = pretrained self.__num_classes = num_classes self.__model_name = backbone_name backbone = build_backbone(backbone_name, pretrained, finetune) anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), )) roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], output_size=7, sampling_ratio=2) self.model = FasterRCNN(backbone=backbone, num_classes=num_classes, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) self.params = [p for p in self.model.parameters() if p.requires_grad] self.optimizer = torch.optim.Adam(params=self.params, lr=0.005, weight_decay=0.0005) self.lr_scheduler = torch.optim.lr_scheduler.StepLR( optimizer=self.optimizer, step_size=3, gamma=0.1)
def train_mask(): torch.manual_seed(1) dataset = unimib_data.UNIMIBDataset(get_transform(train=True)) dataset_test = unimib_data.UNIMIBDataset(get_transform(train=False)) indices = torch.randperm(len(dataset)).tolist() dataset = torch.utils.data.Subset(dataset, indices[:-50]) dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:]) # define training and validation data loaders data_loader = torch.utils.data.DataLoader( dataset, batch_size=1, shuffle=True, collate_fn=utils.collate_fn) data_loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=1, shuffle=False, collate_fn=utils.collate_fn) # device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') device = "cpu" num_classes = 74 model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) backbone = torchvision.models.mobilenet_v2(pretrained=True).features for param in backbone: param.requires_grad = False backbone.out_channels = 1280 model.backbone = backbone model = FasterRCNN(backbone, num_classes=74) model = torch.nn.Sequential( model.rpn, model.roi_heads ) print(model) # for param in model.backbone: # param.requires_grad = False # model = get_instance_segmentation_model(num_classes) model.to(device) # construct an optimizer params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) # and a learning rate scheduler which decreases the learning rate by # 10x every 3 epochs lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1) num_epochs = 1 for epoch in range(num_epochs): # train for one epoch, printing every 10 iterations # train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10) # update the learning rate # lr_scheduler.step() # evaluate on the test dataset evaluate(model, data_loader_test, device=device)
def __init__(self, config, logger): self._logger = logger self._threshold = config['threshold'] modelfile = config['model'] self._device = config['device'] # cpu, cuda, cuda:0 backbone = resnet_fpn_backbone('resnet50', False) self._model = FasterRCNN(backbone, 8) # 8 classes checkpoint = torch.load(modelfile, map_location=self._device) self._model.load_state_dict(checkpoint['model_state_dict']) device = torch.device(self._device) self._model.to(device) self._model.eval()
def fasterrcnn_resnet_fpn(backbone_name, pretrained=False, trainable_backbone_layers=None, num_classes=81, pretrained_backbone=True, **kwargs): """ Construct Faster R-CNN with a ResNet-FPN backbone Arguments: backbone_name(str): Name of the backbone. Refer to torchvision.models.resnet.__dict__ for details pretrained(bool, optional): If True, load weights for the detector pretrained on MS COCO. Only ResNet50-FPN is supported for the moment. trainable_backbone_layers(int, optional): Number of trainable (not frozen) resnet layers starting from final block. num_classes(int, optional): Number of target classes. pretrained_backbone(bool, optional): If True, load weights for backbone pre-trained on ImageNet Refer to torchvision.models.detection.FasterRCNN for kwargs """ trainable_backbone_layers = _validate_trainable_layers( pretrained or pretrained_backbone, trainable_backbone_layers, 5, 3) if pretrained and backbone_name == 'resnet50': # no need to download the backbone if pretrained is set pretrained_backbone = False backbone = resnet_fpn_backbone(backbone_name, pretrained_backbone) model = FasterRCNN(backbone, num_classes, **kwargs) if pretrained and backbone_name == 'resnet50': state_dict = models.utils.load_state_dict_from_url( model_urls['fasterrcnn_resnet50_fpn_coco']) if num_classes == 81: # Remove the parameters for the additional classes state_dict['roi_heads.box_predictor.cls_score.weight'] = \ state_dict['roi_heads.box_predictor.cls_score.weight'][KEEP] state_dict['roi_heads.box_predictor.cls_score.bias'] = \ state_dict['roi_heads.box_predictor.cls_score.bias'][KEEP] state_dict['roi_heads.box_predictor.bbox_pred.weight'] = \ state_dict['roi_heads.box_predictor.bbox_pred.weight'][KEEPX4.flatten()] state_dict['roi_heads.box_predictor.bbox_pred.bias'] = \ state_dict['roi_heads.box_predictor.bbox_pred.bias'][KEEPX4.flatten()] model.load_state_dict(state_dict) elif pretrained: print( "WARNING: No pretrained detector on MS COCO with {}.".format( backbone_name), "Proceed with only pretrained backbone on ImageNet.") return model
def __init__(self, backbone=None, output_channels=2, backbone_output_channels=512): super().__init__() if not backbone: b, _ = remove_backbone_head(resnet18(pretrained=False)) backbone = b # ResNet produces 512-length outputs backbone.out_channels = backbone_output_channels self.segmentation_network = FasterRCNN(backbone, num_classes=output_channels)
def get_model_detection(num_classes, model_name="resnet101", aspect_ratios=(0.5, 1.0, 2.0), anchor_sizes=((32, ), (64, ), (128, ), (256, ), (512, )), backbone_chkpt=None, model_chkpt=None, train_backbone=True, **kwargs): """ """ backbone = load_backbone(model_name, backbone_chkpt) anchors = get_anchors(aspect_ratios=aspect_ratios, anchor_sizes=anchor_sizes) model = FasterRCNN(backbone, num_classes, rpn_anchor_generator=anchors, **kwargs) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) if model_chkpt: backbone.load_state_dict(torch.load(model_chkpt, map_location="cpu")) if not train_backbone: for p in model.backbone.parameters(): p.requires_grad_(False) return model
def Model(self, model_name, use_pretrained=True, use_gpu=True): self.system_dict["model_name"] = model_name self.system_dict["use_pretrained"] = use_pretrained self.system_dict["use_gpu"] = use_gpu if (self.system_dict["model_name"] in self.system_dict["model_set_1"]): first_name, second_name = self.system_dict["model_name"].split("_") if (first_name == "faster-rcnn" and second_name == "mobilenet-v2"): backbone = torchvision.models.mobilenet_v2( pretrained=use_pretrained).features backbone.out_channels = 1280 anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), )) roi_pooler = torchvision.ops.MultiScaleRoIAlign( featmap_names=[0], output_size=7, sampling_ratio=2) self.system_dict["local"]["model"] = FasterRCNN( backbone, num_classes=self.system_dict["local"]["num_classes"], rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) self.set_device(use_gpu=self.system_dict["use_gpu"]) self.system_dict["local"]["model"].to( self.system_dict["local"]["device"])
def FasterRCNN_resnext50_32x4d(): net = torchvision.models.resnext50_32x4d(pretrained=True) modules = list(net.children())[:-1] backbone = nn.Sequential(*modules) # backbone = torchvision.models.resnext50_32x4d(pretrained=True).features # test_backbone(backbone) backbone.out_channels = 2048 anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),)) roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=7, sampling_ratio=2) model = FasterRCNN(backbone, num_classes=config.CLASSES, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) return model
def fasterrcnn_mobilenetv2_fpn(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True, **kwargs): """ Construct a Faster RCNN model with mobilenetv2 backbone """ if pretrained: pretrained_backbone = False backbone = torchvision.models.mobilenet_v2( pretrained=pretrained_backbone).features backbone.out_channels = 1280 anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), )) roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], output_size=7, sampling_ratio=2) model = FasterRCNN(backbone, num_classes=num_classes, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) return model
def fasterrcnn_resnet34(num_classes, **kwargs): backbone = resnet_fpn_backbone("resnet34", pretrained=True) model = FasterRCNN(backbone, num_classes, **kwargs) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) return model
def get_model_instance_segmentation2(num_classes): # COCO 에서 미리 학습된 인스턴스 분할 모델을 읽어옵니다 #model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False) #model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False) backbone = torchvision.models.mobilenet_v2(pretrained=False).features #backbone.out_channels = 1 backbone.out_channels = 1280 anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), )) roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], output_size=1, sampling_ratio=2) model = FasterRCNN(backbone, num_classes=num_classes, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) print("mobilenet_v2 call2 - out_channels :1280, 19,540,921") # 분류를 위한 입력 특징 차원을 얻습니다 #in_features = backbone # 미리 학습된 헤더를 새로운 것으로 바꿉니다 #model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) #in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels #hidden_layer = 1 # and replace the mask predictor with a new one #model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, # hidden_layer, # num_classes) return model
def model_predictions_from_loader( model: FasterRCNN, data_loader: DataLoader, device: torch.device, max_n_batches=1, normalization_transform: Normalize = Normalize([1, 1, 1], [1, 1, 1]) ) -> List[Tuple[BoxedExample, BoxedExample]]: model = model.to(device).eval() input_examples = [] output_examples = [] with torch.no_grad(): for i, (img, labels) in enumerate(data_loader): if i >= max_n_batches: break outputs = model([i.to(device) for i in img]) input_exs = [ model_input_to_boxed_example((im, l), normalization_transform) for (im, l) in zip(img, labels) ] output_exs = [ i.replace(boxes=o["boxes"].cpu().numpy()) for (i, o) in zip(input_exs, outputs) ] input_examples.extend(input_exs) output_examples.extend(output_exs) return list(zip(input_examples, output_examples))
def faster_rcnn_alt_backbone(): # 特徴量のみを取得 backbone = torchvision.models.mobilentet_v2(pretrained=True).features # FasterRCNNはバックボーンからの出力チャネル数を知る必要がある # mobilenet_v2の出力は1280 backbone.out_channels = 1280 # RPN: Resion Proposal Networkに空間ごとにアンカーを生成。 # 5つのサイズと3つのアスペクト比が存在することを意味します。 # 特徴マップごとに異なるサイズとアスペクト比となる可能性があるので, # Tuple[Tuple[int]] という形式で指定します。 anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),), aspect_ratios=((0.5, 1.0, 2.0),)) # 関心領域のトリミングを実行するために使用する特徴マップ(featmap_names)と # 画像の大きさをもとに戻したあとのトリミングのサイズ(output_size)を適宜しましょう。 # バックボーンがテンソルを返す場合、featmap_nameは[0]になっているはずです。 # もう少し一般化して説明すると、バックボーンはOrderedDict[Tensor]を返すことになるので、 # featmap_namesで使用する特徴マップを選択できます。 roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], output_size=7, sampling_ratio=2) model = FasterRCNN(backbone, num_classes=2, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) return model
def __init__(self, pretrained_path): super().__init__() self.pretrained_path = pretrained_path self.output_dim = 800 * 800 # ------------------ # PRE-TRAINED MODEL # ------------------ ae = BasicAE.load_from_checkpoint(pretrained_path) ae.freeze() self.backbone = ae.encoder self.backbone.c3_only = True self.backbone.out_channels = 32 # ------------------ # FAST RCNN # ------------------ anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), )) roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=7, sampling_ratio=2) self.fast_rcnn = FasterRCNN(self.backbone, num_classes=9, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) # for unfreezing encoder later self.frozen = True
def __init__(self, cfg, **kwargs): super(WheatDetector, self).__init__() self.backbone = fpn_backbone(pretrained=True) self.base = FasterRCNN(self.backbone, num_classes=cfg.MODEL.NUM_CLASSES, **kwargs) self.base.roi_heads.fastrcnn_loss = self.fastrcnn_loss
def fasterrcnn_resnet(backbone_name, num_classes=91, pretrained_backbone=True, **kwargs): """ Construct Faster R-CNN with a ResNet backbone Arguments: backbone_name(str): Name of the backbone. Refer to torchvision.models.resnet.__dict__ for details num_classes(int, optional): Number of target classes, default: 91(COCO) pretrained_backbone(bool, optional): If True, load weights for backbone pre-trained on ImageNet Refer to torchvision.models.detection.FasterRCNN for kwargs """ backbone, res5 = resnet_backbone(backbone_name, pretrained_backbone) backbone.out_channels = 1024 box_head = nn.Sequential(res5, nn.AdaptiveAvgPool2d((1, 1))) box_predictor = FastRCNNPredictor(2048, num_classes) model = FasterRCNN(backbone, box_head=box_head, box_predictor=box_predictor, **kwargs) return model
def create_model(num_classes=3): backbone = torchvision.models.mobilenet_v2(pretrained=True).features backbone.out_channels = 1280 anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), )) roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=["0"], output_size=7, sampling_ratio=2) model = FasterRCNN(backbone, num_classes=num_classes, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) # model = torchvision.models.detection.fasterrcnn_resnet50_fpn( # pretrained=False, # pretrained_backbone=True, # num_classes=num_classes # ) # model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) # # get number of input features for the classifier # in_features = model.roi_heads.box_predictor.cls_score.in_features # # replace the pre-trained head with a new one # model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) return model
def __call__(self, classes=3, sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), )): from torchvision.models.detection.rpn import AnchorGenerator import torchvision from torchvision.models.detection import FasterRCNN from torchvision.models.detection.faster_rcnn import TwoMLPHead, FastRCNNPredictor # load a pre-trained model for classification and return # only the features backbone = torchvision.models.squeezenet1_1(pretrained=True).features # FasterRCNN needs to know the number of # output channels in a backbone. For squeezenet1_1, it's 512 # so we need to add it here backbone.out_channels = 512 anchor_generator = AnchorGenerator(sizes=sizes, aspect_ratios=aspect_ratios) roi_out_size = 7 roi_pooler = torchvision.ops.MultiScaleRoIAlign( featmap_names=['0'], output_size=roi_out_size, sampling_ratio=2) representation_size = 256 # Scaled down from 1024 in original implementation. # allows to reduce considerably the number of parameters box_head = TwoMLPHead(backbone.out_channels * roi_out_size**2, representation_size) box_predictor = FastRCNNPredictor(representation_size, classes) model = FasterRCNN(backbone, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler, box_head=box_head, box_predictor=box_predictor) return model
def __init__(self): super(DetectionModel, self).__init__() # Creating the feature extractor self.feature_extractor = DigitModel() self.feature_extractor.out_channels = 512 # The region proposal network will generate 5x3 anchors per spatial location # Each with 5 different sizes and 3 different aspect ratios self.anchor_generator = AnchorGenerator(sizes=((16, 32, 64, 128, 256), ), aspect_ratios=((0.5, 1.0, 2.0), )) # Extract and align features from each proposed region on each scale. # These features would then be sent to the final fully connected layers for classification and bbox regression self.roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], output_size=7, sampling_ratio=2) # We define 11 classes, 1 for each digit 0-9 and an additional class for the background self.num_classes = 11 # Compile these modules within a PyTorch implementation of Faster RCNN self.detector = FasterRCNN(self.feature_extractor, num_classes=self.num_classes, rpn_anchor_generator=self.anchor_generator, box_roi_pool=self.roi_pooler, min_size=384, max_size=512, box_score_thresh=0.5, box_nms_thresh=0.5, box_detections_per_img=5)
def __init__(self, num_classes): super(TrashLiteModel, self).__init__() backbone = torchvision.models.mobilenet_v2(pretrained=True).features backbone.out_channels = 1280 anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), )) roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], output_size=7, sampling_ratio=2) # put the pieces together inside a FasterRCNN model self.model = FasterRCNN(backbone, num_classes=num_classes, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler)
def resnet(n_classes, backbone_n='resnet50', pretrained_backbone=True): # FIX doesn't work """ Faster-RCNN with :param n_classes: number of classes of Fast-RCNN-Predictor :param backbone_n: name of backbone which will extract feature maps :param pretrained_backbone: if True, return pretrained backbone of resnet :return: instance of FasterRCNN """ if backbone_n not in names: raise Exception('Wrong backbone name') backbone = torchvision.models.detection.backbone_utils.resnet_fpn_backbone( backbone_n, pretrained=pretrained_backbone) # set out channels for FasterRCNN backbone.out_channels = 256 # define custom anchors for RPN anchor_generator = AnchorGenerator(sizes=cfg.ANCHOR.SCALES, aspect_ratios=cfg.ANCHOR.RATIOS) roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0], output_size=7, sampling_ratio=2) # define model model = FasterRCNN(backbone, num_classes=n_classes, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler) return model
def fasterrcnn_resnet_fpn(pretrained=False, progress=True, resnet='resnet50', num_classes=91, pretrained_backbone=True, **kwargs): if pretrained: # no need to download the backbone if pretrained is set pretrained_backbone = False backbone = resnet_fpn_backbone(resnet, pretrained_backbone) model = FasterRCNN(backbone, num_classes, **kwargs) if pretrained: target_url = model_urls['fasterrcnn_' + resnet + '_fpn_coco'] state_dict = load_state_dict_from_url(target_url, progress=progress) model.load_state_dict(state_dict) return model