import torch


def get_lastlayers_model_weights(mdl, is_mask_class_specific = False, is_fc_lastlayer = False):
    """
    Returns the weights at the head of the ROI predictor -> classification and bounding box regressor weights + bias
    :returns cls_weights, bbox_pred_weights, bbox_pred_bias
    :rtype: tuple
    """
    if isinstance(mdl, str):
        mdl = torch.load(mdl)['model']
    elif isinstance(mdl, dict):
        if 'model' in mdl:
            mdl = mdl['model']
    else:
        raise TypeError('Expected dict or str')
    cls_weights = mdl['roi_heads.box_predictor.cls_score.weight']
    bbox_pred_weights = mdl['roi_heads.box_predictor.bbox_pred.weight']
    bbox_pred_bias = mdl['roi_heads.box_predictor.bbox_pred.bias']

    return cls_weights, bbox_pred_weights, bbox_pred_bias