class ClassifierLocalizer(nn.Module):
def __init__(self, model_name, num_classes=2):

    super(ClassifierLocalizer, self).__init__()

    self.num_classes = num_classes

  

    # create cnn model

    model = getattr(models, model_name)(pretrained=True)

  

    # remove fc layers and add a new fc layer

    num_features = model.fc.in_features

    model.fc = nn.Linear(num_features, 6) # classifier + localizer

    self.model = model



def forward(self, x):

    x = self.model(x)                    # extract features from CNN

    scores = x[:, :self.num_classes]     # class scores

    coords = x[:, self.num_classes:]     # coordinates

    return scores, coords