def load_model(model_path): model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True) num_classes = 3 in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) new_aspect_ratios = (0.5, 1.0, 2.0) new_aspect_ratios = (new_aspect_ratios, ) * len( model.rpn.anchor_generator.sizes) model.rpn.anchor_generator.aspect_ratios = new_aspect_ratios model.load_state_dict(torch.load(model_path)) print("FFRA model load complete.") return model
def build_model(self) -> nn.Module: model = fasterrcnn_resnet50_fpn(pretrained=False, pretrained_backbone=False) # Replace the classifier with a new two-class classifier. There are # only two "classes": pedestrian and background. in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, self.num_classes) load_exp = self.context.get_hparam("load_from_experiment") if load_exp > 0: os.environ["INFERENCE"] = "True" pretrained = Determined().get_experiment(load_exp).top_checkpoint().load(path=self.download_directory) os.environ["INFERENCE"] = "False" model.load_state_dict(pretrained.state_dict()) return model
def get_model( model_name, num_classes, backbone, fpn, pretrained, pretrained_backbone, trainable_backbone_layers, anchor_generator, **kwargs, ): if backbone is None: # Constructs a model with a ResNet-50-FPN backbone when no backbone is specified. if model_name == "fasterrcnn": model = _models[model_name]( pretrained=pretrained, pretrained_backbone=pretrained_backbone, trainable_backbone_layers=trainable_backbone_layers, ) in_features = model.roi_heads.box_predictor.cls_score.in_features head = FastRCNNPredictor(in_features, num_classes) model.roi_heads.box_predictor = head else: model = _models[model_name](pretrained=pretrained, pretrained_backbone=pretrained_backbone) model.head = RetinaNetHead( in_channels=model.backbone.out_channels, num_anchors=model.head.classification_head.num_anchors, num_classes=num_classes, **kwargs ) else: backbone_model, num_features = backbone_and_num_features( backbone, fpn, pretrained_backbone, trainable_backbone_layers, **kwargs, ) backbone_model.out_channels = num_features if anchor_generator is None: anchor_generator = AnchorGenerator( sizes=((32, 64, 128, 256, 512), ), aspect_ratios=((0.5, 1.0, 2.0), ) ) if not hasattr(backbone_model, "fpn") else None if model_name == "fasterrcnn": model = FasterRCNN(backbone_model, num_classes=num_classes, rpn_anchor_generator=anchor_generator) else: model = RetinaNet(backbone_model, num_classes=num_classes, anchor_generator=anchor_generator) return model
def train(self): # load a model pre-trained pre-trained on COCO model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True) # replace the classifier with a new one, that has # num_classes which is user-defined num_classes = 2 # 1 class (car) + background # 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) model.to(self.device) #self.load_model(self.config['path']['pretrain_model_FRCNN_LR_LR'], model) # 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) data_loader, data_loader_test, _, _, _, _, _, _, _ = self.data_loaders( ) # let's train it for 10 epochs num_epochs = 1000 for epoch in range(num_epochs): # train for one epoch, printing every 10 iterations train_one_epoch(model, optimizer, data_loader, self.device, epoch, print_freq=10) # update the learning rate lr_scheduler.step() # evaluate on the test dataset evaluate_base(model, data_loader_test, device=self.device) if epoch % 1 == 0: self.save_model(model, 'FRCNN_LR_LR', epoch)
def get_model(): # load a model pre-trained pre-trained on COCO model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True) # replace the classifier with a new one, that has num_classes which is user-defined num_classes = 2 # 1 class (person) + background # 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 get_faster_rcnn_model(pretrained=True, trainable_backbone_layers=3, pretrained_backbone=True): if pretrained: model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True, trainable_backbone_layers=trainable_backbone_layers) in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2) return model else: return torchvision.models.detection.fasterrcnn_resnet50_fpn( num_classes=2, trainable_backbone_layers=trainable_backbone_layers, pretrained_backbone=pretrained_backbone)
def __init__(self): backbone = resnet_fpn_backbone('resnet50', False) super(Net, self).__init__(backbone, num_classes=91) self.num_classes = 2 # get the number of input features for the classifier in_features = self.roi_heads.box_predictor.cls_score.in_features # replace the pre-trained head with a new one self.roi_heads.box_predictor = FastRCNNPredictor( in_features, self.num_classes) in_features_mask = self.roi_heads.mask_predictor.conv5_mask.in_channels hidden_layer = 256 self.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, hidden_layer, self.num_classes)
def obtain_model(): # 1 class (wheat) + background num_classes = 2 # define the model 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 __init__(self, num_classes): super(RCNNLoss, self).__init__() import torchvision from torchvision.models.detection.faster_rcnn import FastRCNNPredictor # load an instance segmentation model pre-trained pre-trained on COCO self.model = torchvision.models.detection.fasterrcnn_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, num_classes) self.model.load_state_dict( torch.load("checkpoints/RCNN/rcnn_checkpoint_epoch_40.pt")) print("Finished loading the RCNN")
def get_model(num_classes, pretrained=True): model = maskrcnn_resnet50_fpn(pretrained=pretrained) 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 print("IFM", in_features_mask) 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): # 加载在COCO上预训练的预训练的实例分割模型 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_maskrcnn_model_instance(num_classes): # 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 get_model_instance_segmentation(num_classes): model = torchvision.models.detection.fasterrcnn_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_pretrain_model(model_name, num_classes): if model_name == 'faster_rcnn': model = torchvision.models.detection.fasterrcnn_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) elif model_name == 'mask_rcnn': 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) else: backbone = torchvision.models.mobilenet_v2().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, box_roi_pool=roi_pooler, rpn_anchor_generator=anchor_generator) return model
def _binary_fasterrcnn_resnet50(device: str, model_weights: str): """Loads a pre-trained FasterRCNN with a ResNet50 backbone and 2-class output.""" model = fasterrcnn_resnet50_fpn() # Get the number of input features for the top layer: in_features = model.roi_heads.box_predictor.cls_score.in_features # Replace the head with a new one for just 2 classes: background and logo model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2) # Define the computing device explicitly: checkpoint = torch.load(model_weights, map_location=device) # checkpoint = torch.load(model_weights) model.load_state_dict(checkpoint["state_dict"]) return model.eval().to(device)
def model( num_classes: int, backbone: Optional[nn.Module] = None, remove_internal_transforms: bool = True, pretrained: bool = True, **mask_rcnn_kwargs ) -> nn.Module: """MaskRCNN model implemented by torchvision. # Arguments num_classes: Number of classes. backbone: Backbone model to use. Defaults to a resnet50_fpn model. remove_internal_transforms: The torchvision model internally applies transforms like resizing and normalization, but we already do this at the `Dataset` level, so it's safe to remove those internal transforms. pretrained: Argument passed to `maskrcnn_resnet50_fpn` if `backbone is None`. By default it is set to True: this is generally used when training a new model (transfer learning). `pretrained = False` is used during inference (prediction) for cases where the users have their own pretrained weights. **mask_rcnn_kwargs: Keyword arguments that internally are going to be passed to `torchvision.models.detection.mask_rcnn.MaskRCNN`. # Return A Pytorch `nn.Module`. """ if backbone is None: model = maskrcnn_resnet50_fpn( pretrained=pretrained, pretrained_backbone=pretrained, **mask_rcnn_kwargs ) in_features_box = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features_box, num_classes) in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels model.roi_heads.mask_predictor = MaskRCNNPredictor( in_channels=in_features_mask, dim_reduced=256, num_classes=num_classes ) backbone_param_groups = resnet_fpn.param_groups(model.backbone) else: model = MaskRCNN(backbone, num_classes=num_classes, **mask_rcnn_kwargs) backbone_param_groups = backbone.param_groups() patch_param_groups(model=model, backbone_param_groups=backbone_param_groups) if remove_internal_transforms: remove_internal_model_transforms(model) return model
def predict(): # 准备网络 # load a model; pre-trained on COCO(以下4句为pytorch官方教程例子) model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True) num_classes = 2 # 1 class (wheat) + background # 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) # 指定gpu device = torch.device( 'cuda:1') if torch.cuda.is_available() else torch.device('cpu') model.load_state_dict(torch.load("models_frc_0629/frc_0629_55.pth")) model = model.to(device) print('模型载入成功') save_dir = 'predict_frc_0629/' model.eval() with torch.no_grad(): for image_name_batch, test_images in test_dataloader: images = list(image.to(device) for image in test_images) outputs = model(images) #print(image_name) for i, image in enumerate(images): detection_threshold = 0.5 # (3,800,800) #sample = image.cpu().numpy() sample = image.permute(1, 2, 0).cpu().numpy() #image = image.cpu().numpy() boxes = outputs[i]['boxes'].data.cpu().numpy() scores = outputs[i]['scores'].data.cpu().numpy() boxes = boxes[scores >= detection_threshold].astype(np.int32) fig, ax = plt.subplots(1, 1, figsize=(8, 8)) for box in boxes: cv2.rectangle(sample, (box[0], box[1]), (box[2], box[3]), (220, 0, 0), 2) ax.set_axis_off() ax.imshow(sample) plt.savefig(save_dir + image_name_batch[i]) #cv2.imwrite(save_dir + image_name_batch[i], image) print(image_name_batch[i] + ': 已保存')
def __init__( self, learning_rate: float = 0.0001, num_classes: int = 91, backbone: Optional[str] = None, fpn: bool = True, pretrained: bool = False, pretrained_backbone: bool = True, trainable_backbone_layers: int = 3, **kwargs: Any, ): """ Args: learning_rate: the learning rate num_classes: number of detection classes (including background) backbone: Pretained backbone CNN architecture. fpn: If True, creates a Feature Pyramind Network on top of Resnet based CNNs. pretrained: if true, returns a model pre-trained on COCO train2017 pretrained_backbone: if true, returns a model with backbone pre-trained on Imagenet trainable_backbone_layers: number of trainable resnet layers starting from final block """ if not _TORCHVISION_AVAILABLE: # pragma: no cover raise ModuleNotFoundError("You want to use `torchvision` which is not installed yet.") super().__init__() self.learning_rate = learning_rate self.num_classes = num_classes self.backbone = backbone if backbone is None: self.model = fasterrcnn_resnet50_fpn( pretrained=pretrained, pretrained_backbone=pretrained_backbone, trainable_backbone_layers=trainable_backbone_layers, ) in_features = self.model.roi_heads.box_predictor.cls_score.in_features self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, self.num_classes) else: backbone_model = create_fasterrcnn_backbone( self.backbone, fpn, pretrained_backbone, trainable_backbone_layers, **kwargs, ) self.model = torchvision_FasterRCNN(backbone_model, num_classes=num_classes, **kwargs)
def __init__(self, num_classes, pretrained = True): super(MaskRCNN, self).__init__() self.model = tv.models.detection.maskrcnn_resnet50_fpn(pretrained = pretrained) # get the 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, 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, num_classes)
def prepare_model(): model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=False) num_classes = 2 # 1 class (wheat) + background # 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) model.load_state_dict( torch.load(MODEL_SAVE_PATH, map_location=torch.device('cpu'))) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) model.eval() return model
def get_fast_rcnn_for_fine_tune(num_class: int) -> torch.nn.Module: # load a model pre-trained pre-trained on COCO model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True) # replace the classifier with a new one, that has # num_classes which is user-defined num_classes = num_class # 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 define_modelo_deteccion(): # Definiendo el modelo original y cambiando sus parametros model_FasterRCNN = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True) model_FasterRCNN_in_features = model_FasterRCNN.roi_heads.box_predictor.cls_score.in_features model_FasterRCNN.roi_heads.box_predictor = FastRCNNPredictor( model_FasterRCNN_in_features, 3) model_FasterRCNN.roi_heads.detections_per_img = 500 factor = 1 model_FasterRCNN.rpn.anchor_generator.sizes = ( (16*factor,), (32*factor,), (64*factor,), (128*factor,), (256*factor,)) model_FasterRCNN.transform.min_size = (1232,) model_FasterRCNN.transform.max_size = 1232 return model_FasterRCNN
def _create_model(self): self._net = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True) in_features = self._net.roi_heads.box_predictor.cls_score.in_features self._net.roi_heads.box_predictor = FastRCNNPredictor( in_features, self._num_classes ) in_features_mask = self._net.roi_heads.mask_predictor.conv5_mask.in_channels self._net.roi_heads.mask_predictor = MaskRCNNPredictor( in_features_mask, self._hidden_layer, self._num_classes ) self._net.to(self.device)
def __init__(self, arch, classnames, pretrained=False): super().__init__() assert arch == 'fasterrcnn_resnet50_fpn' pretrained = pretrained == 'coco' if classnames is None: from ..data.coco import COCO_CLASSNAMES self.classnames = COCO_CLASSNAMES else: self.classnames = classnames self.num_classes = len(self.classnames) + 1 self.model = fasterrcnn_resnet50_fpn(pretrained=pretrained) if self.num_classes != 91: in_features = self.model.roi_heads.box_predictor.cls_score.in_features self.model.roi_heads.box_predictor = FastRCNNPredictor( in_features, self.num_classes)
def create_detection_model(rpn_anchor_generator=None): model = torchvision.models.detection.maskrcnn_resnet50_fpn( pretrained=True, pretrained_backbone=True, rpn_anchor_generator=rpn_anchor_generator, image_mean=[0.485, 0.456, 0.406], image_std=[0.229, 0.224, 0.225], ) # replace the classifier with a new one, that has # num_classes which is user-defined num_classes = 2 # 1 class (person) + background # 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 model(): # load the COCO pre-trained model # we will keep the image size to 1024 pixels instead of the original 800, # this will ensure better training and testing results, although it may... # ... increase the training time (a tarde-off) model = torchvision.models.detection.fasterrcnn_resnet50_fpn( pretrained=True, min_size=1024) # one class is pneumonia, and the other is background num_classes = 2 # get the input features for the classifier in_features = model.roi_heads.box_predictor.cls_score.in_features # replace pre-trained head with our features head # the head layer will classify the images based on our data input features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) return model
def prediction(image, XML): #JPEG -> TENSOR if torch.cuda.is_available(): map_location=lambda storage, loc: storage.cuda() else: map_location='cpu' #checkpoint = torch.load('./data/SeaShips/test_model.pth', map_location=map_location) tensor = image_to_tensor(image) #Load pre-trained model model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) model.roi_heads.box_predictor = FastRCNNPredictor(1024,6) model.load_state_dict(torch.load('./data/SeaShips/test_model.pth', map_location=map_location)) model = model.to(device) #get model prediction and XML data in JSON my_list = JSON_output(tensor, XML, model) return my_list
def get_model(name: str = 'fasterrcnn_resnet50_fpn'): if name == 'fasterrcnn_resnet50_fpn': from torchvision.models.detection import fasterrcnn_resnet50_fpn # load a model; pre-trained on COCO model = fasterrcnn_resnet50_fpn(pretrained=True) else: return None # ['Knife', 'Gun', 'Wrench', 'Pliers', 'Scissors'] + background num_classes = 6 # 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 main(): # make folder # Dataset print(f'Preparing Dataset....COCO Dataset') # train_path = r'C:\Users\gjust\Documents\Github\data\COCO\train2017' # train_ann = r'C:\Users\gjust\Documents\Github\data\COCO\annotations\instances_train2017.json' test_path = r'C:\Users\gjust\Documents\Github\data\COCO\val2017' test_ann = r'C:\Users\gjust\Documents\Github\data\COCO\annotations\instances_val2017.json' class_list = ['person', 'dog', 'car'] transform = A.Compose([ A.Resize(width=512, height=512), A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2() ], bbox_params=A.BboxParams(format='coco', label_fields=['class_labels'])) # trainset = MyCocoLimit(root=train_path, annFile=train_ann, class_list=class_list, transform=transform) testset = MyCocoLimit(root=test_path, annFile=test_ann, class_list=class_list, transform=transform) # train_loader = DataLoader(trainset, batch_size=4, shuffle=True, collate_fn=collate_fn) test_loader = DataLoader(testset, batch_size=4, shuffle=False, collate_fn=collate_fn) # GPU device = 'cuda' if torch.cuda.is_available() else 'cpu' if device == 'cuda': torch.backends.cudnn.benchmark = True print(f'Using {device}') # Model print(f'Preparing Model....Faster RCNN') model = fasterrcnn_resnet50_fpn(pretrained=True) num_classes = len(class_list) + 1 in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) model.to(device) # resuming # Optimizer optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9) # Training print('Training....') EPOCH = 50 for e in range(EPOCH): train(model, test_loader, optimizer, device, EPOCH, e) test(model, test_loader, device, EPOCH, e)
def __init__(self, classes=None, device=None): """Initializes a machine learning model for object detection. Models are built on top of PyTorch's `pre-trained models <https://pytorch.org/docs/stable/torchvision/models.html>`_, specifically the Faster R-CNN ResNet-50 FPN, but allow for fine-tuning to predict on custom classes/labels. :param classes: (Optional) A list of classes/labels for the model to predict. If none given, uses the default classes specified `here <https://pytorch.org/docs/stable/torchvision/models.html #object-detection-instance-segmentation-and-person-keypoint-detection>`_. Defaults to None. :type classes: list or None :param device: (Optional) The device on which to run the model, such as the CPU or GPU. See `here <https://pytorch.org/docs/stable/tensor_attributes.html#torch-device>`_ for details on specifying the device. Defaults to the GPU if available and the CPU if not. :type device: torch.device or None **Example**:: >>> from detecto.core import Model >>> model = Model(['dog', 'cat', 'bunny']) """ self._device = device if device else config['default_device'] # Load a model pre-trained on COCO self._model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True) if classes: # Get the 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 (note: +1 because of the __background__ class) self._model.roi_heads.box_predictor = FastRCNNPredictor(in_features, len(classes) + 1) self._disable_normalize = False else: classes = config['default_classes'] self._disable_normalize = True self._model.to(self._device) # Mappings to convert from string labels to ints and vice versa self._classes = ['__background__'] + classes self._int_mapping = {label: index for index, label in enumerate(self._classes)}