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)}
