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
