def __init__(self, dictionary=None): super(MaskRCNN, self).__init__() self.dictionary = dictionary self.input_size = [512, 512] self.dummy_input = torch.zeros(1, 3, self.input_size[0], self.input_size[1]) self.num_classes = len(self.dictionary) self.category = [v for d in self.dictionary for v in d.keys()] self.weight = [ d[v] for d in self.dictionary for v in d.keys() if v in self.category ] # load an instance segmentation model pre-trained pre-trained on COCO self.model = maskrcnn_resnet50_fpn(pretrained=True) # get number of input features for the classifier in_features = self.model.roi_heads.box_predictor.cls_score.in_features # replace the pre-trained head with a new one self.model.roi_heads.box_predictor = FastRCNNPredictor( in_features, self.num_classes) # now get the number of input features for the mask classifier in_features_mask = self.model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # and replace the mask predictor with a new one self.model.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, hidden_layer, self.num_classes)
def __init__(self, n_classes=10, exp_dict=None): super().__init__() self.n_classes = n_classes self.exp_dict = exp_dict self.model = model = torchvision.models.detection.maskrcnn_resnet50_fpn( pretrained_backbone=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, n_classes) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # and replace the mask predictor with a new one model.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, hidden_layer, n_classes) params = [p for p in model.parameters() if p.requires_grad] self.opt = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) self.lr_scheduler = torch.optim.lr_scheduler.StepLR(self.opt, step_size=3, gamma=0.1) warmup_factor = 1. / 1000 warmup_iters = 1000 self.lr_scheduler = warmup_lr_scheduler(self.opt, warmup_iters, warmup_factor)
def get_pretrained_instance_segmentation_maskrcnn( num_categories: int, hidden_layer: int = 256 ) -> Module: """ load an instance segmentation model pre-trained on COCO :param num_categories: :param hidden_layer: :return: """ model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True) # get the number of input features for the classifier replace the pre-trained head with a new one model.roi_heads.box_predictor = FastRCNNPredictor( model.roi_heads.box_predictor.cls_score.in_features, num_categories ) # now get the number of input features for the mask classifier and replace the mask predictor with a new one model.roi_heads.mask_predictor = MaskRCNNPredictor( model.roi_heads.mask_predictor.conv5_mask.in_channels, hidden_layer, num_categories, ) return model
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 get_model_instance_segmentation(num_classes): ''' This is the conventional model which is based on Faster R-CNN Note that to use this model you must install regular pytorch package (instead of from ofekp branch) and use '--model-name faster' in the arguments The correct way to install torchvision will be: pip uninstall torchvision pip install torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html To restore back to EfficientDet use: pip uninstall torchvision pip install git+https://github.com/ofekp/vision.git ''' print("Using Faster-RCNN detection model") # load an instance segmentation model pre-trained pre-trained on COCO model = torchvision.models.detection.maskrcnn_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) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # 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 select_model(self, model_name): self.model_name = model_name if model_name == 'faster_rcnn': 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, self.num_classes) elif model_name == 'mask_rcnn': model = torchvision.models.detection.maskrcnn_resnet50_fpn( pretrained=True) # get number of input features for the classifier in_features = model.roi_heads.box_predictor.cls_score.in_features dim_reduced = 2 # not sure what this should be yet model.roi_heads.box_predictor = MaskRCNNPredictor( in_features, dim_reduced, self.num_classes) else: raise ValueError('Not a valid model name') self.model = model
def get_maskrcnn_model(num_classes=None, num_hidden_layer=256, finetuning=True): """ Get the pretrained mask rcnn model for finetuning or unmodified Args: num_classes ([int], optional): number of class for custom dataset. Defaults to None. num_hidden_layer (int, optional): number of hidden units. Defaults to 256. finetuning (bool, optional): finetuning for the custom dataset. Defaults to True. Returns: [model]: [mask rcnn model] """ # load mask rcnn model pretrained on COCO model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True) if finetuning: # number of input features for the classifier in_features = model.roi_heads.box_predictor.cls_score.in_features # replace the pretrained head model.roi_heads.box_predictor = FastRCNNPredictor( in_features, num_classes) # the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = num_hidden_layer # replace the mask predictor model.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, hidden_layer, num_classes) return model
def get_model(model_name): if model_name == "MaskRCNN": num_classes = 600 # load an instance segmentation model pre-trained pre-trained on COCO model = torchvision.models.detection.maskrcnn_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) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # and replace the mask predictor with a new one model.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, hidden_layer, num_classes) else: print("Model name not recognized.") end(1) return model
def get_mask_rcnn(num_classes: int): """Create model Mask-RCNN Parameters ---------- num_classes: int number of classes Returns ------- torch.nn.Module Mask R-CNN model with backbone of ResNet-50 """ # Load an instance segmentation model pretrained on COCO model = torchvision.models.detection.maskrcnn_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 pretrained head with a new one model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) # Now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # 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 __init__(self, hparams): super().__init__() # Гиперпараметры модели в pytorch_lightning self.hparams = hparams # Mask-RCNN self.mask_rcnn = maskrcnn_resnet50_fpn(pretrained_backbone=True, pretrained=True) # Обновляем выход для предсказания номерного знака num_classes = 2 in_features = self.mask_rcnn.roi_heads.box_predictor.cls_score.in_features self.mask_rcnn.roi_heads.box_predictor = FastRCNNPredictor( in_features, num_classes) self.mask_rcnn.roi_heads.mask_predictor = MaskRCNNPredictor( 256, 256, num_classes) # Разрешаем обновлять только определенные параметры for parameter in self.mask_rcnn.parameters(): parameter.requires_grad = False for parameter in self.mask_rcnn.backbone.fpn.parameters(): parameter.requires_grad = True for parameter in self.mask_rcnn.rpn.parameters(): parameter.requires_grad = True for parameter in self.mask_rcnn.roi_heads.parameters(): parameter.requires_grad = True
def get_instance_segmentation_model(num_classes, model_name='maskrcnn_resnet50_fpn'): # load a pre-trained model for classification # and return only the features if model_name.startswith('efficientnet'): backbone = EfficientNet.from_pretrained(model_name, num_classes=num_classes, include_top=False) # number of output channels backbone.out_channels = int( round_filters(1280, backbone._global_params)) model = MaskRCNN(backbone, num_classes) else: # load an instance segmentation model pre-trained on COCO model = torchvision.models.detection.__dict__[model_name]( pretrained=True) # get the 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) if model_name.startswith('mask') or model_name.startswith('efficientnet'): # now get the number of input features for the mask classifier in_features_mask = \ model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # 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 get_model_instance_segmentation(num_classes): # load an instance segmentation model pre-trained pre-trained on COCO model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True) # FL = FocalLoss(gamma=2, alpha=0.25, magnifier=1) # FL = FocalLoss(gamma=2, alpha=0.5) # FL_wrapped = functools.partial(maskrcnn_loss_focal, focal_loss_func=FL) # RoIHeads_loss_customized.set_customized_loss( # model.roi_heads, maskrcnn_loss_customized=FL_wrapped # ) # RoIHeads_loss_customized.update_forward_func(model.roi_heads) # 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) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # 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 get_model_instance_segmentation(num_classes): # load an instance segmentation model pre-trained pre-trained on COCO model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True) anchor_generator = AnchorGenerator( sizes=tuple([(4, 8, 16, 32, 64, 128, 256, 512) for _ in range(5)]), aspect_ratios=tuple([(0.25, 0.5, 1.0, 2.0) for _ in range(5)])) model.rpn.anchor_generator = anchor_generator # 256 because that's the number of features that FPN returns model.rpn.head = RPNHead(256, anchor_generator.num_anchors_per_location()[0]) # 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) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # 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 get_model_instance_segmentation(num_classes, hidden_layer: int = 256): """ :param num_classes: :type num_classes: :return: :rtype: """ # load an instance segmentation model pre-trained pre-trained on COCO model = torchvision.models.detection.maskrcnn_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) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels # 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 __init__(self, data, backbone=None, pretrained_path=None): super().__init__(data, backbone) self._backbone = models.resnet50 #if not self._check_backbone_support(self._backbone): # raise Exception (f"Enter only compatible backbones from {', '.join(self.supported_backbones)}") self._code = instance_detector_prf model = models.detection.maskrcnn_resnet50_fpn(pretrained=True, min_size=data.chip_size) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, data.c) 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, data.c) self.learn = Learner(data, model, loss_func=mask_rcnn_loss) self.learn.callbacks.append(train_callback(self.learn)) self.learn.model = self.learn.model.to(self._device) # fixes for zero division error when slice is passed self.learn.layer_groups = split_model_idx(self.learn.model, [28]) self.learn.create_opt(lr=3e-3) if pretrained_path is not None: self.load(pretrained_path)
def load_frcnn(object_specific=False): ''' Loads faster rcnn with an RPN and resnet101 backbone network, pretrained on COCO detection dataset. param object_specific: bool, True if want ''' if object_specific: # Read the dataset object files to know how many outputs there should be # num_clases = read file with open(OBJ_LIST, 'r') as f: ids = f.read().splitlines() num_classes = len(ids) else: num_classes = 2 # Load an instance segmentation model pre-trained on COCO model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True) # get the 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) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # 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 ChooseModel(input, n_classes, freeze=False): if input == "faster": print(f"model is Faster RCNN resnet 50") model = getFastRCNNResnet50Fpn(pretrained=True) if freeze: for param in model.parameters(): param.requires_grad = False in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor( in_features, n_classes) elif input == "mask": print(f"model is Mask RCNN") model = getMaskRCNNResnet50Fpn(pretrained=True) if freeze: for param in model.parameters(): param.requires_grad = False in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor( in_features, n_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, n_classes) return model
def get_mask_rcnn(num_classes, max_instances, backbone="resnet101"): # load an instance segmentation model pre-trained pre-trained on COCO if backbone == "resnet50": print("**************Adding Resnet 50 backbone***************") model = torchvision.models.detection.maskrcnn_resnet50_fpn( pretrained=True, box_detections_per_img=max_instances) else: bb = resnet_fpn_backbone(backbone, False) model = MaskRCNN(bb, num_classes=91, box_detections_per_img=max_instances) # 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) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # 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 __init__(self, device, finetune=False, num_classes=2): """ :args device: torch.device, CUDA or CPU :option finetune: If you set it True, you can select the number of the output classes num_classes: the number of output classes. """ torch.cuda.empty_cache() self.device = device # load a model pre-trained pre-trained on COCO self.model = torchvision.models.detection.maskrcnn_resnet50_fpn( pretrained=True) # replace the classifier with a new one, that has # num_classes which is user-defined if finetune: # get number of input features for the classifier in_features = self.model.roi_heads.box_predictor.cls_score.in_features # replace the pre-trained head with a new one self.model.roi_heads.box_predictor = MaskRCNNPredictor( in_features, num_classes) self.model = self.model.to(self.device) self.model.eval()
def get_model_instance_segmentation_custom1(num_classes): # COCO 에서 미리 학습된 인스턴스 분할 모델을 읽어옵니다 model = 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 = 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( "maskrcnn_resnet50_fpn custom call1 - resnet : 24,743,507 mobilenet : 87,366,291 squeezenet : 33,161,683 densnet : 43,702,739, resnet basicblock 3*3 -> 1*1 : 20,549,203 / basic : 20,543,571 / basicblock con1 : 20,195,411 / 채널 : 강제로 128 지정시 13,033,555 / 128 all 변경 : 9,465,555 " ) # 분류를 위한 입력 특징 차원을 얻습니다 in_features = model.roi_heads.box_predictor.cls_score.in_features # 미리 학습된 헤더를 새로운 것으로 바꿉니다 model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 128 # 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 get_detector_model(): model = models.detection.maskrcnn_resnet50_fpn( pretrained=True, pretrained_backbone=True, progress=True, num_classes=91, ) num_classes = 2 in_features = model.roi_heads.box_predictor.cls_score.in_features box_predictor = FastRCNNPredictor(in_features, num_classes) model.roi_heads.box_predictor = box_predictor mask_predictor = MaskRCNNPredictor(256, 256, num_classes) model.roi_heads.mask_predictor = mask_predictor # Заморозим все слои кроме последних for param in model.parameters(): param.requires_grad = True for param in model.backbone.fpn.parameters(): param.requires_grad = True for param in model.rpn.parameters(): param.requires_grad = True for param in model.roi_heads.parameters(): param.requires_grad = True return model
def get_model_instance_segmentation(num_classes, use_my_model=False): """ Load pre-trained model. Replace the box and mask predictors with new predictors with whose outputs are adequate to the `num_classes`. """ # load an instance segmentation model pre-trained pre-trained on COCO if use_my_model is True: model = my_maskrcnn_resnet50_fpn(pretrained=True) else: model = torchvision.models.detection.maskrcnn_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) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # 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 get_instance_segmentation_model(num_classes, backbone, dropout=False): # load an instance segmentation model where backbone is pretrained ImageNet backbone = resnet_fpn_backbone(backbone, pretrained=True) model = MaskRCNN(backbone, num_classes) if dropout: # add drop out after FC layer of box head resolution = model.roi_heads.box_roi_pool.output_size[0] representation_size = 1024 model.roi_heads.box_head = TwoMLPHead( backbone.out_channels * resolution**2, representation_size) # add drop out in mask head mask_layers = (256, 256, 256, 256) mask_dilation = 1 model.roi_heads.mask_head = MaskRCNNHeads(backbone.out_channels, mask_layers, mask_dilation) # get the 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) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # 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 _tune_mask_predictor(model: nn.Module, num_classes: int) -> nn.Module: """ Tune mask predictor in the model. """ # get the number of input features of mask predictor from the pretrained model in_features = model.roi_heads.mask_predictor.conv5_mask.in_channels # replace the mask predictor with a new one model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features, 256, num_classes) return model
def get_model_instance_segmentation(num_classes): model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True) in_featrues = model.roi_heads.box_predictor.cls_score.in_features 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) return model
def _get_model(arch, n_class, pretrained=True, h=256): model = arch(pretrained) # For the classifier in_feats = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_feats, n_class) # For the mask in_chnns = model.roi_heads.mask_predictor.conv5_mask.in_channels model.roi_heads.mask_predictor = MaskRCNNPredictor(in_chnns, 256, n_class) return model
def get_object_detection_model(num_classes, mtype='Resnet50_FPN'): # load a model pre-trained on COCO if mtype == 'Resnet50_FPN': model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True) # get the 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) if mtype == 'MobileNetV3_largeFPN': model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn( pretrained=True) # get the 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) if mtype == 'MobileNetV3_largeFPN_320': # Low resolution network model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_320_fpn( pretrained=True) # get the 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) if mtype == 'MaskRCNN': # load an instance segmentation model pre-trained on COCO model = torchvision.models.detection.maskrcnn_resnet50_fpn( pretrained=True) # get the 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) # now get the number of input features for the mask classifier in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 # and replace the mask predictor with a new one model.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, hidden_layer, num_classes) if mtype == 'YOLOv5x': model = torch.hub.load('ultralytics/yolov5', 'yolov5x') else: raise ValueError('Model type not handled') return model
def __init__(self, num_classes=2, hidden_size=256): super().__init__() self.model_ft = maskrcnn_resnet50_fpn(pretrained=True) in_features = self.model_ft.roi_heads.box_predictor.cls_score.in_features self.model_ft.roi_heads.box_predictor = FastRCNNPredictor( in_features, num_classes) in_features_mask = self.model_ft.roi_heads.mask_predictor.conv5_mask.in_channels self.model_ft.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, hidden_size, num_classes) for param in self.model_ft.parameters(): param.requires_grad = True
def create_model(num_classes): model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 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) return model
def get_model_instance_segmentation(num_classes): model = maskrcnn_resnet50_fpn(pretrained=True) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) mask_predictor_in_channels = model.roi_heads.mask_predictor.conv5_mask.in_channels model.roi_heads.mask_predictor = MaskRCNNPredictor( mask_predictor_in_channels, mask_dim_reduced=256, num_classes=num_classes) return model