import torchvision
from torchvision.models.detection.faster_rcnn import [FastRCNNPredictor](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "torch.nn.Module")
from torchvision.models.detection.mask_rcnn import [MaskRCNNPredictor](https://docs.pytorch.org/docs/stable/generated/torch.nn.Sequential.html#torch.nn.Sequential "torch.nn.Sequential")


def get_model_instance_segmentation(num_classes):
    # load an instance segmentation model pre-trained on COCO
    model = [torchvision.models.detection.maskrcnn_resnet50_fpn](https://docs.pytorch.org/vision/stable/models/generated/torchvision.models.detection.maskrcnn_resnet50_fpn.html#torchvision.models.detection.maskrcnn_resnet50_fpn "torchvision.models.detection.maskrcnn_resnet50_fpn")(weights="DEFAULT")

    # 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](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "torch.nn.Module")(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](https://docs.pytorch.org/docs/stable/generated/torch.nn.Sequential.html#torch.nn.Sequential "torch.nn.Sequential")(
        in_features_mask,
        hidden_layer,
        num_classes
    )

    return model